@latentminds/pi-quotas 0.2.4 → 0.3.1

package/src/extensions/token-status/index.ts

@@ -0,0 +1,241 @@
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+  ThemeColor,
+} from "@mariozechner/pi-coding-agent";
+import {
+  QUOTAS_CONFIG_UPDATED_EVENT,
+  QUOTAS_EXTENSIONS_REGISTER_EVENT,
+  QUOTAS_EXTENSIONS_REQUEST_EVENT,
+  type QuotasConfigUpdatedPayload,
+  configLoader,
+} from "../../config.js";
+import { aggregateAllSessions, formatCost } from "../../lib/session-tokens.js";
+
+const EXTENSION_ID = "pi-quotas-token-status";
+const REFRESH_INTERVAL_MS = 60_000;
+
+/** Go tier limits (approximate, from docs) */
+const GO_LIMITS = {
+  rolling5h: 12, // $12 per 5 hours
+  weekly: 30, // $30 per week
+  monthly: 60, // $60 per month
+} as const;
+
+interface RollingWindowCosts {
+  rolling5h: number;
+  weekly: number;
+  monthly: number;
+}
+
+function computeWindowTimestamps(now: Date): {
+  since5h: Date;
+  sinceWeekly: Date;
+  sinceMonthly: Date;
+} {
+  return {
+    since5h: new Date(now.getTime() - 5 * 60 * 60 * 1000),
+    sinceWeekly: new Date(now.getTime() - 7 * 24 * 60 * 60 * 1000),
+    sinceMonthly: new Date(now.getTime() - 30 * 24 * 60 * 60 * 1000),
+  };
+}
+
+async function computeRollingCosts(cwd?: string): Promise<RollingWindowCosts> {
+  const now = new Date();
+  const { since5h, sinceWeekly, sinceMonthly } = computeWindowTimestamps(now);
+
+  // Run all three aggregations in parallel
+  const [result5h, resultWeekly, resultMonthly] = await Promise.all([
+    aggregateAllSessions({ cwd, since: since5h }),
+    aggregateAllSessions({ cwd, since: sinceWeekly }),
+    aggregateAllSessions({ cwd, since: sinceMonthly }),
+  ]);
+
+  return {
+    rolling5h: result5h.totals.costTotal,
+    weekly: resultWeekly.totals.costTotal,
+    monthly: resultMonthly.totals.costTotal,
+  };
+}
+
+function formatWindowForStatus(
+  theme: ExtensionContext["ui"]["theme"],
+  label: string,
+  cost: number,
+  limit: number,
+): string {
+  const percent = limit > 0 ? Math.round((cost / limit) * 100) : 0;
+  const costStr = formatCost(cost);
+
+  // Color by usage level
+  let color: ThemeColor;
+  if (percent >= 90) {
+    color = "error";
+  } else if (percent >= 70) {
+    color = "warning";
+  } else {
+    color = "dim";
+  }
+
+  return `${theme.fg(color, `${label}:`)}${theme.fg("accent", costStr)}`;
+}
+
+function formatTokenStatus(
+  theme: ExtensionContext["ui"]["theme"],
+  costs: RollingWindowCosts,
+): string {
+  const parts = [
+    formatWindowForStatus(theme, "5h", costs.rolling5h, GO_LIMITS.rolling5h),
+    formatWindowForStatus(theme, "wk", costs.weekly, GO_LIMITS.weekly),
+    formatWindowForStatus(theme, "mo", costs.monthly, GO_LIMITS.monthly),
+  ];
+  return parts.join(" · ");
+}
+
+function isGoProvider(provider: string | undefined): boolean {
+  if (!provider) return false;
+  return provider === "opencode-go" || provider.startsWith("opencode-go/");
+}
+
+function createTokenStatusRefresher() {
+  let refreshTimer: ReturnType<typeof setInterval> | undefined;
+  let activeContext: ExtensionContext | undefined;
+  let lastCosts: RollingWindowCosts | undefined;
+  let inFlight = false;
+  let queued = false;
+
+  async function update(ctx: ExtensionContext): Promise<void> {
+    if (!ctx.hasUI) return;
+    if (inFlight) {
+      queued = true;
+      return;
+    }
+    inFlight = true;
+    try {
+      const costs = await computeRollingCosts(ctx.cwd);
+      if (!ctx.hasUI) return;
+      lastCosts = costs;
+      const status = formatTokenStatus(ctx.ui.theme, costs);
+      ctx.ui.setStatus(EXTENSION_ID, status);
+    } catch {
+      ctx.ui.setStatus(
+        EXTENSION_ID,
+        ctx.ui.theme.fg("warning", "token tracking unavailable"),
+      );
+    } finally {
+      inFlight = false;
+      if (queued) {
+        queued = false;
+        void update(ctx);
+      }
+    }
+  }
+
+  return {
+    async refreshFor(ctx: ExtensionContext): Promise<void> {
+      activeContext = ctx;
+      if (!isGoProvider(ctx.model?.provider)) {
+        ctx.ui.setStatus(EXTENSION_ID, undefined);
+        return;
+      }
+      await update(ctx);
+    },
+    start(): void {
+      if (refreshTimer) clearInterval(refreshTimer);
+      refreshTimer = setInterval(() => {
+        if (activeContext) void update(activeContext);
+      }, REFRESH_INTERVAL_MS);
+      refreshTimer.unref?.();
+    },
+    stop(ctx?: ExtensionContext): void {
+      if (refreshTimer) clearInterval(refreshTimer);
+      refreshTimer = undefined;
+      activeContext = undefined;
+      lastCosts = undefined;
+      ctx?.ui.setStatus(EXTENSION_ID, undefined);
+    },
+    renderLast(ctx: ExtensionContext): boolean {
+      if (!lastCosts || !ctx.hasUI) return false;
+      ctx.ui.setStatus(
+        EXTENSION_ID,
+        formatTokenStatus(ctx.ui.theme, lastCosts),
+      );
+      return true;
+    },
+  };
+}
+
+export default async function (pi: ExtensionAPI) {
+  await configLoader.load();
+  const refresher = createTokenStatusRefresher();
+  let enabled = configLoader.getConfig().tokenStatus;
+  let currentContext: ExtensionContext | undefined;
+
+  function scheduleRefresh(ctx: ExtensionContext): void {
+    void refresher.refreshFor(ctx).catch(() => {
+      if (ctx.hasUI)
+        ctx.ui.setStatus(
+          EXTENSION_ID,
+          ctx.ui.theme.fg("warning", "token tracking unavailable"),
+        );
+    });
+  }
+
+  pi.events.on(QUOTAS_CONFIG_UPDATED_EVENT, (data: unknown) => {
+    const config = (data as QuotasConfigUpdatedPayload).config;
+    enabled = config.tokenStatus;
+    if (!enabled) {
+      refresher.stop(currentContext);
+      return;
+    }
+    if (currentContext) {
+      refresher.start();
+      scheduleRefresh(currentContext);
+    }
+  });
+
+  pi.on("session_start", (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) return;
+    if (!isGoProvider(ctx.model?.provider)) {
+      ctx.ui.setStatus(EXTENSION_ID, undefined);
+      return;
+    }
+    refresher.start();
+    scheduleRefresh(ctx);
+  });
+
+  pi.on("turn_end", (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) return;
+    if (!isGoProvider(ctx.model?.provider)) return;
+    scheduleRefresh(ctx);
+  });
+
+  pi.on("model_select", (_event, ctx) => {
+    currentContext = ctx;
+    if (!enabled) {
+      refresher.stop(ctx);
+      return;
+    }
+    if (!isGoProvider(ctx.model?.provider)) {
+      refresher.stop(ctx);
+      return;
+    }
+    refresher.start();
+    scheduleRefresh(ctx);
+  });
+
+  pi.on("session_shutdown", async (_event, ctx) => {
+    currentContext = undefined;
+    refresher.stop(ctx);
+  });
+
+  pi.events.on(QUOTAS_EXTENSIONS_REQUEST_EVENT, () => {
+    if (configLoader.getConfig().tokenStatus) {
+      pi.events.emit(QUOTAS_EXTENSIONS_REGISTER_EVENT, {
+        feature: "tokenStatus",
+      });
+    }
+  });
+}

