pi update

pi/.pi/agent/extensions/usage-bars/index.ts

@@ -37,7 +37,17 @@ import {
   type UsageData,
 } from "./core";
 
-const CACHE_TTL_MS = 15 * 60 * 1000;           // reuse cached data for 15 min
+// Disk cache TTL for idle/background reads (session start, etc.)
+const CACHE_TTL_MS = 15 * 60 * 1000;
+// Shorter TTL for event-driven polls (after prompt submit / after turn end).
+// With the shared disk cache, only one pi instance per ACTIVE_CACHE_TTL_MS window
+// will actually hit the API — the rest will read from the cached result.
+const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
+// How often to re-poll while the model is actively streaming / running tools.
+// Combined with the shared disk cache this means at most one HTTP request per
+// STREAMING_POLL_INTERVAL_MS regardless of how many pi sessions are open.
+const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
+
 const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000; // 1 hour back-off after 429
 const STATUS_KEY = "usage-bars";
 
@@ -297,6 +307,17 @@ interface UsageState extends UsageByProvider {
   activeProvider: ProviderKey | null;
 }
 
+interface PollOptions {
+  /** Override the disk-cache TTL for this poll (default: CACHE_TTL_MS). */
+  cacheTtl?: number;
+  /**
+   * Skip the shared disk-cache TTL check entirely and always fetch from the
+   * API. Used after an account switch where the cached data belongs to a
+   * different account.
+   */
+  forceFresh?: boolean;
+}
+
 export default function (pi: ExtensionAPI) {
   const endpoints = resolveUsageEndpoints();
   const state: UsageState = {
@@ -311,6 +332,8 @@ export default function (pi: ExtensionAPI) {
 
   let pollInFlight: Promise<void> | null = null;
   let pollQueued = false;
+  /** Timer running during an active agent loop to refresh usage periodically. */
+  let streamingTimer: ReturnType<typeof setInterval> | null = null;
   let ctx: any = null;
 
   function renderPercent(theme: any, value: number): string {
@@ -410,7 +433,7 @@ export default function (pi: ExtensionAPI) {
     return false;
   }
 
-  async function runPoll() {
+  async function runPoll(options: PollOptions = {}) {
     const auth = readAuth();
     const active = state.activeProvider;
 
@@ -420,25 +443,21 @@ export default function (pi: ExtensionAPI) {
       return;
     }
 
-    // --- Shared disk cache check ---
-    // All pi sessions read and write the same cache file so that only one
-    // process hits the API per CACHE_TTL_MS window, no matter how many
-    // sessions are open at once.
     const cache = readUsageCache();
     const now = Date.now();
+    const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
 
     // Respect cross-session rate-limit back-off written by any session.
     const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
     if (now < blockedUntil) {
-      // Use whatever data is in the cache so the bar still shows something.
       if (cache?.data?.[active]) state[active] = cache.data[active]!;
       state.lastPoll = now;
       updateStatus();
       return;
     }
 
-    // If another session already polled recently, use their result.
-    if (cache && now - cache.timestamp < CACHE_TTL_MS && cache.data?.[active]) {
+    // Use shared disk cache unless forceFresh is set (e.g. after account switch).
+    if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
       state[active] = cache.data[active]!;
       state.lastPoll = now;
       updateStatus();
@@ -446,13 +465,6 @@ export default function (pi: ExtensionAPI) {
     }
 
     // --- Proactive token refresh ---
-    // Before hitting the API, check whether the stored access token is expired.
-    // This is the main cause of HTTP 401 errors: switching accounts via
-    // /switch-claude restores a profile whose access token has since expired
-    // (the refresh token is still valid). We use pi's own OAuth resolver so
-    // the new tokens are written back to auth.json and the profile stays in
-    // sync. This is safe at turn_start / session_start because pi hasn't made
-    // any Claude API calls yet, so there's no parallel refresh to conflict with.
     const oauthId = providerToOAuthProviderId(active);
     let effectiveAuth = auth;
     if (oauthId && active !== "zai") {
@@ -506,16 +518,8 @@ export default function (pi: ExtensionAPI) {
 
     state[active] = result;
 
-    // Write result + rate-limit state to shared cache so other sessions
-    // don't need to re-hit the API within CACHE_TTL_MS.
-    //
-    // Error results (other than 429) are NOT cached: they should be retried
-    // on the next input instead of being replayed from cache for 15 minutes.
-    // The most common error is HTTP 401 (expired token after an account switch)
-    // which resolves on the very next poll once the token is refreshed above.
     if (result.error) {
       if (result.error === "HTTP 429") {
-        // Write rate-limit backoff but preserve the last good data in cache.
         const nextCache: import("./core").UsageCache = {
           timestamp: cache?.timestamp ?? now,
           data: { ...(cache?.data ?? {}) },
@@ -541,7 +545,7 @@ export default function (pi: ExtensionAPI) {
     updateStatus();
   }
 
-  async function poll() {
+  async function poll(options: PollOptions = {}) {
     if (pollInFlight) {
       pollQueued = true;
       await pollInFlight;
@@ -550,7 +554,7 @@ export default function (pi: ExtensionAPI) {
 
     do {
       pollQueued = false;
-      pollInFlight = runPoll()
+      pollInFlight = runPoll(options)
         .catch(() => {
           // Never crash extension event handlers on transient polling errors.
         })
@@ -562,6 +566,22 @@ export default function (pi: ExtensionAPI) {
     } while (pollQueued);
   }
 
+  function startStreamingTimer() {
+    if (streamingTimer !== null) return; // already running
+    streamingTimer = setInterval(() => {
+      void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
+    }, STREAMING_POLL_INTERVAL_MS);
+  }
+
+  function stopStreamingTimer() {
+    if (streamingTimer !== null) {
+      clearInterval(streamingTimer);
+      streamingTimer = null;
+    }
+  }
+
+  // ── Session lifecycle ────────────────────────────────────────────────────
+
   pi.on("session_start", async (_event, _ctx) => {
     ctx = _ctx;
     updateProviderFrom(_ctx.model);
@@ -569,19 +589,13 @@ export default function (pi: ExtensionAPI) {
   });
 
   pi.on("session_shutdown", async (_event, _ctx) => {
+    stopStreamingTimer();
     if (_ctx?.hasUI) {
       _ctx.ui.setStatus(STATUS_KEY, undefined);
     }
   });
 
-  // Refresh usage on every turn (like claude-pulse's UserPromptSubmit hook).
-  // The disk cache means the API is only hit at most once per CACHE_TTL_MS
-  // regardless of how many turns or sessions are active.
-  pi.on("turn_start", async (_event, _ctx) => {
-    ctx = _ctx;
-    updateProviderFrom(_ctx.model);
-    await poll();
-  });
+  // ── Model change ─────────────────────────────────────────────────────────
 
   pi.on("model_select", async (event, _ctx) => {
     ctx = _ctx;
@@ -589,6 +603,58 @@ export default function (pi: ExtensionAPI) {
     if (changed) await poll();
   });
 
+  // Keep provider detection up-to-date across turns (cheap, no API call unless
+  // the provider changed).
+  pi.on("turn_start", (_event, _ctx) => {
+    ctx = _ctx;
+    updateProviderFrom(_ctx.model);
+  });
+
+  // ── Agent loop ───────────────────────────────────────────────────────────
+
+  // Poll when the user submits a prompt — captures usage right before the
+  // new turn (useful to see current state before tokens are consumed).
+  pi.on("before_agent_start", async (_event, _ctx) => {
+    ctx = _ctx;
+    await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
+  });
+
+  // Start interval timer so usage stays fresh during long-running agents
+  // (lots of tool calls, extended thinking, etc.).
+  pi.on("agent_start", (_event, _ctx) => {
+    ctx = _ctx;
+    startStreamingTimer();
+  });
+
+  // When the agent finishes, stop the timer and do a final fresh poll to
+  // capture the usage that was just consumed.
+  pi.on("agent_end", async (_event, _ctx) => {
+    ctx = _ctx;
+    stopStreamingTimer();
+    await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
+  });
+
+  // ── Account switch ───────────────────────────────────────────────────────
+
+  // Re-poll immediately when the Claude account is switched via /switch-claude.
+  // We must invalidate the claude entry in the shared disk cache first —
+  // otherwise poll() would serve the previous account's data which is still
+  // within CACHE_TTL_MS.
+  pi.events.on("claude-account:switched", () => {
+    const cache = readUsageCache();
+    if (cache?.data?.claude) {
+      const nextCache: import("./core").UsageCache = {
+        ...cache,
+        data: { ...cache.data },
+      };
+      delete nextCache.data.claude;
+      writeUsageCache(nextCache);
+    }
+    void poll({ forceFresh: true });
+  });
+
+  // ── /usage command ───────────────────────────────────────────────────────
+
   pi.registerCommand("usage", {
     description: "Show API usage for all subscriptions",
     handler: async (_args, _ctx) => {
@@ -609,13 +675,8 @@ export default function (pi: ExtensionAPI) {
           });
         }
       } finally {
-        await poll();
+        await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
       }
     },
   });
-
-  // Re-poll immediately when the Claude account is switched via /switch-claude
-  pi.events.on("claude-account:switched", () => {
-    void poll();
-  });
 }