From c2ba04a85357b8b3578733e7575c7dd69c73eb6a Mon Sep 17 00:00:00 2001 From: t Date: Sat, 5 Sep 2026 09:45:05 +0900 Subject: [PATCH 1/2] fix(oauth): repair proactive-failover policy boundaries Carries the policy half of #3502 onto current dev, rebased, with docs rewritten on top of #3520 rather than replayed. 1. src/oauth/anthropic-routing.ts consulted the pool's proactive strategy even when the pool is disabled, so a disabled pool silently reactivated round-robin/fill-first on the reactive 429 path. Reactive recovery now uses the neutral quota picker there. 2. src/oauth/generic-account-failover.ts honoured only enabled === false per provider, so a provider-specific true could not opt back in when the global default is false. The published narrow-over-broad precedence now applies in both directions, and the typeof guard still lets a malformed value fall through rather than taking a provider out of service. The Kiro continuation half of #3502 is split into the next PR in the stack. Verification: - bun test tests/routing/always-on-429-failover.test.ts tests/oauth/generic-oauth-failover.test.ts tests/adapters/anthropic/anthropic-sidecar-account-failover.test.ts tests/oauth/adapter-event-oauth-failover.test.ts -> 43 pass / 0 fail (both new assertions RED before the source hunks) - bun test tests/test-layout.test.ts tests/test-layout-tooling.test.ts -> 17 pass / 0 fail - bun run typecheck -> exit 0 Co-authored-by: Ingwannu <186453546+Ingwannu@users.noreply.github.com> --- .../docs/reference/configuration/providers.md | 2 +- scripts/test-layout/layout.json | 1 + src/oauth/anthropic-routing.ts | 8 +- src/oauth/generic-account-failover.ts | 15 +- src/types/config.ts | 2 +- src/types/provider.ts | 2 +- structure/04_transports-and-sidecars.md | 9 ++ ...anthropic-sidecar-account-failover.test.ts | 149 ++++++++++++++++++ tests/fixtures/test-layout-expected.json | 1 + .../adapter-event-oauth-failover.test.ts | 5 +- tests/oauth/generic-oauth-failover.test.ts | 16 +- tests/routing/always-on-429-failover.test.ts | 20 ++- 12 files changed, 218 insertions(+), 12 deletions(-) create mode 100644 tests/adapters/anthropic/anthropic-sidecar-account-failover.test.ts diff --git a/docs-site/src/content/docs/reference/configuration/providers.md b/docs-site/src/content/docs/reference/configuration/providers.md index 380c1fb498..a4a3a7b876 100644 --- a/docs-site/src/content/docs/reference/configuration/providers.md +++ b/docs-site/src/content/docs/reference/configuration/providers.md @@ -426,7 +426,7 @@ second account. | Key | Type | Default | Description | | --- | --- | --- | --- | | `oauthAccountFailover.enabled?` | `boolean` | presence-driven | Global override for the **pre-dispatch account preference** only. `false` stops a healthy request being steered toward the account with more known headroom. It does **not** disable 429 rotation. | -| `providers..oauthAccountFailover.enabled?` | `boolean` | inherits | Per-provider override for the same preference; beats the global setting. Only `false` is meaningful — `true` adds nothing over account presence. | +| `providers..oauthAccountFailover.enabled?` | `boolean` | inherits | Per-provider override for the same preference; beats the global setting in either direction. `false` declines the preference for this provider even when the global setting is `true`, and `true` opts this provider in even when the global setting is `false`. Reactive 429 rotation is unaffected either way. | | `providers..oauthAccountFailover.strategy?` | `"quota" \| "round-robin" \| "fill-first"` | — | Declared pool strategy for a generic OAuth provider (#695). Persisted through `ocx account strategy ` or `PUT /api/oauth/accounts/pool`; the generic selector does not act on it yet, so omitted and set behave the same today. | | `providers..oauthAccountFailover.autoSwitchThreshold?` | `number` | — | Declared 0–100 usage percent for a proactive switch on a generic OAuth provider (#695). Set with `ocx account auto-switch threshold `; inert until the selector consumes it. | diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index 0487e8d878..e31d5fe69d 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -203,6 +203,7 @@ "anthropic-pool-toggle-copy.test.ts": "adapters/anthropic", "anthropic-quorum-cache.test.ts": "routing", "anthropic-reasoning.test.ts": "adapters/anthropic", + "anthropic-sidecar-account-failover.test.ts": "adapters/anthropic", "anthropic-stream-hardening.test.ts": "adapters/anthropic", "anthropic-tail-guard.test.ts": "adapters/anthropic", "anthropic-thinking-signature.test.ts": "adapters/anthropic", diff --git a/src/oauth/anthropic-routing.ts b/src/oauth/anthropic-routing.ts index 47206335ae..95197ef1f1 100644 --- a/src/oauth/anthropic-routing.ts +++ b/src/oauth/anthropic-routing.ts @@ -672,7 +672,13 @@ export function rotateAnthropicAccountOn429( // from a count read taken before the failure. quorumCache = null; - const next = pickAlternateAnthropicAccount(config, failedAccountId, now); + // The pool's strategy is a PROACTIVE policy. When the pool is disabled, reactive + // presence-only recovery must not silently reactivate round-robin/fill-first merely + // because those dormant values remain in config. The quota picker is the neutral + // recovery policy already used by the default strategy. + const next = isAnthropicAccountPoolEnabled(config) + ? pickAlternateAnthropicAccount(config, failedAccountId, now) + : pickLowestUsage(config, failedAccountId, now); if (!next) { console.warn("[anthropic-pool] all eligible Anthropic OAuth accounts are in cooldown; returning 429"); return null; diff --git a/src/oauth/generic-account-failover.ts b/src/oauth/generic-account-failover.ts index 9a7a2566f9..321b5f92a1 100644 --- a/src/oauth/generic-account-failover.ts +++ b/src/oauth/generic-account-failover.ts @@ -178,15 +178,20 @@ export function isGenericOAuthFailoverEnabled( * Whether the pre-dispatch account PREFERENCE may run for this provider. * * Unlike reactive rotation, this moves a request that upstream has not refused, so it stays - * refusable: an explicit `false` — per provider first, then global — turns it off. `true` adds - * nothing over presence, so only `false` is honoured; that keeps the predicate identical to the - * old behaviour for every operator who never wrote the key, and a malformed value falls through - * rather than taking a provider out of service. + * refusable: an explicit provider value wins over the global default, and a global `false` + * turns it off only when the provider has no override. A malformed value falls through rather + * than taking a provider out of service. */ function isProactivePreferenceEnabled(config: OcxConfig, providerName: string, now: number): boolean { const provider = config.providers?.[providerName]; if (!provider || !isGenericFailoverProvider(providerName, provider)) return false; - if (provider.oauthAccountFailover?.enabled === false) return false; + const perProvider = provider.oauthAccountFailover?.enabled; + // Preserve the published narrow-over-broad precedence. A provider-specific true may + // opt this provider into proactive preference even when the global default is false; + // a provider-specific false refuses it even when the global setting is true. + if (typeof perProvider === "boolean") { + return perProvider && hasFailoverAccountQuorum(providerName, now); + } if (config.oauthAccountFailover?.enabled === false) return false; return hasFailoverAccountQuorum(providerName, now); } diff --git a/src/types/config.ts b/src/types/config.ts index fdbec5ba25..114a743980 100644 --- a/src/types/config.ts +++ b/src/types/config.ts @@ -791,7 +791,7 @@ export interface OcxConfig { * What `enabled: false` still refuses is the PRE-DISPATCH preference: steering a request * upstream has not refused toward the account with more known headroom. That moves a healthy * request, so it stays a real choice. `providers..oauthAccountFailover` overrides this - * per provider, and only `false` is meaningful — `true` adds nothing over presence. + * per provider in either direction; reactive 429 rotation remains presence-driven. */ oauthAccountFailover?: { enabled?: boolean; diff --git a/src/types/provider.ts b/src/types/provider.ts index 459634fd16..46082c9620 100644 --- a/src/types/provider.ts +++ b/src/types/provider.ts @@ -430,7 +430,7 @@ export interface OcxProviderConfig { * activate it, and a 429 with an idle second account is a defect rather than a preference. * What an explicit `false` still refuses is the pre-dispatch preference that steers a HEALTHY * request toward the account with more known headroom. It beats the global - * `oauthAccountFailover`; only `false` is meaningful, since `true` adds nothing over presence. + * `oauthAccountFailover` in either direction; reactive 429 rotation remains presence-driven. */ oauthAccountFailover?: { enabled?: boolean; diff --git a/structure/04_transports-and-sidecars.md b/structure/04_transports-and-sidecars.md index eb6cbc82a5..e1068b68ae 100644 --- a/structure/04_transports-and-sidecars.md +++ b/structure/04_transports-and-sidecars.md @@ -1412,9 +1412,18 @@ surface is listed here so a maintainer can find the owner without grepping: | Image/video generation loop | `src/images/loop.ts`, `src/images/plan.ts`, `src/images/fulfill.ts`, `src/images/xai-client.ts`, `src/images/xai-video-client.ts`, `src/images/artifacts.ts` | A provider-returned image URL is downloaded into a local artifact once, then served locally; warnings stay URL-free because provider CDN URLs may embed credentials. | | GitHub Copilot | `src/providers/xai-transport.ts` (`resolveProviderTransport`), `src/providers/github-copilot-transport.ts` | `resolveProviderTransport` selects the Copilot transport when the routed provider name is `github-copilot`; the Copilot module then resolves its headers and base URL, and the registry seeds the provider row and model fallback. | | API-key pools | `src/providers/key-failover.ts` | A 429 rotates the active key and records a cooldown; `provider.apiKey` keeps mirroring the active entry so routing stays single-key. | +| OAuth account failover | `src/oauth/generic-account-failover.ts`, `src/oauth/anthropic-routing.ts` | Reactive pre-output 429 recovery is presence-driven with 2+ eligible accounts. Pool and `oauthAccountFailover` flags govern proactive routing, not the reactive retry: a disabled Anthropic pool recovers through quota ordering rather than its dormant strategy, and a per-provider `enabled` beats the global default in either direction. | | Alibaba regions | `src/providers/alibaba-region-backup.ts`, `src/providers/alibaba-region-migration.ts`, `src/providers/alibaba-region-startup.ts` | Region migration backs up before rewriting and is idempotent across restarts. | | Discovery and quota | `src/providers/model-discovery.ts`, `src/providers/quota.ts` | Discovery rejects a response over 4 MiB or past 2,000 raw rows before caching it. | +[Decision Log] +- 목적과 의도: Keep reactive OAuth 429 recovery available without silently enabling proactive account-routing policy the operator switched off. +- 기존 구현 및 제약 조건: #3495 made reactive recovery presence-driven, but a disabled Anthropic pool still consulted its dormant strategy on the reactive path, and a per-provider `oauthAccountFailover.enabled: true` could no longer beat a global `false`. +- 검토한 주요 대안: Restore the old all-or-nothing enable flag; leave the merged behavior and document the gaps; or keep the reactive/proactive split and repair the exact policy boundaries. +- 선택한 방식: Keep presence-driven reactive recovery, apply proactive precedence only before dispatch, and use quota ordering for disabled-pool Anthropic recovery. +- 다른 대안 대신 이 방식을 선택한 이유: This preserves the merged product decision without letting disabled proactive settings influence a retry, and it restores the published narrow-over-broad precedence in both directions. +- 장점, 단점 및 영향: 429 recovery stays automatic for operators with multiple eligible accounts; operators who require no automatic account switch must keep one eligible account, which the GUI and public docs state explicitly. + ## Sidecars Web search and vision sidecars run only when the main request needs that capability and a usable diff --git a/tests/adapters/anthropic/anthropic-sidecar-account-failover.test.ts b/tests/adapters/anthropic/anthropic-sidecar-account-failover.test.ts new file mode 100644 index 0000000000..31b2389fb8 --- /dev/null +++ b/tests/adapters/anthropic/anthropic-sidecar-account-failover.test.ts @@ -0,0 +1,149 @@ +/** + * The web-search sidecar is a rotation site of its own, and it reaches Anthropic through the + * shared 429 hook rather than the main response loop. A pool that is switched off must still + * recover there: `anthropicAccountPool.enabled: false` declines PROACTIVE routing, not the + * reactive retry that runs only after upstream has already refused the request. + */ +import { afterAll, afterEach, beforeAll, beforeEach, expect, mock, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { ProviderAdapter } from "../../../src/adapters/base"; +import { clearAnthropicAccountPoolState } from "../../../src/oauth/anthropic-routing"; +import { clearGenericFailoverHealth } from "../../../src/oauth/generic-account-failover"; +import { getAccountSet, saveCredential, setActiveAccount } from "../../../src/oauth/store"; +import type { OcxConfig, OcxParsedRequest, OcxProviderConfig } from "../../../src/types"; +import { removeTreeWithRetry } from "../../helpers/remove-tree"; + +const previousHome = process.env.OPENCODEX_HOME; +let testHome = ""; +let handleResponses: typeof import("../../../src/server/responses")["handleResponses"]; +let observedKeys: string[] = []; +let sidecarMode = false; + +function fixtureAdapter(provider: OcxProviderConfig): ProviderAdapter { + return { + name: "anthropic", + buildRequest() { + return { + url: provider.baseUrl, + method: "POST", + headers: { authorization: `Bearer ${provider.apiKey ?? ""}` }, + body: "{}", + }; + }, + async *parseStream() { + yield { type: "done" as const }; + }, + }; +} + +beforeAll(async () => { + const actualResolver = await import("../../../src/server/adapter-resolve"); + const actualResolveAdapter = actualResolver.resolveAdapter; + mock.module("../../../src/server/adapter-resolve", () => ({ + ...actualResolver, + resolveAdapter(provider: OcxProviderConfig, cacheRetention?: "none" | "short" | "long") { + if (provider.adapter === "test-anthropic-sidecar") return fixtureAdapter(provider); + return actualResolveAdapter(provider, cacheRetention); + }, + })); + + mock.module("../../../src/web-search", () => ({ + buildWebSearchTool: () => ({ + name: "web_search", + parameters: { type: "object", properties: {} }, + }), + planWebSearch: () => sidecarMode + ? { + backend: "anthropic", + hostedTool: { type: "web_search" }, + settings: { model: "claude-haiku-4-5", reasoning: "low", timeoutMs: 1_000 }, + maxSearches: 1, + } + : undefined, + shouldResolveOpenAiWebSearchSidecar: () => false, + runWithWebSearch: async (args: { + parsed: OcxParsedRequest; + adapter: ProviderAdapter; + on429?: (retryAfter: string | null) => Promise; + }) => { + const first = await args.adapter.buildRequest(args.parsed); + observedKeys.push(new Headers(first.headers).get("authorization") ?? ""); + const rotated = await args.on429?.("30"); + if (!rotated) throw new Error("Anthropic sidecar did not rotate after 429"); + const second = await rotated.buildRequest(args.parsed); + observedKeys.push(new Headers(second.headers).get("authorization") ?? ""); + return new Response("sidecar-ok", { status: 200 }); + }, + })); + + ({ handleResponses } = await import("../../../src/server/responses")); +}); + +beforeEach(() => { + testHome = mkdtempSync(join(tmpdir(), "ocx-oauth-429-boundaries-")); + process.env.OPENCODEX_HOME = testHome; + observedKeys = []; + sidecarMode = false; + clearAnthropicAccountPoolState(); + clearGenericFailoverHealth(); +}); + +afterEach(() => { + clearAnthropicAccountPoolState(); + clearGenericFailoverHealth(); + removeTreeWithRetry(testHome); +}); + +afterAll(() => { + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + mock.restore(); +}); + +test("Anthropic web-search sidecar rotates on 429 when proactive pooling is disabled", async () => { + sidecarMode = true; + for (let index = 0; index < 2; index += 1) { + await saveCredential("anthropic", { + access: `anthropic-access-${index}`, + refresh: `anthropic-refresh-${index}`, + expires: Date.now() + 3_600_000, + accountId: `anthropic-account-${index}`, + } as never, { addAccount: true }); + } + const ids = getAccountSet("anthropic")!.accounts.map(account => account.id); + await setActiveAccount("anthropic", ids[0]!); + + const config = { + port: 0, + defaultProvider: "anthropic", + anthropicAccountPool: { enabled: false, strategy: "round-robin" }, + providers: { + anthropic: { + adapter: "test-anthropic-sidecar", + baseUrl: "https://anthropic-sidecar.test/v1", + authMode: "oauth", + models: ["model"], + }, + }, + } as unknown as OcxConfig; + + const response = await handleResponses(new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "anthropic/model", + input: "search", + stream: true, + tools: [{ type: "web_search" }], + }), + }), config, { model: "", provider: "" }); + + expect(response.status).toBe(200); + expect(await response.text()).toBe("sidecar-ok"); + expect(observedKeys).toEqual([ + "Bearer anthropic-access-0", + "Bearer anthropic-access-1", + ]); +}); diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index 42e3b74834..c1848c7d44 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -40,6 +40,7 @@ "anthropic-pool-toggle-copy.test.ts": "adapters/anthropic", "anthropic-quorum-cache.test.ts": "routing", "anthropic-reasoning.test.ts": "adapters/anthropic", + "anthropic-sidecar-account-failover.test.ts": "adapters/anthropic", "anthropic-stream-hardening.test.ts": "adapters/anthropic", "anthropic-tail-guard.test.ts": "adapters/anthropic", "anthropic-thinking-signature.test.ts": "adapters/anthropic", diff --git a/tests/oauth/adapter-event-oauth-failover.test.ts b/tests/oauth/adapter-event-oauth-failover.test.ts index 8d5b4dc17d..8309d9ffd3 100644 --- a/tests/oauth/adapter-event-oauth-failover.test.ts +++ b/tests/oauth/adapter-event-oauth-failover.test.ts @@ -137,9 +137,12 @@ describe("#2568 adapter-event OAuth failover", () => { [{ type: "text", text: "ok" }], ]; - const body = await (await handleResponses(request(true), config(false), { model: "", provider: "" })).text(); + const response = await handleResponses(request(true), config(false), { model: "", provider: "" }); + const body = await response.text(); + expect(response.status).toBe(200); expect(attemptKeys).toEqual(["cursor-access-1", "cursor-access-0"]); + expect(body).toContain("ok"); expect(body).not.toContain("rate_limit_exceeded"); }); diff --git a/tests/oauth/generic-oauth-failover.test.ts b/tests/oauth/generic-oauth-failover.test.ts index 0bb219d07e..76fe18fd1a 100644 --- a/tests/oauth/generic-oauth-failover.test.ts +++ b/tests/oauth/generic-oauth-failover.test.ts @@ -12,7 +12,7 @@ import { preferredInitialAccount, rotateGenericOAuthAccountOn429, } from "../../src/oauth/generic-account-failover"; -import { getAccountSet, markAccountNeedsReauth, saveCredential } from "../../src/oauth/store"; +import { getAccountSet, markAccountNeedsReauth, saveCredential, setActiveAccount } from "../../src/oauth/store"; import { clearAccountQuotaCache, setCachedProviderAccountQuotaForTests } from "../../src/providers/quota"; import { resolveCopilotApiBaseUrl } from "../../src/oauth/github-copilot"; import { resolveProviderTransport } from "../../src/providers/xai-transport"; @@ -137,6 +137,20 @@ describe("#2568 generic OAuth account failover", () => { expect(preferredInitialAccount(config(true, false), "xai")).toBeNull(); }); + test("a provider-level true overrides a global proactive opt-out", async () => { + // The documented precedence is narrow-over-broad in BOTH directions. A per-provider false + // refuses the preference under a global true (pinned above); the mirror case is an operator + // who declines steering globally and opts one provider back in. Honouring only `false` + // silently drops that opt-in. + const ids = await seed(2); + await setActiveAccount("xai", ids[0]!); + clearGenericFailoverHealth("xai"); + setCachedProviderAccountQuotaForTests("xai", ids[0]!, { fiveHourPercent: 99 }); + setCachedProviderAccountQuotaForTests("xai", ids[1]!, { fiveHourPercent: 1 }); + + expect(preferredInitialAccount(config(false, true), "xai")).toBe(ids[1]); + }); + test("a second account flagged for reauth is not a quorum", async () => { // A revoked account cannot serve the replay, so counting it would arm the failover machinery // for a user who still has exactly one usable credential. diff --git a/tests/routing/always-on-429-failover.test.ts b/tests/routing/always-on-429-failover.test.ts index b4d825983d..2068c7a58a 100644 --- a/tests/routing/always-on-429-failover.test.ts +++ b/tests/routing/always-on-429-failover.test.ts @@ -24,7 +24,7 @@ import { } from "../../src/oauth/anthropic-routing"; import { clearPoolRotationState } from "../../src/codex/pool-rotation"; import { getAccountSet, saveCredential, setActiveAccount } from "../../src/oauth/store"; -import { clearAccountQuotaCache } from "../../src/providers/quota"; +import { clearAccountQuotaCache, setCachedProviderAccountQuotaForTests } from "../../src/providers/quota"; import type { OcxConfig } from "../../src/types"; import { removeTreeWithRetry } from "../helpers/remove-tree"; @@ -99,6 +99,24 @@ describe("Anthropic reactive 429 failover without the pool flag", () => { expect(rotateAnthropicAccountOn429(poolDisabled(), ids[0]!, null)).toBe(ids[1]); }); + test("a disabled pool does not apply its dormant proactive strategy to reactive recovery", async () => { + // The pool flag buys PROACTIVE routing: affinity, quota ranking, and the declared strategy. + // Leaving `strategy: "round-robin"` in a config whose pool is off is not an opt-in to + // round-robin -- it is dormant configuration. Reactive recovery must therefore fall back to + // the neutral quota picker rather than reactivating the strategy the operator switched off. + const ids = await seedAccounts(3); + setCachedProviderAccountQuotaForTests("anthropic", ids[1]!, { fiveHourPercent: 90 }); + setCachedProviderAccountQuotaForTests("anthropic", ids[2]!, { fiveHourPercent: 10 }); + const disabledRoundRobin = { + ...poolAbsent(), + anthropicAccountPool: { enabled: false, strategy: "round-robin" }, + } as OcxConfig; + + // Round-robin would hand back ids[1] (the next account in order); quota ordering picks the + // account with the most headroom instead. + expect(rotateAnthropicAccountOn429(disabledRoundRobin, ids[0]!, null)).toBe(ids[2]); + }); + test("a single account is still a strict no-op", async () => { // Rotating to itself would replay the same 429 on the same credential, and cooling the only // account would take the provider out of service for nothing. From 49c48662f5f3bad49f5034d544e168ca41427e0c Mon Sep 17 00:00:00 2001 From: t Date: Sat, 5 Sep 2026 09:48:04 +0900 Subject: [PATCH 2/2] fix(responses): carry rotated Kiro auth context into the terminal continuation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second half of the #3502 split. The bounded terminal-guard continuation dispatches a shallow clone of the parsed request, so a 429 rotation that wrote _kiroAuthContext onto the outer request only left the clone carrying the FAILED account's region and profile ARN — a rotated bearer paired with an old identity, which is the mixed-identity failure applyFailoverSnapshot exists to prevent. applyFailoverSnapshot now takes the request being retried as a defaulted second parameter and synchronizes both owners, so every other call site is unchanged. The continuation call site passes nextParsed. Verification: - bun test tests/providers/kiro/kiro-auth-context-continuation.test.ts tests/oauth/generic-oauth-failover.test.ts tests/routing/always-on-429-failover.test.ts tests/test-layout.test.ts tests/test-layout-tooling.test.ts -> 53 pass / 0 fail (RED without the core.ts hunk: the third build carried kiro-access-1 with account 0's us-east-1 profile) - bun run typecheck -> exit 0 Co-authored-by: Ingwannu <186453546+Ingwannu@users.noreply.github.com> --- scripts/test-layout/layout.json | 1 + src/server/responses/core.ts | 16 +- tests/fixtures/test-layout-expected.json | 1 + tests/oauth/generic-oauth-failover.test.ts | 6 +- .../kiro-auth-context-continuation.test.ts | 187 ++++++++++++++++++ 5 files changed, 207 insertions(+), 4 deletions(-) create mode 100644 tests/providers/kiro/kiro-auth-context-continuation.test.ts diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index e31d5fe69d..0d04e08ac5 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -684,6 +684,7 @@ "kimi-oauth-identity.test.ts": "providers", "kiro-account-quota.test.ts": "providers/kiro", "kiro-adapter.test.ts": "providers/kiro", + "kiro-auth-context-continuation.test.ts": "providers/kiro", "kiro-builder-id-profile.test.ts": "providers/kiro", "kiro-calibration.test.ts": "providers/kiro", "kiro-images.test.ts": "providers/kiro", diff --git a/src/server/responses/core.ts b/src/server/responses/core.ts index 0c67a834b1..beb88fe7f2 100644 --- a/src/server/responses/core.ts +++ b/src/server/responses/core.ts @@ -3364,7 +3364,10 @@ async function handleResponsesInner( * tolerates project discovery failing, so a stored account can legitimately have no project; * sending that account's bearer with the FAILED account's project is worse than not rotating. */ - const applyFailoverSnapshot = (snapshot: OAuthAccessSnapshot): boolean => { + const applyFailoverSnapshot = ( + snapshot: OAuthAccessSnapshot, + retryParsed: OcxParsedRequest = parsed, + ): boolean => { if (route.provider.googleMode === "cloud-code-assist" && !snapshot.projectId) return false; let rotatedProvider: OcxProviderConfig = { ...route.provider, apiKey: snapshot.accessToken }; if (route.providerName === "github-copilot") { @@ -3377,7 +3380,14 @@ async function handleResponsesInner( } if (snapshot.projectId) rotatedProvider = { ...rotatedProvider, project: snapshot.projectId }; route.provider = rotatedProvider; - if (route.providerName === "kiro") parsed._kiroAuthContext = { ...(snapshot.kiro ?? {}) }; + if (route.providerName === "kiro") { + const kiroContext = { ...(snapshot.kiro ?? {}) }; + // Terminal-guard continuations are rebuilt from a shallow clone. Updating only the + // outer request pairs the new bearer with the failed account's region/profile on + // the retry. Keep both owners synchronized; for ordinary paths they are identical. + parsed._kiroAuthContext = kiroContext; + if (retryParsed !== parsed) retryParsed._kiroAuthContext = { ...kiroContext }; + } // Re-stamp: a request that rotated accounts must be attributed to the account that actually // served it. All three rotation sites funnel through here, so this is the only re-stamp // needed -- and putting it anywhere else would let one of the three drift. @@ -6686,7 +6696,7 @@ async function handleResponsesInner( const snapshot = await failoverAccountSnapshot(route.providerName, nextAccountId); genericFailoverAccountId = nextAccountId; genericFailovers += 1; - if (applyFailoverSnapshot(snapshot)) { + if (applyFailoverSnapshot(snapshot, nextParsed)) { invalidateSameTargetRequest(); activeAdapter = resolveAdapter( resolveWireProtocolOverride(route.providerName, route.modelId, route.provider, inboundWire), diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index c1848c7d44..ef52b3720f 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -521,6 +521,7 @@ "kimi-oauth-identity.test.ts": "providers", "kiro-account-quota.test.ts": "providers/kiro", "kiro-adapter.test.ts": "providers/kiro", + "kiro-auth-context-continuation.test.ts": "providers/kiro", "kiro-builder-id-profile.test.ts": "providers/kiro", "kiro-calibration.test.ts": "providers/kiro", "kiro-images.test.ts": "providers/kiro", diff --git a/tests/oauth/generic-oauth-failover.test.ts b/tests/oauth/generic-oauth-failover.test.ts index 76fe18fd1a..e5bd26e014 100644 --- a/tests/oauth/generic-oauth-failover.test.ts +++ b/tests/oauth/generic-oauth-failover.test.ts @@ -303,7 +303,7 @@ describe("sidecar on429 wiring", () => { // Kiro's routing metadata) live in exactly one place. A fourth rotation site that swaps the // bearer by hand would reintroduce the mixed-identity bug this helper exists to prevent. const snapshotUses = coreSource.match(/failoverAccountSnapshot\(/g) ?? []; - const helperUses = coreSource.match(/applyFailoverSnapshot\(snapshot\)/g) ?? []; + const helperUses = coreSource.match(/applyFailoverSnapshot\(snapshot(?:, nextParsed)?\)/g) ?? []; // Four since the continuation loop gained its own generic-OAuth arm: the streaming loop grew // one with #2568 and the continuation loop did not, so an xAI/Cursor continuation 429 stayed // terminal. Bumping this count is the deliberate act of admitting a fourth rotation site -- @@ -375,6 +375,10 @@ describe("sidecar on429 wiring", () => { // Kiro routing metadata still travels with its own token. expect(body).toContain("_kiroAuthContext"); + // ...and reaches the object actually retried. The terminal-guard continuation dispatches a + // shallow clone, so writing only the outer request pairs the rotated bearer with the failed + // account's region/profile. + expect(coreSource).toContain("applyFailoverSnapshot(snapshot, nextParsed)"); }); test("pre-dispatch selection replaces the CCA project instead of inheriting one", () => { diff --git a/tests/providers/kiro/kiro-auth-context-continuation.test.ts b/tests/providers/kiro/kiro-auth-context-continuation.test.ts new file mode 100644 index 0000000000..83e22d0b48 --- /dev/null +++ b/tests/providers/kiro/kiro-auth-context-continuation.test.ts @@ -0,0 +1,187 @@ +/** + * A Kiro bearer is useless without the region and profile ARN that were issued with it. The + * bounded terminal continuation dispatches a SHALLOW CLONE of the parsed request, so a rotation + * that writes `_kiroAuthContext` onto the outer request only leaves the clone carrying the + * FAILED account's routing metadata — a new token paired with an old identity, which is the + * exact mixed-identity failure the shared snapshot helper exists to prevent. + */ +import { afterAll, afterEach, beforeAll, beforeEach, expect, mock, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { ProviderAdapter } from "../../../src/adapters/base"; +import { clearAnthropicAccountPoolState } from "../../../src/oauth/anthropic-routing"; +import { clearGenericFailoverHealth } from "../../../src/oauth/generic-account-failover"; +import { getAccountSet, saveCredential, setActiveAccount } from "../../../src/oauth/store"; +import type { AdapterEvent, OcxConfig, OcxParsedRequest, OcxProviderConfig } from "../../../src/types"; +import { removeTreeWithRetry } from "../../helpers/remove-tree"; + +const previousHome = process.env.OPENCODEX_HOME; +let testHome = ""; +let handleResponses: typeof import("../../../src/server/responses")["handleResponses"]; +let kiroBuilds: Array<{ key: string; profileArn?: string; apiRegion?: string }> = []; + +function kiroContinuationEvents(phase: string): AdapterEvent[] { + if (phase === "plan") { + return [ + { type: "text_delta", text: "I will modify the file now." }, + { type: "done", stopReason: "end_turn" }, + ]; + } + if (phase === "complete") { + return [ + { type: "tool_call_start", id: "call_read", name: "read_file" }, + { type: "tool_call_delta", arguments: "{}" }, + { type: "tool_call_end" }, + { type: "done", stopReason: "tool_use" }, + ]; + } + throw new Error(`unexpected phase: ${phase}`); +} + +function kiroFixtureAdapter(provider: OcxProviderConfig): ProviderAdapter { + return { + // Anthropic enables the bounded terminal continuation, while the provider id remains + // Kiro so generic OAuth snapshot pairing is exercised. + name: "anthropic", + buildRequest(parsed: OcxParsedRequest) { + kiroBuilds.push({ + key: provider.apiKey ?? "", + ...(parsed._kiroAuthContext?.profileArn + ? { profileArn: parsed._kiroAuthContext.profileArn } + : {}), + ...(parsed._kiroAuthContext?.apiRegion + ? { apiRegion: parsed._kiroAuthContext.apiRegion } + : {}), + }); + return { + url: provider.baseUrl, + method: "POST", + headers: { authorization: `Bearer ${provider.apiKey ?? ""}` }, + body: "{}", + }; + }, + async *parseStream(response: Response): AsyncGenerator { + yield* kiroContinuationEvents(response.headers.get("x-test-phase") ?? ""); + }, + }; +} + +beforeAll(async () => { + const actualResolver = await import("../../../src/server/adapter-resolve"); + const actualResolveAdapter = actualResolver.resolveAdapter; + mock.module("../../../src/server/adapter-resolve", () => ({ + ...actualResolver, + resolveAdapter(provider: OcxProviderConfig, cacheRetention?: "none" | "short" | "long") { + if ( + provider.adapter === "test-kiro-continuation" + || (provider.adapter === "kiro" && provider.apiKey?.startsWith("kiro-access-")) + ) return kiroFixtureAdapter(provider); + return actualResolveAdapter(provider, cacheRetention); + }, + })); + + ({ handleResponses } = await import("../../../src/server/responses")); +}); + +beforeEach(() => { + testHome = mkdtempSync(join(tmpdir(), "ocx-kiro-continuation-auth-")); + process.env.OPENCODEX_HOME = testHome; + kiroBuilds = []; + clearAnthropicAccountPoolState(); + clearGenericFailoverHealth(); +}); + +afterEach(() => { + clearAnthropicAccountPoolState(); + clearGenericFailoverHealth(); + removeTreeWithRetry(testHome); +}); + +afterAll(() => { + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + mock.restore(); +}); + +test("Kiro continuation 429 keeps the rotated bearer and routing metadata together", async () => { + const profiles = [ + "arn:aws:codewhisperer:us-east-1:123456789012:profile/account-a", + "arn:aws:codewhisperer:eu-west-1:123456789012:profile/account-b", + ]; + const regions = ["us-east-1", "eu-west-1"]; + for (let index = 0; index < 2; index += 1) { + await saveCredential("kiro", { + access: `kiro-access-${index}`, + refresh: `kiro-refresh-${index}`, + expires: Date.now() + 3_600_000, + accountId: `kiro-account-${index}`, + kiro: { + profileArn: profiles[index], + apiRegion: regions[index], + ssoRegion: regions[index], + }, + } as never, { addAccount: true }); + } + const ids = getAccountSet("kiro")!.accounts.map(account => account.id); + await setActiveAccount("kiro", ids[0]!); + + const config = { + port: 0, + defaultProvider: "kiro", + providers: { + kiro: { + adapter: "test-kiro-continuation", + baseUrl: "https://kiro-continuation.test/v1", + authMode: "oauth", + models: ["model"], + }, + }, + } as unknown as OcxConfig; + + const phases = ["plan", "rate-limit", "complete"]; + const originalFetch = globalThis.fetch; + globalThis.fetch = (async () => { + const phase = phases.shift(); + if (phase === "rate-limit") { + return Response.json( + { error: { message: "rate limited" } }, + { status: 429, headers: { "retry-after": "30" } }, + ); + } + if (!phase) throw new Error("unexpected extra request"); + return new Response("", { status: 200, headers: { "x-test-phase": phase } }); + }) as typeof fetch; + + try { + const response = await handleResponses(new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "kiro/model", + input: "Please modify the file", + stream: true, + tools: [{ + type: "function", + name: "read_file", + description: "Read one file", + parameters: { type: "object", properties: {} }, + }], + }), + }), config, { model: "", provider: "" }); + + expect(response.status).toBe(200); + expect(await response.text()).toContain("read_file"); + } finally { + globalThis.fetch = originalFetch; + } + + // The third build is the continuation retry: it must carry the ROTATED account's bearer and + // that same account's region/profile, not account 0's routing metadata. + expect(kiroBuilds).toEqual([ + { key: "kiro-access-0", profileArn: profiles[0], apiRegion: regions[0] }, + { key: "kiro-access-0", profileArn: profiles[0], apiRegion: regions[0] }, + { key: "kiro-access-1", profileArn: profiles[1], apiRegion: regions[1] }, + ]); + expect(phases).toEqual([]); +});