fix: stabilize Claude usage cache across token refreshes

src/claude-usage.test.ts

@@ -1,6 +1,10 @@
 import { describe, expect, it } from 'vitest';
 
-import type { ClaudeUsageData } from './claude-usage.js';
+import {
+  getUsageCacheReadKeys,
+  getUsageCacheWriteKey,
+  type ClaudeUsageData,
+} from './claude-usage.js';
 
 describe('ClaudeUsageData', () => {
   it('represents the API response structure correctly', () => {
@@ -28,4 +32,23 @@ describe('ClaudeUsageData', () => {
     expect(data.five_hour?.utilization).toBe(10);
     expect(data.seven_day).toBeUndefined();
   });
+
+  it('prefers account-based cache keys and falls back to token suffixes', () => {
+    expect(
+      getUsageCacheReadKeys(
+        'sk-ant-oat01-currentHNCJmAAA',
+        0,
+        'sk-ant-oat01-credsEqnOPAAA',
+      ),
+    ).toEqual(['account-0', 'EqnOPAAA', 'HNCJmAAA']);
+  });
+
+  it('writes cache under account key when account index is known', () => {
+    expect(getUsageCacheWriteKey('sk-ant-oat01-currentHNCJmAAA', 0)).toBe(
+      'account-0',
+    );
+    expect(getUsageCacheWriteKey('sk-ant-oat01-currentHNCJmAAA')).toBe(
+      'HNCJmAAA',
+    );
+  });
 });

src/claude-usage.ts

@@ -71,11 +71,43 @@ function saveUsageDiskCache(): void {
   }
 }
 
-function cacheKey(token: string): string {
+function legacyTokenCacheKey(token: string): string {
   // Use last 8 chars — prefix is always "sk-ant-oat01-" for all tokens
   return token.slice(-8);
 }
 
+function accountCacheKey(accountIndex: number): string {
+  return `account-${accountIndex}`;
+}
+
+export function getUsageCacheWriteKey(
+  token: string,
+  accountIndex?: number,
+): string {
+  return accountIndex != null
+    ? accountCacheKey(accountIndex)
+    : legacyTokenCacheKey(token);
+}
+
+export function getUsageCacheReadKeys(
+  token: string,
+  accountIndex?: number,
+  credentialsAccessToken?: string | null,
+): string[] {
+  const keys: string[] = [];
+  if (accountIndex != null) keys.push(accountCacheKey(accountIndex));
+
+  if (credentialsAccessToken) {
+    const credsKey = legacyTokenCacheKey(credentialsAccessToken);
+    if (!keys.includes(credsKey)) keys.push(credsKey);
+  }
+
+  const tokenKey = legacyTokenCacheKey(token);
+  if (!keys.includes(tokenKey)) keys.push(tokenKey);
+
+  return keys;
+}
+
 // Rate limit: at most one API call per token per 5 minutes
 const MIN_FETCH_INTERVAL_MS = 300_000;
 
@@ -86,8 +118,26 @@ async function fetchUsageForToken(
   loadUsageDiskCache();
 
   // Return cached data if attempted recently (avoid API rate-limit)
-  const key = cacheKey(token);
+  const writeKey = getUsageCacheWriteKey(token, accountIndex);
-  const cached = usageDiskCache[key];
+  const readKeys = getUsageCacheReadKeys(
+    token,
+    accountIndex,
+    accountIndex != null ? readCredentialsAccessToken(accountIndex) : null,
+  );
+  let cachedKey: string | null = null;
+  let cached: UsageCacheEntry | undefined;
+  for (const key of readKeys) {
+    const entry = usageDiskCache[key];
+    if (!entry) continue;
+    cachedKey = key;
+    cached = entry;
+    break;
+  }
+  if (cached && cachedKey && cachedKey !== writeKey && !usageDiskCache[writeKey]) {
+    usageDiskCache[writeKey] = { ...cached };
+    cached = usageDiskCache[writeKey];
+    saveUsageDiskCache();
+  }
   const lastAttempt = cached?.lastAttemptAt ?? cached?.fetchedAt ?? 0;
   if (cached && Date.now() - lastAttempt < MIN_FETCH_INTERVAL_MS) {
     return cached.usage;
@@ -110,7 +160,8 @@ async function fetchUsageForToken(
       logger.warn(
         {
           account: accountIndex != null ? accountIndex + 1 : '?',
-          tokenKey: key,
+          tokenKey: legacyTokenCacheKey(token),
+          cacheKey: writeKey,
         },
         'Claude usage API: token expired or invalid (401)',
       );
@@ -121,7 +172,8 @@ async function fetchUsageForToken(
       logger.warn(
         {
           account: accountIndex != null ? accountIndex + 1 : '?',
-          tokenKey: key,
+          tokenKey: legacyTokenCacheKey(token),
+          cacheKey: writeKey,
           staleMinutes: Math.round(staleMs / 60_000),
         },
         'Claude usage API: rate limited (429), returning cached data',
@@ -138,7 +190,8 @@ async function fetchUsageForToken(
         {
           account: accountIndex != null ? accountIndex + 1 : '?',
           status: res.status,
-          tokenKey: key,
+          tokenKey: legacyTokenCacheKey(token),
+          cacheKey: writeKey,
         },
         `Claude usage API: unexpected status ${res.status}`,
       );
@@ -160,12 +213,18 @@ async function fetchUsageForToken(
 
     // Persist to disk cache — success: update both fetchedAt and lastAttemptAt
     const now = Date.now();
-    usageDiskCache[key] = { usage: result, fetchedAt: now, lastAttemptAt: now };
+    usageDiskCache[writeKey] = {
+      usage: result,
+      fetchedAt: now,
+      lastAttemptAt: now,
+    };
     saveUsageDiskCache();
 
     logger.debug(
       {
-        tokenKey: key,
+        account: accountIndex != null ? accountIndex + 1 : '?',
+        tokenKey: legacyTokenCacheKey(token),
+        cacheKey: writeKey,
         h5: result.five_hour?.utilization,
         d7: result.seven_day?.utilization,
       },
@@ -175,9 +234,24 @@ async function fetchUsageForToken(
     return result;
   } catch (err) {
     if ((err as Error).name === 'AbortError') {
-      logger.warn({ tokenKey: key }, 'Claude usage API: request timed out');
+      logger.warn(
+        {
+          account: accountIndex != null ? accountIndex + 1 : '?',
+          tokenKey: legacyTokenCacheKey(token),
+          cacheKey: writeKey,
+        },
+        'Claude usage API: request timed out',
+      );
     } else {
-      logger.warn({ err, tokenKey: key }, 'Claude usage API: fetch failed');
+      logger.warn(
+        {
+          err,
+          account: accountIndex != null ? accountIndex + 1 : '?',
+          tokenKey: legacyTokenCacheKey(token),
+          cacheKey: writeKey,
+        },
+        'Claude usage API: fetch failed',
+      );
     }
     // Record attempt time so we back off on persistent failures
     if (cached) {
@@ -236,6 +310,27 @@ function readCredentialsPlanType(accountIndex: number): string | null {
   }
 }
 
+function readCredentialsAccessToken(accountIndex: number): string | null {
+  try {
+    const credsPath =
+      accountIndex === 0
+        ? path.join(os.homedir(), '.claude', '.credentials.json')
+        : path.join(
+            os.homedir(),
+            '.claude-accounts',
+            String(accountIndex),
+            '.credentials.json',
+          );
+    if (!fs.existsSync(credsPath)) return null;
+    const data = readJsonFile<{
+      claudeAiOauth?: { accessToken?: string };
+    }>(credsPath);
+    return data?.claudeAiOauth?.accessToken || null;
+  } catch {
+    return null;
+  }
+}
+
 async function fetchProfileForToken(
   token: string,
 ): Promise<ClaudeAccountProfile | null> {

src/token-refresh.test.ts

@@ -1,6 +1,9 @@
 import { describe, expect, it } from 'vitest';
 
-import { shouldStartTokenRefreshLoop } from './token-refresh.js';
+import {
+  applyUpdatedTokensToEnvContent,
+  shouldStartTokenRefreshLoop,
+} from './token-refresh.js';
 
 describe('shouldStartTokenRefreshLoop', () => {
   it('starts refresh for the Claude service', () => {
@@ -10,4 +13,32 @@ describe('shouldStartTokenRefreshLoop', () => {
   it('skips refresh for the Codex service', () => {
     expect(shouldStartTokenRefreshLoop('codex')).toBe(false);
   });
+
+  it('updates both multi-token and single-token env vars when present', () => {
+    const next = applyUpdatedTokensToEnvContent(
+      [
+        'CLAUDE_CODE_OAUTH_TOKEN=old-primary',
+        'CLAUDE_CODE_OAUTH_TOKENS=old-primary,old-secondary',
+        'OTHER=value',
+      ].join('\n'),
+      ['new-primary', 'new-secondary'],
+    );
+
+    expect(next).toContain('CLAUDE_CODE_OAUTH_TOKEN=new-primary');
+    expect(next).toContain(
+      'CLAUDE_CODE_OAUTH_TOKENS=new-primary,new-secondary',
+    );
+  });
+
+  it('adds the multi-token env var when it is missing', () => {
+    const next = applyUpdatedTokensToEnvContent(
+      ['CLAUDE_CODE_OAUTH_TOKEN=old-primary', 'OTHER=value'].join('\n'),
+      ['new-primary', 'new-secondary'],
+    );
+
+    expect(next).toContain('CLAUDE_CODE_OAUTH_TOKEN=new-primary');
+    expect(next).toContain(
+      'CLAUDE_CODE_OAUTH_TOKENS=new-primary,new-secondary',
+    );
+  });
 });

src/token-refresh.ts

@@ -137,37 +137,45 @@ function syncToSessionDirs(credsPath: string): void {
 /**
  * Update CLAUDE_CODE_OAUTH_TOKENS in .env so refreshed tokens survive restarts.
  */
+export function applyUpdatedTokensToEnvContent(
+  content: string,
+  tokens: string[],
+): string {
+  if (tokens.length === 0) return content;
+  const multiValue = tokens.join(',');
+  const multiLineRe = /^CLAUDE_CODE_OAUTH_TOKENS=.*/m;
+  const singleLineRe = /^CLAUDE_CODE_OAUTH_TOKEN=.*/m;
+
+  let next = content;
+  if (multiLineRe.test(next)) {
+    next = next.replace(multiLineRe, `CLAUDE_CODE_OAUTH_TOKENS=${multiValue}`);
+  } else {
+    next = `${next.replace(/\s*$/, '')}\nCLAUDE_CODE_OAUTH_TOKENS=${multiValue}\n`;
+  }
+
+  if (singleLineRe.test(next)) {
+    next = next.replace(singleLineRe, `CLAUDE_CODE_OAUTH_TOKEN=${tokens[0]}`);
+  }
+
+  return next;
+}
+
 function updateEnvTokens(): void {
   const envFile = path.join(process.cwd(), '.env');
   try {
     if (!fs.existsSync(envFile)) return;
-    let content = fs.readFileSync(envFile, 'utf-8');
+    const content = fs.readFileSync(envFile, 'utf-8');
 
     const allTokens = getAllTokens();
     if (allTokens.length === 0) return;
 
-    const newValue = allTokens.map((t) => t.token).join(',');
+    const nextContent = applyUpdatedTokensToEnvContent(
-
+      content,
-    // Replace existing CLAUDE_CODE_OAUTH_TOKENS line, or append if missing
+      allTokens.map((t) => t.token),
-    const tokenLineRe = /^CLAUDE_CODE_OAUTH_TOKENS=.*/m;
-    if (tokenLineRe.test(content)) {
-      content = content.replace(
-        tokenLineRe,
-        `CLAUDE_CODE_OAUTH_TOKENS=${newValue}`,
     );
-    } else {
-      // Also update the single-token var if present
-      const singleLineRe = /^CLAUDE_CODE_OAUTH_TOKEN=.*/m;
-      if (singleLineRe.test(content)) {
-        content = content.replace(
-          singleLineRe,
-          `CLAUDE_CODE_OAUTH_TOKEN=${allTokens[0].token}`,
-        );
-      }
-    }
 
     const tempPath = `${envFile}.tmp`;
-    fs.writeFileSync(tempPath, content, { mode: 0o600 });
+    fs.writeFileSync(tempPath, nextContent, { mode: 0o600 });
     fs.renameSync(tempPath, envFile);
     logger.debug('Updated .env with refreshed tokens');
   } catch (err) {