@latentminds/pi-quotas 0.2.6 → 0.3.1

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);
+  });
+});

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

@@ -28,6 +28,21 @@ import { formatWindowStatus, type WindowStatus } from "./format-status.js";
 
 const EXTENSION_ID = "pi-quotas-usage";
 const REFRESH_INTERVAL_MS = 60_000;
+const STALE_CONTEXT_MESSAGE = "This extension ctx is stale";
+
+function isStaleContextError(error: unknown): boolean {
+  return error instanceof Error && error.message.includes(STALE_CONTEXT_MESSAGE);
+}
+
+function getContextProvider(ctx: ExtensionContext | undefined): string | undefined {
+  if (!ctx) return undefined;
+  try {
+    return ctx.model?.provider;
+  } catch (error) {
+    if (isStaleContextError(error)) return undefined;
+    throw error;
+  }
+}
 
 function formatFooterResetTime(resetsAt: string): string {
   const remaining = formatTimeRemaining(new Date(resetsAt));
@@ -93,30 +108,75 @@ function createStatusRefresher() {
   let inFlight = false;
   let queued = false;
 
-  async function update(ctx: ExtensionContext): Promise<void> {
-    if (!ctx.hasUI || !activeProvider || !isSupportedProvider(activeProvider)) return;
+  // Bumped whenever the active ctx/provider is replaced or the refresher stops.
+  // This prevents an old async fetch from writing to a replacement session.
+  let generation = 0;
+
+  function deactivate(): void {
+    if (refreshTimer) clearInterval(refreshTimer);
+    refreshTimer = undefined;
+    activeContext = undefined;
+    activeProvider = undefined;
+    lastStatus = undefined;
+    queued = false;
+    generation++;
+  }
+
+  function setStatusSafely(
+    ctx: ExtensionContext | undefined,
+    text: string | undefined | ((ctx: ExtensionContext) => string | undefined),
+  ): boolean {
+    if (!ctx) return false;
+    try {
+      if (!ctx.hasUI) return true;
+      ctx.ui.setStatus(EXTENSION_ID, typeof text === "function" ? text(ctx) : text);
+      return true;
+    } catch (error) {
+      if (isStaleContextError(error) && activeContext === ctx) deactivate();
+      return false;
+    }
+  }
+
+  async function update(ctx: ExtensionContext, requestGeneration = generation): Promise<void> {
     if (inFlight) {
       queued = true;
       return;
     }
     inFlight = true;
     try {
-      const result = await fetchProviderQuotas(ctx.modelRegistry.authStorage, activeProvider);
+      if (requestGeneration !== generation || activeContext !== ctx) return;
+      if (!ctx.hasUI || !activeProvider || !isSupportedProvider(activeProvider)) return;
+
+      const provider = activeProvider;
+      const result = await fetchProviderQuotas(ctx.modelRegistry.authStorage, provider);
+      if (requestGeneration !== generation || activeContext !== ctx) return;
+
       if (!result.success) {
-        ctx.ui.setStatus(EXTENSION_ID, ctx.ui.theme.fg("warning", "usage unavailable"));
+        // A "not applicable" result (e.g. a direct Anthropic API key with no
+        // OAuth subscription usage) is expected, not a failure — show nothing
+        // rather than a persistent "usage unavailable" warning.
+        if (result.error.kind === "not_applicable") {
+          setStatusSafely(ctx, undefined);
+          return;
+        }
+        setStatusSafely(ctx, (ctx) => ctx.ui.theme.fg("warning", "usage unavailable"));
         return;
       }
       const windows: WindowStatus[] = toStatusWindows(result.data.windows);
       const status = formatStatusForFooter(ctx, windows);
       lastStatus = status === undefined ? undefined : windows;
-      ctx.ui.setStatus(EXTENSION_ID, status);
-    } catch {
-      ctx.ui.setStatus(EXTENSION_ID, ctx.ui.theme.fg("warning", "usage unavailable"));
+      setStatusSafely(ctx, status);
+    } catch (error) {
+      if (isStaleContextError(error)) {
+        if (activeContext === ctx) deactivate();
+        return;
+      }
+      setStatusSafely(ctx, (ctx) => ctx.ui.theme.fg("warning", "usage unavailable"));
     } finally {
       inFlight = false;
-      if (queued) {
+      if (queued && activeContext) {
         queued = false;
-        void update(ctx);
+        void update(activeContext, generation).catch(() => undefined);
       }
     }
   }
@@ -124,32 +184,29 @@ function createStatusRefresher() {
   return {
     async refreshFor(ctx: ExtensionContext): Promise<void> {
       activeContext = ctx;
-      activeProvider = ctx.model?.provider;
+      activeProvider = getContextProvider(ctx);
+      generation++;
+      const requestGeneration = generation;
       if (!activeProvider || !isSupportedProvider(activeProvider)) {
-        ctx.ui.setStatus(EXTENSION_ID, undefined);
+        setStatusSafely(ctx, undefined);
         return;
       }
-      await update(ctx);
+      await update(ctx, requestGeneration);
     },
     start(): void {
       if (refreshTimer) clearInterval(refreshTimer);
       refreshTimer = setInterval(() => {
-        if (activeContext) void update(activeContext);
+        if (activeContext) void update(activeContext, generation).catch(() => undefined);
       }, REFRESH_INTERVAL_MS);
       refreshTimer.unref?.();
     },
     stop(ctx?: ExtensionContext): void {
-      if (refreshTimer) clearInterval(refreshTimer);
-      refreshTimer = undefined;
-      activeContext = undefined;
-      activeProvider = undefined;
-      lastStatus = undefined;
-      ctx?.ui.setStatus(EXTENSION_ID, undefined);
+      deactivate();
+      setStatusSafely(ctx, undefined);
     },
     renderLast(ctx: ExtensionContext): boolean {
-      if (!lastStatus || !ctx.hasUI) return false;
-      ctx.ui.setStatus(EXTENSION_ID, formatStatusForFooter(ctx, lastStatus));
-      return true;
+      if (!lastStatus) return false;
+      return setStatusSafely(ctx, (ctx) => formatStatusForFooter(ctx, lastStatus ?? []));
     },
   };
 }
@@ -157,6 +214,7 @@ function createStatusRefresher() {
 export default async function (pi: ExtensionAPI) {
   await configLoader.load();
   const refresher = createStatusRefresher();
+  const unsubscribeEventBusListeners: Array<() => void> = [];
   let enabled = configLoader.getConfig().usageStatus;
   let deferToSynthetic = configLoader.getConfig().deferToSynthetic;
   let currentContext: ExtensionContext | undefined;
@@ -164,25 +222,22 @@ export default async function (pi: ExtensionAPI) {
   /** Whether pi-synthetic's usage footer is active in this session. */
   let syntheticUsageActive = false;
 
-  pi.events.on(SYNTHETIC_EXTENSIONS_REGISTER_EVENT, (data: unknown) => {
+  unsubscribeEventBusListeners.push(pi.events.on(SYNTHETIC_EXTENSIONS_REGISTER_EVENT, (data: unknown) => {
     const { feature } = data as SyntheticExtensionsRegisterPayload;
     if (feature === "usageStatus") {
       syntheticUsageActive = true;
       // If currently showing synthetic data, clear our footer
-      if (currentContext && enabled && deferToSynthetic && currentContext.model?.provider === "synthetic") {
-        currentContext.ui.setStatus(EXTENSION_ID, undefined);
-        refresher.stop();
+      if (currentContext && enabled && deferToSynthetic && getContextProvider(currentContext) === "synthetic") {
+        refresher.stop(currentContext);
       }
     }
-  });
+  }));
 
   function scheduleRefresh(ctx: ExtensionContext): void {
-    void refresher.refreshFor(ctx).catch(() => {
-      if (ctx.hasUI) ctx.ui.setStatus(EXTENSION_ID, ctx.ui.theme.fg("warning", "usage unavailable"));
-    });
+    void refresher.refreshFor(ctx).catch(() => undefined);
   }
 
-  pi.events.on(QUOTAS_CONFIG_UPDATED_EVENT, (data: unknown) => {
+  unsubscribeEventBusListeners.push(pi.events.on(QUOTAS_CONFIG_UPDATED_EVENT, (data: unknown) => {
     const config = (data as QuotasConfigUpdatedPayload).config;
     enabled = config.usageStatus;
     deferToSynthetic = config.deferToSynthetic;
@@ -194,7 +249,7 @@ export default async function (pi: ExtensionAPI) {
       refresher.start();
       scheduleRefresh(currentContext);
     }
-  });
+  }));
 
   /**
    * Whether to suppress our footer because pi-synthetic is showing
@@ -206,9 +261,12 @@ export default async function (pi: ExtensionAPI) {
 
   pi.on("session_start", (_event, ctx) => {
     currentContext = ctx;
-    if (!enabled) return;
-    if (shouldDeferToSynthetic(ctx.model?.provider)) {
-      ctx.ui.setStatus(EXTENSION_ID, undefined);
+    if (!enabled) {
+      refresher.stop(ctx);
+      return;
+    }
+    if (shouldDeferToSynthetic(getContextProvider(ctx))) {
+      refresher.stop(ctx);
       return;
     }
     refresher.start();
@@ -218,7 +276,10 @@ export default async function (pi: ExtensionAPI) {
   pi.on("turn_end", (_event, ctx) => {
     currentContext = ctx;
     if (!enabled) return;
-    if (shouldDeferToSynthetic(ctx.model?.provider)) return;
+    if (shouldDeferToSynthetic(getContextProvider(ctx))) {
+      refresher.stop(ctx);
+      return;
+    }
     scheduleRefresh(ctx);
   });
 
@@ -228,8 +289,8 @@ export default async function (pi: ExtensionAPI) {
       refresher.stop(ctx);
       return;
     }
-    if (shouldDeferToSynthetic(ctx.model?.provider)) {
-      ctx.ui.setStatus(EXTENSION_ID, undefined);
+    if (shouldDeferToSynthetic(getContextProvider(ctx))) {
+      refresher.stop(ctx);
       return;
     }
     scheduleRefresh(ctx);
@@ -239,11 +300,14 @@ export default async function (pi: ExtensionAPI) {
     currentContext = undefined;
     syntheticUsageActive = false;
     refresher.stop(ctx);
+    for (const unsubscribe of unsubscribeEventBusListeners.splice(0)) {
+      unsubscribe();
+    }
   });
 
-  pi.events.on(QUOTAS_EXTENSIONS_REQUEST_EVENT, () => {
+  unsubscribeEventBusListeners.push(pi.events.on(QUOTAS_EXTENSIONS_REQUEST_EVENT, () => {
     if (configLoader.getConfig().usageStatus) {
       pi.events.emit(QUOTAS_EXTENSIONS_REGISTER_EVENT, { feature: "usageStatus" });
     }
-  });
+  }));
 }

package/src/lib/quotas.ts

@@ -8,6 +8,8 @@ export const SUPPORTED_PROVIDERS: SupportedQuotaProvider[] = [
   "github-copilot",
   "openrouter",
   "synthetic",
+  "zai",
+  "opencode-go",
 ];
 
 export const PROVIDER_LABELS: Record<SupportedQuotaProvider, string> = {
@@ -16,6 +18,8 @@ export const PROVIDER_LABELS: Record<SupportedQuotaProvider, string> = {
   "github-copilot": "GitHub Copilot",
   openrouter: "OpenRouter",
   synthetic: "Synthetic",
+  zai: "Z.ai",
+  "opencode-go": "OpenCode Go",
 };
 
 const PROVIDER_TTLS_MS: Record<SupportedQuotaProvider, number> = {
@@ -24,6 +28,8 @@ const PROVIDER_TTLS_MS: Record<SupportedQuotaProvider, number> = {
   "github-copilot": 5 * 60_000,
   openrouter: 60_000,
   synthetic: 60_000,
+  zai: 60_000,
+  "opencode-go": 60_000,
 };
 
 type CacheEntry = {
@@ -54,7 +60,12 @@ export async function fetchProviderQuotas(
   const now = Date.now();
   const ttl = PROVIDER_TTLS_MS[provider];
 
-  if (!options?.force && entry.result && entry.fetchedAt && now - entry.fetchedAt < ttl) {
+  if (
+    !options?.force &&
+    entry.result &&
+    entry.fetchedAt &&
+    now - entry.fetchedAt < ttl
+  ) {
     return entry.result;
   }
   if (!options?.force && entry.inFlight) return entry.inFlight;

package/src/lib/session-tokens.test.ts

@@ -0,0 +1,137 @@
+import { describe, expect, it } from "vitest";
+import {
+  addBuckets,
+  emptyBuckets,
+  formatCost,
+  formatNumber,
+  formatTokenSummary,
+  type TokenBuckets,
+} from "./session-tokens.js";
+
+describe("emptyBuckets", () => {
+  it("returns all zeros", () => {
+    const b = emptyBuckets();
+    expect(b.input).toBe(0);
+    expect(b.output).toBe(0);
+    expect(b.cacheRead).toBe(0);
+    expect(b.cacheWrite).toBe(0);
+    expect(b.totalTokens).toBe(0);
+    expect(b.costTotal).toBe(0);
+  });
+});
+
+describe("addBuckets", () => {
+  it("sums two bucket sets", () => {
+    const a: TokenBuckets = {
+      input: 100,
+      output: 200,
+      cacheRead: 50,
+      cacheWrite: 25,
+      totalTokens: 375,
+      costTotal: 0.05,
+    };
+    const b: TokenBuckets = {
+      input: 300,
+      output: 400,
+      cacheRead: 100,
+      cacheWrite: 75,
+      totalTokens: 875,
+      costTotal: 0.15,
+    };
+    const result = addBuckets(a, b);
+    expect(result.input).toBe(400);
+    expect(result.output).toBe(600);
+    expect(result.cacheRead).toBe(150);
+    expect(result.cacheWrite).toBe(100);
+    expect(result.totalTokens).toBe(1250);
+    expect(result.costTotal).toBe(0.2);
+  });
+
+  it("handles empty buckets", () => {
+    const a = emptyBuckets();
+    const b: TokenBuckets = {
+      input: 100,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: 100,
+      costTotal: 0.01,
+    };
+    const result = addBuckets(a, b);
+    expect(result).toEqual(b);
+  });
+});
+
+describe("formatNumber", () => {
+  it("formats small numbers directly", () => {
+    expect(formatNumber(0)).toBe("0");
+    expect(formatNumber(42)).toBe("42");
+    expect(formatNumber(999)).toBe("999");
+  });
+
+  it("formats thousands with commas", () => {
+    expect(formatNumber(1000)).toBe("1,000");
+    expect(formatNumber(12345)).toBe("12.3K");
+  });
+
+  it("formats millions", () => {
+    expect(formatNumber(1_000_000)).toBe("1.0M");
+    expect(formatNumber(5_500_000)).toBe("5.5M");
+  });
+});
+
+describe("formatCost", () => {
+  it("formats zero", () => {
+    expect(formatCost(0)).toBe("$0.00");
+  });
+
+  it("formats small costs", () => {
+    expect(formatCost(0.005)).toBe("<$0.01");
+    expect(formatCost(0.01)).toBe("$0.01");
+    expect(formatCost(1.23)).toBe("$1.23");
+  });
+
+  it("formats large costs", () => {
+    expect(formatCost(1000)).toBe("$1.0K");
+    expect(formatCost(1234.56)).toBe("$1.2K");
+  });
+});
+
+describe("formatTokenSummary", () => {
+  it("formats full token breakdown", () => {
+    const tokens: TokenBuckets = {
+      input: 50000,
+      output: 10000,
+      cacheRead: 5000,
+      cacheWrite: 2000,
+      totalTokens: 67000,
+      costTotal: 0.45,
+    };
+    const result = formatTokenSummary(tokens);
+    expect(result).toContain("50.0K in");
+    expect(result).toContain("10.0K out");
+    expect(result).toContain("5,000 cached");
+    expect(result).toContain("$0.45");
+  });
+
+  it("omits zero fields", () => {
+    const tokens: TokenBuckets = {
+      input: 1000,
+      output: 500,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: 1500,
+      costTotal: 0,
+    };
+    const result = formatTokenSummary(tokens);
+    expect(result).toContain("in");
+    expect(result).toContain("out");
+    expect(result).not.toContain("cached");
+    expect(result).not.toContain("$");
+  });
+
+  it("handles empty buckets", () => {
+    const result = formatTokenSummary(emptyBuckets());
+    expect(result).toBe("");
+  });
+});