@aliou/pi-synthetic 0.15.0 → 0.17.0

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

@@ -1,7 +1,7 @@
 import type {
   ExtensionAPI,
   ExtensionContext,
-} from "@mariozechner/pi-coding-agent";
+} from "@earendil-works/pi-coding-agent";
 import {
   configLoader,
   SYNTHETIC_CONFIG_UPDATED_EVENT,
@@ -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/extensions/web-search/index.ts

@@ -1,4 +1,4 @@
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import {
   configLoader,
   SYNTHETIC_CONFIG_UPDATED_EVENT,

package/src/extensions/web-search/tool.ts

@@ -5,9 +5,9 @@ import type {
   ExtensionContext,
   Theme,
   ToolRenderResultOptions,
-} from "@mariozechner/pi-coding-agent";
-import { getMarkdownTheme, keyHint } from "@mariozechner/pi-coding-agent";
-import { Container, Markdown, Text } from "@mariozechner/pi-tui";
+} from "@earendil-works/pi-coding-agent";
+import { getMarkdownTheme, keyHint } from "@earendil-works/pi-coding-agent";
+import { Container, Markdown, Text } from "@earendil-works/pi-tui";
 import { type Static, Type } from "typebox";
 import { configLoader } from "../../config";
 import { getSyntheticApiKey } from "../../lib/env";

package/src/lib/env.ts

@@ -1,4 +1,4 @@
-import type { AuthStorage } from "@mariozechner/pi-coding-agent";
+import type { AuthStorage } from "@earendil-works/pi-coding-agent";
 
 const PROVIDER_ID = "synthetic";
 

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

package/src/services/quota-store.ts

@@ -0,0 +1,112 @@
+import type { QuotaSource, QuotasResponse } from "../types/quotas";
+
+export interface QuotaSnapshot {
+  quotas: QuotasResponse;
+  source: QuotaSource;
+  updatedAt: number; // epoch ms
+}
+
+type Listener = (snapshot: QuotaSnapshot) => void;
+
+/**
+ * Pi-agnostic in-memory quota store.
+ *
+ * Ingests quota data from headers or API, handles throttling of
+ * header-sourced updates, and notifies subscribers on change.
+ *
+ * Usage:
+ *   const store = new QuotaStore();
+ *   store.subscribe((snap) => { ... });
+ *   store.ingest(quotas, "header");
+ *   store.ingest(quotas, "api");
+ */
+export class QuotaStore {
+  private snapshot: QuotaSnapshot | undefined;
+  private listeners = new Set<Listener>();
+  private lastHeaderIngestAt = 0;
+  private inFlightRefresh: Promise<QuotaSnapshot | undefined> | undefined;
+  private inFlightId = 0;
+
+  /** Throttle header ingestion: skip if last header ingest was within this window. */
+  headerThrottleMs = 5_000;
+
+  /** Current snapshot (may be undefined if no data has been ingested yet). */
+  getSnapshot(): QuotaSnapshot | undefined {
+    return this.snapshot;
+  }
+
+  /** Subscribe to snapshot updates. Returns unsubscribe function. */
+  subscribe(listener: Listener): () => void {
+    this.listeners.add(listener);
+    return () => {
+      this.listeners.delete(listener);
+    };
+  }
+
+  private emit(snapshot: QuotaSnapshot): void {
+    for (const l of this.listeners) l(snapshot);
+  }
+
+  /**
+   * Ingest quota data. Returns true if the snapshot was updated
+   * (i.e. not throttled).
+   *
+   * Header-sourced data is throttled: if the last header ingest was
+   * within `headerThrottleMs`, it is silently dropped.
+   * API-sourced data always goes through.
+   */
+  ingest(quotas: QuotasResponse, source: QuotaSource): boolean {
+    const now = Date.now();
+
+    if (source === "header") {
+      if (now - this.lastHeaderIngestAt < this.headerThrottleMs) return false;
+      this.lastHeaderIngestAt = now;
+    }
+
+    this.snapshot = { quotas, source, updatedAt: now };
+    this.emit(this.snapshot);
+    return true;
+  }
+
+  /**
+   * Refresh quotas by calling the provided fetcher.
+   * Deduplicates concurrent calls — only one fetch runs at a time.
+   */
+  async refreshFromApi(
+    fetcher: () => Promise<QuotasResponse | undefined>,
+  ): Promise<QuotaSnapshot | undefined> {
+    if (this.inFlightRefresh) return this.inFlightRefresh;
+
+    this.inFlightId++;
+    const id = this.inFlightId;
+
+    this.inFlightRefresh = (async () => {
+      try {
+        const quotas = await fetcher();
+        if (quotas && id === this.inFlightId) {
+          this.ingest(quotas, "api");
+        }
+        return this.snapshot;
+      } finally {
+        if (id === this.inFlightId) {
+          this.inFlightRefresh = undefined;
+        }
+      }
+    })();
+
+    return this.inFlightRefresh;
+  }
+
+  /** Returns true if a refresh is currently in flight. */
+  get isRefreshing(): boolean {
+    return !!this.inFlightRefresh;
+  }
+
+  /** Clear all state. Call on session shutdown or reset. */
+  clear(): void {
+    this.inFlightId++; // Invalidates in-flight refresh
+    this.snapshot = undefined;
+    this.lastHeaderIngestAt = 0;
+    this.inFlightRefresh = undefined;
+  }
+}