pi-statusline/index.ts

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");
    },
  });
}