package/src/extensions/token-status/provider-detection.test.ts

@@ -0,0 +1,30 @@
+import { describe, expect, it } from "vitest";
+
+// Extract isGoProvider logic for testing
+function isGoProvider(provider: string | undefined): boolean {
+  if (!provider) return false;
+  return provider === "opencode-go" || provider.startsWith("opencode-go/");
+}
+
+describe("isGoProvider", () => {
+  it("detects opencode-go provider", () => {
+    expect(isGoProvider("opencode-go")).toBe(true);
+  });
+
+  it("detects opencode-go/ prefixed providers", () => {
+    expect(isGoProvider("opencode-go/anthropic")).toBe(true);
+    expect(isGoProvider("opencode-go/openai")).toBe(true);
+  });
+
+  it("rejects non-go providers", () => {
+    expect(isGoProvider("opencode")).toBe(false);
+    expect(isGoProvider("anthropic")).toBe(false);
+    expect(isGoProvider("openai")).toBe(false);
+    expect(isGoProvider("github-copilot")).toBe(false);
+  });
+
+  it("rejects undefined/empty", () => {
+    expect(isGoProvider(undefined)).toBe(false);
+    expect(isGoProvider("")).toBe(false);
+  });
+});

package/src/extensions/usage-status/index.test.ts

