From 0ffc43bfc09b1d784c2e9afa261d32ee8d75dcfd Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 21:36:35 +0000 Subject: [PATCH] fix(ai): preserved fresh codex blocks - Delayed healthy-usage reconciliation for newly-set local Codex blocks so lagging /usage responses cannot immediately undo a real 429 backoff. - Added regression coverage for fresh usage-limit blocks that see healthy usage during selection. - Kept broker stale-block reconciliation covered by seeding a persisted-only block. Fixes #4980 --- packages/ai/src/auth-storage.ts | 20 +++ .../test/auth-storage-codex-selection.test.ts | 114 +++++++++++++----- 2 files changed, 103 insertions(+), 31 deletions(-) diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index bd7b42fb6..acdaff5e7 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -965,6 +965,8 @@ export class AuthStorage { #sessionLastCredential: Map> = new Map(); /** Maps provider:type -> credentialIndex -> blockedUntilMs for temporary backoff. */ #credentialBackoff: Map> = new Map(); + /** Earliest time a freshly-set in-memory block may be cleared by live usage reconciliation. */ + #credentialBackoffProbeAfter: Map> = new Map(); #usageProviderResolver?: (provider: Provider) => UsageProvider | undefined; #rankingStrategyResolver?: (provider: Provider) => CredentialRankingStrategy | undefined; #usageCache: UsageCache; @@ -1345,6 +1347,9 @@ export class AuthStorage { if (backoffMap.size === 0) { this.#credentialBackoff.delete(backoffKey); } + const probeAfterMap = this.#credentialBackoffProbeAfter.get(backoffKey); + probeAfterMap?.delete(credentialIndex); + if (probeAfterMap?.size === 0) this.#credentialBackoffProbeAfter.delete(backoffKey); return undefined; } return blockedUntil; @@ -1437,6 +1442,9 @@ export class AuthStorage { const nextBlockedUntil = Math.max(existing, blockedUntilMs); backoffMap.set(credentialIndex, nextBlockedUntil); this.#credentialBackoff.set(backoffKey, backoffMap); + const probeAfterMap = this.#credentialBackoffProbeAfter.get(backoffKey) ?? new Map(); + probeAfterMap.set(credentialIndex, Math.min(nextBlockedUntil, Date.now() + USAGE_REPORT_TTL_MS)); + this.#credentialBackoffProbeAfter.set(backoffKey, probeAfterMap); this.#invalidateUsageReportCache(provider); const upsertCredentialBlock = this.#store.upsertCredentialBlock?.bind(this.#store); @@ -4385,6 +4393,11 @@ export class AuthStorage { backoffMap.delete(index); if (backoffMap.size === 0) this.#credentialBackoff.delete(key); } + for (const [key, probeAfterMap] of this.#credentialBackoffProbeAfter) { + if (key !== providerKey && !key.startsWith(scopedPrefix)) continue; + probeAfterMap.delete(index); + if (probeAfterMap.size === 0) this.#credentialBackoffProbeAfter.delete(key); + } } /** @@ -4410,6 +4423,13 @@ export class AuthStorage { const blockScope = this.#rankingStrategyResolver?.(provider)?.blockScope?.({}); const blockedUntilMs = this.#getCredentialBlockedUntil(provider, providerKey, credentialIndex, blockScope); if (blockedUntilMs === undefined) return; + // `/usage` can lag the request path that just returned 429. Fresh local + // blocks get one usage-cache window before healthy reports may clear them. + const nowMs = Date.now(); + const scopedBackoffKey = this.#toScopedBackoffKey(providerKey, blockScope); + const globalProbeAfterMs = this.#credentialBackoffProbeAfter.get(providerKey)?.get(credentialIndex) ?? 0; + const scopedProbeAfterMs = this.#credentialBackoffProbeAfter.get(scopedBackoffKey)?.get(credentialIndex) ?? 0; + if (Math.max(globalProbeAfterMs, scopedProbeAfterMs) > nowMs) return; this.#clearCredentialBlocks(provider, credentialId); logger.info("Cleared stale Codex usage-limit block after healthy live usage report", { credentialId, diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 14d0e7c8a..3963880e0 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -551,6 +551,73 @@ describe("AuthStorage codex oauth ranking", () => { expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBe(blockedUntilMs); }); + test("keeps a fresh Codex usage-limit block when selection sees healthy usage", async () => { + if (!authStorage || !store?.getCredentialBlock) { + throw new Error("test setup failed"); + } + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-fresh-blocked", "fresh-blocked@example.com") }, + { type: "oauth", ...createCredential("acct-fresh-healthy", "fresh-healthy@example.com") }, + ]); + + usageByAccount.set( + "acct-fresh-blocked", + createCodexUsageReport({ + accountId: "acct-fresh-blocked", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "fresh-blocked@example.com", + accountId: "acct-fresh-blocked", + }, + }), + ); + usageByAccount.set( + "acct-fresh-healthy", + createCodexUsageReport({ + accountId: "acct-fresh-healthy", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "fresh-healthy@example.com", + accountId: "acct-fresh-healthy", + }, + }), + ); + + const blockedRow = store.listAuthCredentials("openai-codex").find(row => { + const credential = row.credential; + return credential.type === "oauth" && credential.accountId === "acct-fresh-blocked"; + }); + if (!blockedRow) throw new Error("expected blocked credential row"); + + let blockedSessionId: string | undefined; + for (let index = 0; index < 100; index += 1) { + const sessionId = `codex-fresh-block-selected-${index}`; + if ((await authStorage.getApiKey("openai-codex", sessionId)) === "api-acct-fresh-blocked") { + blockedSessionId = sessionId; + break; + } + } + if (!blockedSessionId) throw new Error("expected a session selecting the soon-blocked account"); + + const markResult = await authStorage.markUsageLimitReached("openai-codex", blockedSessionId, { + retryAfterMs: 6 * 24 * HOUR_MS, + }); + + expect(markResult.switched).toBe(true); + const selectionAfterBlock = await authStorage.getApiKey("openai-codex", blockedSessionId); + expect(selectionAfterBlock).not.toBe("api-acct-fresh-blocked"); + expect(selectionAfterBlock).toBe("api-acct-fresh-healthy"); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + }); test("an older in-flight healthy Codex usage report does not clear a newer usage-limit block", async () => { if (!authStorage || !store?.getCredentialBlock) { throw new Error("test setup failed"); @@ -639,7 +706,7 @@ describe("AuthStorage codex oauth ranking", () => { }); test("broker-sourced healthy Codex usage clears remote gateway backoff", async () => { - if (!authStorage || !store?.getCredentialBlock) { + if (!authStorage || !store?.getCredentialBlock || !store.upsertCredentialBlock) { throw new Error("test setup failed"); } @@ -679,6 +746,18 @@ describe("AuthStorage codex oauth ranking", () => { }), ); + const staleBlockedRow = store.listAuthCredentials("openai-codex").find(row => { + const credential = row.credential; + return credential.type === "oauth" && credential.accountId === "acct-broker-blocked"; + }); + if (!staleBlockedRow) throw new Error("expected stale blocked credential row"); + store.upsertCredentialBlock({ + credentialId: staleBlockedRow.id, + providerKey: "openai-codex:oauth", + blockScope: "shared", + blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, + }); + const token = "codex-broker-reconcile"; const handle = startAuthBroker({ storage: authStorage, @@ -688,18 +767,6 @@ describe("AuthStorage codex oauth ranking", () => { }); try { const brokerClient = new AuthBrokerClient({ url: handle.url, token }); - const originalUpsertCredentialBlock = brokerClient.upsertCredentialBlock.bind(brokerClient); - const blockPersisted = Promise.withResolvers(); - vi.spyOn(brokerClient, "upsertCredentialBlock").mockImplementation(async (id, block, signal) => { - try { - const response = await originalUpsertCredentialBlock(id, block, signal); - blockPersisted.resolve(); - return response; - } catch (error) { - blockPersisted.reject(error); - throw error; - } - }); const initialResult = await brokerClient.fetchSnapshot(); if (initialResult.status !== 200) throw new Error("expected broker snapshot"); const blockedRow = initialResult.snapshot.credentials.find(entry => { @@ -715,31 +782,16 @@ describe("AuthStorage codex oauth ranking", () => { const clientStorage = new AuthStorage(remoteStore); await clientStorage.reload(); try { - let blockedSessionId: string | undefined; - for (let index = 0; index < 100; index += 1) { - const sessionId = `broker-codex-local-block-${index}`; - const apiKey = await clientStorage.getApiKey("openai-codex", sessionId); - if (apiKey === "api-acct-broker-blocked") { - blockedSessionId = sessionId; - break; - } - } - if (!blockedSessionId) throw new Error("expected a session selecting the blocked account"); - - const markResult = await clientStorage.markUsageLimitReached("openai-codex", blockedSessionId, { - retryAfterMs: 6 * 24 * HOUR_MS, - }); - - expect(markResult.switched).toBe(true); expect(remoteStore.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); - await blockPersisted.promise; expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); await clientStorage.fetchUsageReports(); expect(remoteStore.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeUndefined(); expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeUndefined(); - expect(await clientStorage.getApiKey("openai-codex", blockedSessionId)).toBe("api-acct-broker-blocked"); + expect(await clientStorage.getApiKey("openai-codex", "broker-codex-reconciled")).toBe( + "api-acct-broker-blocked", + ); } finally { clientStorage.close(); remoteStore.close();