@latentminds/pi-quotas 0.2.4 → 0.3.1

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,19 +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);
       }
     }
-  });
+  }));
 
-  pi.events.on(QUOTAS_CONFIG_UPDATED_EVENT, (data: unknown) => {
+  function scheduleRefresh(ctx: ExtensionContext): void {
+    void refresher.refreshFor(ctx).catch(() => undefined);
+  }
+
+  unsubscribeEventBusListeners.push(pi.events.on(QUOTAS_CONFIG_UPDATED_EVENT, (data: unknown) => {
     const config = (data as QuotasConfigUpdatedPayload).config;
     enabled = config.usageStatus;
     deferToSynthetic = config.deferToSynthetic;
@@ -186,9 +247,9 @@ export default async function (pi: ExtensionAPI) {
     }
     if (currentContext) {
       refresher.start();
-      void refresher.refreshFor(currentContext);
+      scheduleRefresh(currentContext);
     }
-  });
+  }));
 
   /**
    * Whether to suppress our footer because pi-synthetic is showing
@@ -198,46 +259,55 @@ export default async function (pi: ExtensionAPI) {
     return deferToSynthetic && syntheticUsageActive && provider === "synthetic";
   }
 
-  pi.on("session_start", async (_event, ctx) => {
+  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();
-    await refresher.refreshFor(ctx);
+    scheduleRefresh(ctx);
   });
 
-  pi.on("turn_end", async (_event, ctx) => {
+  pi.on("turn_end", (_event, ctx) => {
     currentContext = ctx;
     if (!enabled) return;
-    if (shouldDeferToSynthetic(ctx.model?.provider)) return;
-    await refresher.refreshFor(ctx);
+    if (shouldDeferToSynthetic(getContextProvider(ctx))) {
+      refresher.stop(ctx);
+      return;
+    }
+    scheduleRefresh(ctx);
   });
 
-  pi.on("model_select", async (_event, ctx) => {
+  pi.on("model_select", (_event, ctx) => {
     currentContext = ctx;
     if (!enabled) {
       refresher.stop(ctx);
       return;
     }
-    if (shouldDeferToSynthetic(ctx.model?.provider)) {
-      ctx.ui.setStatus(EXTENSION_ID, undefined);
+    if (shouldDeferToSynthetic(getContextProvider(ctx))) {
+      refresher.stop(ctx);
       return;
     }
-    await refresher.refreshFor(ctx);
+    scheduleRefresh(ctx);
   });
 
   pi.on("session_shutdown", async (_event, ctx) => {
     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("");
+  });
+});