@aliou/pi-synthetic 0.14.0 → 0.16.0

package/src/extensions/sub-bar-integration/index.ts

@@ -1,4 +1,4 @@
-import type { AuthStorage, ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import {
   configLoader,
   SYNTHETIC_CONFIG_UPDATED_EVENT,
@@ -6,9 +6,13 @@ import {
   SYNTHETIC_EXTENSIONS_REQUEST_EVENT,
   type SyntheticConfigUpdatedPayload,
 } from "../../config";
-import { getSyntheticApiKey } from "../../lib/env";
-import type { QuotasResponse } from "../../types/quotas";
-import { fetchQuotas, formatResetTime } from "../../utils/quotas";
+import {
+  type QuotasResponse,
+  SYNTHETIC_QUOTAS_REQUEST_EVENT,
+  SYNTHETIC_QUOTAS_UPDATED_EVENT,
+  type SyntheticQuotasUpdatedPayload,
+} from "../../types/quotas";
+import { formatResetTime } from "../../utils/quotas";
 
 interface RateWindow {
   label: string;
@@ -24,14 +28,6 @@ interface UsageSnapshot {
   lastSuccessAt?: number;
 }
 
-interface SubCoreSettingsPayload {
-  settings?: {
-    behavior?: {
-      refreshInterval: number;
-    };
-  };
-}
-
 function toUsageSnapshot(quotas: QuotasResponse): UsageSnapshot {
   const windows: RateWindow[] = [];
 
@@ -105,64 +101,42 @@ function toUsageSnapshot(quotas: QuotasResponse): UsageSnapshot {
   };
 }
 
-async function emitCurrentUsage(
-  pi: ExtensionAPI,
-  authStorage: AuthStorage,
-): Promise<void> {
-  const apiKey = await getSyntheticApiKey(authStorage);
-  if (!apiKey) return;
-  const result = await fetchQuotas(apiKey);
-  if (!result.success) return;
-  pi.events.emit("sub-core:update-current", {
-    state: {
-      provider: "synthetic",
-      usage: toUsageSnapshot(result.data.quotas),
-    },
-  });
-}
-
 export function registerSubBarIntegration(pi: ExtensionAPI): void {
-  let interval: NodeJS.Timeout | undefined;
-  let refreshMs = 60000;
   let subCoreReady = false;
   let currentProvider: string | undefined;
-  let currentAuthStorage: AuthStorage | undefined;
   let enabled = configLoader.getConfig().subBarIntegration;
 
   function isSynthetic(): boolean {
     return enabled && currentProvider === "synthetic";
   }
 
-  function stop(): void {
-    if (interval) {
-      clearInterval(interval);
-      interval = undefined;
-    }
+  function emitUsage(quotas: QuotasResponse): void {
+    pi.events.emit("sub-core:update-current", {
+      state: {
+        provider: "synthetic",
+        usage: toUsageSnapshot(quotas),
+      },
+    });
   }
 
-  function startPolling(authStorage: AuthStorage): void {
-    stop();
-    currentAuthStorage = authStorage;
-    void emitCurrentUsage(pi, authStorage);
-    const ms = Math.max(10000, refreshMs);
-    interval = setInterval(() => {
-      if (isSynthetic() && currentAuthStorage) {
-        void emitCurrentUsage(pi, currentAuthStorage);
-      }
-    }, ms);
-    interval.unref?.();
+  function requestQuotas(): void {
+    pi.events.emit(SYNTHETIC_QUOTAS_REQUEST_EVENT, undefined);
   }
 
+  // Receive quota updates from the provider extension
+  pi.events.on(SYNTHETIC_QUOTAS_UPDATED_EVENT, (data: unknown) => {
+    if (!isSynthetic() || !subCoreReady) return;
+    const { quotas } = data as SyntheticQuotasUpdatedPayload;
+    emitUsage(quotas);
+  });
+
   pi.events.on(SYNTHETIC_CONFIG_UPDATED_EVENT, (data: unknown) => {
     enabled = (data as SyntheticConfigUpdatedPayload).config.subBarIntegration;
 
-    if (!enabled) {
-      stop();
-      return;
-    }
+    if (!enabled) return;
 
-    if (subCoreReady && currentAuthStorage && currentProvider === "synthetic") {
-      startPolling(currentAuthStorage);
+    if (subCoreReady && currentProvider === "synthetic") {
+      requestQuotas();
     }
   });
 
@@ -170,48 +144,24 @@ export function registerSubBarIntegration(pi: ExtensionAPI): void {
     subCoreReady = true;
   });
 
-  pi.events.on("sub-core:settings:updated", (data: unknown) => {
-    const payload = data as SubCoreSettingsPayload;
-    if (payload.settings?.behavior?.refreshInterval) {
-      refreshMs = payload.settings.behavior.refreshInterval * 1000;
-      if (interval && isSynthetic() && currentAuthStorage) {
-        startPolling(currentAuthStorage);
-      }
-    }
+  pi.on("session_start", async (_event, ctx) => {
+    currentProvider = ctx.model?.provider;
   });
 
-  pi.on("session_start", async (_event, ctx) => {
+  pi.on("model_select", async (_event, ctx) => {
     currentProvider = ctx.model?.provider;
-    currentAuthStorage = ctx.modelRegistry.authStorage;
 
     if (subCoreReady && isSynthetic()) {
-      const apiKey = await getSyntheticApiKey(currentAuthStorage);
-      if (apiKey) {
-        startPolling(currentAuthStorage);
-      }
+      requestQuotas();
     }
   });
 
-  pi.on("model_select", async (event, ctx) => {
-    currentProvider = event.model?.provider;
-    currentAuthStorage = ctx.modelRegistry.authStorage;
-
-    if (subCoreReady && isSynthetic()) {
-      const apiKey = await getSyntheticApiKey(currentAuthStorage);
-      if (apiKey) {
-        startPolling(currentAuthStorage);
-      } else {
-        stop();
-      }
-    } else {
-      stop();
-    }
+  pi.on("session_before_switch", (_event, ctx) => {
+    currentProvider = ctx.model?.provider;
   });
 
   pi.on("session_shutdown", () => {
     currentProvider = undefined;
-    currentAuthStorage = undefined;
-    stop();
   });
 }
 

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

@@ -9,9 +9,15 @@ import {
   SYNTHETIC_EXTENSIONS_REQUEST_EVENT,
   type SyntheticConfigUpdatedPayload,
 } from "../../config";
-import { getSyntheticApiKey } from "../../lib/env";
-import type { QuotasResponse } from "../../types/quotas";
-import { fetchQuotas, formatResetTime } from "../../utils/quotas";
+import {
+  type QuotasResponse,
+  SYNTHETIC_QUOTAS_READ_EVENT,
+  SYNTHETIC_QUOTAS_REQUEST_EVENT,
+  type SyntheticQuotasReadPayload,
+  type SyntheticQuotasRequestPayload,
+  type SyntheticQuotasSnapshotPayload,
+} from "../../types/quotas";
+import { formatResetTime } from "../../utils/quotas";
 import {
   assessWindow,
   getSeverityColor,
@@ -20,7 +26,6 @@ import {
 } from "../../utils/quotas-severity";
 
 const EXTENSION_ID = "synthetic-usage";
-const REFRESH_INTERVAL_MS = 60_000;
 
 type WindowStatus = {
   label: string;
@@ -73,175 +78,95 @@ function formatStatus(ctx: ExtensionContext, windows: WindowStatus[]): string {
   return parts.join(" ");
 }
 
-function createStatusRefresher() {
-  let refreshTimer: ReturnType<typeof setInterval> | undefined;
-  let activeContext: ExtensionContext | undefined;
-  let isRefreshInFlight = false;
-  let queuedRefresh = false;
-  let lastSnapshot: WindowStatus[] | undefined;
+export default async function (pi: ExtensionAPI) {
+  await configLoader.load();
 
-  async function updateFooterStatus(ctx: ExtensionContext): Promise<void> {
-    if (!ctx.hasUI) return;
-    if (isRefreshInFlight) {
-      queuedRefresh = true;
-      return;
-    }
-    isRefreshInFlight = true;
-    try {
-      const apiKey = await getSyntheticApiKey(ctx.modelRegistry.authStorage);
-      if (!apiKey) {
-        lastSnapshot = undefined;
-        ctx.ui.setStatus(EXTENSION_ID, undefined);
-        return;
-      }
-      const result = await fetchQuotas(apiKey);
-      if (!result.success) {
-        ctx.ui.setStatus(
-          EXTENSION_ID,
-          ctx.ui.theme.fg("warning", "usage unavailable"),
-        );
-        return;
-      }
-      const windows = parseSnapshot(result.data.quotas);
-      lastSnapshot = windows;
-      if (windows.length === 0) {
-        ctx.ui.setStatus(EXTENSION_ID, undefined);
-        return;
-      }
-      ctx.ui.setStatus(EXTENSION_ID, formatStatus(ctx, windows));
-    } catch {
-      ctx.ui.setStatus(
-        EXTENSION_ID,
-        ctx.ui.theme.fg("warning", "usage unavailable"),
-      );
-    } finally {
-      isRefreshInFlight = false;
-      if (queuedRefresh) {
-        queuedRefresh = false;
-        void updateFooterStatus(ctx);
-      }
-    }
-  }
+  let enabled = configLoader.getConfig().usageStatus;
 
-  function refreshFor(ctx: ExtensionContext): Promise<void> {
-    activeContext = ctx;
-    return updateFooterStatus(ctx);
+  function requestQuotas(
+    respond: (snapshot: SyntheticQuotasSnapshotPayload | undefined) => void,
+  ): void {
+    pi.events.emit(SYNTHETIC_QUOTAS_REQUEST_EVENT, {
+      respond,
+    } satisfies SyntheticQuotasRequestPayload);
   }
 
-  function startAutoRefresh(): void {
-    if (refreshTimer) clearInterval(refreshTimer);
-    refreshTimer = setInterval(() => {
-      if (!activeContext) return;
-      void updateFooterStatus(activeContext);
-    }, REFRESH_INTERVAL_MS);
-    refreshTimer.unref?.();
+  function readQuotas(
+    respond: (snapshot: SyntheticQuotasSnapshotPayload | undefined) => void,
+  ): void {
+    pi.events.emit(SYNTHETIC_QUOTAS_READ_EVENT, {
+      respond,
+    } satisfies SyntheticQuotasReadPayload);
   }
 
-  function stopAutoRefresh(ctx?: ExtensionContext): void {
-    if (refreshTimer) {
-      clearInterval(refreshTimer);
-      refreshTimer = undefined;
+  function renderSnapshot(
+    ctx: ExtensionContext,
+    snapshot: SyntheticQuotasSnapshotPayload | undefined,
+  ): void {
+    if (!ctx.hasUI) return;
+    if (!snapshot) {
+      ctx.ui.setStatus(
+        EXTENSION_ID,
+        ctx.ui.theme.fg("dim", "loading usage..."),
+      );
+      return;
     }
-    ctx?.ui.setStatus(EXTENSION_ID, undefined);
-  }
 
-  async function setLoadingStatus(ctx: ExtensionContext): Promise<void> {
-    if (!ctx.hasUI) return;
-    const apiKey = await getSyntheticApiKey(
-      ctx.modelRegistry.authStorage,
-    ).catch(() => undefined);
-    if (!apiKey) {
+    const windows = parseSnapshot(snapshot.quotas);
+    if (windows.length === 0) {
       ctx.ui.setStatus(EXTENSION_ID, undefined);
       return;
     }
-    ctx.ui.setStatus(EXTENSION_ID, ctx.ui.theme.fg("dim", "loading usage..."));
-  }
 
-  function renderFromLastSnapshot(ctx: ExtensionContext): boolean {
-    if (!ctx.hasUI || !lastSnapshot) return false;
-    ctx.ui.setStatus(EXTENSION_ID, formatStatus(ctx, lastSnapshot));
-    return true;
+    ctx.ui.setStatus(EXTENSION_ID, formatStatus(ctx, windows));
   }
 
-  return {
-    refreshFor,
-    startAutoRefresh,
-    stopAutoRefresh,
-    setLoadingStatus,
-    renderFromLastSnapshot,
-  };
-}
-
-export default async function (pi: ExtensionAPI) {
-  await configLoader.load();
+  function clearStatus(ctx: ExtensionContext): void {
+    if (!ctx.hasUI) return;
+    ctx.ui.setStatus(EXTENSION_ID, undefined);
+  }
 
-  const refresher = createStatusRefresher();
-  let enabled = configLoader.getConfig().usageStatus;
-  let currentContext: ExtensionContext | undefined;
-  let currentProvider: string | undefined;
+  function renderFromStoreOrRefresh(ctx: ExtensionContext): void {
+    if (!enabled || ctx.model?.provider !== "synthetic") {
+      clearStatus(ctx);
+      return;
+    }
+    readQuotas((snapshot) => {
+      if (snapshot) {
+        renderSnapshot(ctx, snapshot);
+      } else {
+        renderSnapshot(ctx, undefined); // show loading
+        requestQuotas((refreshed) => renderSnapshot(ctx, refreshed));
+      }
+    });
+  }
 
   pi.events.on(SYNTHETIC_CONFIG_UPDATED_EVENT, (data: unknown) => {
     enabled = (data as SyntheticConfigUpdatedPayload).config.usageStatus;
+  });
 
-    if (!enabled) {
-      refresher.stopAutoRefresh(currentContext);
-      return;
-    }
-
-    if (currentContext && currentProvider === "synthetic") {
-      refresher.startAutoRefresh();
-      void refresher.refreshFor(currentContext);
-    }
+  pi.on("session_start", (_event, ctx) => {
+    renderFromStoreOrRefresh(ctx);
   });
 
-  pi.on("session_start", async (_event, ctx) => {
-    currentContext = ctx;
-    currentProvider = ctx.model?.provider;
-    if (!enabled || ctx.model?.provider !== "synthetic") return;
-    refresher.startAutoRefresh();
-    await refresher.setLoadingStatus(ctx);
-    await refresher.refreshFor(ctx);
+  pi.on("model_select", (_event, ctx) => {
+    renderFromStoreOrRefresh(ctx);
   });
 
-  pi.on("turn_end", (_event, ctx) => {
-    currentContext = ctx;
-    currentProvider = ctx.model?.provider;
-    if (!enabled || ctx.model?.provider !== "synthetic") return;
-    void refresher.refreshFor(ctx);
+  pi.on("agent_end", (_event, ctx) => {
+    renderFromStoreOrRefresh(ctx);
   });
 
-  pi.on("session_start", (event, ctx) => {
-    // Handle session switches (model_select handles mid-session provider changes)
-    if (
-      event.reason === "new" ||
-      event.reason === "resume" ||
-      event.reason === "fork"
-    ) {
-      currentContext = ctx;
-      currentProvider = ctx.model?.provider;
-      if (enabled && ctx.model?.provider === "synthetic") {
-        void refresher.refreshFor(ctx);
-      } else {
-        refresher.stopAutoRefresh(ctx);
-      }
-    }
+  pi.on("turn_end", (_event, ctx) => {
+    renderFromStoreOrRefresh(ctx);
   });
 
-  pi.on("model_select", (_event, ctx) => {
-    currentContext = ctx;
-    currentProvider = ctx.model?.provider;
-    if (enabled && ctx.model?.provider === "synthetic") {
-      refresher.startAutoRefresh();
-      void refresher.refreshFor(ctx);
-    } else {
-      refresher.stopAutoRefresh(ctx);
-    }
+  pi.on("session_before_switch", (_event, ctx) => {
+    clearStatus(ctx);
   });
 
   pi.on("session_shutdown", (_event, ctx) => {
-    currentContext = undefined;
-    currentProvider = undefined;
-    refresher.stopAutoRefresh(ctx);
+    clearStatus(ctx);
   });
 
   pi.events.on(SYNTHETIC_EXTENSIONS_REQUEST_EVENT, () => {

package/src/services/quota-store.test.ts

@@ -0,0 +1,211 @@
+import {
+  afterEach,
+  assert,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  vi,
+} from "vitest";
+import type { QuotasResponse } from "../types/quotas";
+import { QuotaStore } from "./quota-store";
+
+beforeEach(() => {
+  vi.useFakeTimers();
+});
+
+afterEach(() => {
+  vi.useRealTimers();
+});
+
+describe("QuotaStore", () => {
+  const sampleQuotas: QuotasResponse = {
+    subscription: { limit: 100, requests: 5, renewsAt: "2026-01-01T00:00:00Z" },
+  };
+
+  describe("ingest", () => {
+    it("stores and emits API-sourced data", () => {
+      const store = new QuotaStore();
+      const received: QuotasResponse[] = [];
+      store.subscribe((snap) => received.push(snap.quotas));
+
+      const result = store.ingest(sampleQuotas, "api");
+
+      expect(result).toBe(true);
+      expect(store.getSnapshot()?.quotas).toBe(sampleQuotas);
+      expect(store.getSnapshot()?.source).toBe("api");
+      expect(received).toHaveLength(1);
+    });
+
+    it("stores and emits header-sourced data", () => {
+      const store = new QuotaStore();
+      const result = store.ingest(sampleQuotas, "header");
+
+      expect(result).toBe(true);
+      expect(store.getSnapshot()?.source).toBe("header");
+    });
+
+    it("throttles header ingestion within throttle window", () => {
+      const store = new QuotaStore();
+      store.ingest(sampleQuotas, "header");
+
+      // Within throttle window — should be dropped
+      const result = store.ingest(sampleQuotas, "header");
+      expect(result).toBe(false);
+
+      // Advance past throttle
+      vi.advanceTimersByTime(store.headerThrottleMs + 1);
+      const result2 = store.ingest(sampleQuotas, "header");
+      expect(result2).toBe(true);
+    });
+
+    it("does NOT throttle API ingestion", () => {
+      const store = new QuotaStore();
+      store.ingest(sampleQuotas, "api");
+
+      // API is never throttled
+      const result = store.ingest(sampleQuotas, "api");
+      expect(result).toBe(true);
+    });
+
+    it("header after API emit always goes through", () => {
+      const store = new QuotaStore();
+      store.ingest(sampleQuotas, "api");
+
+      // Header 1ms after API should not be blocked
+      vi.advanceTimersByTime(1);
+      const result = store.ingest(sampleQuotas, "header");
+      expect(result).toBe(true);
+    });
+
+    it("updates timestamp on each successful ingest", () => {
+      const store = new QuotaStore();
+      store.ingest(sampleQuotas, "api");
+      const snap1 = store.getSnapshot();
+      assert(snap1);
+      const t1 = snap1.updatedAt;
+
+      vi.advanceTimersByTime(10_000);
+      store.ingest(sampleQuotas, "api");
+      const snap2 = store.getSnapshot();
+      assert(snap2);
+      const t2 = snap2.updatedAt;
+
+      expect(t2).toBeGreaterThan(t1);
+    });
+  });
+
+  describe("subscribe", () => {
+    it("notifies subscribers on ingest", () => {
+      const store = new QuotaStore();
+      const calls: QuotasResponse[] = [];
+      store.subscribe((snap) => calls.push(snap.quotas));
+
+      store.ingest(sampleQuotas, "api");
+      expect(calls).toHaveLength(1);
+      expect(calls[0]).toBe(sampleQuotas);
+    });
+
+    it("does not notify on throttled ingest", () => {
+      const store = new QuotaStore();
+      const calls: QuotasResponse[] = [];
+      store.subscribe((snap) => calls.push(snap.quotas));
+
+      store.ingest(sampleQuotas, "header");
+      store.ingest(sampleQuotas, "header"); // throttled
+
+      expect(calls).toHaveLength(1);
+    });
+
+    it("unsubscribes when unsubscribe function is called", () => {
+      const store = new QuotaStore();
+      const calls: QuotasResponse[] = [];
+      const unsub = store.subscribe((snap) => calls.push(snap.quotas));
+
+      unsub();
+      store.ingest(sampleQuotas, "api");
+
+      expect(calls).toHaveLength(0);
+    });
+
+    it("supports multiple subscribers", () => {
+      const store = new QuotaStore();
+      const calls1: QuotasResponse[] = [];
+      const calls2: QuotasResponse[] = [];
+      store.subscribe((snap) => calls1.push(snap.quotas));
+      store.subscribe((snap) => calls2.push(snap.quotas));
+
+      store.ingest(sampleQuotas, "api");
+
+      expect(calls1).toHaveLength(1);
+      expect(calls2).toHaveLength(1);
+    });
+  });
+
+  describe("refreshFromApi", () => {
+    it("calls the fetcher and ingests the result", async () => {
+      const store = new QuotaStore();
+      const fetcher = vi.fn().mockResolvedValue(sampleQuotas);
+
+      const result = await store.refreshFromApi(fetcher);
+
+      assert(result);
+      expect(result.quotas).toBe(sampleQuotas);
+      expect(result.source).toBe("api");
+      expect(fetcher).toHaveBeenCalledOnce();
+    });
+
+    it("deduplicates concurrent calls", async () => {
+      const store = new QuotaStore();
+      let resolveFirst!: (v: QuotasResponse) => void;
+      const first = new Promise<QuotasResponse>((r) => (resolveFirst = r));
+      const fetcher = vi.fn().mockImplementation(() => first);
+
+      // Start two concurrent refreshes
+      const p1 = store.refreshFromApi(fetcher);
+      const p2 = store.refreshFromApi(fetcher);
+
+      // Only one fetcher call
+      expect(fetcher).toHaveBeenCalledOnce();
+      expect(store.isRefreshing).toBe(true);
+
+      // Resolve the fetch
+      resolveFirst(sampleQuotas);
+      await p1;
+      await p2;
+
+      expect(store.isRefreshing).toBe(false);
+    });
+
+    it("handles fetcher returning undefined", async () => {
+      const store = new QuotaStore();
+      const fetcher = vi.fn().mockResolvedValue(undefined);
+
+      const result = await store.refreshFromApi(fetcher);
+
+      expect(result).toBeUndefined();
+      expect(store.getSnapshot()).toBeUndefined();
+    });
+  });
+
+  describe("clear", () => {
+    it("resets all state", () => {
+      const store = new QuotaStore();
+      store.ingest(sampleQuotas, "api");
+
+      store.clear();
+
+      expect(store.getSnapshot()).toBeUndefined();
+    });
+
+    it("resets header throttle after clear", () => {
+      const store = new QuotaStore();
+      store.ingest(sampleQuotas, "header");
+      expect(store.ingest(sampleQuotas, "header")).toBe(false);
+
+      store.clear();
+
+      expect(store.ingest(sampleQuotas, "header")).toBe(true);
+    });
+  });
+});