pi extension fixes

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

@@ -20,6 +20,7 @@ import {
   clampPercent,
   colorForPercent,
   detectProvider,
+  ensureFreshAuthForProviders,
   fetchAllUsages,
   fetchClaudeUsage,
   fetchCodexUsage,
@@ -30,6 +31,7 @@ import {
   readUsageCache,
   resolveUsageEndpoints,
   writeUsageCache,
+  type OAuthProviderId,
   type ProviderKey,
   type UsageByProvider,
   type UsageData,
@@ -443,34 +445,60 @@ export default function (pi: ExtensionAPI) {
       return;
     }
 
-    // --- Actually hit the API ---
-    // Skip independent token refresh — pi manages OAuth tokens and refreshes
-    // them in memory. A parallel refresh here would cause token rotation
-    // conflicts (Anthropic invalidates the old refresh token on use).
+    // --- 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") {
+      const creds = auth[oauthId as keyof typeof auth] as
+        | { access?: string; refresh?: string; expires?: number }
+        | undefined;
+      const expires = typeof creds?.expires === "number" ? creds.expires : 0;
+      const tokenExpiredOrMissing =
+        !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
+      if (tokenExpiredOrMissing && creds?.refresh) {
+        try {
+          const refreshed = await ensureFreshAuthForProviders([oauthId as OAuthProviderId], {
+            auth,
+            persist: true,
+          });
+          if (refreshed.auth) effectiveAuth = refreshed.auth;
+        } catch {
+          // Ignore refresh errors — fall through with existing auth
+        }
+      }
+    }
+
     let result: UsageData;
 
     if (active === "codex") {
-      const access = auth["openai-codex"]?.access;
+      const access = effectiveAuth["openai-codex"]?.access;
       result = access
         ? await fetchCodexUsage(access)
         : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
     } else if (active === "claude") {
-      const access = auth.anthropic?.access;
+      const access = effectiveAuth.anthropic?.access;
       result = access
         ? await fetchClaudeUsage(access)
         : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
     } else if (active === "zai") {
-      const token = auth.zai?.access || auth.zai?.key;
+      const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
       result = token
         ? await fetchZaiUsage(token, { endpoints })
         : { session: 0, weekly: 0, error: "missing token (try /login again)" };
     } else if (active === "gemini") {
-      const creds = auth["google-gemini-cli"];
+      const creds = effectiveAuth["google-gemini-cli"];
       result = creds?.access
         ? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
         : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
     } else {
-      const creds = auth["google-antigravity"];
+      const creds = effectiveAuth["google-antigravity"];
       result = creds?.access
         ? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
         : { session: 0, weekly: 0, error: "missing access token (try /login again)" };
@@ -479,18 +507,35 @@ export default function (pi: ExtensionAPI) {
     state[active] = result;
 
     // Write result + rate-limit state to shared cache so other sessions
-    // (and our own next timer tick) don't need to re-hit the API.
-    const nextCache: import("./core").UsageCache = {
-      timestamp: now,
-      data: { ...(cache?.data ?? {}), [active]: result },
-      rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
-    };
-    if (result.error === "HTTP 429") {
-      nextCache.rateLimitedUntil![active] = now + RATE_LIMITED_BACKOFF_MS;
+    // 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 ?? {}) },
+          rateLimitedUntil: {
+            ...(cache?.rateLimitedUntil ?? {}),
+            [active]: now + RATE_LIMITED_BACKOFF_MS,
+          },
+        };
+        writeUsageCache(nextCache);
+      }
+      // All other errors: don't update cache — next turn will retry from scratch.
     } else {
+      const nextCache: import("./core").UsageCache = {
+        timestamp: now,
+        data: { ...(cache?.data ?? {}), [active]: result },
+        rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
+      };
       delete nextCache.rateLimitedUntil![active];
+      writeUsageCache(nextCache);
     }
-    writeUsageCache(nextCache);
 
     state.lastPoll = now;
     updateStatus();