fix: stabilize Claude usage cache across token refreshes

This commit is contained in:
Eyejoker
2026-03-26 16:37:00 +09:00
parent 34c9ecd10b
commit d865a6b07f
4 changed files with 190 additions and 33 deletions

View File

@@ -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',
);
});
});

View File

@@ -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<ClaudeAccountProfile | null> {

View File

@@ -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',
);
});
});

View File

@@ -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) {