diff --git a/packages/ai/src/auth-retry.ts b/packages/ai/src/auth-retry.ts index 87da67210..22a8f8e55 100644 --- a/packages/ai/src/auth-retry.ts +++ b/packages/ai/src/auth-retry.ts @@ -1,6 +1,7 @@ import type { OAuthAccess } from "./auth-storage"; import * as AIError from "./error"; import { isAuthRetryableError } from "./error/auth-classify"; +import { isUsageLimit } from "./error/flags"; /** * Context passed to an {@link ApiKeyResolver} on each resolution attempt. @@ -92,7 +93,8 @@ export async function resolveRetryKey( previousKey?: string, ): Promise { try { - return (await resolver({ lastChance, error, signal, previousKey })) || undefined; + const rotateSibling = lastChance || (!lastChance && isUsageLimit(error)); + return (await resolver({ lastChance: rotateSibling, error, signal, previousKey })) || undefined; } catch { return undefined; } diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 3d0ef1262..67b698838 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -1390,12 +1390,14 @@ export class AuthStorage { const credentialId = this.#getStoredCredentials(provider)[credentialIndex]?.id; if (credentialId === undefined) return blockedUntil; - const persistedGlobalBlockedUntil = this.#readPersistedCredentialBlock(credentialId, providerKey, ""); - if ( - persistedGlobalBlockedUntil !== undefined && - (blockedUntil === undefined || persistedGlobalBlockedUntil > blockedUntil) - ) { - blockedUntil = persistedGlobalBlockedUntil; + if (!blockScope || provider !== "openai-codex") { + const persistedGlobalBlockedUntil = this.#readPersistedCredentialBlock(credentialId, providerKey, ""); + if ( + persistedGlobalBlockedUntil !== undefined && + (blockedUntil === undefined || persistedGlobalBlockedUntil > blockedUntil) + ) { + blockedUntil = persistedGlobalBlockedUntil; + } } if (blockScope) { const persistedScopedBlockedUntil = this.#readPersistedCredentialBlock(credentialId, providerKey, blockScope); diff --git a/packages/ai/src/usage/openai-codex.ts b/packages/ai/src/usage/openai-codex.ts index 355744208..c7adebb4c 100644 --- a/packages/ai/src/usage/openai-codex.ts +++ b/packages/ai/src/usage/openai-codex.ts @@ -505,6 +505,9 @@ export const openaiCodexUsageProvider: UsageProvider = { const FIVE_HOUR_MS = 5 * 60 * 60 * 1000; export const codexRankingStrategy: CredentialRankingStrategy = { + blockScope() { + return "shared"; + }, findWindowLimits(report) { const findLimit = (key: "primary" | "secondary"): UsageLimit | undefined => { const direct = report.limits.find(l => l.id === `openai-codex:${key}`); diff --git a/packages/ai/test/auth-retry.test.ts b/packages/ai/test/auth-retry.test.ts index 73b1f66d1..44a4a9f89 100644 --- a/packages/ai/test/auth-retry.test.ts +++ b/packages/ai/test/auth-retry.test.ts @@ -99,6 +99,28 @@ describe("withAuth", () => { ]); }); + it("switches accounts before refreshing the same account on usage limits", async () => { + const keys: string[] = []; + const contexts: ApiKeyResolveContext[] = []; + const result = await withAuth( + ctx => { + contexts.push(ctx); + return ctx.error === undefined ? "k0" : ctx.lastChance ? "k2" : "k1"; + }, + async key => { + keys.push(key); + if (key === "k2") return "success"; + throw usageLimitError(); + }, + ); + expect(result).toBe("success"); + expect(keys).toEqual(["k0", "k2"]); + expect(contexts.map(ctx => ({ lastChance: ctx.lastChance, hasError: ctx.error !== undefined }))).toEqual([ + { lastChance: false, hasError: false }, + { lastChance: true, hasError: true }, + ]); + }); + it("stops retrying when the resolver returns undefined", async () => { const keys: string[] = []; const original = authError(); diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index 3dd0a169f..b6aa79847 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -643,6 +643,52 @@ describe("AuthStorage codex oauth ranking", () => { expect(await authStorage.getApiKey("openai-codex", sessionId)).toBe("api-acct-plus"); }); + + test("ignores legacy global Codex blocks when a scoped quota window has fresh siblings", async () => { + if (!authStorage || !store) throw new Error("test setup failed"); + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-k12", "k12@example.com") }, + { type: "oauth", ...createCredential("acct-plus", "plus@example.com") }, + ]); + usageByAccount.set( + "acct-k12", + createCodexUsageReport({ + accountId: "acct-k12", + primary: { usedFraction: 1, resetInMs: FIVE_HOUR_MS }, + secondary: { usedFraction: 1, resetInMs: WEEK_MS }, + }), + ); + usageByAccount.set( + "acct-plus", + createCodexUsageReport({ + accountId: "acct-plus", + primary: { usedFraction: 0.2, resetInMs: FIVE_HOUR_MS }, + secondary: { usedFraction: 0.74, resetInMs: WEEK_MS }, + }), + ); + const plus = store + .listAuthCredentials("openai-codex") + .find(row => row.credential.type === "oauth" && row.credential.accountId === "acct-plus"); + if (!plus || !store.upsertCredentialBlock) throw new Error("missing plus credential row"); + store.upsertCredentialBlock({ + credentialId: plus.id, + providerKey: "openai-codex:oauth", + blockScope: "", + blockedUntilMs: Date.now() + WEEK_MS, + }); + const k12 = store + .listAuthCredentials("openai-codex") + .find(row => row.credential.type === "oauth" && row.credential.accountId === "acct-k12"); + if (!k12 || !store.upsertCredentialBlock) throw new Error("missing k12 credential row"); + store.upsertCredentialBlock({ + credentialId: k12.id, + providerKey: "openai-codex:oauth", + blockScope: "shared", + blockedUntilMs: Date.now() + HOUR_MS, + }); + + expect(await authStorage.getApiKey("openai-codex", "session-with-legacy-global-block")).toBe("api-acct-plus"); + }); }); // ─────────────────────────────────────────────────────────────────────────────