feat: add configurable pi statusbar extension
This commit is contained in:
238
index.ts
Normal file
238
index.ts
Normal file
@@ -0,0 +1,238 @@
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import { readFileSync } from "node:fs";
|
||||
import { homedir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { statusbarConfig } from "./config";
|
||||
import { contextModule, costModule, directoryModule, modelModule, thinkingModule } from "./modules/basic";
|
||||
import { usageModule } from "./modules/usage";
|
||||
import { moduleType, renderRow } from "./render";
|
||||
import { claudeUsageProvider } from "./usage/claude";
|
||||
import { openaiCodexUsageProvider } from "./usage/openai-codex";
|
||||
import { zaiUsageProvider } from "./usage/zai";
|
||||
import type { ModuleContext, ModuleSpec, RenderedModule, StatusbarState } from "./types";
|
||||
import type { Provider, UsageCredential, UsageProvider, UsageReport } from "./usage";
|
||||
|
||||
const AUTH_PATH = join(homedir(), ".pi", "agent", "auth.json");
|
||||
const REFRESH_MS = 60_000;
|
||||
const ANTHROPIC_REFRESH_MS = 15 * 60_000;
|
||||
|
||||
const usageProviders: Record<Provider, UsageProvider> = {
|
||||
"openai-codex": openaiCodexUsageProvider,
|
||||
zai: zaiUsageProvider,
|
||||
anthropic: claudeUsageProvider,
|
||||
};
|
||||
|
||||
function isRecord(value: unknown): value is Record<string, unknown> {
|
||||
return typeof value === "object" && value !== null && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function activeProvider(ctx: any): Provider | undefined {
|
||||
const provider = ctx.model?.provider as string | undefined;
|
||||
if (provider === "openai-codex" || provider === "zai" || provider === "anthropic") return provider;
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function readPiCredential(provider: Provider): UsageCredential | undefined {
|
||||
try {
|
||||
const auth = JSON.parse(readFileSync(AUTH_PATH, "utf8"));
|
||||
if (!isRecord(auth) || !isRecord(auth[provider])) return undefined;
|
||||
|
||||
const raw = auth[provider];
|
||||
if (raw.type === "oauth") {
|
||||
return {
|
||||
type: "oauth",
|
||||
accessToken: typeof raw.access === "string" ? raw.access : undefined,
|
||||
refreshToken: typeof raw.refresh === "string" ? raw.refresh : undefined,
|
||||
expiresAt: typeof raw.expires === "number" ? raw.expires : undefined,
|
||||
accountId: typeof raw.accountId === "string" ? raw.accountId : undefined,
|
||||
email: typeof raw.email === "string" ? raw.email : undefined,
|
||||
metadata: raw,
|
||||
};
|
||||
}
|
||||
|
||||
if (raw.type === "api_key") {
|
||||
return {
|
||||
type: "api_key",
|
||||
apiKey: typeof raw.key === "string" ? raw.key : undefined,
|
||||
accountId: typeof raw.accountId === "string" ? raw.accountId : undefined,
|
||||
email: typeof raw.email === "string" ? raw.email : undefined,
|
||||
metadata: raw,
|
||||
};
|
||||
}
|
||||
} catch {}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function renderModule(moduleCtx: ModuleContext, spec: ModuleSpec): RenderedModule {
|
||||
switch (moduleType(spec)) {
|
||||
case "directory":
|
||||
return directoryModule(moduleCtx);
|
||||
case "context":
|
||||
return contextModule(moduleCtx);
|
||||
case "model":
|
||||
return modelModule(moduleCtx);
|
||||
case "thinking":
|
||||
return thinkingModule(moduleCtx);
|
||||
case "cost":
|
||||
return costModule(moduleCtx);
|
||||
case "usage":
|
||||
return usageModule(moduleCtx, spec);
|
||||
default:
|
||||
return { text: "" };
|
||||
}
|
||||
}
|
||||
|
||||
function renderFooter(ctx: any, footerData: any, state: StatusbarState, width: number, theme?: any): string[] {
|
||||
const moduleCtx: ModuleContext = { ctx, footerData, state, config: statusbarConfig };
|
||||
return statusbarConfig.rows.map(row => renderRow({
|
||||
left: (row.left ?? []).map(spec => renderModule(moduleCtx, spec)),
|
||||
center: (row.center ?? []).map(spec => renderModule(moduleCtx, spec)),
|
||||
right: (row.right ?? []).map(spec => renderModule(moduleCtx, spec)),
|
||||
}, width, statusbarConfig.separator, theme));
|
||||
}
|
||||
|
||||
function usageSummary(report: UsageReport | undefined, error: string | undefined): string {
|
||||
if (error) return `usage: ${error}`;
|
||||
if (!report) return "usage loading";
|
||||
return "usage refreshed";
|
||||
}
|
||||
|
||||
export default function piStatusbarExtension(pi: ExtensionAPI) {
|
||||
let timer: ReturnType<typeof setInterval> | undefined;
|
||||
let inFlight: AbortController | undefined;
|
||||
let lastProvider: Provider | undefined;
|
||||
let latestCtx: any;
|
||||
let requestRender: (() => void) | undefined;
|
||||
const statusbarState: StatusbarState = {};
|
||||
|
||||
function updateThinkingLevel() {
|
||||
try {
|
||||
statusbarState.thinkingLevel = pi.getThinkingLevel?.() ?? "off";
|
||||
} catch {
|
||||
statusbarState.thinkingLevel = "off";
|
||||
}
|
||||
}
|
||||
|
||||
function rerender(ctx: any) {
|
||||
latestCtx = ctx;
|
||||
requestRender?.();
|
||||
if (ctx.hasUI) ctx.ui.setStatus("pi-statusbar", undefined);
|
||||
}
|
||||
|
||||
function installFooter(ctx: any) {
|
||||
if (!ctx.hasUI) return;
|
||||
latestCtx = ctx;
|
||||
ctx.ui.setFooter((tui: any, theme: any, footerData: any) => {
|
||||
requestRender = () => tui.requestRender();
|
||||
return {
|
||||
invalidate() {},
|
||||
render(width: number): string[] {
|
||||
return renderFooter(latestCtx ?? ctx, footerData, statusbarState, width, theme);
|
||||
},
|
||||
dispose() {
|
||||
requestRender = undefined;
|
||||
},
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
async function refresh(ctx: any, force = false) {
|
||||
latestCtx = ctx;
|
||||
const provider = activeProvider(ctx);
|
||||
lastProvider = provider;
|
||||
if (!provider) {
|
||||
statusbarState.report = undefined;
|
||||
statusbarState.error = undefined;
|
||||
rerender(ctx);
|
||||
return;
|
||||
}
|
||||
|
||||
const minRefreshMs = provider === "anthropic" ? ANTHROPIC_REFRESH_MS : REFRESH_MS;
|
||||
if (!force && statusbarState.report?.provider === provider && Date.now() - statusbarState.report.fetchedAt < minRefreshMs) {
|
||||
rerender(ctx);
|
||||
return;
|
||||
}
|
||||
|
||||
const credential = readPiCredential(provider);
|
||||
if (!credential) {
|
||||
statusbarState.report = { provider, fetchedAt: Date.now(), limits: [] };
|
||||
statusbarState.error = "not logged in";
|
||||
rerender(ctx);
|
||||
return;
|
||||
}
|
||||
|
||||
inFlight?.abort();
|
||||
inFlight = new AbortController();
|
||||
statusbarState.report = undefined;
|
||||
statusbarState.error = undefined;
|
||||
rerender(ctx);
|
||||
|
||||
try {
|
||||
const report = await usageProviders[provider].fetchUsage({
|
||||
provider,
|
||||
credential,
|
||||
baseUrl: ctx.model?.baseUrl,
|
||||
signal: inFlight.signal,
|
||||
}, {
|
||||
fetch: globalThis.fetch.bind(globalThis),
|
||||
logger: {
|
||||
debug: () => undefined,
|
||||
warn: () => undefined,
|
||||
},
|
||||
});
|
||||
|
||||
statusbarState.report = report ?? { provider, fetchedAt: Date.now(), limits: [] };
|
||||
statusbarState.error = report ? undefined : "unavailable";
|
||||
} catch (error) {
|
||||
statusbarState.report = { provider, fetchedAt: Date.now(), limits: [] };
|
||||
statusbarState.error = error instanceof Error ? error.message : String(error);
|
||||
}
|
||||
|
||||
if (lastProvider === provider) rerender(ctx);
|
||||
}
|
||||
|
||||
pi.on("session_start", async (_event, ctx) => {
|
||||
updateThinkingLevel();
|
||||
installFooter(ctx);
|
||||
await refresh(ctx, true);
|
||||
|
||||
// Usage Refresh Timer
|
||||
timer = setInterval(() => void refresh(latestCtx ?? ctx), REFRESH_MS);
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async (_event, ctx) => {
|
||||
if (timer) clearInterval(timer);
|
||||
timer = undefined;
|
||||
inFlight?.abort();
|
||||
inFlight = undefined;
|
||||
requestRender = undefined;
|
||||
if (ctx.hasUI) {
|
||||
ctx.ui.setStatus("pi-statusbar", undefined);
|
||||
ctx.ui.setFooter(undefined);
|
||||
}
|
||||
});
|
||||
|
||||
pi.on("model_select", async (_event, ctx) => {
|
||||
updateThinkingLevel();
|
||||
await refresh(ctx, true);
|
||||
});
|
||||
|
||||
pi.on("thinking_level_select", async (event, ctx) => {
|
||||
statusbarState.thinkingLevel = event.level;
|
||||
rerender(ctx);
|
||||
});
|
||||
|
||||
pi.on("agent_end", async (_event, ctx) => {
|
||||
updateThinkingLevel();
|
||||
await refresh(ctx);
|
||||
});
|
||||
|
||||
pi.registerCommand("refresh-usage", {
|
||||
description: "Refresh account usage for the active provider",
|
||||
handler: async (_args, ctx) => {
|
||||
await refresh(ctx, true);
|
||||
ctx.ui.notify(usageSummary(statusbarState.report, statusbarState.error), statusbarState.error ? "warning" : "info");
|
||||
},
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user