diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 0e3af1a2d..02c36dea5 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -6,6 +6,10 @@ - Fixed API-key `/login` providers replacing sibling credentials instead of appending new keys for the same provider. ([#3265](https://github.com/can1357/oh-my-pi/issues/3265)) +### Fixed + +- Fixed OpenAI Codex OAuth account rotation for quota failures that surface as bare HTTP 429 or `insufficient_quota`, so pre-content failures temporarily block only the exhausted credential and retry a healthy sibling. The 429 status-only fallback applies only to absent/opaque bodies; informative transient bodies (`Too many requests`, `Service overloaded 529`, `Please retry in 5s`, …) defer to `parseRateLimitReason` and stay in the provider's own backoff layer instead of burning sibling credentials. ([#3231](https://github.com/can1357/oh-my-pi/issues/3231)) + ## [16.1.14] - 2026-06-22 ### Added diff --git a/packages/ai/src/auth-gateway/server.ts b/packages/ai/src/auth-gateway/server.ts index e41f91884..2003726ae 100644 --- a/packages/ai/src/auth-gateway/server.ts +++ b/packages/ai/src/auth-gateway/server.ts @@ -19,14 +19,14 @@ */ import { Effort } from "@oh-my-pi/pi-catalog/effort"; -import { extractRetryHint, logger } from "@oh-my-pi/pi-utils"; +import { extractHttpStatusFromError, extractRetryHint, logger } from "@oh-my-pi/pi-utils"; import type { ApiKeyResolver } from "../auth-retry"; import type { AuthStorage } from "../auth-storage"; import * as anthropicMessages from "../providers/anthropic-messages-server"; import * as openaiChat from "../providers/openai-chat-server"; import * as openaiResponses from "../providers/openai-responses-server"; import * as piNative from "../providers/pi-native-server"; -import { isUsageLimitError } from "../rate-limit-utils"; +import { isUsageLimitError, isUsageLimitOutcome } from "../rate-limit-utils"; import { streamSimple } from "../stream"; import type { Api, AssistantMessageEventStream, Context, Model, SimpleStreamOptions } from "../types"; import { deterministicUuid } from "../utils/deterministic-id"; @@ -315,7 +315,7 @@ async function refreshGatewayApiKeyAfterAuthError( peer: string, ): Promise { const message = error instanceof Error ? error.message : String(error); - if (isUsageLimitError(message)) { + if (isUsageLimitOutcome(extractHttpStatusFromError(error), message)) { const retryAfterMs = extractRetryHint(undefined, message); const { switched, retryAtMs } = await storage.markUsageLimitReached(provider, sessionId, { retryAfterMs, diff --git a/packages/ai/src/auth-retry.ts b/packages/ai/src/auth-retry.ts index 12bc14cc7..d4a45bf87 100644 --- a/packages/ai/src/auth-retry.ts +++ b/packages/ai/src/auth-retry.ts @@ -1,6 +1,6 @@ import { extractHttpStatusFromError } from "@oh-my-pi/pi-utils"; import type { OAuthAccess } from "./auth-storage"; -import { isUsageLimitError } from "./rate-limit-utils"; +import { isUsageLimitOutcome } from "./rate-limit-utils"; /** * Context passed to an {@link ApiKeyResolver} on each resolution attempt. @@ -72,15 +72,20 @@ export function seedApiKeyResolver(seed: string | undefined, resolver: ApiKeyRes /** * Classifies whether an error should trigger a credential refresh/rotation - * retry: a hard `401`, or a rotatable usage-limit ("usage_limit_reached", - * Codex's "you have hit your ChatGPT usage limit", etc.). + * retry: a hard `401`, body-classified usage limit (Codex + * `usage_limit_reached`, Anthropic account rate-limit, Google + * `resource_exhausted`, OpenAI `insufficient_quota`, …), or a bare `429` + * whose payload did not preserve a richer quota code. Transient 429s + * (`Too many requests`, per-minute caps) classify as `RATE_LIMIT_EXCEEDED` + * via {@link parseRateLimitReason} and stay in the upstream-backoff lane. */ export function isAuthRetryableError(error: unknown): boolean { - if (extractHttpStatusFromError(error) === 401) return true; + const status = extractHttpStatusFromError(error); + if (status === 401) return true; const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; - if (!message) return false; - if (extractHttpStatusFromError({ message }) === 401) return true; - return isUsageLimitError(message); + const embeddedStatus = message ? extractHttpStatusFromError({ message }) : undefined; + if (embeddedStatus === 401) return true; + return isUsageLimitOutcome(status ?? embeddedStatus, message); } /** diff --git a/packages/ai/src/auth-storage.ts b/packages/ai/src/auth-storage.ts index 7ecf412ad..45a909167 100644 --- a/packages/ai/src/auth-storage.ts +++ b/packages/ai/src/auth-storage.ts @@ -10,9 +10,9 @@ import { Database, type Statement } from "bun:sqlite"; import * as fs from "node:fs/promises"; import * as path from "node:path"; -import { getAgentDbPath, logger } from "@oh-my-pi/pi-utils"; +import { extractHttpStatusFromError, getAgentDbPath, logger } from "@oh-my-pi/pi-utils"; import type { ApiKeyResolver } from "./auth-retry"; -import { isUsageLimitError } from "./rate-limit-utils"; +import { isUsageLimitOutcome } from "./rate-limit-utils"; import { getProviderDefinition } from "./registry"; import { getOAuthApiKey, getOAuthProvider, refreshOAuthToken } from "./registry/oauth"; import type { OAuthController, OAuthCredentials, OAuthProvider, OAuthProviderId } from "./registry/oauth/types"; @@ -4084,8 +4084,9 @@ export class AuthStorage { if (!sessionCredential) return false; const error = options?.error; + const status = extractHttpStatusFromError(error); const message = error instanceof Error ? error.message : typeof error === "string" ? error : undefined; - if (message && isUsageLimitError(message)) { + if (isUsageLimitOutcome(status, message)) { return ( await this.markUsageLimitReached(provider, sessionId, { modelId: options?.modelId, diff --git a/packages/ai/src/rate-limit-utils.ts b/packages/ai/src/rate-limit-utils.ts index fed913b47..b2772cec1 100644 --- a/packages/ai/src/rate-limit-utils.ts +++ b/packages/ai/src/rate-limit-utils.ts @@ -100,7 +100,55 @@ export function calculateRateLimitBackoffMs(reason: RateLimitReason): number { /** Detect usage/quota limit errors in error messages (persistent, requires credential switch). */ const USAGE_LIMIT_PATTERN = - /usage.?limit|usage_limit_reached|usage_not_included|limit_reached|quota.?exceeded|quota.?reached|resource.?exhausted|exhausted your capacity|quota will reset|insufficient.?balance/i; + /usage.?limit|usage_limit_reached|usage_not_included|limit_reached|quota.?exceeded|quota.?reached|resource.?exhausted|exhausted your capacity|quota will reset|insufficient.?(?:balance|quota)/i; + +/** + * HTTP status codes that, absent richer body classification, represent an + * account-local usage cap rather than a bad credential or a transient blip. + * Always combine with {@link isUsageLimitOutcome} when a message is available + * — a 429 carrying transient rate-limit wording is NOT a usage cap. + */ +export function isUsageLimitStatus(status: number | undefined): boolean { + return status === 429; +} + +/** + * Returns true for failures that should burn one credential and rotate to a + * sibling account. Decision tree: + * + * 1. Body matches {@link isUsageLimitError} (Codex `usage_limit_reached`, + * Anthropic account rate-limit, Google `resource_exhausted`, OpenAI + * `insufficient_quota`, …) → rotate. + * 2. Status is not 429 → backoff (caller's domain). + * 3. Body is absent or {@link isOpaqueStatusBody opaque} (just the status, + * empty JSON, HTTP framing only) → rotate conservatively: the server + * gave us nothing else to go on. + * 4. Body has content → defer to {@link parseRateLimitReason}. Only + * `QUOTA_EXHAUSTED` rotates; `RATE_LIMIT_EXCEEDED` (`Too many requests`, + * per-minute caps), `MODEL_CAPACITY_EXHAUSTED` (`Service overloaded`), + * `SERVER_ERROR`, and `UNKNOWN` (`Please retry in 5s`) stay in the + * provider's own backoff layer so transient 429s don't burn sibling + * credentials. + */ +export function isUsageLimitOutcome(status: number | undefined, message: string | undefined): boolean { + if (message && isUsageLimitError(message)) return true; + if (!isUsageLimitStatus(status)) return false; + if (!message || isOpaqueStatusBody(message)) return true; + return parseRateLimitReason(message) === "QUOTA_EXHAUSTED"; +} + +/** + * A 429 body is opaque when it carries no signal beyond the status itself — + * empty, whitespace-only, the status digits with HTTP/JSON framing, or + * generic punctuation. Anything else (retry hints, capacity wording, error + * descriptions) is informative enough to defer to the classifier. + */ +function isOpaqueStatusBody(message: string): boolean { + const cleaned = message + .replace(/\b429\b/g, "") + .replace(/\b(?:http|https|status|error|code|response|message)\b/gi, ""); + return !/[a-z\d]{3,}/i.test(cleaned); +} export function isUsageLimitError(errorMessage: string): boolean { return USAGE_LIMIT_PATTERN.test(errorMessage) || ACCOUNT_RATE_LIMIT_PATTERN.test(errorMessage); diff --git a/packages/ai/src/stream.ts b/packages/ai/src/stream.ts index 2dce95b30..14ff913ac 100644 --- a/packages/ai/src/stream.ts +++ b/packages/ai/src/stream.ts @@ -48,7 +48,7 @@ import { streamOpenAIResponses, } from "./providers/register-builtins"; import { isSyntheticModel, streamSynthetic } from "./providers/synthetic"; -import { isUsageLimitError } from "./rate-limit-utils"; +import { isUsageLimitOutcome } from "./rate-limit-utils"; import { PROVIDER_REGISTRY } from "./registry"; import type { Api, @@ -384,13 +384,18 @@ function extractStatusFromAssistantError(message: AssistantMessage): number | un function isRetryableUpstreamError(error: unknown, status: number | undefined, message: string | undefined): boolean { // 401 means the credential is bad. Usage-limit phrasing (Codex's // "You have hit your ChatGPT usage limit", Anthropic's "usage_limit_reached", - // Google's "resource_exhausted") means this account is parked but a - // sibling credential can usually pick the request up. Both are - // rotatable via `onAuthError` — the auth-gateway maps the former to - // `invalidateCredentialMatching` and the latter to `markUsageLimitReached`. + // Google's "resource_exhausted", OpenAI's "insufficient_quota") and 429s + // without transient rate-limit wording mean this account is parked but a + // sibling credential can usually pick the request up. Both are rotatable + // via `onAuthError` — the auth-gateway maps the former to + // `invalidateCredentialMatching` and the latter to + // `markUsageLimitReached`. Transient 429s ("Too many requests", + // per-minute caps) classify as RATE_LIMIT_EXCEEDED in + // `parseRateLimitReason` and stay in the provider's own backoff layer + // instead of burning siblings. if (status === 401) return true; void error; - return !!message && isUsageLimitError(message); + return isUsageLimitOutcome(status, message); } function createAssistantAuthError(message: AssistantMessage): Error { diff --git a/packages/ai/test/auth-storage-codex-selection.test.ts b/packages/ai/test/auth-storage-codex-selection.test.ts index e98fd6443..2f61b1ac6 100644 --- a/packages/ai/test/auth-storage-codex-selection.test.ts +++ b/packages/ai/test/auth-storage-codex-selection.test.ts @@ -267,6 +267,59 @@ describe("AuthStorage codex oauth ranking", () => { expect(apiKey).toBe("api-acct-healthy"); }); + test("temporarily blocks only the exhausted Codex OAuth credential after a quota 429", async () => { + if (!authStorage) throw new Error("test setup failed"); + + await authStorage.set("openai-codex", [ + { type: "oauth", ...createCredential("acct-A", "a@example.com") }, + { type: "oauth", ...createCredential("acct-B", "b@example.com") }, + ]); + usageByAccount.set( + "acct-A", + createCodexUsageReport({ + accountId: "acct-A", + primary: { usedFraction: 0.1, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.1, resetInMs: WEEK_MS }, + }), + ); + usageByAccount.set( + "acct-B", + createCodexUsageReport({ + accountId: "acct-B", + primary: { usedFraction: 0.1, resetInMs: HOUR_MS }, + secondary: { usedFraction: 0.1, resetInMs: WEEK_MS }, + }), + ); + + const sessionId = "session-codex-quota-429"; + const firstKey = await authStorage.getApiKey("openai-codex", sessionId); + if (!firstKey) throw new Error("expected initial Codex credential"); + const exhaustedAccount = firstKey.replace(/^api-/, ""); + const healthyAccount = exhaustedAccount === "acct-A" ? "acct-B" : "acct-A"; + usageByAccount.set( + exhaustedAccount, + createCodexUsageReport({ + accountId: exhaustedAccount, + primary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 }, + secondary: { usedFraction: 1, resetInMs: 5 * 60 * 1000 }, + }), + ); + + const usageLimitSpy = vi.spyOn(authStorage, "markUsageLimitReached"); + const switched = await authStorage.rotateSessionCredential("openai-codex", sessionId, { + error: Object.assign(new Error("insufficient_quota"), { status: 429 }), + }); + + expect(switched).toBe(true); + expect(usageLimitSpy).toHaveBeenCalledTimes(1); + expect(await authStorage.getApiKey("openai-codex", sessionId)).toBe(`api-${healthyAccount}`); + const activeAccounts = (await authStorage.checkCredentials()) + .map(result => result.accountId) + .filter((accountId): accountId is string => accountId !== undefined) + .sort(); + expect(activeAccounts).toEqual(["acct-A", "acct-B"]); + }); + test("falls back to earliest-unblocking account when all exhausted", async () => { if (!authStorage) throw new Error("test setup failed"); diff --git a/packages/ai/test/auth-storage-force-refresh-rotate.test.ts b/packages/ai/test/auth-storage-force-refresh-rotate.test.ts index 62a2865ec..c6eb5d117 100644 --- a/packages/ai/test/auth-storage-force-refresh-rotate.test.ts +++ b/packages/ai/test/auth-storage-force-refresh-rotate.test.ts @@ -22,6 +22,14 @@ function usageLimitError(): Error & { status: number } { }); } +function quotaPayloadError(message: string, status?: number): Error & { status?: number } { + return status === undefined ? new Error(message) : Object.assign(new Error(message), { status }); +} + +function invalidRequestError(): Error & { status: number } { + return Object.assign(new Error("400 invalid_request_error: model unsupported"), { status: 400 }); +} + describe("AuthStorage forceRefresh + rotateSessionCredential", () => { let tempDir = ""; let store: AuthCredentialStore | undefined; @@ -141,6 +149,84 @@ describe("AuthStorage forceRefresh + rotateSessionCredential", () => { expect(second).not.toBe(first); }); + test("rotateSessionCredential treats quota payloads as temporary usage blocks", async () => { + if (!authStorage) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "acc-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "acc-B", refresh: "ref-B", expires: farExpiry() }, + { type: "oauth", access: "acc-C", refresh: "ref-C", expires: farExpiry() }, + { type: "oauth", access: "acc-D", refresh: "ref-D", expires: farExpiry() }, + { type: "oauth", access: "acc-E", refresh: "ref-E", expires: farExpiry() }, + ]); + + for (const [index, error] of [ + [0, quotaPayloadError("429", 429)], + [1, quotaPayloadError("insufficient_quota")], + [2, quotaPayloadError("usage_limit_exceeded")], + [3, quotaPayloadError("usage_limit_reached")], + ] as const) { + const sessionId = `quota-payload-${index}`; + const first = await authStorage.getApiKey(PROVIDER, sessionId); + const usageLimitSpy = vi.spyOn(authStorage, "markUsageLimitReached"); + + const rotated = await authStorage.rotateSessionCredential(PROVIDER, sessionId, { error }); + + expect(rotated).toBe(true); + expect(usageLimitSpy).toHaveBeenCalledTimes(1); + expect(await authStorage.getApiKey(PROVIDER, sessionId)).not.toBe(first); + usageLimitSpy.mockRestore(); + } + }); + + test("rotateSessionCredential does not treat invalid requests as quota blocks", async () => { + if (!authStorage) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "acc-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "acc-B", refresh: "ref-B", expires: farExpiry() }, + ]); + + await authStorage.getApiKey(PROVIDER, "invalid-request"); + const usageLimitSpy = vi.spyOn(authStorage, "markUsageLimitReached"); + + await authStorage.rotateSessionCredential(PROVIDER, "invalid-request", { error: invalidRequestError() }); + + expect(usageLimitSpy).not.toHaveBeenCalled(); + }); + + test("rotateSessionCredential leaves informative transient 429s out of the quota block path", async () => { + if (!authStorage) throw new Error("test setup failed"); + registerProvider(); + await authStorage.set(PROVIDER, [ + { type: "oauth", access: "acc-A", refresh: "ref-A", expires: farExpiry() }, + { type: "oauth", access: "acc-B", refresh: "ref-B", expires: farExpiry() }, + ]); + + const transient429Bodies = [ + "Cloud Code Assist API error (429): Too many requests", + "Please retry in 5s", + "Service overloaded 529", + ]; + + for (const [index, body] of transient429Bodies.entries()) { + const sessionId = `transient-429-${index}`; + await authStorage.getApiKey(PROVIDER, sessionId); + const usageLimitSpy = vi.spyOn(authStorage, "markUsageLimitReached"); + + await authStorage.rotateSessionCredential(PROVIDER, sessionId, { + error: Object.assign(new Error(body), { status: 429 }), + }); + + // `Too many requests`, server retry hints, and capacity overload are + // owned by the provider's own retry layer — burning a sibling + // credential here would orphan a healthy account for the default + // backoff window. + expect(usageLimitSpy).not.toHaveBeenCalled(); + usageLimitSpy.mockRestore(); + } + }); + test("rotateSessionCredential reports no sibling for a single-credential setup", async () => { if (!authStorage) throw new Error("test setup failed"); registerProvider(); diff --git a/packages/ai/test/rate-limit-utils.test.ts b/packages/ai/test/rate-limit-utils.test.ts index 387e00168..4b77a215e 100644 --- a/packages/ai/test/rate-limit-utils.test.ts +++ b/packages/ai/test/rate-limit-utils.test.ts @@ -1,5 +1,11 @@ import { describe, expect, it } from "bun:test"; -import { calculateRateLimitBackoffMs, isUsageLimitError, parseRateLimitReason } from "@oh-my-pi/pi-ai/rate-limit-utils"; +import { + calculateRateLimitBackoffMs, + isUsageLimitError, + isUsageLimitOutcome, + isUsageLimitStatus, + parseRateLimitReason, +} from "@oh-my-pi/pi-ai/rate-limit-utils"; describe("parseRateLimitReason", () => { it("classifies Google Quota exceeded as QUOTA_EXHAUSTED", () => { @@ -118,6 +124,61 @@ describe("isUsageLimitError", () => { expect(isUsageLimitError("quota reached")).toBe(true); expect(isUsageLimitError("quota_reached")).toBe(true); }); + + it("detects OpenAI quota payload codes as credential-rotatable usage limits", () => { + for (const message of ["insufficient_quota", "usage_limit_exceeded", "usage_limit_reached"]) { + expect(isUsageLimitError(message)).toBe(true); + } + expect(isUsageLimitStatus(429)).toBe(true); + expect(isUsageLimitStatus(400)).toBe(false); + }); +}); + +describe("isUsageLimitOutcome", () => { + it("rotates on bare/opaque 429 bodies (status-only fallback)", () => { + expect(isUsageLimitOutcome(429, undefined)).toBe(true); + expect(isUsageLimitOutcome(429, "")).toBe(true); + expect(isUsageLimitOutcome(429, "429")).toBe(true); + expect(isUsageLimitOutcome(429, "HTTP 429")).toBe(true); + expect(isUsageLimitOutcome(429, "Error 429")).toBe(true); + expect(isUsageLimitOutcome(429, "{}")).toBe(true); + }); + + it("rotates on 429 carrying quota payload codes", () => { + for (const message of ["insufficient_quota", "usage_limit_exceeded", "usage_limit_reached"]) { + expect(isUsageLimitOutcome(429, message)).toBe(true); + } + }); + + it("keeps informative transient 429s in the upstream-backoff lane", () => { + // RATE_LIMIT_EXCEEDED — generic throttling. + expect(isUsageLimitOutcome(429, "Cloud Code Assist API error (429): Too many requests")).toBe(false); + expect(isUsageLimitOutcome(429, "Requests per minute limit reached")).toBe(false); + // MODEL_CAPACITY_EXHAUSTED — provider overload, not account quota. + expect(isUsageLimitOutcome(429, "Service overloaded 529")).toBe(false); + // UNKNOWN but carries a transient retry hint — body is informative, + // so we defer to parseRateLimitReason and stay out of the quota lane. + expect(isUsageLimitOutcome(429, "Please retry in 5s")).toBe(false); + }); + + it("still rotates on 429 with explicit account rate-limit framing", () => { + expect( + isUsageLimitOutcome( + 429, + '{"type":"error","error":{"type":"rate_limit_error","message":"This request would exceed your account\'s rate limit. Please try again later."}}', + ), + ).toBe(true); + }); + + it("rotates on usage-limit message regardless of status", () => { + expect(isUsageLimitOutcome(undefined, "usage_limit_reached")).toBe(true); + expect(isUsageLimitOutcome(500, "insufficient_quota")).toBe(true); + }); + + it("does not rotate on auth/invalid-request statuses with unrelated bodies", () => { + expect(isUsageLimitOutcome(401, "Invalid API key")).toBe(false); + expect(isUsageLimitOutcome(400, "invalid_request_error: model unsupported")).toBe(false); + }); }); describe("calculateRateLimitBackoffMs", () => { diff --git a/packages/ai/test/stream-auth-retry.test.ts b/packages/ai/test/stream-auth-retry.test.ts index 4d27dfbd4..307b37412 100644 --- a/packages/ai/test/stream-auth-retry.test.ts +++ b/packages/ai/test/stream-auth-retry.test.ts @@ -350,6 +350,113 @@ describe("streamSimple resolver auth retry", () => { expect(keys).toEqual(["old-key", "new-key"]); }); + it("rotates before emitting content for Codex quota payloads", async () => { + const payloads: Array<{ message: string; status?: number }> = [ + { message: "429", status: 429 }, + { message: '{"error":{"code":"insufficient_quota","message":"quota exhausted"}}' }, + { message: '{"error":{"code":"usage_limit_exceeded","message":"usage limit exceeded"}}' }, + { message: '{"error":{"code":"usage_limit_reached","message":"usage limit reached"}}' }, + ]; + let activePayload = payloads[0]!; + let keys: unknown[] = []; + registerCustomApi( + API, + (_model: Model, _context: Context, options?: SimpleStreamOptions) => { + pushKey(keys, options); + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + if (options?.apiKey === "credential-B") { + ok(stream); + return; + } + stream.push({ type: "start", partial: assistant() }); + stream.push({ + type: "error", + reason: "error", + error: assistantError(activePayload.message, activePayload.status), + }); + }); + return stream; + }, + SOURCE_ID, + ); + + for (const payload of payloads) { + activePayload = payload; + keys = []; + const eventTypes: string[] = []; + const retryContexts: ApiKeyResolveContext[] = []; + const stream = streamSimple(model(), context, { + apiKey: async ctx => { + if (ctx.error !== undefined) retryContexts.push(ctx); + return ctx.error === undefined ? "credential-A" : ctx.lastChance ? "credential-B" : "credential-A"; + }, + }); + for await (const event of stream) { + eventTypes.push(event.type); + } + + expect((await stream.result()).content).toEqual([{ type: "text", text: "ok" }]); + expect(keys).toEqual(["credential-A", "credential-B"]); + expect(eventTypes).toEqual(["start", "done"]); + expect(retryContexts.map(ctx => ctx.lastChance)).toEqual([false, true]); + } + }); + + it("does not rotate or refresh on informative transient 429 bodies", async () => { + const transient429Bodies = [ + "Cloud Code Assist API error (429): Too many requests", + "Please retry in 5s", + "Service overloaded 529", + ]; + let active = transient429Bodies[0]!; + const keys: unknown[] = []; + const retryResolves: ApiKeyResolveContext[] = []; + registerCustomApi( + API, + (_model: Model, _context: Context, options?: SimpleStreamOptions) => { + pushKey(keys, options); + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + stream.push({ type: "start", partial: assistant() }); + stream.push({ + type: "error", + reason: "error", + error: assistantError(active, 429), + }); + }); + return stream; + }, + SOURCE_ID, + ); + + for (const body of transient429Bodies) { + active = body; + keys.length = 0; + retryResolves.length = 0; + const eventTypes: string[] = []; + const stream = streamSimple(model(), context, { + apiKey: async ctx => { + if (ctx.error !== undefined) retryResolves.push(ctx); + return ctx.error === undefined ? "credential-A" : "credential-B"; + }, + }); + + for await (const event of stream) { + eventTypes.push(event.type); + } + const result = await stream.result(); + + // The provider's own retry/backoff layer owns these — the auth + // retry loop must NOT capture, refresh, or burn a sibling. + expect(retryResolves).toEqual([]); + expect(keys).toEqual(["credential-A"]); + expect(eventTypes).toEqual(["start", "error"]); + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toContain(body); + } + }); + it("rotates on the exact Google Resource exhausted 429 error before content", async () => { const keys: unknown[] = []; const retryContexts: ApiKeyResolveContext[] = [];