diff --git a/src/utils/__tests__/usage-fetch.test.ts b/src/utils/__tests__/usage-fetch.test.ts index 57e836c6..1067c25c 100644 --- a/src/utils/__tests__/usage-fetch.test.ts +++ b/src/utils/__tests__/usage-fetch.test.ts @@ -1,4 +1,5 @@ import type * as childProcess from 'child_process'; +import { createHash } from 'crypto'; import * as fs from 'fs'; import { createRequire } from 'module'; import * as os from 'os'; @@ -690,6 +691,137 @@ describe('fetchUsageData error handling', () => { } }); + it('refetches a fresh cache when the token fingerprint changes (account switch)', () => { + const harness = createProbeHarness(); + + try { + const home = harness.createTokenHome('account-switch'); + const cacheDir = path.join(home.home, '.cache', 'ccstatusline'); + fs.mkdirSync(cacheDir, { recursive: true }); + const cacheFile = path.join(cacheDir, 'usage.json'); + // A complete cache written under a different account's token. + fs.writeFileSync(cacheFile, JSON.stringify({ sessionUsage: 5, tokenHash: 'deadbeefdeadbeef' })); + const seededMtimeMs = fs.statSync(cacheFile).mtimeMs; + + const result = harness.runProbe({ + claudeConfigDir: home.claudeConfig, + home: home.home, + mode: 'success', + nowMs: seededMtimeMs + 5000, + pathDir: home.bin, + requiredFields: ['sessionUsage'], + responseBody: successResponseBody + }); + + // The cached fingerprint mismatches the live token, so the still-fresh + // cache is rejected and the API is hit once for the new account. + expect(result.requestCount).toBe(1); + expect(result.first.sessionUsage).toBe(42); + } finally { + harness.cleanup(); + } + }); + + it('does not serve a mismatched account cache during an active lock', () => { + const harness = createProbeHarness(); + + try { + const home = harness.createTokenHome('account-switch-active-lock'); + const cacheDir = path.join(home.home, '.cache', 'ccstatusline'); + fs.mkdirSync(cacheDir, { recursive: true }); + const cacheFile = path.join(cacheDir, 'usage.json'); + const lockFile = path.join(cacheDir, 'usage.lock'); + fs.writeFileSync(cacheFile, JSON.stringify({ sessionUsage: 5, tokenHash: 'deadbeefdeadbeef' })); + + const seededMtimeMs = fs.statSync(cacheFile).mtimeMs; + const lockedNowMs = seededMtimeMs + 5000; + fs.writeFileSync(lockFile, JSON.stringify({ + blockedUntil: Math.floor(lockedNowMs / 1000) + 30, + error: 'timeout' + })); + + const result = harness.runProbe({ + claudeConfigDir: home.claudeConfig, + home: home.home, + mode: 'unexpected', + nowMs: lockedNowMs, + pathDir: home.bin, + requiredFields: ['sessionUsage'] + }); + + expect(result.first).toEqual({ error: 'timeout' }); + expect(result.second).toEqual({ error: 'timeout' }); + expect(result.requestCount).toBe(0); + } finally { + harness.cleanup(); + } + }); + + it('does not serve a mismatched account cache during a rate-limit backoff', () => { + const harness = createProbeHarness(); + + try { + const home = harness.createTokenHome('account-switch-rate-limit'); + const cacheDir = path.join(home.home, '.cache', 'ccstatusline'); + fs.mkdirSync(cacheDir, { recursive: true }); + const cacheFile = path.join(cacheDir, 'usage.json'); + fs.writeFileSync(cacheFile, JSON.stringify({ sessionUsage: 5, tokenHash: 'deadbeefdeadbeef' })); + + const seededMtimeMs = fs.statSync(cacheFile).mtimeMs; + const rateLimitedNowMs = seededMtimeMs + 5000; + const result = harness.runProbe({ + claudeConfigDir: home.claudeConfig, + home: home.home, + mode: 'status', + nowMs: rateLimitedNowMs, + pathDir: home.bin, + requiredFields: ['sessionUsage'], + responseBody: rateLimitedResponseBody, + responseHeaders: { 'retry-after': '3600' }, + statusCode: 429 + }); + + expect(result.first).toEqual({ error: 'rate-limited' }); + expect(result.second).toEqual({ error: 'rate-limited' }); + expect(result.requestCount).toBe(1); + expect(parseLockContents(result.lockContents)).toEqual({ + blockedUntil: Math.floor(rateLimitedNowMs / 1000) + 3600, + error: 'rate-limited' + }); + } finally { + harness.cleanup(); + } + }); + + it('serves a fresh cache whose token fingerprint matches (same account)', () => { + const harness = createProbeHarness(); + + try { + const home = harness.createTokenHome('account-same'); + const cacheDir = path.join(home.home, '.cache', 'ccstatusline'); + fs.mkdirSync(cacheDir, { recursive: true }); + const cacheFile = path.join(cacheDir, 'usage.json'); + const matchingHash = createHash('sha256').update('test-token').digest('hex').slice(0, 16); + fs.writeFileSync(cacheFile, JSON.stringify({ sessionUsage: 5, tokenHash: matchingHash })); + const seededMtimeMs = fs.statSync(cacheFile).mtimeMs; + + const result = harness.runProbe({ + claudeConfigDir: home.claudeConfig, + home: home.home, + mode: 'unexpected', + nowMs: seededMtimeMs + 5000, + pathDir: home.bin, + requiredFields: ['sessionUsage'] + }); + + // Fingerprint matches and the cache is fresh, so it is served with no API call. + expect(result.requestCount).toBe(0); + expect(result.first.sessionUsage).toBe(5); + } finally { + harness.cleanup(); + } + }); + it('treats enabled extra usage without a monthly limit as complete for extra usage widget fields', () => { const harness = createProbeHarness(); diff --git a/src/utils/usage-fetch.ts b/src/utils/usage-fetch.ts index dd13691a..d9da4c63 100644 --- a/src/utils/usage-fetch.ts +++ b/src/utils/usage-fetch.ts @@ -1,4 +1,5 @@ import { execFileSync } from 'child_process'; +import { createHash } from 'crypto'; import * as fs from 'fs'; import * as https from 'https'; import { HttpsProxyAgent } from 'https-proxy-agent'; @@ -68,6 +69,8 @@ const CachedUsageDataSchema = z.object({ error: z.string().nullable().optional() }); +const CachedTokenHashSchema = z.object({ tokenHash: z.string().optional() }); + const UsageApiBucketSchema = z.looseObject({ utilization: z.number().nullable().optional(), resets_at: z.string().nullable().optional() @@ -133,6 +136,27 @@ function parseCachedUsageData(rawJson: string): UsageData | null { }; } +// One-way fingerprint of the usage token, persisted alongside the cache so a +// login switch (e.g. enterprise<->personal, a different token) invalidates the +// cache immediately instead of waiting out the TTL. A truncated SHA-256 is a +// stable identifier, not the token itself, so it is safe to write to disk. +function fingerprintUsageToken(token: string): string { + return createHash('sha256').update(token).digest('hex').slice(0, 16); +} + +function readCachedTokenHash(rawJson: string): string | undefined { + return parseJsonWithSchema(rawJson, CachedTokenHashSchema)?.tokenHash; +} + +function tokenHashMatches(cachedHash: string | undefined, currentHash: string | null): boolean { + // With no current token we cannot fingerprint-gate, so fall through to the + // existing no-token handling rather than discarding an otherwise usable cache. + if (currentHash === null) { + return true; + } + return cachedHash === currentHash; +} + function parseUsageApiResponse(rawJson: string): UsageData | null { const parsed = parseJsonWithSchema(rawJson, UsageApiResponseSchema); if (!parsed) { @@ -214,10 +238,11 @@ function hasRequiredUsageFields(data: UsageData, requiredFields: readonly UsageD function getStaleUsageOrError( error: UsageError, now: number, + currentTokenHash: string | null, errorCacheMaxAge = LOCK_MAX_AGE, requiredFields: readonly UsageDataField[] = [] ): UsageData { - const stale = readStaleUsageCache(); + const stale = readStaleUsageCache(currentTokenHash); if (stale && !stale.error && hasRequiredUsageFields(stale, requiredFields)) { return cacheUsageData(stale, now); } @@ -382,9 +407,13 @@ export function getUsageToken(): string | null { ?? readUsageTokenFromCredentialsFile(); } -function readStaleUsageCache(): UsageData | null { +function readStaleUsageCache(currentTokenHash: string | null): UsageData | null { try { - return parseCachedUsageData(fs.readFileSync(CACHE_FILE, 'utf8')); + const rawCache = fs.readFileSync(CACHE_FILE, 'utf8'); + if (!tokenHashMatches(readCachedTokenHash(rawCache), currentTokenHash)) { + return null; + } + return parseCachedUsageData(rawCache); } catch { return null; } @@ -561,13 +590,23 @@ export async function fetchUsageData(options: FetchUsageDataOptions = {}): Promi } } + // Resolve the token up front (before lock/rate-limit checks so auth + // failures are not masked as timeout) and fingerprint it so the file cache + // can be invalidated on an account switch: a different token, written by a + // logout/login, no longer matches the cached fingerprint. + const token = getUsageToken(); + const currentTokenHash = token ? fingerprintUsageToken(token) : null; + // Check file cache try { const stat = fs.statSync(CACHE_FILE); const fileAge = now - Math.floor(stat.mtimeMs / 1000); if (fileAge < CACHE_MAX_AGE) { - const fileData = parseCachedUsageData(fs.readFileSync(CACHE_FILE, 'utf8')); - if (fileData && !fileData.error && hasRequiredUsageFields(fileData, requiredFields)) { + const rawCache = fs.readFileSync(CACHE_FILE, 'utf8'); + const fileData = parseCachedUsageData(rawCache); + if (fileData && !fileData.error + && tokenHashMatches(readCachedTokenHash(rawCache), currentTokenHash) + && hasRequiredUsageFields(fileData, requiredFields)) { return cacheUsageData(fileData, now); } } @@ -575,10 +614,8 @@ export async function fetchUsageData(options: FetchUsageDataOptions = {}): Promi // File doesn't exist or read error - continue to API call } - // Get token before lock/rate-limit checks so auth failures are not masked as timeout. - const token = getUsageToken(); if (!token) { - return getStaleUsageOrError('no-credentials', now, LOCK_MAX_AGE, requiredFields); + return getStaleUsageOrError('no-credentials', now, currentTokenHash, LOCK_MAX_AGE, requiredFields); } const activeLock = readActiveUsageLock(now); @@ -586,6 +623,7 @@ export async function fetchUsageData(options: FetchUsageDataOptions = {}): Promi return getStaleUsageOrError( activeLock.error, now, + currentTokenHash, Math.max(1, activeLock.blockedUntil - now), requiredFields ); @@ -599,29 +637,29 @@ export async function fetchUsageData(options: FetchUsageDataOptions = {}): Promi if (response.kind === 'rate-limited') { writeUsageLock(now + response.retryAfterSeconds, 'rate-limited'); - return getStaleUsageOrError('rate-limited', now, response.retryAfterSeconds, requiredFields); + return getStaleUsageOrError('rate-limited', now, currentTokenHash, response.retryAfterSeconds, requiredFields); } if (response.kind === 'error') { - return getStaleUsageOrError('api-error', now, LOCK_MAX_AGE, requiredFields); + return getStaleUsageOrError('api-error', now, currentTokenHash, LOCK_MAX_AGE, requiredFields); } const usageData = parseUsageApiResponse(response.body); if (!usageData) { writeUsageLock(now + LOCK_MAX_AGE, 'parse-error'); - return getStaleUsageOrError('parse-error', now, LOCK_MAX_AGE, requiredFields); + return getStaleUsageOrError('parse-error', now, currentTokenHash, LOCK_MAX_AGE, requiredFields); } // Validate we got actual data if (usageData.sessionUsage === undefined && usageData.weeklyUsage === undefined) { writeUsageLock(now + LOCK_MAX_AGE, 'parse-error'); - return getStaleUsageOrError('parse-error', now, LOCK_MAX_AGE, requiredFields); + return getStaleUsageOrError('parse-error', now, currentTokenHash, LOCK_MAX_AGE, requiredFields); } // Save to cache try { ensureCacheDirExists(); - fs.writeFileSync(CACHE_FILE, JSON.stringify(usageData)); + fs.writeFileSync(CACHE_FILE, JSON.stringify({ ...usageData, tokenHash: currentTokenHash ?? undefined })); } catch { // Ignore cache write errors } @@ -629,6 +667,6 @@ export async function fetchUsageData(options: FetchUsageDataOptions = {}): Promi return cacheUsageData(usageData, now); } catch { writeUsageLock(now + LOCK_MAX_AGE, 'parse-error'); - return getStaleUsageOrError('parse-error', now, LOCK_MAX_AGE, requiredFields); + return getStaleUsageOrError('parse-error', now, currentTokenHash, LOCK_MAX_AGE, requiredFields); } }