@aliou/pi-synthetic 0.15.0 → 0.16.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aliou/pi-synthetic",
-  "version": "0.15.0",
+  "version": "0.16.0",
   "license": "MIT",
   "type": "module",
   "private": false,

package/src/extensions/provider/index.ts

@@ -1,4 +1,4 @@
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type { AuthStorage, ExtensionAPI } from "@mariozechner/pi-coding-agent";
 import {
   configLoader,
   emitSyntheticConfigUpdated,
@@ -12,6 +12,18 @@ import {
   type SyntheticFeatureId,
   seedSyntheticConfigIfMissing,
 } from "../../config";
+import { getSyntheticApiKey } from "../../lib/env";
+import { QuotaStore } from "../../services/quota-store";
+import {
+  parseQuotaHeader,
+  type QuotasResponse,
+  SYNTHETIC_QUOTAS_READ_EVENT,
+  SYNTHETIC_QUOTAS_REQUEST_EVENT,
+  SYNTHETIC_QUOTAS_UPDATED_EVENT,
+  type SyntheticQuotasReadPayload,
+  type SyntheticQuotasRequestPayload,
+} from "../../types/quotas";
+import { fetchQuotas } from "../../utils/quotas";
 import { SYNTHETIC_MODELS } from "./models";
 
 export function buildSyntheticProviderModels(includeProxiedModels: boolean) {
@@ -71,6 +83,54 @@ export default async function (pi: ExtensionAPI) {
     getLoadedFeatures: () => loadedFeatures,
   });
 
+  const quotaStore = new QuotaStore();
+  let currentAuthStorage: AuthStorage | undefined;
+
+  async function fetchQuotasFromAuth(): Promise<QuotasResponse | undefined> {
+    if (!currentAuthStorage) return undefined;
+    const apiKey = await getSyntheticApiKey(currentAuthStorage);
+    if (!apiKey) return undefined;
+    const result = await fetchQuotas(apiKey);
+    return result.success ? result.data.quotas : undefined;
+  }
+
+  quotaStore.subscribe((snapshot) => {
+    pi.events.emit(SYNTHETIC_QUOTAS_UPDATED_EVENT, {
+      quotas: snapshot.quotas,
+      source: snapshot.source,
+      updatedAt: snapshot.updatedAt,
+    });
+  });
+
+  pi.on("after_provider_response", (event, ctx) => {
+    if (ctx.model?.provider !== "synthetic") return;
+    const quotas = parseQuotaHeader(event.headers);
+    if (quotas) quotaStore.ingest(quotas, "header");
+  });
+
+  pi.events.on(SYNTHETIC_QUOTAS_REQUEST_EVENT, async (data: unknown) => {
+    const payload = data as SyntheticQuotasRequestPayload | undefined;
+    const snapshot = await quotaStore.refreshFromApi(fetchQuotasFromAuth);
+    if (payload?.respond) {
+      payload.respond(snapshot);
+    }
+  });
+
+  pi.events.on(SYNTHETIC_QUOTAS_READ_EVENT, (data: unknown) => {
+    const { respond } = data as SyntheticQuotasReadPayload;
+    respond(quotaStore.getSnapshot());
+  });
+
+  pi.on("session_before_switch", () => {
+    quotaStore.clear();
+    currentAuthStorage = undefined;
+  });
+
+  pi.on("session_shutdown", () => {
+    quotaStore.clear();
+    currentAuthStorage = undefined;
+  });
+
   pi.on("session_start", async (_event, ctx) => {
     const messages = pendingMessages.splice(0).map((m) => `- ${m}`);
     if (messages.length > 0) {
@@ -81,7 +141,13 @@ export default async function (pi: ExtensionAPI) {
     }
 
     loadedFeatures.clear();
+    quotaStore.clear();
+    currentAuthStorage = ctx.modelRegistry.authStorage;
     pi.events.emit(SYNTHETIC_EXTENSIONS_REQUEST_EVENT, undefined);
     emitSyntheticConfigUpdated(pi);
+
+    if (ctx.model?.provider === "synthetic") {
+      await quotaStore.refreshFromApi(fetchQuotasFromAuth);
+    }
   });
 }

package/src/extensions/quota-warnings/index.ts

@@ -1,4 +1,7 @@
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
+import type {
+  ExtensionAPI,
+  ExtensionContext,
+} from "@mariozechner/pi-coding-agent";
 import {
   configLoader,
   SYNTHETIC_CONFIG_UPDATED_EVENT,
@@ -6,48 +9,100 @@ import {
   SYNTHETIC_EXTENSIONS_REQUEST_EVENT,
   type SyntheticConfigUpdatedPayload,
 } from "../../config";
-import { clearAlertState, triggerCheck } from "./notifier";
+import { QuotaWarningNotifier } from "../../services/quota-warnings";
+import {
+  SYNTHETIC_QUOTAS_READ_EVENT,
+  SYNTHETIC_QUOTAS_REQUEST_EVENT,
+  type SyntheticQuotasReadPayload,
+  type SyntheticQuotasRequestPayload,
+  type SyntheticQuotasSnapshotPayload,
+} from "../../types/quotas";
 
 export default async function (pi: ExtensionAPI) {
   await configLoader.load();
 
   let enabled = configLoader.getConfig().quotaWarnings;
-  let currentModel: { provider: string; id: string } | undefined;
-  let currentContext: Parameters<typeof triggerCheck>[0] | undefined;
+
+  const notifier = new QuotaWarningNotifier();
+
+  function requestQuotas(
+    respond?: (snapshot: SyntheticQuotasSnapshotPayload | undefined) => void,
+  ): void {
+    pi.events.emit(SYNTHETIC_QUOTAS_REQUEST_EVENT, {
+      respond,
+    } satisfies SyntheticQuotasRequestPayload);
+  }
+
+  function readQuotas(
+    respond: (snapshot: SyntheticQuotasSnapshotPayload | undefined) => void,
+  ): void {
+    pi.events.emit(SYNTHETIC_QUOTAS_READ_EVENT, {
+      respond,
+    } satisfies SyntheticQuotasReadPayload);
+  }
+
+  function evaluateFromStoreOrRefresh(ctx: ExtensionContext): void {
+    if (!enabled || ctx.model?.provider !== "synthetic") return;
+    readQuotas((snapshot) => {
+      if (snapshot) {
+        notifier.evaluate(
+          snapshot.quotas,
+          snapshot.source === "header",
+          (message, level) => {
+            ctx.ui.notify(message, level);
+          },
+        );
+      } else {
+        requestQuotas((refreshed) => {
+          if (!refreshed) return;
+          notifier.evaluate(
+            refreshed.quotas,
+            refreshed.source === "header",
+            (message, level) => {
+              ctx.ui.notify(message, level);
+            },
+          );
+        });
+      }
+    });
+  }
 
   pi.events.on(SYNTHETIC_CONFIG_UPDATED_EVENT, (data: unknown) => {
     enabled = (data as SyntheticConfigUpdatedPayload).config.quotaWarnings;
 
     if (!enabled) {
-      clearAlertState();
+      notifier.clearAlertState();
       return;
     }
 
-    if (currentContext && currentModel?.provider === "synthetic") {
-      clearAlertState();
-      triggerCheck(currentContext, currentModel, false);
-    }
+    notifier.clearAlertState();
+    // In config updates we don't have ctx, so we just clear. The next lifecycle event will refresh.
   });
 
-  pi.on("session_start", async (_event, ctx) => {
-    currentContext = ctx;
-    currentModel = ctx.model;
-    if (!enabled || ctx.model?.provider !== "synthetic") return;
-    clearAlertState();
-    triggerCheck(ctx, ctx.model, false);
+  pi.on("session_start", (_event, ctx) => {
+    notifier.clearAlertState();
+    evaluateFromStoreOrRefresh(ctx);
   });
 
-  pi.on("agent_end", async (_event, ctx) => {
-    currentContext = ctx;
-    currentModel = ctx.model;
-    if (!enabled || ctx.model?.provider !== "synthetic") return;
-    triggerCheck(ctx, ctx.model, true);
+  pi.on("model_select", (_event, ctx) => {
+    notifier.clearAlertState();
+    evaluateFromStoreOrRefresh(ctx);
+  });
+
+  pi.on("agent_end", (_event, ctx) => {
+    evaluateFromStoreOrRefresh(ctx);
+  });
+
+  pi.on("turn_end", (_event, ctx) => {
+    evaluateFromStoreOrRefresh(ctx);
+  });
+
+  pi.on("session_before_switch", () => {
+    notifier.clearAlertState();
   });
 
-  pi.on("session_shutdown", async () => {
-    currentContext = undefined;
-    currentModel = undefined;
-    clearAlertState();
+  pi.on("session_shutdown", () => {
+    notifier.clearAlertState();
   });
 
   pi.events.on(SYNTHETIC_EXTENSIONS_REQUEST_EVENT, () => {

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, () => {