diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 4420c3286..50c78e065 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed `omp usage invalidate` to discard stale OAuth and API-key usage snapshots, then force a cache-bypassing, per-provider serialized refresh so upgraded subscriptions do not silently retain pre-change quota data. + ## [17.3.3] - 2026-08-14 ### Fixed diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index bab6c6c54..43a0f9211 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -677,6 +677,7 @@ const DEFAULT_USAGE_PROVIDER_MAP = new Map( ); const USAGE_CACHE_PREFIX = "usage_cache:"; +const USAGE_FORCE_REFRESH_CACHE_PREFIX = "force-refresh:"; const USAGE_HEADER_INGEST_INTERVAL_MS = 60_000; const USAGE_LAST_GOOD_RETENTION_MS = 24 * 60 * 60_000; /** @@ -686,6 +687,12 @@ const USAGE_LAST_GOOD_RETENTION_MS = 24 * 60 * 60_000; * on the next poll. */ const USAGE_FAILURE_BACKOFF_MS = 10_000; +/** + * A manual invalidation persists across the next CLI process and serializes + * same-provider probes, avoiding a cold-account burst against IP-limited + * upstream usage endpoints. + */ +const USAGE_FORCE_REFRESH_TTL_MS = 5 * 60_000; // Bumped from 3s — Claude usage retries up to 3 times with exponential backoff // (~3.5s total worst case); a tight per-request budget aborts retries mid-cycle. const DEFAULT_USAGE_REQUEST_TIMEOUT_MS = 10_000; @@ -792,6 +799,11 @@ type UsageRequestDescriptor = { baseUrl?: string; }; +type ForcedUsageRefresh = { + all: boolean; + providers: Set; +}; + type AuthApiKeyOptions = { baseUrl?: string; modelId?: string; @@ -2967,6 +2979,47 @@ export class AuthStorage { return versionOverride === undefined ? provider : `${versionOverride}:${provider}`; } + #usageForceRefreshCacheKey(provider?: Provider): string { + return provider + ? `${USAGE_FORCE_REFRESH_CACHE_PREFIX}provider:${provider}` + : `${USAGE_FORCE_REFRESH_CACHE_PREFIX}all`; + } + + #markUsageForceRefresh(provider?: Provider): void { + this.#usageCache.set(this.#usageForceRefreshCacheKey(provider), { + value: true, + expiresAt: Date.now() + USAGE_FORCE_REFRESH_TTL_MS, + }); + } + + #hasUsageForceRefresh(provider?: Provider): boolean { + const key = this.#usageForceRefreshCacheKey(provider); + const entry = this.#usageCache.get(key); + if (entry?.value !== true) return false; + if (entry.expiresAt > Date.now()) return true; + this.#usageCache.set(key, { value: null, expiresAt: 0 }); + return false; + } + + #usageForceRefresh(requests: readonly UsageRequestDescriptor[]): ForcedUsageRefresh { + const all = this.#hasUsageForceRefresh(); + const providers = new Set(); + for (const request of requests) providers.add(request.provider); + if (!all) { + for (const provider of providers) { + if (!this.#hasUsageForceRefresh(provider)) providers.delete(provider); + } + } + return { all, providers }; + } + + #clearUsageForceRefresh(refresh: ForcedUsageRefresh): void { + if (refresh.all) this.#usageCache.set(this.#usageForceRefreshCacheKey(), { value: null, expiresAt: 0 }); + for (const provider of refresh.providers) { + this.#usageCache.set(this.#usageForceRefreshCacheKey(provider), { value: null, expiresAt: 0 }); + } + } + #buildUsageReportCacheKey(request: UsageRequestDescriptor): string { const baseUrl = this.#normalizeUsageBaseUrl(request.baseUrl) || "default"; const identity = this.#buildUsageCacheIdentity(request.credential); @@ -3219,19 +3272,23 @@ export class AuthStorage { } } - async #fetchUsageCached(request: UsageRequestDescriptor, timeoutMs?: number): Promise { + async #fetchUsageCached( + request: UsageRequestDescriptor, + timeoutMs?: number, + forceRefresh = false, + ): Promise { const cacheKey = this.#buildUsageReportCacheKey(request); const now = Date.now(); - const cached = this.#usageCache.get(cacheKey); + const cached = forceRefresh ? undefined : this.#usageCache.get(cacheKey); // Fresh cache hit: return whatever's there (success or null fallback). if (cached && cached.expiresAt > now) { return cached.value; } - const inFlight = this.#usageRequestInFlight.get(cacheKey); - if (inFlight) return inFlight; - const usageCacheEpoch = this.#usageCacheEpoch; + const inFlightKey = `${cacheKey}\0${usageCacheEpoch}`; + const inFlight = this.#usageRequestInFlight.get(inFlightKey); + if (inFlight) return inFlight; const promise = (async () => { const report = await this.#fetchUsageUncached(request, timeoutMs); if (usageCacheEpoch !== this.#usageCacheEpoch) return report; @@ -3251,7 +3308,8 @@ export class AuthStorage { // re-hit the endpoint on every poll. Most providers serve the last good // value through transient failures. Session-cookie providers can opt out // so an expired login does not display stale quota indefinitely. - const retainLastGood = this.#usageProviderResolver?.(request.provider)?.retainLastGoodOnFailure !== false; + const retainLastGood = + !forceRefresh && this.#usageProviderResolver?.(request.provider)?.retainLastGoodOnFailure !== false; const lastGood = retainLastGood ? (this.#usageCache.getStale(cacheKey)?.value ?? null) : null; @@ -3260,10 +3318,10 @@ export class AuthStorage { this.#usageCache.set(cacheKey, { value: lastGood, expiresAt: coolDown }); return lastGood; })().finally(() => { - this.#usageRequestInFlight.delete(cacheKey); + this.#usageRequestInFlight.delete(inFlightKey); }); - this.#usageRequestInFlight.set(cacheKey, promise); + this.#usageRequestInFlight.set(inFlightKey, promise); return promise; } @@ -3936,6 +3994,35 @@ export class AuthStorage { return true; } + /** + * Fetch every requested report, keeping normal polls parallel while a + * manually invalidated provider probes accounts one at a time. + */ + #fetchUsageRequests( + requests: readonly UsageRequestDescriptor[], + serializedProviders: ReadonlySet, + ): Promise> { + const tails = new Map>(); + return Promise.all( + requests.map(request => { + const forceRefresh = serializedProviders.has(request.provider); + if (!forceRefresh) { + return this.#fetchUsageCached(request, this.#usageRequestTimeoutMs); + } + const tail = tails.get(request.provider) ?? Promise.resolve(); + const current = tail.then(() => this.#fetchUsageCached(request, this.#usageRequestTimeoutMs, true)); + tails.set( + request.provider, + current.then( + () => undefined, + () => undefined, + ), + ); + return current; + }), + ); + } + async fetchUsageReports(options?: { baseUrlResolver?: (provider: Provider) => string | undefined; /** Caller's cancel signal; only rejects this caller, never the shared upstream fetch. */ @@ -3954,15 +4041,15 @@ export class AuthStorage { // dispatch + credential selection) coalesce into one upstream call. // Each caller's `signal` only cancels THAT caller's await; the // shared upstream fetch runs to completion so peers aren't punished. - const OVERRIDE_KEY = "__override__"; - let shared = this.#usageReportsInFlight.get(OVERRIDE_KEY); + const overrideKey = `__override__\0${this.#usageCacheEpoch}`; + let shared = this.#usageReportsInFlight.get(overrideKey); if (!shared) { // Don't forward the caller signal into the shared fetch — first caller's // abort would otherwise cancel the upstream for every peer. shared = override().finally(() => { - this.#usageReportsInFlight.delete(OVERRIDE_KEY); + this.#usageReportsInFlight.delete(overrideKey); }); - this.#usageReportsInFlight.set(OVERRIDE_KEY, shared); + this.#usageReportsInFlight.set(overrideKey, shared); } const reports = await raceUsageWithSignal(shared, options?.signal); if (shouldReconcileStoreHookReports && reports) this.#reconcileCodexUsageBlocksFromReports(reports); @@ -3982,7 +4069,8 @@ export class AuthStorage { // a single decorrelation snapshot for 30s, defeating the jitter (some // accounts can be missing from one fetch and present in the next; the // aggregate cache freezes whichever set landed first). - const cacheKey = this.#buildUsageReportsCacheKey(requests); + const forcedRefresh = this.#usageForceRefresh(requests); + const cacheKey = `${this.#buildUsageReportsCacheKey(requests)}\0${this.#usageCacheEpoch}`; const inFlight = this.#usageReportsInFlight.get(cacheKey); if (inFlight) return inFlight; @@ -3998,9 +4086,7 @@ export class AuthStorage { }); } - const results = await Promise.all( - requests.map(request => this.#fetchUsageCached(request, this.#usageRequestTimeoutMs)), - ); + const results = await this.#fetchUsageRequests(requests, forcedRefresh.providers); const reports = results.filter((report): report is UsageReport => report !== null); const deduped = this.#dedupeUsageReports(reports); // no outer cache write — see comment above. @@ -4021,6 +4107,7 @@ export class AuthStorage { }; }), }); + this.#clearUsageForceRefresh(forcedRefresh); return resolved; })().finally(() => { this.#usageReportsInFlight.delete(cacheKey); @@ -5739,26 +5826,28 @@ export class AuthStorage { /** * Drop report snapshots for a user-requested refresh so a failed probe - * cannot replay the pre-invalidation last-good value. + * cannot replay the pre-invalidation last-good value. The persisted marker + * makes the next same-provider refresh serial rather than a cold fan-out. */ async #clearUsageReportCache(provider?: string): Promise { this.#usageCacheEpoch += 1; const prefix = provider ? `report:${this.#usageCacheProviderKey(provider)}:` : "report:"; - if (this.#usageCache.deletePrefix(prefix)) return; - - // Third-party stores may not support prefix deletion. Clear every active - // request key instead, including API-key and environment credentials. - const requests = await this.#collectUsageRequests(); - for (const request of requests) { - if (provider && request.provider !== provider) continue; - this.#usageCache.set(this.#buildUsageReportCacheKey(request), { value: null, expiresAt: 0 }); + if (!this.#usageCache.deletePrefix(prefix)) { + // Third-party stores may not support prefix deletion. Clear every active + // request key instead, including API-key and environment credentials. + const requests = await this.#collectUsageRequests(); + for (const request of requests) { + if (provider && request.provider !== provider) continue; + this.#usageCache.set(this.#buildUsageReportCacheKey(request), { value: null, expiresAt: 0 }); + } } + if (!this.#fetchUsageReportsOverride && !this.#store.fetchUsageReports) this.#markUsageForceRefresh(provider); } /** * Discard cached usage reports before a user-requested refresh. The next - * read probes upstream; a failure reports no fresh usage instead of replaying - * an invalidated last-good snapshot. + * read probes upstream serially per provider; a failure reports no fresh + * usage instead of replaying an invalidated last-good snapshot. */ async invalidateUsageCache(provider?: string, signal?: AbortSignal): Promise { await this.#clearUsageReportCache(provider); diff --git a/packages/ai/test/auth-storage-usage-cache.test.ts b/packages/ai/test/auth-storage-usage-cache.test.ts index 13840ee23..893ee0bc4 100644 --- a/packages/ai/test/auth-storage-usage-cache.test.ts +++ b/packages/ai/test/auth-storage-usage-cache.test.ts @@ -342,7 +342,6 @@ describe("AuthStorage usage cache: last-good failure fallback", () => { }); }); - describe("AuthStorage usage cache: explicit invalidation", () => { it("clears cached API-key reports before the next usage read", async () => { const store = makeStore([ @@ -391,6 +390,138 @@ describe("AuthStorage usage cache: explicit invalidation", () => { storage.close(); } }); + + it("serializes a persisted Codex refresh and returns an upgraded plan", async () => { + const store = makeStore([ + { + id: 1, + provider: "openai-codex", + credential: { + type: "oauth", + access: "access-free", + refresh: "refresh-free", + expires: Date.now() + 3_600_000, + accountId: "account-free", + email: "free@example.com", + }, + disabledCause: null, + }, + { + id: 2, + provider: "openai-codex", + credential: { + type: "oauth", + access: "access-upgraded", + refresh: "refresh-upgraded", + expires: Date.now() + 3_600_000, + accountId: "account-upgraded", + email: "upgraded@example.com", + }, + disabledCause: null, + }, + { + id: 3, + provider: "openai-codex", + credential: { + type: "oauth", + access: "access-other", + refresh: "refresh-other", + expires: Date.now() + 3_600_000, + accountId: "account-other", + email: "other@example.com", + }, + disabledCause: null, + }, + ]); + let upgraded = false; + const refreshStarted = new Map>(); + const refreshReleases = new Map>(); + for (const accountId of ["account-free", "account-upgraded", "account-other"]) { + refreshStarted.set(accountId, Promise.withResolvers()); + refreshReleases.set(accountId, Promise.withResolvers()); + } + const startedAccounts: string[] = []; + const usageProvider: UsageProvider = { + id: "openai-codex", + supports: params => params.provider === "openai-codex" && params.credential.type === "oauth", + async fetchUsage(params) { + const accountId = params.credential.accountId; + if (!accountId) return null; + if (upgraded) { + const started = refreshStarted.get(accountId); + const release = refreshReleases.get(accountId); + if (!started || !release) throw new Error(`unexpected account ${accountId}`); + startedAccounts.push(accountId); + started.resolve(); + await release.promise; + } + return { + provider: "openai-codex", + fetchedAt: Date.now(), + limits: [ + { + id: "openai-codex:7d", + label: "7 days", + scope: { provider: "openai-codex", windowId: "7d" }, + amount: { used: 0, limit: 100, unit: "percent" }, + status: "ok", + }, + ], + metadata: { + accountId, + email: params.credential.email, + planType: upgraded && accountId === "account-upgraded" ? "pro" : "free", + }, + }; + }, + }; + const initialStorage = new AuthStorage(store, { + usageProviderResolver: provider => (provider === "openai-codex" ? usageProvider : undefined), + }); + await initialStorage.reload(); + try { + expect(await initialStorage.fetchUsageReports()).toHaveLength(3); + await initialStorage.invalidateUsageCache(); + } finally { + initialStorage.close(); + } + + upgraded = true; + const refreshedStorage = new AuthStorage(store, { + usageProviderResolver: provider => (provider === "openai-codex" ? usageProvider : undefined), + }); + await refreshedStorage.reload(); + try { + const refresh = refreshedStorage.fetchUsageReports(); + const freeStarted = refreshStarted.get("account-free"); + const freeRelease = refreshReleases.get("account-free"); + if (!freeStarted || !freeRelease) throw new Error("missing free-account refresh gates"); + await freeStarted.promise; + expect(startedAccounts).toEqual(["account-free"]); + freeRelease.resolve(); + + const upgradedStarted = refreshStarted.get("account-upgraded"); + const upgradedRelease = refreshReleases.get("account-upgraded"); + if (!upgradedStarted || !upgradedRelease) throw new Error("missing upgraded-account refresh gates"); + await upgradedStarted.promise; + expect(startedAccounts).toEqual(["account-free", "account-upgraded"]); + upgradedRelease.resolve(); + + const otherStarted = refreshStarted.get("account-other"); + const otherRelease = refreshReleases.get("account-other"); + if (!otherStarted || !otherRelease) throw new Error("missing other-account refresh gates"); + await otherStarted.promise; + expect(startedAccounts).toEqual(["account-free", "account-upgraded", "account-other"]); + otherRelease.resolve(); + + const reports = await refresh; + expect(reports?.find(report => report.metadata?.accountId === "account-upgraded")?.metadata?.planType).toBe( + "pro", + ); + } finally { + refreshedStorage.close(); + } + }); }); describe("AuthStorage usage cache: provider failure policy", () => { it("drops stale QwenCloud quota after the optional console session expires", async () => { diff --git a/packages/ai/test/remote-auth-store.test.ts b/packages/ai/test/remote-auth-store.test.ts index 3cc7c2492..56b203664 100644 --- a/packages/ai/test/remote-auth-store.test.ts +++ b/packages/ai/test/remote-auth-store.test.ts @@ -16,6 +16,7 @@ import { import { snapshotResponseSchema } from "@oh-my-pi/pi-ai/auth-broker/wire-schemas"; import * as oauthUtils from "@oh-my-pi/pi-ai/registry/oauth"; import type { UsageLimit, UsageReport } from "@oh-my-pi/pi-ai/usage"; +import * as claudeUsage from "@oh-my-pi/pi-ai/usage/claude"; import { removeWithRetries } from "../../utils/src/temp"; function requireLimit(report: UsageReport, id: string): UsageLimit { @@ -48,7 +49,9 @@ describe("RemoteAuthCredentialStore + AuthStorage integration", () => { accountId: "account-1", email: "a@example.com", }); - serverStorage = new AuthStorage(serverStore); + serverStorage = new AuthStorage(serverStore, { + usageProviderResolver: provider => (provider === "anthropic" ? claudeUsage.claudeUsageProvider : undefined), + }); await serverStorage.reload(); handle = startAuthBroker({ storage: serverStorage, @@ -1105,6 +1108,53 @@ describe("RemoteAuthCredentialStore + AuthStorage integration", () => { clientStorage.close(); }); + test("broker invalidation drops server-side last-good usage reports", async () => { + const credential = serverStore!.listAuthCredentials("anthropic")[0]; + if (!credential || credential.credential.type !== "oauth") throw new Error("expected OAuth credential"); + serverStore!.updateAuthCredential(credential.id, { + ...credential.credential, + expires: Date.now() + 3_600_000, + }); + await serverStorage!.reload(); + + let calls = 0; + const fetchSpy = vi.spyOn(claudeUsage.claudeUsageProvider, "fetchUsage").mockImplementation(async () => { + calls += 1; + if (calls > 1) return null; + return { + provider: "anthropic", + fetchedAt: Date.now(), + limits: [ + { + id: "anthropic:5h", + label: "Claude 5 Hour", + scope: { provider: "anthropic", windowId: "5h" }, + amount: { used: 80, limit: 100, unit: "percent" }, + status: "ok", + }, + ], + metadata: { accountId: "account-1", email: "a@example.com" }, + }; + }); + const brokerClient = new AuthBrokerClient({ url: handle!.url, token }); + const initialResult = await brokerClient.fetchSnapshot(); + if (initialResult.status !== 200) throw new Error("expected snapshot"); + const remoteStore = new RemoteAuthCredentialStore({ + client: brokerClient, + initialSnapshot: initialResult.snapshot, + }); + const clientStorage = new AuthStorage(remoteStore); + await clientStorage.reload(); + try { + expect(await clientStorage.fetchUsageReports()).toHaveLength(1); + await clientStorage.invalidateUsageCache(); + expect(await clientStorage.fetchUsageReports()).toEqual([]); + expect(fetchSpy).toHaveBeenCalledTimes(2); + } finally { + clientStorage.close(); + } + }); + test("account pool exposes only qualified usage reports for visible OAuth identities", async () => { const brokerClient = new AuthBrokerClient({ url: "http://127.0.0.1:9", token: "unused" }); const now = Date.now();