fix: stabilize Claude usage cache across token refreshes
This commit is contained in:
@@ -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',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user