From c143185c018e98812d44d365c0710fa5006f5d20 Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 21:03:25 +0000 Subject: [PATCH 1/7] fix(ai): rechecked codex blocks during selection - Re-fetched usage for blocked Codex OAuth candidates during ranking so fresh recovered windows can clear stale persisted blocks. - Relaxed Codex block reconciliation to trust live allowed/limitReached metadata with an available primary window. - Added regression coverage for selection-path stale block recovery. Fixes #4980 --- packages/ai/CHANGELOG.md | 4 + packages/ai/src/auth-storage.ts | 51 +++++++---- .../test/auth-storage-codex-selection.test.ts | 91 +++++++++++++++---- 3 files changed, 112 insertions(+), 34 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index e6d81932f..a74b0c89c 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed Codex OAuth credential selection to re-check blocked accounts during ranking and clear stale usage-limit blocks when live usage shows the 5-hour window recovered ([#4980](https://github.com/can1357/oh-my-pi/issues/4980)). + ## [16.3.14] - 2026-07-09 ### Changed diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 315d43c48..886052d4a 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -3316,18 +3316,36 @@ export class AuthStorage { args.order.map(async idx => { const selection = args.credentials[idx]; if (!selection) return null; - const blockedUntil = this.#getCredentialBlockedUntil( + let blockedUntil = this.#getCredentialBlockedUntil( args.provider, args.providerKey, selection.index, args.blockScope, ); - if (blockedUntil !== undefined) return { selection, usage: null, usageChecked: false, blockedUntil }; - const usage = await this.#getUsageReport(args.provider, selection.credential, { - ...args.options, - timeoutMs: this.#usageRequestTimeoutMs, - }); - return { selection, usage, usageChecked: true, blockedUntil: undefined as number | undefined }; + let usage: UsageReport | null = null; + let usageChecked = false; + if (blockedUntil !== undefined && args.provider === "openai-codex") { + usage = await this.#getUsageReport(args.provider, selection.credential, { + ...args.options, + timeoutMs: this.#usageRequestTimeoutMs, + }); + usageChecked = true; + blockedUntil = this.#getCredentialBlockedUntil( + args.provider, + args.providerKey, + selection.index, + args.blockScope, + ); + } + if (blockedUntil !== undefined) return { selection, usage, usageChecked, blockedUntil }; + if (!usageChecked) { + usage = await this.#getUsageReport(args.provider, selection.credential, { + ...args.options, + timeoutMs: this.#usageRequestTimeoutMs, + }); + usageChecked = true; + } + return { selection, usage, usageChecked, blockedUntil: undefined as number | undefined }; }), ); const timeoutSignal = Promise.withResolvers(); @@ -4371,19 +4389,18 @@ export class AuthStorage { /** * Self-heal a stale Codex usage-limit block: when a fresh live usage report - * shows the account is allowed and below every limit, drop the persisted and - * in-memory `openai-codex:oauth` blocks so the balancer re-includes it. Only - * Codex — its ranking strategy uses the single model-independent `"shared"` - * scope, so clearing every block for the credential id is exact. + * says the account is allowed and its primary rolling window is available, + * drop the persisted and in-memory `openai-codex:oauth` blocks so credential + * selection can re-include seats whose short window recovered before a + * longer persisted block naturally expires. */ #isHealthyCodexUsageReport(report: UsageReport): boolean { + if (report.provider !== "openai-codex") return false; const metadata = report.metadata; - return ( - report.provider === "openai-codex" && - metadata?.allowed === true && - metadata.limitReached === false && - !this.#isUsageLimitReached(report.limits) - ); + if (metadata?.allowed !== true || metadata.limitReached !== false) return false; + const primary = codexRankingStrategy.findWindowLimits(report).primary; + if (primary) return !this.#isUsageLimitExhausted(primary); + return !this.#isUsageLimitReached(report.limits); } #reconcileCodexUsageBlockForCredential(provider: Provider, credentialId: number, report: UsageReport): void { diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 6bf0ff584..a63ac005e 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -385,15 +385,6 @@ describe("AuthStorage codex oauth ranking", () => { }), ); - const blockedSelectionCounts = await countApiKeySelections( - authStorage, - "openai-codex", - "stale-codex-block-before-fetch", - 40, - ); - expect(countFor(blockedSelectionCounts, "api-acct-blocked")).toBe(0); - expect(countFor(blockedSelectionCounts, "api-acct-healthy")).toBeGreaterThan(0); - const generationBeforeFetch = authStorage.getGeneration(); await authStorage.fetchUsageReports(); @@ -411,6 +402,80 @@ describe("AuthStorage codex oauth ranking", () => { expect(countFor(reconciledSelectionCounts, "api-acct-healthy")).toBeGreaterThan(0); }); + test("re-evaluates a stale persisted Codex block during selection when the 5h window recovered", async () => { + if (!authStorage || !store?.upsertCredentialBlock || !store.getCredentialBlock) { + throw new Error("test setup failed"); + } + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-recovered-blocked", "recovered-blocked@example.com") }, + { type: "oauth", ...createCredential("acct-recovered-sibling", "recovered-sibling@example.com") }, + ]); + + const blockedRow = store.listAuthCredentials("openai-codex").find(row => { + const credential = row.credential; + return credential.type === "oauth" && credential.accountId === "acct-recovered-blocked"; + }); + if (!blockedRow) throw new Error("expected blocked credential row"); + + store.upsertCredentialBlock({ + credentialId: blockedRow.id, + providerKey: "openai-codex:oauth", + blockScope: "shared", + blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, + }); + + const fiveHourWindow: UsageWindowConfig = { + windowId: "5h", + windowLabel: "5 Hours", + durationMs: FIVE_HOUR_MS, + }; + + usageByAccount.set( + "acct-recovered-blocked", + createCodexUsageReport({ + accountId: "acct-recovered-blocked", + primary: { usedFraction: 0.2, resetInMs: 4 * HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: 6 * 24 * HOUR_MS }, + primaryWindow: fiveHourWindow, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "recovered-blocked@example.com", + accountId: "acct-recovered-blocked", + }, + }), + ); + usageByAccount.set( + "acct-recovered-sibling", + createCodexUsageReport({ + accountId: "acct-recovered-sibling", + primary: { usedFraction: 0.2, resetInMs: 4 * HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: 6 * 24 * HOUR_MS }, + primaryWindow: fiveHourWindow, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "recovered-sibling@example.com", + accountId: "acct-recovered-sibling", + }, + }), + ); + + const selectionCounts = await countApiKeySelections( + authStorage, + "openai-codex", + "codex-stale-block-selection-recovered", + 150, + ); + + expect(countFor(selectionCounts, "api-acct-recovered-blocked")).toBeGreaterThan(0); + expect(countFor(selectionCounts, "api-acct-recovered-sibling")).toBeGreaterThan(0); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeUndefined(); + }); + 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"); @@ -496,14 +561,6 @@ describe("AuthStorage codex oauth ranking", () => { await inFlightReports; expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); - const selectionCounts = await countApiKeySelections( - authStorage, - "openai-codex", - "codex-inflight-race-after-resolve", - 40, - ); - expect(countFor(selectionCounts, "api-acct-race-blocked")).toBe(0); - expect(countFor(selectionCounts, "api-acct-race-healthy")).toBeGreaterThan(0); }); test("broker-sourced healthy Codex usage clears remote gateway backoff", async () => { From 5b0d670d6b901a0140b74e6f9a8532807f9e9c44 Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 21:17:08 +0000 Subject: [PATCH 2/7] fix(ai): kept exhausted codex blocks - Restored all-limit Codex block clearing so recovered primary windows do not clear blocks while another reported quota remains exhausted. - Added regression coverage for primary-recovered/secondary-exhausted selection-path reconciliation. Fixes #4980 --- packages/ai/CHANGELOG.md | 2 +- packages/ai/src/auth-storage.ts | 9 +-- .../test/auth-storage-codex-selection.test.ts | 75 +++++++++++++++++++ 3 files changed, 79 insertions(+), 7 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index a74b0c89c..417ccf1ff 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -4,7 +4,7 @@ ### Fixed -- Fixed Codex OAuth credential selection to re-check blocked accounts during ranking and clear stale usage-limit blocks when live usage shows the 5-hour window recovered ([#4980](https://github.com/can1357/oh-my-pi/issues/4980)). +- Fixed Codex OAuth credential selection to re-check blocked accounts during ranking and clear stale usage-limit blocks when live usage shows all reported windows recovered ([#4980](https://github.com/can1357/oh-my-pi/issues/4980)). ## [16.3.14] - 2026-07-09 diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 886052d4a..bd7b42fb6 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -4389,17 +4389,14 @@ export class AuthStorage { /** * Self-heal a stale Codex usage-limit block: when a fresh live usage report - * says the account is allowed and its primary rolling window is available, - * drop the persisted and in-memory `openai-codex:oauth` blocks so credential - * selection can re-include seats whose short window recovered before a - * longer persisted block naturally expires. + * says the account is allowed and below every reported limit, drop the + * persisted and in-memory `openai-codex:oauth` blocks so credential selection + * can re-include recovered seats before a stale block naturally expires. */ #isHealthyCodexUsageReport(report: UsageReport): boolean { if (report.provider !== "openai-codex") return false; const metadata = report.metadata; if (metadata?.allowed !== true || metadata.limitReached !== false) return false; - const primary = codexRankingStrategy.findWindowLimits(report).primary; - if (primary) return !this.#isUsageLimitExhausted(primary); return !this.#isUsageLimitReached(report.limits); } diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index a63ac005e..14d0e7c8a 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -476,6 +476,81 @@ describe("AuthStorage codex oauth ranking", () => { expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeUndefined(); }); + test("keeps a stale Codex block when the 5h window recovered but the 7d window remains exhausted", async () => { + if (!authStorage || !store?.upsertCredentialBlock || !store.getCredentialBlock) { + throw new Error("test setup failed"); + } + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-secondary-exhausted", "secondary-exhausted@example.com") }, + { type: "oauth", ...createCredential("acct-secondary-healthy", "secondary-healthy@example.com") }, + ]); + + const blockedRow = store.listAuthCredentials("openai-codex").find(row => { + const credential = row.credential; + return credential.type === "oauth" && credential.accountId === "acct-secondary-exhausted"; + }); + if (!blockedRow) throw new Error("expected blocked credential row"); + + const blockedUntilMs = Date.now() + 6 * 24 * HOUR_MS; + store.upsertCredentialBlock({ + credentialId: blockedRow.id, + providerKey: "openai-codex:oauth", + blockScope: "shared", + blockedUntilMs, + }); + + const fiveHourWindow: UsageWindowConfig = { + windowId: "5h", + windowLabel: "5 Hours", + durationMs: FIVE_HOUR_MS, + }; + + usageByAccount.set( + "acct-secondary-exhausted", + createCodexUsageReport({ + accountId: "acct-secondary-exhausted", + primary: { usedFraction: 0.2, resetInMs: 4 * HOUR_MS }, + secondary: { usedFraction: 1, resetInMs: 6 * 24 * HOUR_MS }, + primaryWindow: fiveHourWindow, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "secondary-exhausted@example.com", + accountId: "acct-secondary-exhausted", + }, + }), + ); + usageByAccount.set( + "acct-secondary-healthy", + createCodexUsageReport({ + accountId: "acct-secondary-healthy", + primary: { usedFraction: 0.2, resetInMs: 4 * HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: 6 * 24 * HOUR_MS }, + primaryWindow: fiveHourWindow, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "secondary-healthy@example.com", + accountId: "acct-secondary-healthy", + }, + }), + ); + + const selectionCounts = await countApiKeySelections( + authStorage, + "openai-codex", + "codex-stale-block-secondary-exhausted", + 150, + ); + + expect(countFor(selectionCounts, "api-acct-secondary-exhausted")).toBe(0); + expect(countFor(selectionCounts, "api-acct-secondary-healthy")).toBeGreaterThan(0); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBe(blockedUntilMs); + }); + 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"); From 0ffc43bfc09b1d784c2e9afa261d32ee8d75dcfd Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 21:36:35 +0000 Subject: [PATCH 3/7] 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(); From 04214a8c8504dd1e391745dfe2599314126fabc9 Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 22:07:08 +0000 Subject: [PATCH 4/7] fix(ai): guarded broker codex blocks - Preserved fresh broker-sourced credential blocks by exposing store-level reconciliation delays to AuthStorage. - Tracked fresh block observations in SQLite and remote broker stores without protecting initial stale snapshots. - Added broker sibling regression coverage for healthy usage lag after a shared 429 block. Fixes #4980 --- packages/ai/src/auth-broker/remote-store.ts | 52 +++++++- packages/ai/src/auth-storage.ts | 30 ++++- .../test/auth-storage-codex-selection.test.ts | 120 ++++++++++++++++++ 3 files changed, 195 insertions(+), 7 deletions(-) diff --git a/packages/ai/src/auth-broker/remote-store.ts b/packages/ai/src/auth-broker/remote-store.ts index f56ba6b3a..a2c30eb61 100644 --- a/packages/ai/src/auth-broker/remote-store.ts +++ b/packages/ai/src/auth-broker/remote-store.ts @@ -39,6 +39,7 @@ import type { * one broker call instead of N. */ const USAGE_CACHE_TTL_MS = 15_000; +const CREDENTIAL_BLOCK_RECONCILE_DELAY_MS = 5 * 60_000; const WAIT_THRESHOLD_MS = 1_000; const MAX_WAIT_MS = 5_000; const BACKGROUND_WAIT_MS = 30_000; @@ -208,6 +209,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { #cache: Map = new Map(); #usageCache?: UsageCacheEntry; #usageInflight?: Promise; + #credentialBlockReconcileAfter: Map = new Map(); #usageCacheEpoch = 0; #closed = false; /** @@ -223,7 +225,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { constructor(opts: RemoteAuthCredentialStoreOptions) { this.#client = opts.client; this.#streamSnapshots = opts.streamSnapshots ?? true; - this.#applySnapshot(opts.initialSnapshot ?? emptySnapshot(), opts.initialSnapshot?.generation ?? 0); + this.#applySnapshot(opts.initialSnapshot ?? emptySnapshot(), opts.initialSnapshot?.generation ?? 0, false); this.#onSnapshot = opts.onSnapshot; void this.#runBackground(); } @@ -236,10 +238,12 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { return this.#snapshot; } - #applySnapshot(snapshot: SnapshotResponse, generation: number): void { + #applySnapshot(snapshot: SnapshotResponse, generation: number, protectNewBlocks = true): void { const nowMs = Date.now(); + const previousCredentials = this.#snapshot.credentials; const credentials = snapshot.credentials.map(entry => this.#normalizeSnapshotEntryBlocks(entry, nowMs)); - if (snapshotBlocksChanged(this.#snapshot.credentials, credentials)) this.#invalidateUsageCache(); + if (snapshotBlocksChanged(previousCredentials, credentials)) this.#invalidateUsageCache(); + if (protectNewBlocks) this.#protectNewSnapshotBlocks(previousCredentials, credentials, nowMs); this.#snapshot = { ...snapshot, credentials }; this.#generation = generation; this.#snapshotReceivedAt = nowMs; @@ -251,6 +255,29 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { logger.debug("auth-broker snapshot callback failed", { error: String(error) }); } } + #protectNewSnapshotBlocks(previous: readonly SnapshotEntry[], next: readonly SnapshotEntry[], nowMs: number): void { + const previousBlocksByKey = new Map(); + for (const entry of previous) { + for (const block of entry.blocks ?? []) { + previousBlocksByKey.set(`${entry.id}\0${block.providerKey}\0${block.blockScope}`, block.blockedUntilMs); + } + } + const activeKeys = new Set(); + for (const entry of next) { + for (const block of entry.blocks ?? []) { + const key = `${entry.id}\0${block.providerKey}\0${block.blockScope}`; + activeKeys.add(key); + if (previousBlocksByKey.get(key) === block.blockedUntilMs) continue; + this.#credentialBlockReconcileAfter.set( + key, + Math.min(block.blockedUntilMs, nowMs + CREDENTIAL_BLOCK_RECONCILE_DELAY_MS), + ); + } + } + for (const key of this.#credentialBlockReconcileAfter.keys()) { + if (!activeKeys.has(key)) this.#credentialBlockReconcileAfter.delete(key); + } + } async #runBackground(): Promise { let backoffMs = BACKGROUND_BACKOFF_INITIAL_MS; @@ -340,11 +367,13 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { const incoming = this.#normalizeSnapshotEntryBlocks(entry, Date.now()); const index = this.#snapshot.credentials.findIndex(candidate => candidate.id === incoming.id); const previousBlocks = index === -1 ? undefined : this.#snapshot.credentials[index]?.blocks; - if (!credentialBlockSnapshotsEqual(previousBlocks, incoming.blocks)) this.#invalidateUsageCache(); + const blocksChanged = !credentialBlockSnapshotsEqual(previousBlocks, incoming.blocks); + if (blocksChanged) this.#invalidateUsageCache(); const credentials = index === -1 ? [...this.#snapshot.credentials, incoming] : this.#snapshot.credentials.map((candidate, i) => (i === index ? incoming : candidate)); + if (blocksChanged) this.#protectNewSnapshotBlocks(this.#snapshot.credentials, credentials, Date.now()); this.#snapshot = { ...this.#snapshot, generation, serverNowMs, refresher, credentials }; this.#generation = generation; this.#snapshotReceivedAt = Date.now(); @@ -392,6 +421,11 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { return block.blockedUntilMs; } + getCredentialBlockReconcileAfter(credentialId: number, providerKey: string, blockScope: string): number | undefined { + if (this.getCredentialBlock(credentialId, providerKey, blockScope) === undefined) return undefined; + return this.#credentialBlockReconcileAfter.get(`${credentialId}\0${providerKey}\0${blockScope}`); + } + listCredentialBlocks(credentialIds: readonly number[]): StoredCredentialBlock[] { const nowMs = Date.now(); this.cleanExpiredCredentialBlocks(nowMs); @@ -416,6 +450,10 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { upsertCredentialBlock(block: StoredCredentialBlock): void { this.#upsertSnapshotBlock(block); this.#invalidateUsageCache(); + this.#credentialBlockReconcileAfter.set( + `${block.credentialId}\0${block.providerKey}\0${block.blockScope}`, + Math.min(block.blockedUntilMs, Date.now() + CREDENTIAL_BLOCK_RECONCILE_DELAY_MS), + ); const body = toCredentialBlockSnapshot(block); void this.#client .upsertCredentialBlock(block.credentialId, body) @@ -434,6 +472,9 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { deleteCredentialBlocks(credentialId: number): void { this.#deleteSnapshotBlocks(credentialId); + for (const key of this.#credentialBlockReconcileAfter.keys()) { + if (key.startsWith(`${credentialId}\0`)) this.#credentialBlockReconcileAfter.delete(key); + } this.#invalidateUsageCache(); void this.#client .deleteCredentialBlocks(credentialId) @@ -450,6 +491,9 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { cleanExpiredCredentialBlocks(nowMs: number): void { this.#pruneExpiredCredentialBlocks(nowMs); + for (const [key, reconcileAfterMs] of this.#credentialBlockReconcileAfter) { + if (reconcileAfterMs <= nowMs) this.#credentialBlockReconcileAfter.delete(key); + } } /** diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index acdaff5e7..36eb052d9 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -320,6 +320,8 @@ export interface AuthCredentialStore { cleanExpiredCache(): void; /** Non-expired block for one (credential, providerKey, scope) key, or undefined. */ getCredentialBlock?(credentialId: number, providerKey: string, blockScope: string): number | undefined; + /** Earliest time a shared-store block should be eligible for live-usage reconciliation. */ + getCredentialBlockReconcileAfter?(credentialId: number, providerKey: string, blockScope: string): number | undefined; /** Upsert with MAX semantics: keep the later blockedUntilMs on conflict. */ upsertCredentialBlock?(block: StoredCredentialBlock): void; /** Drop every block row for a credential (all providerKeys/scopes). */ @@ -4423,13 +4425,19 @@ 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. + // `/usage` can lag the request path that just returned 429. Fresh local or + // broker-sourced 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; + const getStoreReconcileAfter = this.#store.getCredentialBlockReconcileAfter?.bind(this.#store); + const storeGlobalProbeAfterMs = getStoreReconcileAfter?.(credentialId, providerKey, "") ?? 0; + const storeScopedProbeAfterMs = getStoreReconcileAfter?.(credentialId, providerKey, blockScope ?? "") ?? 0; + if (Math.max(globalProbeAfterMs, scopedProbeAfterMs, storeGlobalProbeAfterMs, storeScopedProbeAfterMs) > nowMs) { + return; + } this.#clearCredentialBlocks(provider, credentialId); logger.info("Cleared stale Codex usage-limit block after healthy live usage report", { credentialId, @@ -5148,6 +5156,7 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore { #upsertCredentialBlockStmt: Statement; #deleteCredentialBlocksStmt: Statement; #deleteExpiredCredentialBlocksStmt: Statement; + #credentialBlockReconcileAfter: Map = new Map(); #insertUsageHistoryStmt: Statement; #insertUsageCostStmt: Statement; #listUsageCostsStmt: Statement; @@ -5816,6 +5825,11 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore { return typeof row?.blocked_until_ms === "number" ? row.blocked_until_ms : undefined; } + getCredentialBlockReconcileAfter(credentialId: number, providerKey: string, blockScope: string): number | undefined { + if (this.getCredentialBlock(credentialId, providerKey, blockScope) === undefined) return undefined; + return this.#credentialBlockReconcileAfter.get(`${credentialId}\0${providerKey}\0${blockScope}`); + } + upsertCredentialBlock(block: StoredCredentialBlock): void { this.#upsertCredentialBlockStmt.run( block.credentialId, @@ -5823,14 +5837,24 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore { block.blockScope, block.blockedUntilMs, ); + this.#credentialBlockReconcileAfter.set( + `${block.credentialId}\0${block.providerKey}\0${block.blockScope}`, + Math.min(block.blockedUntilMs, Date.now() + USAGE_REPORT_TTL_MS), + ); } deleteCredentialBlocks(credentialId: number): void { this.#deleteCredentialBlocksStmt.run(credentialId); + for (const key of this.#credentialBlockReconcileAfter.keys()) { + if (key.startsWith(`${credentialId}\0`)) this.#credentialBlockReconcileAfter.delete(key); + } } cleanExpiredCredentialBlocks(nowMs: number): void { this.#deleteExpiredCredentialBlocksStmt.run(nowMs); + for (const [key, reconcileAfterMs] of this.#credentialBlockReconcileAfter) { + if (reconcileAfterMs <= nowMs) this.#credentialBlockReconcileAfter.delete(key); + } } listCredentialBlocks(credentialIds: readonly number[]): StoredCredentialBlock[] { diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 3963880e0..880074ad3 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -13,6 +13,7 @@ const WEEK_MS = 7 * 24 * 60 * 60 * 1000; const HOUR_MS = 60 * 60 * 1000; const FIVE_HOUR_MS = 5 * HOUR_MS; +const STALE_BLOCK_GUARD_MS = 5 * 60_000 + 1; type UsageWindowSpec = { usedFraction: number; @@ -353,6 +354,7 @@ describe("AuthStorage codex oauth ranking", () => { blockScope: "shared", blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, }); + store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_MS); usageByAccount.set( "acct-blocked", @@ -424,6 +426,7 @@ describe("AuthStorage codex oauth ranking", () => { blockScope: "shared", blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, }); + store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_MS); const fiveHourWindow: UsageWindowConfig = { windowId: "5h", @@ -618,6 +621,122 @@ describe("AuthStorage codex oauth ranking", () => { expect(selectionAfterBlock).toBe("api-acct-fresh-healthy"); expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); }); + + test("keeps broker-sourced fresh Codex block when sibling selection sees healthy usage", async () => { + if (!authStorage || !store?.getCredentialBlock) { + throw new Error("test setup failed"); + } + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-broker-fresh-blocked", "broker-fresh-blocked@example.com") }, + { type: "oauth", ...createCredential("acct-broker-fresh-healthy", "broker-fresh-healthy@example.com") }, + ]); + + usageByAccount.set( + "acct-broker-fresh-blocked", + createCodexUsageReport({ + accountId: "acct-broker-fresh-blocked", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "broker-fresh-blocked@example.com", + accountId: "acct-broker-fresh-blocked", + }, + }), + ); + usageByAccount.set( + "acct-broker-fresh-healthy", + createCodexUsageReport({ + accountId: "acct-broker-fresh-healthy", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "broker-fresh-healthy@example.com", + accountId: "acct-broker-fresh-healthy", + }, + }), + ); + + const token = "codex-broker-fresh-block"; + const handle = startAuthBroker({ + storage: authStorage, + bind: "127.0.0.1:0", + bearerTokens: [token], + disableRefresher: true, + }); + try { + const clientA = new AuthBrokerClient({ url: handle.url, token }); + const clientB = new AuthBrokerClient({ url: handle.url, token }); + const initialResult = await clientB.fetchSnapshot(); + if (initialResult.status !== 200) throw new Error("expected initial broker snapshot"); + const blockedRow = initialResult.snapshot.credentials.find(entry => { + const credential = entry.credential; + return credential.type === "oauth" && credential.accountId === "acct-broker-fresh-blocked"; + }); + if (!blockedRow) throw new Error("expected blocked credential row"); + + const remoteStoreA = new RemoteAuthCredentialStore({ + client: clientA, + initialSnapshot: initialResult.snapshot, + streamSnapshots: false, + }); + const remoteStoreB = new RemoteAuthCredentialStore({ + client: clientB, + initialSnapshot: initialResult.snapshot, + streamSnapshots: false, + }); + const clientStorageA = new AuthStorage(remoteStoreA); + const clientStorageB = new AuthStorage(remoteStoreB); + await clientStorageA.reload(); + await clientStorageB.reload(); + try { + let blockedSessionId: string | undefined; + for (let index = 0; index < 100; index += 1) { + const sessionId = `codex-broker-fresh-block-selected-${index}`; + if ((await clientStorageA.getApiKey("openai-codex", sessionId)) === "api-acct-broker-fresh-blocked") { + blockedSessionId = sessionId; + break; + } + } + if (!blockedSessionId) throw new Error("expected client A to select the soon-blocked account"); + + const markResult = await clientStorageA.markUsageLimitReached("openai-codex", blockedSessionId, { + retryAfterMs: 6 * 24 * HOUR_MS, + }); + expect(markResult.switched).toBe(true); + + const updatedSnapshot = await clientB.fetchSnapshot({ + ifGenerationGt: initialResult.generation, + waitMs: 1000, + }); + if (updatedSnapshot.status !== 200) throw new Error("expected broker snapshot containing fresh block"); + + await remoteStoreB.refreshSnapshot(); + expect(remoteStoreB.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + + expect(await clientStorageB.getApiKey("openai-codex", "codex-broker-fresh-block-sibling")).toBe( + "api-acct-broker-fresh-healthy", + ); + expect(remoteStoreB.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + } finally { + clientStorageA.close(); + clientStorageB.close(); + remoteStoreA.close(); + remoteStoreB.close(); + } + } finally { + await handle.close(); + } + }); + 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"); @@ -757,6 +876,7 @@ describe("AuthStorage codex oauth ranking", () => { blockScope: "shared", blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, }); + store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_MS); const token = "codex-broker-reconcile"; const handle = startAuthBroker({ From d8745b62b417be9df135acf47c9706554faf1705 Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 22:23:07 +0000 Subject: [PATCH 5/7] fix(ai): protected initial broker blocks - Protected Codex blocks present in RemoteAuthCredentialStore initial snapshots from immediate healthy-usage reconciliation. - Added regression coverage for clients that start after a broker peer already persisted a fresh 429 block. - Kept stale broker reconciliation explicit by expiring the test guard before the stale-block refresh path. Fixes #4980 --- packages/ai/src/auth-broker/remote-store.ts | 2 +- .../test/auth-storage-codex-selection.test.ts | 142 ++++++++++++++++++ 2 files changed, 143 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/auth-broker/remote-store.ts b/packages/ai/src/auth-broker/remote-store.ts index a2c30eb61..c1552870d 100644 --- a/packages/ai/src/auth-broker/remote-store.ts +++ b/packages/ai/src/auth-broker/remote-store.ts @@ -225,7 +225,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { constructor(opts: RemoteAuthCredentialStoreOptions) { this.#client = opts.client; this.#streamSnapshots = opts.streamSnapshots ?? true; - this.#applySnapshot(opts.initialSnapshot ?? emptySnapshot(), opts.initialSnapshot?.generation ?? 0, false); + this.#applySnapshot(opts.initialSnapshot ?? emptySnapshot(), opts.initialSnapshot?.generation ?? 0); this.#onSnapshot = opts.onSnapshot; void this.#runBackground(); } diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 880074ad3..19f227eb3 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -737,6 +737,147 @@ describe("AuthStorage codex oauth ranking", () => { } }); + test("protects fresh Codex blocks present in the initial broker snapshot from healthy selection reconciliation", async () => { + if (!authStorage || !store?.getCredentialBlock) { + throw new Error("test setup failed"); + } + + await authStorage.set("openai-codex", [ + { + type: "oauth", + ...createCredential("acct-broker-initial-snapshot-blocked", "broker-initial-snapshot-blocked@example.com"), + }, + { + type: "oauth", + ...createCredential("acct-broker-initial-snapshot-healthy", "broker-initial-snapshot-healthy@example.com"), + }, + ]); + + usageByAccount.set( + "acct-broker-initial-snapshot-blocked", + createCodexUsageReport({ + accountId: "acct-broker-initial-snapshot-blocked", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "broker-initial-snapshot-blocked@example.com", + accountId: "acct-broker-initial-snapshot-blocked", + }, + }), + ); + usageByAccount.set( + "acct-broker-initial-snapshot-healthy", + createCodexUsageReport({ + accountId: "acct-broker-initial-snapshot-healthy", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "broker-initial-snapshot-healthy@example.com", + accountId: "acct-broker-initial-snapshot-healthy", + }, + }), + ); + + const token = "codex-broker-initial-snapshot-block"; + const handle = startAuthBroker({ + storage: authStorage, + bind: "127.0.0.1:0", + bearerTokens: [token], + disableRefresher: true, + }); + try { + const clientA = new AuthBrokerClient({ url: handle.url, token }); + const clientB = new AuthBrokerClient({ url: handle.url, token }); + const clientAInitial = await clientA.fetchSnapshot(); + if (clientAInitial.status !== 200) throw new Error("expected client A broker snapshot"); + + const remoteStoreA = new RemoteAuthCredentialStore({ + client: clientA, + initialSnapshot: clientAInitial.snapshot, + streamSnapshots: false, + }); + const clientStorageA = new AuthStorage(remoteStoreA); + await clientStorageA.reload(); + try { + let blockedSessionId: string | undefined; + let blockedAccountId: string | undefined; + for (let index = 0; index < 100; index += 1) { + const sessionId = `codex-broker-initial-snapshot-block-selected-${index}`; + const apiKey = await clientStorageA.getApiKey("openai-codex", sessionId); + if ( + apiKey === "api-acct-broker-initial-snapshot-blocked" || + apiKey === "api-acct-broker-initial-snapshot-healthy" + ) { + blockedSessionId = sessionId; + blockedAccountId = apiKey.replace(/^api-/, ""); + break; + } + } + if (!blockedSessionId || !blockedAccountId) { + throw new Error("expected client A to select a Codex account to block"); + } + const healthyAccountId = + blockedAccountId === "acct-broker-initial-snapshot-blocked" + ? "acct-broker-initial-snapshot-healthy" + : "acct-broker-initial-snapshot-blocked"; + + const markResult = await clientStorageA.markUsageLimitReached("openai-codex", blockedSessionId, { + retryAfterMs: 6 * 24 * HOUR_MS, + }); + expect(markResult.switched).toBe(true); + + const snapshotWithBlock = await clientB.fetchSnapshot({ + ifGenerationGt: clientAInitial.generation, + waitMs: 1000, + }); + if (snapshotWithBlock.status !== 200) throw new Error("expected broker snapshot containing initial block"); + + const blockedRow = snapshotWithBlock.snapshot.credentials.find(entry => { + const credential = entry.credential; + return credential.type === "oauth" && credential.accountId === blockedAccountId; + }); + if (!blockedRow) throw new Error("expected blocked credential row"); + expect( + blockedRow.blocks?.some( + block => block.providerKey === "openai-codex:oauth" && block.blockScope === "shared", + ), + ).toBe(true); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + + const remoteStoreB = new RemoteAuthCredentialStore({ + client: clientB, + initialSnapshot: snapshotWithBlock.snapshot, + streamSnapshots: false, + }); + const clientStorageB = new AuthStorage(remoteStoreB); + await clientStorageB.reload(); + try { + expect(remoteStoreB.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + + expect(await clientStorageB.getApiKey("openai-codex", "codex-broker-initial-snapshot-sibling")).toBe( + `api-${healthyAccountId}`, + ); + expect(remoteStoreB.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + } finally { + clientStorageB.close(); + remoteStoreB.close(); + } + } finally { + clientStorageA.close(); + remoteStoreA.close(); + } + } finally { + await handle.close(); + } + }); + 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"); @@ -904,6 +1045,7 @@ describe("AuthStorage codex oauth ranking", () => { try { expect(remoteStore.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + remoteStore.cleanExpiredCredentialBlocks(Date.now() + STALE_BLOCK_GUARD_MS); await clientStorage.fetchUsageReports(); From 4911ce02e85e3b8456eddba7795ae0488f1c2e0e Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 22:50:33 +0000 Subject: [PATCH 6/7] fix(ai): persisted codex block guard - Derived SQLite credential-block reconciliation delays from persisted updated_at so fresh blocks survive process restarts and sibling local stores. - Added reopened-SQLite regression coverage for healthy usage lag after a persisted Codex 429 block. - Aged explicit stale-block fixtures by moving updated_at outside the guard window. Fixes #4980 --- packages/ai/src/auth-storage.ts | 21 +++- .../test/auth-storage-codex-selection.test.ts | 103 +++++++++++++++++- 2 files changed, 118 insertions(+), 6 deletions(-) diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 36eb052d9..478d540f4 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -4941,6 +4941,7 @@ type CredentialBlockRow = { provider_key: string; block_scope: string; blocked_until_ms: number; + updated_at: number; }; type SerializedCredentialRecord = { @@ -5203,10 +5204,10 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore { ); this.#deleteExpiredCacheStmt = this.#db.prepare(`DELETE FROM cache WHERE expires_at <= ${SQLITE_NOW_EPOCH}`); this.#getCredentialBlockStmt = this.#db.prepare( - "SELECT blocked_until_ms FROM auth_credential_blocks WHERE credential_id = ? AND provider_key = ? AND block_scope = ? AND blocked_until_ms > ?", + "SELECT blocked_until_ms, updated_at FROM auth_credential_blocks WHERE credential_id = ? AND provider_key = ? AND block_scope = ? AND blocked_until_ms > ?", ); this.#listCredentialBlocksByCredentialStmt = this.#db.prepare( - "SELECT credential_id, provider_key, block_scope, blocked_until_ms FROM auth_credential_blocks WHERE credential_id = ? AND blocked_until_ms > ? ORDER BY provider_key ASC, block_scope ASC", + "SELECT credential_id, provider_key, block_scope, blocked_until_ms, updated_at FROM auth_credential_blocks WHERE credential_id = ? AND blocked_until_ms > ? ORDER BY provider_key ASC, block_scope ASC", ); this.#upsertCredentialBlockStmt = this.#db.prepare( `INSERT INTO auth_credential_blocks (credential_id, provider_key, block_scope, blocked_until_ms, updated_at) @@ -5820,14 +5821,24 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore { const nowMs = Date.now(); this.#deleteExpiredCredentialBlocksStmt.run(nowMs); const row = this.#getCredentialBlockStmt.get(credentialId, providerKey, blockScope, nowMs) as - | { blocked_until_ms?: number } + | { blocked_until_ms?: number; updated_at?: number } | undefined; return typeof row?.blocked_until_ms === "number" ? row.blocked_until_ms : undefined; } getCredentialBlockReconcileAfter(credentialId: number, providerKey: string, blockScope: string): number | undefined { - if (this.getCredentialBlock(credentialId, providerKey, blockScope) === undefined) return undefined; - return this.#credentialBlockReconcileAfter.get(`${credentialId}\0${providerKey}\0${blockScope}`); + const nowMs = Date.now(); + this.#deleteExpiredCredentialBlocksStmt.run(nowMs); + const row = this.#getCredentialBlockStmt.get(credentialId, providerKey, blockScope, nowMs) as + | { blocked_until_ms?: number; updated_at?: number } + | undefined; + if (typeof row?.blocked_until_ms !== "number") return undefined; + const memoryReconcileAfter = + this.#credentialBlockReconcileAfter.get(`${credentialId}\0${providerKey}\0${blockScope}`) ?? 0; + const persistedReconcileAfter = + typeof row.updated_at === "number" ? row.updated_at * 1000 + USAGE_REPORT_TTL_MS : 0; + const reconcileAfter = Math.max(memoryReconcileAfter, persistedReconcileAfter); + return reconcileAfter > nowMs ? Math.min(row.blocked_until_ms, reconcileAfter) : undefined; } upsertCredentialBlock(block: StoredCredentialBlock): void { diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 19f227eb3..85c7c822c 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -1,3 +1,4 @@ +import { Database } from "bun:sqlite"; import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; import * as fs from "node:fs/promises"; import * as os from "node:os"; @@ -15,6 +16,17 @@ const HOUR_MS = 60 * 60 * 1000; const FIVE_HOUR_MS = 5 * HOUR_MS; const STALE_BLOCK_GUARD_MS = 5 * 60_000 + 1; +function ageCredentialBlockRows(dbPath: string): void { + const db = new Database(dbPath); + try { + db.prepare("UPDATE auth_credential_blocks SET updated_at = ?").run( + Math.floor((Date.now() - STALE_BLOCK_GUARD_MS) / 1000), + ); + } finally { + db.close(); + } +} + type UsageWindowSpec = { usedFraction: number; resetInMs: number; @@ -145,6 +157,7 @@ function expectWeightedPreference(counts: Map, preferred: string describe("AuthStorage codex oauth ranking", () => { let tempDir = ""; let store: AuthCredentialStore | null = null; + let dbPath = ""; let authStorage: AuthStorage | null = null; const usageByAccount = new Map(); @@ -159,7 +172,8 @@ describe("AuthStorage codex oauth ranking", () => { beforeEach(async () => { tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "pi-ai-auth-codex-selection-")); - store = await SqliteAuthCredentialStore.open(path.join(tempDir, "agent.db")); + dbPath = path.join(tempDir, "agent.db"); + store = await SqliteAuthCredentialStore.open(dbPath); authStorage = new AuthStorage(store, { usageProviderResolver: provider => (provider === "openai-codex" ? usageProvider : undefined), }); @@ -354,6 +368,7 @@ describe("AuthStorage codex oauth ranking", () => { blockScope: "shared", blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, }); + ageCredentialBlockRows(dbPath); store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_MS); usageByAccount.set( @@ -426,6 +441,7 @@ describe("AuthStorage codex oauth ranking", () => { blockScope: "shared", blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, }); + ageCredentialBlockRows(dbPath); store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_MS); const fiveHourWindow: UsageWindowConfig = { @@ -622,6 +638,90 @@ describe("AuthStorage codex oauth ranking", () => { expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); }); + test("protects a fresh Codex block after reopening SQLite storage", async () => { + if (!authStorage || !store?.getCredentialBlock) { + throw new Error("test setup failed"); + } + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-reopened-blocked", "reopened-blocked@example.com") }, + { type: "oauth", ...createCredential("acct-reopened-healthy", "reopened-healthy@example.com") }, + ]); + + usageByAccount.set( + "acct-reopened-blocked", + createCodexUsageReport({ + accountId: "acct-reopened-blocked", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "reopened-blocked@example.com", + accountId: "acct-reopened-blocked", + }, + }), + ); + usageByAccount.set( + "acct-reopened-healthy", + createCodexUsageReport({ + accountId: "acct-reopened-healthy", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "reopened-healthy@example.com", + accountId: "acct-reopened-healthy", + }, + }), + ); + + const firstSelectionSessionId = "codex-reopened-fresh-block-initial"; + const firstSelection = await authStorage.getApiKey("openai-codex", firstSelectionSessionId); + if (!firstSelection) throw new Error("expected initial Codex credential"); + + const blockedAccountId = firstSelection.replace(/^api-/, ""); + const healthyAccountId = + blockedAccountId === "acct-reopened-blocked" ? "acct-reopened-healthy" : "acct-reopened-blocked"; + + const blockedRow = store.listAuthCredentials("openai-codex").find(row => { + const credential = row.credential; + return credential.type === "oauth" && credential.accountId === blockedAccountId; + }); + if (!blockedRow) throw new Error("expected blocked credential row"); + + const markResult = await authStorage.markUsageLimitReached("openai-codex", firstSelectionSessionId, { + retryAfterMs: 6 * 24 * HOUR_MS, + }); + + expect(markResult.switched).toBe(true); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + + authStorage.close(); + authStorage = null; + store = null; + + const reopenedStore = await SqliteAuthCredentialStore.open(dbPath); + const reopenedAuthStorage = new AuthStorage(reopenedStore, { + usageProviderResolver: provider => (provider === "openai-codex" ? usageProvider : undefined), + }); + try { + await reopenedAuthStorage.reload(); + + const selectionAfterReopen = await reopenedAuthStorage.getApiKey( + "openai-codex", + "codex-reopened-fresh-block-sibling", + ); + expect(selectionAfterReopen).toBe(`api-${healthyAccountId}`); + expect(reopenedStore.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBeDefined(); + } finally { + reopenedAuthStorage.close(); + } + }); + test("keeps broker-sourced fresh Codex block when sibling selection sees healthy usage", async () => { if (!authStorage || !store?.getCredentialBlock) { throw new Error("test setup failed"); @@ -1017,6 +1117,7 @@ describe("AuthStorage codex oauth ranking", () => { blockScope: "shared", blockedUntilMs: Date.now() + 6 * 24 * HOUR_MS, }); + ageCredentialBlockRows(dbPath); store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_MS); const token = "codex-broker-reconcile"; From 6d7341b139f61f3bce4573eedab7a7396991e03a Mon Sep 17 00:00:00 2001 From: roboomp Date: Thu, 9 Jul 2026 23:18:36 +0000 Subject: [PATCH 7/7] fix(ai): refreshed broker block guards - Sent credential-block updatedAtMs through broker snapshots so same-deadline block refreshes are observable by remote clients. - Refreshed RemoteAuthCredentialStore reconciliation guards when updatedAtMs changes even if blockedUntilMs is unchanged. - Added broker regression coverage for same-deadline Codex block re-upserts after a local guard expires. Fixes #4980 --- packages/ai/src/auth-broker/remote-store.ts | 23 ++- packages/ai/src/auth-broker/server.ts | 1 + packages/ai/src/auth-broker/wire-schemas.ts | 1 + packages/ai/src/auth-storage.ts | 3 + .../test/auth-storage-codex-selection.test.ts | 144 ++++++++++++++++++ 5 files changed, 166 insertions(+), 6 deletions(-) diff --git a/packages/ai/src/auth-broker/remote-store.ts b/packages/ai/src/auth-broker/remote-store.ts index c1552870d..4351fe616 100644 --- a/packages/ai/src/auth-broker/remote-store.ts +++ b/packages/ai/src/auth-broker/remote-store.ts @@ -51,7 +51,9 @@ function compareCredentialBlockSnapshots(a: CredentialBlockSnapshot, b: Credenti if (provider !== 0) return provider; const scope = a.blockScope.localeCompare(b.blockScope); if (scope !== 0) return scope; - return a.blockedUntilMs - b.blockedUntilMs; + const blockedUntil = a.blockedUntilMs - b.blockedUntilMs; + if (blockedUntil !== 0) return blockedUntil; + return (a.updatedAtMs ?? 0) - (b.updatedAtMs ?? 0); } function toCredentialBlockSnapshot(block: StoredCredentialBlock): CredentialBlockSnapshot { @@ -59,6 +61,7 @@ function toCredentialBlockSnapshot(block: StoredCredentialBlock): CredentialBloc providerKey: block.providerKey, blockScope: block.blockScope, blockedUntilMs: block.blockedUntilMs, + ...(block.updatedAtMs !== undefined ? { updatedAtMs: block.updatedAtMs } : {}), }; } @@ -75,7 +78,8 @@ function credentialBlockSnapshotsEqual( if ( leftBlock.providerKey !== rightBlock.providerKey || leftBlock.blockScope !== rightBlock.blockScope || - leftBlock.blockedUntilMs !== rightBlock.blockedUntilMs + leftBlock.blockedUntilMs !== rightBlock.blockedUntilMs || + leftBlock.updatedAtMs !== rightBlock.updatedAtMs ) { return false; } @@ -256,10 +260,13 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { } } #protectNewSnapshotBlocks(previous: readonly SnapshotEntry[], next: readonly SnapshotEntry[], nowMs: number): void { - const previousBlocksByKey = new Map(); + const previousBlocksByKey = new Map(); for (const entry of previous) { for (const block of entry.blocks ?? []) { - previousBlocksByKey.set(`${entry.id}\0${block.providerKey}\0${block.blockScope}`, block.blockedUntilMs); + previousBlocksByKey.set( + `${entry.id}\0${block.providerKey}\0${block.blockScope}`, + `${block.blockedUntilMs}\0${block.updatedAtMs ?? ""}`, + ); } } const activeKeys = new Set(); @@ -267,10 +274,12 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { for (const block of entry.blocks ?? []) { const key = `${entry.id}\0${block.providerKey}\0${block.blockScope}`; activeKeys.add(key); - if (previousBlocksByKey.get(key) === block.blockedUntilMs) continue; + const signature = `${block.blockedUntilMs}\0${block.updatedAtMs ?? ""}`; + if (previousBlocksByKey.get(key) === signature) continue; + const updatedAtMs = block.updatedAtMs ?? nowMs; this.#credentialBlockReconcileAfter.set( key, - Math.min(block.blockedUntilMs, nowMs + CREDENTIAL_BLOCK_RECONCILE_DELAY_MS), + Math.min(block.blockedUntilMs, updatedAtMs + CREDENTIAL_BLOCK_RECONCILE_DELAY_MS), ); } } @@ -440,6 +449,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { providerKey: block.providerKey, blockScope: block.blockScope, blockedUntilMs: block.blockedUntilMs, + updatedAtMs: block.updatedAtMs, }); } } @@ -691,6 +701,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { providerKey: block.providerKey, blockScope: block.blockScope, blockedUntilMs: block.blockedUntilMs, + ...(block.updatedAtMs !== undefined ? { updatedAtMs: block.updatedAtMs } : {}), })) .sort(compareCredentialBlockSnapshots); if (blocks.length > 0) return { ...entry, blocks }; diff --git a/packages/ai/src/auth-broker/server.ts b/packages/ai/src/auth-broker/server.ts index 820ea8bd6..bf5b8afb3 100644 --- a/packages/ai/src/auth-broker/server.ts +++ b/packages/ai/src/auth-broker/server.ts @@ -290,6 +290,7 @@ function buildCredentialBlockGroups( providerKey: block.providerKey, blockScope: block.blockScope, blockedUntilMs: block.blockedUntilMs, + updatedAtMs: block.updatedAtMs, }; const existing = byCredentialId.get(block.credentialId); if (existing) { diff --git a/packages/ai/src/auth-broker/wire-schemas.ts b/packages/ai/src/auth-broker/wire-schemas.ts index 741a9ea11..35180ebf1 100644 --- a/packages/ai/src/auth-broker/wire-schemas.ts +++ b/packages/ai/src/auth-broker/wire-schemas.ts @@ -74,6 +74,7 @@ export const credentialBlockSnapshotSchema = type({ providerKey: type("string").atLeastLength(1), blockScope: "string", blockedUntilMs: "number", + "updatedAtMs?": "number", }); export const snapshotEntrySchema = type({ diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 478d540f4..c0f2b6eda 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -134,6 +134,8 @@ export interface StoredCredentialBlock { blockScope: string; /** Epoch milliseconds. */ blockedUntilMs: number; + /** Last row update timestamp in epoch milliseconds, when provided by the backing store. */ + updatedAtMs?: number; } /** @@ -5884,6 +5886,7 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore { providerKey: row.provider_key, blockScope: row.block_scope, blockedUntilMs: row.blocked_until_ms, + updatedAtMs: row.updated_at * 1000, }); } } diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 85c7c822c..d66c7afe7 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -837,6 +837,150 @@ describe("AuthStorage codex oauth ranking", () => { } }); + test("refreshes broker-sourced Codex block protection when the same deadline is re-upserted", async () => { + if (!authStorage || !store?.getCredentialBlock) { + throw new Error("test setup failed"); + } + + await authStorage.set("openai-codex", [ + { + type: "oauth", + ...createCredential("acct-broker-same-deadline-blocked", "broker-same-deadline-blocked@example.com"), + }, + { + type: "oauth", + ...createCredential("acct-broker-same-deadline-healthy", "broker-same-deadline-healthy@example.com"), + }, + ]); + + usageByAccount.set( + "acct-broker-same-deadline-blocked", + createCodexUsageReport({ + accountId: "acct-broker-same-deadline-blocked", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "broker-same-deadline-blocked@example.com", + accountId: "acct-broker-same-deadline-blocked", + }, + }), + ); + usageByAccount.set( + "acct-broker-same-deadline-healthy", + createCodexUsageReport({ + accountId: "acct-broker-same-deadline-healthy", + primary: { usedFraction: 0.2, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.3, resetInMs: WEEK_MS }, + metadata: { + allowed: true, + limitReached: false, + planType: "pro", + email: "broker-same-deadline-healthy@example.com", + accountId: "acct-broker-same-deadline-healthy", + }, + }), + ); + + const token = "codex-broker-same-deadline-block"; + const handle = startAuthBroker({ + storage: authStorage, + bind: "127.0.0.1:0", + bearerTokens: [token], + disableRefresher: true, + }); + try { + const clientA = new AuthBrokerClient({ url: handle.url, token }); + const clientB = new AuthBrokerClient({ url: handle.url, token }); + const initialResult = await clientB.fetchSnapshot(); + if (initialResult.status !== 200) throw new Error("expected initial broker snapshot"); + const blockedRow = initialResult.snapshot.credentials.find(entry => { + const credential = entry.credential; + return credential.type === "oauth" && credential.accountId === "acct-broker-same-deadline-blocked"; + }); + if (!blockedRow) throw new Error("expected blocked credential row"); + + const blockedUntilMs = Date.now() + 6 * 24 * HOUR_MS; + await clientA.upsertCredentialBlock(blockedRow.id, { + providerKey: "openai-codex:oauth", + blockScope: "shared", + blockedUntilMs, + }); + + const initialUpdatedAtSec = Math.floor(Date.now() / 1000) - 1; + const db = new Database(dbPath); + try { + const result = db + .prepare( + "UPDATE auth_credential_blocks SET updated_at = ? WHERE credential_id = ? AND provider_key = ? AND block_scope = ?", + ) + .run(initialUpdatedAtSec, blockedRow.id, "openai-codex:oauth", "shared") as { changes: number }; + if (result.changes !== 1) throw new Error("expected to age the broker block update timestamp"); + } finally { + db.close(); + } + + const snapshotWithBlock = await clientB.fetchSnapshot({ + ifGenerationGt: initialResult.generation, + waitMs: 1000, + }); + if (snapshotWithBlock.status !== 200) + throw new Error("expected broker snapshot containing same-deadline block"); + const initialSnapshotBlock = snapshotWithBlock.snapshot.credentials + .find(entry => entry.id === blockedRow.id) + ?.blocks?.find(block => block.providerKey === "openai-codex:oauth" && block.blockScope === "shared"); + expect(initialSnapshotBlock?.blockedUntilMs).toBe(blockedUntilMs); + expect(initialSnapshotBlock?.updatedAtMs).toBe(initialUpdatedAtSec * 1000); + + const remoteStoreB = new RemoteAuthCredentialStore({ + client: clientB, + initialSnapshot: snapshotWithBlock.snapshot, + streamSnapshots: false, + }); + const clientStorageB = new AuthStorage(remoteStoreB); + await clientStorageB.reload(); + try { + expect(remoteStoreB.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBe(blockedUntilMs); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBe(blockedUntilMs); + + remoteStoreB.cleanExpiredCredentialBlocks(Date.now() + STALE_BLOCK_GUARD_MS); + + await clientA.upsertCredentialBlock(blockedRow.id, { + providerKey: "openai-codex:oauth", + blockScope: "shared", + blockedUntilMs, + }); + const refreshedSnapshot = await clientB.fetchSnapshot({ + ifGenerationGt: snapshotWithBlock.generation, + waitMs: 1000, + }); + if (refreshedSnapshot.status !== 200) { + throw new Error("expected broker snapshot containing refreshed same-deadline block"); + } + + await remoteStoreB.refreshSnapshot(); + const refreshedBlock = remoteStoreB.snapshot.credentials + .find(entry => entry.id === blockedRow.id) + ?.blocks?.find(block => block.providerKey === "openai-codex:oauth" && block.blockScope === "shared"); + expect(refreshedBlock?.blockedUntilMs).toBe(blockedUntilMs); + expect(refreshedBlock?.updatedAtMs).toBeGreaterThan(initialSnapshotBlock!.updatedAtMs!); + + expect(await clientStorageB.getApiKey("openai-codex", "codex-broker-same-deadline-sibling")).toBe( + "api-acct-broker-same-deadline-healthy", + ); + expect(remoteStoreB.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBe(blockedUntilMs); + expect(store.getCredentialBlock(blockedRow.id, "openai-codex:oauth", "shared")).toBe(blockedUntilMs); + } finally { + clientStorageB.close(); + remoteStoreB.close(); + } + } finally { + await handle.close(); + } + }); + test("protects fresh Codex blocks present in the initial broker snapshot from healthy selection reconciliation", async () => { if (!authStorage || !store?.getCredentialBlock) { throw new Error("test setup failed");