diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 9925e111815..b00be853cf9 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -938,6 +938,7 @@ "gemini-web-search.test.ts": "adapters/google", "generic-session-affinity.test.ts": "oauth", "server-generic-oauth-affinity.test.ts": "server", + "generic-oauth-combo-account-cooldown.test.ts": "oauth", "generic-oauth-failover.test.ts": "oauth", "generic-oauth-refusal-composition.test.ts": "oauth", "github-copilot-auto.test.ts": "providers/github-copilot", "github-copilot-auto-runtime.test.ts": "providers/github-copilot", "github-copilot-catalog-hints.test.ts": "providers/github-copilot", "github-copilot-account-origin.test.ts": "providers/github-copilot", diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index c8cbd8e042c..9fca39601c3 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -483,7 +483,8 @@ export function noteGenericPoolSelection( * Cool the account that actually 429'd and name the next eligible one, or null. * * Returns the id only; the caller mints the credential so a failed refresh does not leave the - * cooldown applied to an account we then could not use. + * cooldown applied to an account we then could not use. `commit: false` still cools the failed + * account but only peeks the round-robin ring, for callers that will not dispatch the answer. */ export function rotateGenericOAuthAccountOn429( config: OcxConfig, @@ -492,9 +493,10 @@ export function rotateGenericOAuthAccountOn429( retryAfterHeader: string | null | undefined, now = Date.now(), requestedModelId?: string | null, + commit = true, ): string | null { return rotateGenericOAuthAccountOnRefusal(config, providerName, failedAccountId, "rate", - retryAfterHeader, now, requestedModelId); + retryAfterHeader, now, requestedModelId, undefined, commit); } /** Kiro adds account-scoped refusal classes; all other providers retain their 429 policy. */ @@ -507,6 +509,7 @@ export function rotateGenericOAuthAccountOnRefusal( now = Date.now(), requestedModelId?: string | null, monthlyCooldownMs?: number, + commit = true, ): string | null { if (kind === "other" || (providerName !== "kiro" && kind !== "rate")) return null; if (!isGenericOAuthFailoverEnabled(config, providerName)) return null; @@ -566,8 +569,8 @@ export function rotateGenericOAuthAccountOnRefusal( } if (strategy === "round-robin") { // PICK here, not peek: the failure already happened and this answer is the one being used, - // so the ring genuinely advances. - return pickRoundRobinAccount( + // so the ring genuinely advances. A caller that only asks whether a sibling exists peeks. + return (commit ? pickRoundRobinAccount : peekRoundRobinAccount)( genericPoolKey(providerName), candidates, genericStickyLimit(config, providerName), diff --git a/src/server/responses/core-combo.ts b/src/server/responses/core-combo.ts index f4714f0c9b7..50b5eac6fb5 100644 --- a/src/server/responses/core-combo.ts +++ b/src/server/responses/core-combo.ts @@ -52,14 +52,15 @@ import { import { hasUnreadableEncryptedAgentTask } from "./encrypted-payload"; import { routeConcreteModel, comboRouteDecisionTrace } from "../../router"; import { memoryModelRouteReason } from "./memory-models"; -import { poolAccountProviderLabel } from "../../providers/label"; +import { baseProviderLabel, poolAccountProviderLabel } from "../../providers/label"; import { getAccountSet } from "../../oauth/store"; import { configuredAnthropicInstance } from "../../providers/anthropic-instance"; import { anthropicRoutingFor, formatAnthropicProviderForLog, } from "../../oauth/anthropic-routing"; -import { codexAccountLogLabel } from "../../codex/account-label"; +import { codexAccountLogLabel, oauthAccountLogLabel } from "../../codex/account-label"; +import { isGenericFailoverProvider, rotateGenericOAuthAccountOn429 } from "../../oauth/generic-account-failover"; import { codexQuotaScopeForModel, getCodexQuotaHealthSnapshot } from "../../codex/routing"; import { isCanonicalOpenAiForwardProvider } from "../../providers/openai-tiers"; import type { AgentTaskRecoveryFailureReason } from "./agent-task-recovery"; @@ -135,6 +136,27 @@ function cooledPoolAccountLabel(config: OcxConfig, providerName: string, modelId ? label : undefined; } +/** + * The generic OAuth pool's counterpart of `cooledPoolAccountLabel`. Its 429 names one account + * through `accountLogLabel`, but the per-request account budget can end the pool's own rotation + * before the last refused account is recorded. Hand that account to the pool's rotator, which + * cools it under the pool's own Retry-After policy, and spare the target only when the pool still + * has an eligible sibling. The rotator only peeks round-robin here: nothing is sent to that sibling. + * A lone account or an all-cooled pool returns null and cools the target. + */ +function cooledGenericOAuthAccountLabel( + config: OcxConfig, providerName: string, modelId: string, label: string | undefined, retryAfter: string | undefined, +): string | undefined { + const provider = config.providers[providerName]; + if (!label || !provider || !isGenericFailoverProvider(providerName, provider)) return undefined; + const base = baseProviderLabel(providerName); + const matches = getAccountSet(providerName)?.accounts.filter(account => + oauthAccountLogLabel(account.id, base) === label) ?? []; + return matches.length === 1 + && rotateGenericOAuthAccountOn429(config, providerName, matches[0]!.id, retryAfter, Date.now(), modelId, false) + ? label : undefined; +} + /** * True only for the Anthropic OAuth pool's OWN "all accounts cooled" 429. Its Retry-After is the @@ -1140,7 +1162,9 @@ export async function executeComboResponses( pick.target.provider, pick.target.model, poolAccountProviderLabel(childLog.provider, pick.target.provider), - ); + ) ?? (failure.response.status === 429 && failureCooldownScope === "target" + ? cooledGenericOAuthAccountLabel(config, pick.target.provider, pick.target.model, childLog.accountLogLabel, failure.retryAfter) + : undefined); // The target still cools on the pool's own refusal (the anti-hammer guard), but for the local // fallback rather than the pool's restated per-account Retry-After. const poolLocalRefusal = isAnthropicPoolLocalRefusal( diff --git a/structure/runtime.md b/structure/runtime.md index dc720c31f82..c281555d49d 100644 --- a/structure/runtime.md +++ b/structure/runtime.md @@ -504,7 +504,7 @@ This is also why the classifier cannot duplicate visible output. Native byte str Regression coverage: `tests/responses/responses-forward-prompt-envelope.test.ts`, `tests/routing/router-combo-failover-classification.test.ts`, `tests/routing/routing-policy-fallback.test.ts`, `tests/helpers/combo-context-overflow-cases.ts`, and `tests/server/server-combo-failover-e2e.test.ts`. -`src/combos/failover.ts` uses a 10-minute fallback for a spent account usage window (codes `usage_limit_exceeded`, `usage_limit_reached`, `1308`, or `usage limit reached` / `usage limit has been reached` / `token-plan quota has been exhausted` prose, including HTTP 502) and for provider-scoped credential or billing failure codes such as `invalid_api_key` and `insufficient_quota`. This duration does not change failure classification or cooldown scope; upstream retry/reset signals and configured durations retain precedence. It caps explicit upstream `Retry-After` target cooldowns at 24 hours while reset-derived, configured, and fallback cooldowns remain capped at 10 minutes. `src/combos/resolve.ts` omits a target cooldown for a 429 only when the failed child identifies one account in an Anthropic OAuth or canonical Codex pool and that account has a live pool cooldown. `src/server/responses/core-combo.ts` checks the pool's health record before passing that attribution; unknown accounts, ambiguous labels, API-key and generic OAuth providers, and other failures keep their ordinary target cooldown. The Anthropic OAuth pool's own all-accounts-cooled 429 is the one unattributed exception: `isAnthropicPoolLocalRefusal` in `core-combo.ts` requires the `anthropic` provider with `authMode: "oauth"`, status 429, no identified failed account, no eligible pool account and a live pool retry-after. The target then still cools, but for the local fallback instead of the restated per-account `Retry-After`, so a resumed or newly added account is not ignored until a 24-hour target cooldown expires. An upstream 429 on an API-key provider or on an identified pool account keeps its full `Retry-After`. Regression coverage: `tests/routing/always-on-429-failover.test.ts`. Cooldown entries in `src/combos/failover.ts` retain the failure status. The quota-only snapshot query records target/expiry pairs for unexpired 429/402 cooldowns using only cooldown map lookups. Its config-aware wrapper in `src/combos/resolve.ts` applies the picker's provider availability and cached quota checks before dispatch; request eligibility runs lazily only when exhaustion considers a 503, with throwing checks treated as ineligible. Responses exhaustion uses this evidence under the [failover contract](transports/responses-failover.md); ordinary cooldown selection and timing stay unchanged. +`src/combos/failover.ts` uses a 10-minute fallback for a spent account usage window (codes `usage_limit_exceeded`, `usage_limit_reached`, `1308`, or `usage limit reached` / `usage limit has been reached` / `token-plan quota has been exhausted` prose, including HTTP 502) and for provider-scoped credential or billing failure codes such as `invalid_api_key` and `insufficient_quota`. This duration does not change failure classification or cooldown scope; upstream retry/reset signals and configured durations retain precedence. It caps explicit upstream `Retry-After` target cooldowns at 24 hours while reset-derived, configured, and fallback cooldowns remain capped at 10 minutes. `src/combos/resolve.ts` omits a target cooldown for a 429 only when the failed child identifies one account in an Anthropic OAuth or canonical Codex pool and that account has a live pool cooldown, or identifies one account in a generic OAuth pool that still has an eligible sibling. `src/server/responses/core-combo.ts` checks the pool's health record before passing that attribution; for a generic OAuth pool it first hands the account (matched through its `accountLogLabel`) to `rotateGenericOAuthAccountOn429`, which cools it under the pool's own Retry-After policy even when the per-request account budget ended rotation before that account was recorded, and attributes the failure only when the rotator names a sibling. Unknown accounts, ambiguous labels, a lone or fully cooled generic OAuth pool, API-key providers, and other failures keep their ordinary target cooldown. An API-key pool needs no attribution: its 429 rotation walks every key inside the attempt, so a 429 reaches the combo only once every key is cooled. Regression coverage: `tests/oauth/generic-oauth-combo-account-cooldown.test.ts`. The Anthropic OAuth pool's own all-accounts-cooled 429 is the one unattributed exception: `isAnthropicPoolLocalRefusal` in `core-combo.ts` requires the `anthropic` provider with `authMode: "oauth"`, status 429, no identified failed account, no eligible pool account and a live pool retry-after. The target then still cools, but for the local fallback instead of the restated per-account `Retry-After`, so a resumed or newly added account is not ignored until a 24-hour target cooldown expires. An upstream 429 on an API-key provider or on an identified pool account keeps its full `Retry-After`. Regression coverage: `tests/routing/always-on-429-failover.test.ts`. Cooldown entries in `src/combos/failover.ts` retain the failure status. The quota-only snapshot query records target/expiry pairs for unexpired 429/402 cooldowns using only cooldown map lookups. Its config-aware wrapper in `src/combos/resolve.ts` applies the picker's provider availability and cached quota checks before dispatch; request eligibility runs lazily only when exhaustion considers a 503, with throwing checks treated as ineligible. Responses exhaustion uses this evidence under the [failover contract](transports/responses-failover.md); ordinary cooldown selection and timing stay unchanged. ## Combo default effort precedence diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index f5047535934..432cc842a7a 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -661,7 +661,7 @@ "gemini-inline.test.ts": "images", "gemini-web-search.test.ts": "adapters/google", "generic-session-affinity.test.ts": "oauth", "server-generic-oauth-affinity.test.ts": "server", - "generic-oauth-failover.test.ts": "oauth", "generic-oauth-refusal-composition.test.ts": "oauth", + "generic-oauth-combo-account-cooldown.test.ts": "oauth", "generic-oauth-failover.test.ts": "oauth", "generic-oauth-refusal-composition.test.ts": "oauth", "github-copilot-auto.test.ts": "providers/github-copilot", "github-copilot-auto-runtime.test.ts": "providers/github-copilot", "github-copilot-catalog-hints.test.ts": "providers/github-copilot", diff --git a/tests/oauth/generic-oauth-combo-account-cooldown.test.ts b/tests/oauth/generic-oauth-combo-account-cooldown.test.ts new file mode 100644 index 00000000000..b8ec163c75d --- /dev/null +++ b/tests/oauth/generic-oauth-combo-account-cooldown.test.ts @@ -0,0 +1,149 @@ +/** + * A generic OAuth pool's spent account must not park the whole combo target. + * + * The per-request account budget stops the pool's own rotation before it reaches every stored + * account, so a 429 can reach the combo layer while a sibling is still healthy. Cooling the target + * for that account's `Retry-After` (up to the 24h server-delay ceiling) would answer zero-attempt + * 503s, or pay for a fallback target, while the pool still has an account that answers. + */ +import { afterEach, beforeEach, expect, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { clearComboSelectionState, clearComboTargetCooldowns, isComboTargetInCooldown } from "../../src/combos"; +import { clearGenericFailoverHealth, eligibleFailoverAccounts } from "../../src/oauth/generic-account-failover"; +import { genericPoolKey, peekRoundRobinAccount } from "../../src/oauth/pool-kernel"; +import { getAccountSet, saveCredential, setActiveAccount } from "../../src/oauth/store"; +import { clearAccountQuotaCache } from "../../src/providers/quota"; +import { clearResponseStateForTests, flushResponseState } from "../../src/responses/state"; +import { handleResponses } from "../../src/server/responses"; +import type { OcxConfig } from "../../src/types"; +import { chatSuccess } from "../helpers/combo-failover-upstream"; +import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; + +const TARGET = { provider: "xai", model: "m1" }; +const originalHome = process.env.OPENCODEX_HOME; +const originalFetch = globalThis.fetch; +let home = ""; +let releaseSpendHome: (() => void) | undefined; +/** Account index per upstream send (`B` for the fallback target). */ +let sends: string[] = []; + +async function seed(count: number): Promise { + for (let index = 1; index <= count; index++) { + await saveCredential("xai", { + access: `synthetic-xai-access-${index}`, refresh: `synthetic-xai-refresh-${index}`, + expires: Date.now() + 3_600_000, accountId: `synthetic-xai-${index}`, + } as never, { addAccount: true }); + } + await setActiveAccount("xai", getAccountSet("xai")!.accounts[0]!.id); +} + +/** Accounts at or below `spentThrough` answer a one-hour 429; the rest answer 200. */ +function upstream(spentThrough: number): void { + globalThis.fetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input instanceof Request ? input.url : input); + if (url.startsWith("https://fallback.test/")) { + sends.push("B"); + return chatSuccess("fallback answer", "m2"); + } + const index = Number(new Headers(init?.headers).get("authorization")?.split("-").pop()); + sends.push(String(index)); + return index <= spentThrough + ? Response.json({ error: { message: "synthetic rate limit" } }, { status: 429, headers: { "retry-after": "3600" } }) + : chatSuccess("pool answer", "m1"); + }) as typeof fetch; +} + +const config: OcxConfig = { + port: 0, defaultProvider: "xai", + providers: { + xai: { adapter: "openai-chat", baseUrl: "https://api.x.ai/v1", authMode: "oauth", models: ["m1"] }, + fallback: { adapter: "openai-chat", baseUrl: "https://fallback.test/v1", authMode: "key", apiKey: "synthetic-fallback-key", models: ["m2"] }, + }, + combos: { pooled: { strategy: "failover", targets: [TARGET, { provider: "fallback", model: "m2" }] } }, +}; + +async function post(): Promise { + const response = await handleResponses(new Request("http://localhost/v1/responses", { + method: "POST", headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "combo/pooled", input: "Answer briefly", stream: false }), + }), config, { model: "", provider: "" }); + expect(response.status).toBe(200); + return response.text(); +} + +beforeEach(() => { + home = mkdtempSync(join(tmpdir(), "ocx-generic-combo-cooldown-")); + process.env.OPENCODEX_HOME = home; + sends = []; + clearComboSelectionState(); + clearComboTargetCooldowns(); + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + clearResponseStateForTests(); + releaseSpendHome = acquireOwnedSpendHome(); +}); + +afterEach(async () => { + await flushResponseState(); + releaseSpendHome?.(); + releaseSpendHome = undefined; + globalThis.fetch = originalFetch; + clearComboSelectionState(); + clearComboTargetCooldowns(); + clearGenericFailoverHealth(); + clearAccountQuotaCache(); + clearResponseStateForTests(); + if (originalHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = originalHome; + if (home) removeTreeWithRetry(home); +}); + +test("a spent generic OAuth account cools the account, not the combo target", async () => { + await seed(5); + upstream(4); + + expect(await post()).toContain("fallback answer"); + // Precondition: the request budget ended the pool's rotation before the healthy fifth account. + expect(sends).not.toContain("5"); + expect(sends.at(-1)).toBe("B"); + expect(isComboTargetInCooldown("pooled", TARGET)).toBe(false); + // Every account this request reached is cooled by the pool, the last one included. + expect(eligibleFailoverAccounts("xai")).toEqual([getAccountSet("xai")!.accounts[4]!.id]); + + // The pool, not the fallback target, serves the next request. Without proactive preference the + // pool may still open on its cooled active account once before rotating; that is the pool's own + // contract and stays bounded by its rotation. + sends = []; + expect(await post()).toContain("pool answer"); + expect(sends).not.toContain("B"); + expect(sends.at(-1)).toBe("5"); +}, 20_000); + +test("every generic OAuth account spent still cools the combo target", async () => { + await seed(2); + upstream(2); + + expect(await post()).toContain("fallback answer"); + expect(sends).toEqual(["1", "2", "B"]); + expect(isComboTargetInCooldown("pooled", TARGET)).toBe(true); +}, 20_000); + +test("the combo's sibling check does not advance a round-robin pool onto an unsent account", async () => { + await seed(5); + upstream(4); + const roundRobin: OcxConfig = { ...config, pool: { kernel: true }, providers: { ...config.providers, + xai: { ...config.providers.xai!, oauthAccountFailover: { strategy: "round-robin", stickyLimit: 3 } } } }; + const ids = getAccountSet("xai")!.accounts.map(account => account.id); + const response = await handleResponses(new Request("http://localhost/v1/responses", { + method: "POST", headers: { "content-type": "application/json" }, + body: JSON.stringify({ model: "combo/pooled", input: "Answer briefly", stream: false }), + }), roundRobin, { model: "", provider: "" }); + expect(await response.text()).toContain("fallback answer"); + expect(sends).not.toContain("5"); + expect(isComboTargetInCooldown("pooled", TARGET)).toBe(false); + // Account 5 was only probed for eligibility, never sent; the ring must not stick to it. + expect(peekRoundRobinAccount(genericPoolKey("xai"), ids, 3)).not.toBe(ids[4]); +}, 20_000); diff --git a/tests/oauth/generic-oauth-failover.test.ts b/tests/oauth/generic-oauth-failover.test.ts index bb4150200f5..63a253f5c60 100644 --- a/tests/oauth/generic-oauth-failover.test.ts +++ b/tests/oauth/generic-oauth-failover.test.ts @@ -559,8 +559,11 @@ describe("sidecar on429 wiring", () => { // The counts differ by rotator because the recovery sites differ, and each number is a // statement about which providers can recover where: // - // generic = 5: streaming loop, continuation loop, sidecar hook, runTurn preflight, - // native Responses passthrough. The Antigravity 403 verify-account + // generic = 6: streaming loop, continuation loop, sidecar hook, runTurn preflight, + // native Responses passthrough, plus the combo failure attribution in + // core-combo.ts, which sends nothing: it records the refused account the + // request budget left uncooled and spares the combo target only when the + // rotator names a sibling. The Antigravity 403 verify-account // arm deliberately does NOT use this rotator: a verification refusal // must not record rate-limit cooldown semantics, so it moves via // rotateAntigravityAccountOnAuthRefusal instead. The new default @@ -581,7 +584,7 @@ describe("sidecar on429 wiring", () => { // // Adding a recovery site means deciding, deliberately, which rotators it needs and // updating the matching number. That decision is the thing this test exists to force. - expect(counts.generic).toBe(5); + expect(counts.generic).toBe(6); expect(counts.anthropic).toBe(3); expect(counts.key).toBe(7); });