diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index e90b9e3db..ab234c9ad 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -49,6 +49,7 @@ - Refined account selection logic to correctly identify plan types from account metadata - Fixed OpenAI Codex multi-account routing for GPT-5.6: Sol and Luna requests now prefer Plus-or-higher accounts while Terra remains available to Free/Go accounts; local pro-mode aliases inherit their base model's Codex plan eligibility. - Fixed xAI Grok OAuth login to use xAI's device authorization flow: `/login` now opens the verification URL, displays the device code, and polls for approval instead of asking for a pasted redirect or linking to Hermes Agent documentation. +- 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-broker/remote-store.ts b/packages/ai/src/auth-broker/remote-store.ts index f56ba6b3a..4351fe616 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; @@ -50,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 { @@ -58,6 +61,7 @@ function toCredentialBlockSnapshot(block: StoredCredentialBlock): CredentialBloc providerKey: block.providerKey, blockScope: block.blockScope, blockedUntilMs: block.blockedUntilMs, + ...(block.updatedAtMs !== undefined ? { updatedAtMs: block.updatedAtMs } : {}), }; } @@ -74,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; } @@ -208,6 +213,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { #cache: Map = new Map(); #usageCache?: UsageCacheEntry; #usageInflight?: Promise; + #credentialBlockReconcileAfter: Map = new Map(); #usageCacheEpoch = 0; #closed = false; /** @@ -236,10 +242,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 +259,34 @@ 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}\0${block.updatedAtMs ?? ""}`, + ); + } + } + 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); + 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, updatedAtMs + 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 +376,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 +430,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); @@ -406,6 +449,7 @@ export class RemoteAuthCredentialStore implements AuthCredentialStore { providerKey: block.providerKey, blockScope: block.blockScope, blockedUntilMs: block.blockedUntilMs, + updatedAtMs: block.updatedAtMs, }); } } @@ -416,6 +460,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 +482,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 +501,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); + } } /** @@ -647,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 fb9020f0d..417d1e1c8 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; } /** @@ -320,6 +322,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). */ @@ -1024,6 +1028,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; @@ -1404,6 +1410,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; @@ -1496,6 +1505,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); @@ -3373,18 +3385,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(); @@ -4433,23 +4463,24 @@ 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); + } } /** * 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 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; - return ( - report.provider === "openai-codex" && - metadata?.allowed === true && - metadata.limitReached === false && - !this.#isUsageLimitReached(report.limits) - ); + if (metadata?.allowed !== true || metadata.limitReached !== false) return false; + return !this.#isUsageLimitReached(report.limits); } #reconcileCodexUsageBlockForCredential(provider: Provider, credentialId: number, report: UsageReport): void { @@ -4462,6 +4493,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 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; + 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, @@ -4965,6 +5009,7 @@ type CredentialBlockRow = { provider_key: string; block_scope: string; blocked_until_ms: number; + updated_at: number; }; type SerializedCredentialRecord = { @@ -5180,6 +5225,7 @@ export class SqliteAuthCredentialStore implements AuthCredentialStore { #upsertCredentialBlockStmt: Statement; #deleteCredentialBlocksStmt: Statement; #deleteExpiredCredentialBlocksStmt: Statement; + #credentialBlockReconcileAfter: Map = new Map(); #insertUsageHistoryStmt: Statement; #insertUsageCostStmt: Statement; #listUsageCostsStmt: Statement; @@ -5226,10 +5272,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) @@ -5843,11 +5889,26 @@ 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 { + 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 { this.#upsertCredentialBlockStmt.run( block.credentialId, @@ -5855,14 +5916,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[] { @@ -5881,6 +5952,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 a0ff1e114..50d859786 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"; @@ -13,6 +14,18 @@ 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; + +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; @@ -144,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(); @@ -158,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), }); @@ -353,6 +368,8 @@ 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( "acct-blocked", @@ -385,15 +402,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 +419,709 @@ 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, + }); + ageCredentialBlockRows(dbPath); + store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_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("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("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("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"); + } + + 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("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"); + } + + 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"); @@ -496,18 +1207,10 @@ 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 () => { - if (!authStorage || !store?.getCredentialBlock) { + if (!authStorage || !store?.getCredentialBlock || !store.upsertCredentialBlock) { throw new Error("test setup failed"); } @@ -547,6 +1250,20 @@ 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, + }); + ageCredentialBlockRows(dbPath); + store.cleanExpiredCredentialBlocks?.(Date.now() + STALE_BLOCK_GUARD_MS); + const token = "codex-broker-reconcile"; const handle = startAuthBroker({ storage: authStorage, @@ -556,18 +1273,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 => { @@ -583,31 +1288,17 @@ 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(); + remoteStore.cleanExpiredCredentialBlocks(Date.now() + STALE_BLOCK_GUARD_MS); 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();