From d865a6b07f3f575099701fe82c8948e1addfc060 Mon Sep 17 00:00:00 2001 From: Eyejoker Date: Thu, 26 Mar 2026 16:37:00 +0900 Subject: [PATCH] fix: stabilize Claude usage cache across token refreshes --- src/claude-usage.test.ts | 25 ++++++++- src/claude-usage.ts | 115 ++++++++++++++++++++++++++++++++++---- src/token-refresh.test.ts | 33 ++++++++++- src/token-refresh.ts | 50 ++++++++++------- 4 files changed, 190 insertions(+), 33 deletions(-) diff --git a/src/claude-usage.test.ts b/src/claude-usage.test.ts index 884a7bb..2639d63 100644 --- a/src/claude-usage.test.ts +++ b/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', + ); + }); }); diff --git a/src/claude-usage.ts b/src/claude-usage.ts index a07bf4a..c52e5ba 100644 --- a/src/claude-usage.ts +++ b/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 cached = usageDiskCache[key]; + const writeKey = getUsageCacheWriteKey(token, accountIndex); + 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 { diff --git a/src/token-refresh.test.ts b/src/token-refresh.test.ts index 8a11308..46699bc 100644 --- a/src/token-refresh.test.ts +++ b/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', + ); + }); }); diff --git a/src/token-refresh.ts b/src/token-refresh.ts index 14cf51d..f0b847d 100644 --- a/src/token-refresh.ts +++ b/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(','); - - // Replace existing CLAUDE_CODE_OAUTH_TOKENS line, or append if missing - 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 nextContent = applyUpdatedTokensToEnvContent( + content, + allTokens.map((t) => t.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) {