@@ -0,0 +1,182 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import usageStatusExtension from "./index.js";
+import { fetchProviderQuotas } from "../../lib/quotas.js";
+
+vi.mock("../../config.js", () => ({
+  QUOTAS_CONFIG_UPDATED_EVENT: "quotas:config:updated",
+  QUOTAS_EXTENSIONS_REGISTER_EVENT: "quotas:extensions:register",
+  QUOTAS_EXTENSIONS_REQUEST_EVENT: "quotas:extensions:request",
+  configLoader: {
+    load: vi.fn(async () => undefined),
+    getConfig: vi.fn(() => ({
+      configVersion: "test",
+      quotasCommand: true,
+      providerCommands: true,
+      usageStatus: true,
+      quotaWarnings: true,
+      deferToSynthetic: true,
+    })),
+  },
+}));
+
+vi.mock("../../lib/quotas.js", () => ({
+  isSupportedProvider: (provider: string | undefined) => provider === "anthropic",
+  fetchProviderQuotas: vi.fn(async () => ({
+    success: true,
+    data: { provider: "anthropic", windows: [] },
+  })),
+}));
+
+const STALE_CONTEXT_ERROR =
+  "This extension ctx is stale after session replacement or reload.";
+
+type EventHandler = (event: unknown, ctx: ExtensionContext) => unknown;
+
+function createFakePi() {
+  const extensionHandlers = new Map<string, EventHandler[]>();
+  const eventBusHandlers = new Map<string, Array<(data: unknown) => void>>();
+
+  const pi = {
+    on(event: string, handler: EventHandler) {
+      const handlers = extensionHandlers.get(event) ?? [];
+      handlers.push(handler);
+      extensionHandlers.set(event, handlers);
+    },
+    events: {
+      on(channel: string, handler: (data: unknown) => void) {
+        const handlers = eventBusHandlers.get(channel) ?? [];
+        handlers.push(handler);
+        eventBusHandlers.set(channel, handlers);
+        return () => {
+          const current = eventBusHandlers.get(channel) ?? [];
+          eventBusHandlers.set(channel, current.filter((entry) => entry !== handler));
+        };
+      },
+      emit(channel: string, data: unknown) {
+        for (const handler of eventBusHandlers.get(channel) ?? []) handler(data);
+      },
+    },
+  } as unknown as ExtensionAPI;
+
+  return {
+    pi,
+    async emitExtensionEvent(event: string, ctx: ExtensionContext) {
+      for (const handler of extensionHandlers.get(event) ?? []) {
+        await handler({ type: event, reason: "test" }, ctx);
+      }
+    },
+    emitBusEvent(channel: string, data: unknown) {
+      pi.events.emit(channel, data);
+    },
+    listenerCount(channel: string) {
+      return eventBusHandlers.get(channel)?.length ?? 0;
+    },
+  };
+}
+
+function createContext(provider: string) {
+  let stale = false;
+  const setStatus = vi.fn(() => {
+    if (stale) throw new Error(STALE_CONTEXT_ERROR);
+  });
+
+  const ctx = {
+    get hasUI() {
+      if (stale) throw new Error(STALE_CONTEXT_ERROR);
+      return true;
+    },
+    get model() {
+      if (stale) throw new Error(STALE_CONTEXT_ERROR);
+      return { provider };
+    },
+    modelRegistry: { authStorage: {} },
+    ui: {
+      theme: { fg: (_color: string, text: string) => text },
+      setStatus,
+    },
+  } as unknown as ExtensionContext;
+
+  return {
+    ctx,
+    setStale() {
+      stale = true;
+    },
+    setStatus,
+  };
+}
+
+afterEach(() => {
+  vi.useRealTimers();
+  vi.clearAllMocks();
+});
+
+describe("usage-status extension lifecycle", () => {
+  it("ignores interval refreshes for stale session contexts", async () => {
+    vi.useFakeTimers();
+    const { pi, emitExtensionEvent } = createFakePi();
+    const { ctx, setStale } = createContext("unsupported-provider");
+
+    await usageStatusExtension(pi);
+    await emitExtensionEvent("session_start", ctx);
+
+    setStale();
+
+    expect(() => vi.advanceTimersByTime(60_000)).not.toThrow();
+    await vi.runOnlyPendingTimersAsync();
+  });
+
+  it("does not throw when event-bus callbacks see a stale session context", async () => {
+    const { pi, emitExtensionEvent, emitBusEvent } = createFakePi();
+    const { ctx, setStale } = createContext("synthetic");
+
+    await usageStatusExtension(pi);
+    await emitExtensionEvent("session_start", ctx);
+
+    setStale();
+
+    expect(() => {
+      emitBusEvent("synthetic:extensions:register", { feature: "usageStatus" });
+      emitBusEvent("quotas:config:updated", {
+        config: { usageStatus: true, deferToSynthetic: true },
+      });
+    }).not.toThrow();
+  });
+
+  it("unsubscribes event-bus listeners during session shutdown", async () => {
+    const { pi, emitExtensionEvent, listenerCount } = createFakePi();
+    const { ctx } = createContext("unsupported-provider");
+
+    await usageStatusExtension(pi);
+    expect(listenerCount("quotas:config:updated")).toBe(1);
+    expect(listenerCount("synthetic:extensions:register")).toBe(1);
+    expect(listenerCount("quotas:extensions:request")).toBe(1);
+
+    await emitExtensionEvent("session_shutdown", ctx);
+
+    expect(listenerCount("quotas:config:updated")).toBe(0);
+    expect(listenerCount("synthetic:extensions:register")).toBe(0);
+    expect(listenerCount("quotas:extensions:request")).toBe(0);
+  });
+
+  it("clears the footer silently for not_applicable credentials instead of warning", async () => {
+    vi.useFakeTimers();
+    vi.mocked(fetchProviderQuotas).mockResolvedValueOnce({
+      success: false,
+      error: { kind: "not_applicable", message: "Direct API key" },
+    } as any);
+
+    const { pi, emitExtensionEvent } = createFakePi();
+    const { ctx, setStatus } = createContext("anthropic");
+
+    await usageStatusExtension(pi);
+    await emitExtensionEvent("session_start", ctx);
+    await vi.runOnlyPendingTimersAsync();
+    await vi.advanceTimersByTimeAsync(0);
+
+    const calls = setStatus.mock.calls as unknown as Array<[string, string | undefined]>;
+    const last = calls[calls.length - 1]?.[1];
+    expect(last).toBeUndefined();
+    expect(calls.some((c) => c[1] === "usage unavailable")).toBe(false);
+  });
+});