pi-statusline/index.ts

import { AuthStorage, type AuthCredential, type ExtensionAPI, type ExtensionContext, type ExtensionCommandContext, type ReadonlyFooterDataProvider, type SessionStartEvent, type SessionShutdownEvent, type AgentEndEvent } from "@mariozechner/pi-coding-agent";
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 REFRESH_MS = 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: ExtensionContext): Provider | undefined {
  const provider = ctx.model?.provider as string | undefined;
  if (provider === "openai-codex" || provider === "zai" || provider === "anthropic") return provider;
  return undefined;
}

function credentialString(raw: unknown, key: string): string | undefined {
  return isRecord(raw) && typeof raw[key] === "string" ? raw[key] : undefined;
}

function buildUsageCredential(raw: AuthCredential | undefined, apiKey: string): UsageCredential {
  if (raw?.type === "oauth") {
    return {
      type: "oauth",
      accessToken: apiKey,
      refreshToken: raw.refresh,
      expiresAt: raw.expires,
      accountId: credentialString(raw, "accountId"),
      email: credentialString(raw, "email"),
      metadata: raw,
    };
  }

  return {
    type: "api_key",
    apiKey,
    accountId: credentialString(raw, "accountId"),
    email: credentialString(raw, "email"),
    metadata: raw,
  };
}

async function readPiCredential(authStorage: AuthStorage, provider: Provider): Promise<UsageCredential | undefined> {
  authStorage.reload();
  const apiKey = await authStorage.getApiKey(provider);
  if (!apiKey) return undefined;
  return buildUsageCredential(authStorage.get(provider), apiKey);
}

async function forceRefreshPiCredential(authStorage: AuthStorage, provider: Provider): Promise<UsageCredential> {
  authStorage.reload();
  const raw = authStorage.get(provider);
  const oauthProvider = authStorage.getOAuthProviders().find(candidate => candidate.id === provider);
  if (raw?.type !== "oauth" || !oauthProvider) throw new Error("login expired");

  // Refresh Provider OAuth Token
  const refreshed = await oauthProvider.refreshToken(raw);
  authStorage.set(provider, { type: "oauth", ...refreshed });
  return buildUsageCredential(authStorage.get(provider), oauthProvider.getApiKey(refreshed));
}

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: ExtensionContext, footerData: ReadonlyFooterDataProvider, 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: ExtensionContext | undefined;
  let requestRender: (() => void) | undefined;
  const statusbarState: StatusbarState = {};
  const authStorage = AuthStorage.create();

  function updateThinkingLevel() {
    try {
      statusbarState.thinkingLevel = pi.getThinkingLevel?.() ?? "off";
    } catch {
      statusbarState.thinkingLevel = "off";
    }
  }

  let rerenderScheduled = false;
  function rerender(ctx: ExtensionContext) {
    latestCtx = ctx;
    if (rerenderScheduled) return;
    rerenderScheduled = true;

    // Debounce Render - Coalesce rapid rerender calls into a single pass.
    queueMicrotask(() => {
      rerenderScheduled = false;
      requestRender?.();
      const current = latestCtx;
      if (current?.hasUI) current.ui.setStatus("pi-statusbar", undefined);
    });
  }

  function installFooter(ctx: ExtensionContext) {
    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: ExtensionContext, force = false) {
    latestCtx = ctx;
    const provider = activeProvider(ctx);
    lastProvider = provider;
    if (!provider) {
      statusbarState.report = undefined;
      statusbarState.error = undefined;
      rerender(ctx);
      return;
    }

    const minRefreshMs = REFRESH_MS;
    if (!force && statusbarState.report?.provider === provider && Date.now() - statusbarState.report.fetchedAt < minRefreshMs) {
      rerender(ctx);
      return;
    }

    const credential = await readPiCredential(authStorage, provider);
    if (!credential) {
      statusbarState.report = { provider, fetchedAt: Date.now(), limits: [] };
      statusbarState.error = "not logged in";
      rerender(ctx);
      return;
    }

    inFlight?.abort();
    const controller = new AbortController();
    inFlight = controller;

    try {
      let activeCredential = credential;
      const usageCtx = {
        fetch: globalThis.fetch.bind(globalThis),
        logger: {
          debug: () => undefined,
          warn: () => undefined,
        },
      };
      const fetchParams = () => ({
        provider,
        credential: activeCredential,
        baseUrl: ctx.model?.baseUrl,
        signal: controller.signal,
      });

      let report: UsageReport | null;
      try {
        report = await usageProviders[provider].fetchUsage(fetchParams(), usageCtx);
      } catch (error) {
        if (!(error instanceof Error) || error.message !== "unauthorized") throw error;
        activeCredential = await forceRefreshPiCredential(authStorage, provider);
        report = await usageProviders[provider].fetchUsage(fetchParams(), usageCtx);
      }

      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: SessionStartEvent, ctx: ExtensionContext) => {
    updateThinkingLevel();
    installFooter(ctx);
    await refresh(ctx, true);

    // Usage Refresh Timer
    timer = setInterval(() => void refresh(latestCtx ?? ctx), REFRESH_MS);
  });

  pi.on("session_shutdown", async (_event: SessionShutdownEvent, ctx: ExtensionContext) => {
    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: AgentEndEvent, ctx: ExtensionContext) => {
    updateThinkingLevel();
    await refresh(ctx);
  });

  pi.registerCommand("refresh-usage", {
    description: "Refresh account usage for the active provider",
    handler: async (_args: string, ctx: ExtensionCommandContext) => {
      await refresh(ctx, true);
      ctx.ui.notify(usageSummary(statusbarState.report, statusbarState.error), statusbarState.error ? "warning" : "info");
    },
  });
}