diff --git a/src/providers/key-failover.ts b/src/providers/key-failover.ts index 4e9f2e60a4..9bc8f8d210 100644 --- a/src/providers/key-failover.ts +++ b/src/providers/key-failover.ts @@ -8,8 +8,10 @@ * * Modelled after src/codex/routing.ts cooldown logic but scoped to plain API-key pools. */ -import { saveConfigPreservingClaudeCode } from "../config"; +import { mutatePersistedConfig } from "../config"; +import { routedProviderConfig } from "../router"; import type { OcxConfig, OcxProviderConfig, RateLimitRetryPolicy, TransientRetryPolicy } from "../types"; +import { OPENCODE_GO_SESSION_HEADER } from "./opencode-go-transport"; import { resolveProviderTransport, type OcxProviderTransport } from "./xai-transport"; import { sweepExpiredOnWrite } from "../lib/state-store-sweeper"; @@ -174,7 +176,8 @@ export function rateLimitRetryDelayMs( * The returned object is a snapshot of the PERSISTED config — it carries none of the * registry backfills `routedProviderConfig` merges in at request time. Request paths must * not assign it to an active route wholesale; use `rotateProviderTransportOn429`, which - * takes only the swapped key and keeps the routed provider intact. + * rebuilds from this committed row, reapplies registry metadata, and retains only explicit + * runtime transport state (`fetch` and generated OpenCode session affinity). */ export function rotateKeyOn429( config: OcxConfig, @@ -187,50 +190,69 @@ export function rotateKeyOn429( if (!provider) return null; if (provider.authMode === "oauth" || provider.authMode === "forward") return null; - const pool = provider.apiKeyPool; - if (!pool || pool.length < 2) return null; - - // Cool the key that ACTUALLY failed. Under concurrent 429s another request may already have - // rotated provider.apiKey — cooling the live key would punish an innocent replacement and can - // exhaust a 2-key pool from a single bad key. CAS semantics: callers pass the key they used. const failedKey = attemptedKey ?? provider.apiKey; - const currentEntry = pool.find(e => e.key === failedKey); - if (currentEntry) { - const cooldownMs = parseRetryAfterMs(retryAfterHeader, now) ?? DEFAULT_COOLDOWN_MS; - keyCooldowns.set(cooldownKey(providerName, currentEntry.id), { - cooldownUntil: now + cooldownMs, - }); - sweepExpiredOnWrite(now); - } + type Rotation = + | { provider: OcxProviderConfig; failedId?: string; candidateId?: string } + | { exhaustedCount: number; failedId?: string }; + const outcome = mutatePersistedConfig(fresh => { + const freshProvider = fresh.providers[providerName]; + if (!freshProvider || freshProvider.authMode === "oauth" || freshProvider.authMode === "forward") { + return { changed: false, value: null }; + } + const pool = freshProvider.apiKeyPool; + if (!pool || pool.length < 2) return { changed: false, value: null }; + + // The callback can be rerun after rebasing, so identify the failed key here but + // defer the in-memory cooldown side effect until persistence has succeeded. + const failedEntry = pool.find(entry => entry.key === failedKey); - // Lost the race: someone already rotated away from the failed key. If the live key is healthy, - // retry with it as-is instead of rotating a second time. - if (attemptedKey !== undefined && provider.apiKey !== attemptedKey) { - const liveEntry = pool.find(e => e.key === provider.apiKey); - if (liveEntry && !isKeyInCooldown(providerName, liveEntry.id, now)) { - return { ...provider }; + if (freshProvider.apiKey !== failedKey) { + const activeEntry = pool.find(entry => entry.key === freshProvider.apiKey); + if (activeEntry && !isKeyInCooldown(providerName, activeEntry.id, now)) { + return { + changed: false, + value: { provider: structuredClone(freshProvider), failedId: failedEntry?.id }, + }; + } } - } - // Pick the next key that is NOT in cooldown - const currentIndex = currentEntry ? pool.indexOf(currentEntry) : -1; - for (let i = 1; i < pool.length; i++) { - const candidate = pool[(currentIndex + i) % pool.length]!; - if (!isKeyInCooldown(providerName, candidate.id, now)) { - // Swap active key - provider.apiKey = candidate.key; - saveConfigPreservingClaudeCode(config); - console.warn( - // Log ids only — labels are user-supplied free text and could carry secret material. - `[key-failover] ${providerName}: 429 on key ${currentEntry?.id ?? "?"}; rotating to key ${candidate.id}`, - ); - return { ...provider }; + const currentIndex = failedEntry ? pool.indexOf(failedEntry) : -1; + const candidateCount = failedEntry ? pool.length - 1 : pool.length; + for (let offset = 1; offset <= candidateCount; offset += 1) { + const candidate = pool[(currentIndex + offset) % pool.length]!; + if (isKeyInCooldown(providerName, candidate.id, now)) continue; + freshProvider.apiKey = candidate.key; + return { + changed: true, + value: { + provider: structuredClone(freshProvider), + failedId: failedEntry?.id, + candidateId: candidate.id, + }, + }; } + return { changed: false, value: { exhaustedCount: pool.length, failedId: failedEntry?.id } }; + }); + if (outcome.status === "unavailable" || outcome.value === null) return null; + if (outcome.value.failedId) { + const cooldownMs = parseRetryAfterMs(retryAfterHeader, now) ?? DEFAULT_COOLDOWN_MS; + keyCooldowns.set(cooldownKey(providerName, outcome.value.failedId), { cooldownUntil: now + cooldownMs }); + sweepExpiredOnWrite(now); + } + if ("exhaustedCount" in outcome.value) { + console.warn(`[key-failover] ${providerName}: all ${outcome.value.exhaustedCount} keys in cooldown; returning 429 to client`); + return null; } - // All keys in cooldown - console.warn(`[key-failover] ${providerName}: all ${pool.length} keys in cooldown; returning 429 to client`); - return null; + const committed = structuredClone(outcome.value.provider); + config.providers[providerName] = committed; + if (outcome.value.candidateId) { + console.warn( + // Log ids only — labels are user-supplied free text and could carry secret material. + `[key-failover] ${providerName}: 429 on key ${outcome.value.failedId ?? "?"}; rotating to key ${outcome.value.candidateId}`, + ); + } + return structuredClone(committed); } export function sweepExpiredApiKeyCooldowns(now = Date.now()): number { @@ -253,13 +275,8 @@ interface RotateProviderTransportOptions { /** * Rotate a failed key and re-apply provider-specific transport metadata to the replacement. * - * `routedProvider` is the request's active provider (the `routedProviderConfig` output the - * route was built with). The result inherits it and swaps ONLY the API key: the persisted - * config that `rotateKeyOn429` snapshots predates registry backfill, so building the retry - * provider from that snapshot would silently drop every field the registry merged in at - * routing time (scalar flags like `promptCacheKey`/`parallelToolCalls`, merged model - * metadata such as `noTemperatureModels`, a pinned baseUrl). Mirrors the OAuth-401 replay - * path in src/server/responses/core.ts, which spreads `route.provider` for the same reason. + * Route the authoritative committed row again so concurrent provider edits take effect, then + * restore only transport-only state that can never come from persisted configuration. */ export function rotateProviderTransportOn429( config: OcxConfig, @@ -274,13 +291,23 @@ export function rotateProviderTransportOn429( options.now, options.attemptedKey, ); - return rotated - ? resolveProviderTransport( - providerName, - { ...routedProvider, apiKey: rotated.apiKey }, - options.promptCacheKey, - ) - : null; + if (!rotated) return null; + + const committedRoute = routedProviderConfig(providerName, rotated); + const routedSession = routedProvider.headers?.[OPENCODE_GO_SESSION_HEADER]; + const retryProvider: OcxProviderTransport = { + ...committedRoute, + ...(routedProvider.fetch !== undefined ? { fetch: routedProvider.fetch } : {}), + ...(routedSession !== undefined + ? { + headers: { + ...(committedRoute.headers ?? {}), + [OPENCODE_GO_SESSION_HEADER]: routedSession, + }, + } + : {}), + }; + return resolveProviderTransport(providerName, retryProvider, options.promptCacheKey); } /** Clear cooldown state for a provider (e.g. after manual key management). */ diff --git a/structure/04_transports-and-sidecars.md b/structure/04_transports-and-sidecars.md index 4c47dd0233..eb6cbc82a5 100644 --- a/structure/04_transports-and-sidecars.md +++ b/structure/04_transports-and-sidecars.md @@ -1022,9 +1022,11 @@ its own: it forwards what the request already carries — Codex's session key on (metadata.user_id hash, else the system+tools cohort hash) — and a request with no key stays keyless. An explicit provider-level `promptCacheKey: false` continues to opt out, and the flag is persisted through `providerConfigSeed`/`enrichProviderFromRegistry` for new configs; key-pool 429 -rotation keeps it — along with every other registry backfill — because the retry inherits the -request's routed provider and swaps only the API key (`rotateProviderTransportOn429` in -src/providers/key-failover.ts). If an opted-in upstream rejects the field, OpenCodex does not strip it and retry or mutate the +rotation keeps it — along with every other registry backfill — because the retry starts from the +fresh committed provider row and routes it again (`rotateProviderTransportOn429` in +src/providers/key-failover.ts). Stale request-time config fields are deliberately discarded so a +concurrent deletion stays authoritative; only runtime `fetch` state and generated OpenCode session +affinity survive the rebuild. If an opted-in upstream rejects the field, OpenCodex does not strip it and retry or mutate the saved configuration. Other OpenAI-compatible providers remain deny-by-default because strict backends may reject the OpenAI-specific field. diff --git a/tests/adapters/key-failover.test.ts b/tests/adapters/key-failover.test.ts index 5cff50912d..e2595a84f0 100644 --- a/tests/adapters/key-failover.test.ts +++ b/tests/adapters/key-failover.test.ts @@ -1,8 +1,14 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; -import { mkdtempSync} from "node:fs"; +import { mkdtempSync, unlinkSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { createOpenAIChatAdapter } from "../../src/adapters/openai-chat"; +import { + getConfigPath, + loadConfig, + mutatePersistedConfig, + saveConfig, +} from "../../src/config"; import { clearKeyCooldowns, getKeyCooldownUntil, @@ -19,7 +25,7 @@ import { removeTreeWithRetry } from "../helpers/remove-tree"; let home: string; function makeConfig(provider: Partial): OcxConfig { - return { + const config = { port: 10199, defaultProvider: "p", providers: { @@ -30,6 +36,8 @@ function makeConfig(provider: Partial): OcxConfig { } as OcxProviderConfig, }, } as OcxConfig; + saveConfig(config); + return config; } function pool3(): OcxProviderConfig["apiKeyPool"] { @@ -106,6 +114,16 @@ describe("rotateKeyOn429", () => { expect(rotateKeyOn429(makeConfig({}), "missing", null)).toBeNull(); }); + test("unavailable persistence does not publish a tentative cooldown", () => { + const config = makeConfig({ apiKey: "key-alpha-000111222333", apiKeyPool: pool3() }); + const now = 1_000_000; + unlinkSync(getConfigPath()); + + expect(rotateKeyOn429(config, "p", null, now, "key-alpha-000111222333")).toBeNull(); + expect(getKeyCooldownUntil("p", "k1", now)).toBeNull(); + expect(config.providers.p.apiKey).toBe("key-alpha-000111222333"); + }); + test("clearKeyCooldowns scoped to a provider", () => { const config = makeConfig({ apiKey: "key-alpha-000111222333", apiKeyPool: pool3() }); const now = 1_000_000; @@ -131,6 +149,40 @@ describe("rotateKeyOn429", () => { // A REAL beta failure afterwards still rotates to gamma. expect(rotateKeyOn429(config, "p", null, now, "key-beta-444555666777")?.apiKey).toBe("key-gamma-888999000111"); }); + + test("two stale handlers adopt one committed rotation without rotating twice", () => { + const first = makeConfig({ apiKey: "key-alpha-000111222333", apiKeyPool: pool3() }); + const second = loadConfig(); + const now = 1_000_000; + + expect(rotateKeyOn429(first, "p", null, now, "key-alpha-000111222333")?.apiKey) + .toBe("key-beta-444555666777"); + expect(rotateKeyOn429(second, "p", null, now, "key-alpha-000111222333")?.apiKey) + .toBe("key-beta-444555666777"); + expect(second.providers.p.apiKey).toBe("key-beta-444555666777"); + expect(getKeyCooldownUntil("p", "k2", now)).toBeNull(); + }); + + test("rebases over a concurrent pool edit without resurrecting a removed key", () => { + const stale = makeConfig({ + apiKey: "key-alpha-000111222333", + apiKeyPool: pool3(), + note: "stale", + }); + const added = { id: "k4", key: "key-delta-222333444555", addedAt: 4 }; + const edit = mutatePersistedConfig(fresh => { + fresh.providers.p.apiKeyPool = [fresh.providers.p.apiKeyPool![0]!, fresh.providers.p.apiKeyPool![2]!, added]; + fresh.providers.p.note = "concurrent"; + return { changed: true, value: undefined }; + }); + expect(edit.status).toBe("committed"); + + const rotated = rotateKeyOn429(stale, "p", null, 1_000_000, "key-alpha-000111222333"); + expect(rotated?.apiKey).toBe("key-gamma-888999000111"); + expect(rotated?.apiKeyPool?.map(entry => entry.id)).toEqual(["k1", "k3", "k4"]); + expect(rotated?.note).toBe("concurrent"); + expect(stale.providers.p).toEqual(loadConfig().providers.p); + }); }); describe("rotateProviderTransportOn429", () => { @@ -146,6 +198,7 @@ describe("rotateProviderTransportOn429", () => { baseUrl: "https://opencode.ai/zen/go/v1", }; delete config.providers.p; + writeFileSync(getConfigPath(), `${JSON.stringify(config, null, 2)}\n`); const initial = resolveOpenCodeGoTransport( config.providers["opencode-go"], @@ -177,6 +230,7 @@ describe("rotateProviderTransportOn429", () => { baseUrl: "https://api.kimi.com/coding/v1", }; delete config.providers.p; + writeFileSync(getConfigPath(), `${JSON.stringify(config, null, 2)}\n`); expect(config.providers["kimi-code"].promptCacheKey).toBeUndefined(); const parsed: OcxParsedRequest = { @@ -196,14 +250,14 @@ describe("rotateProviderTransportOn429", () => { }); expect(rotated?.apiKey).toBe("key-beta-444555666777"); expect(rotated?.promptCacheKey).toBe(true); + expect(rotated?.modelContextWindows).toBeDefined(); const retryBody = JSON.parse(createOpenAIChatAdapter(rotated!).buildRequest(parsed).body); expect(retryBody.prompt_cache_key).toBe(promptCacheKey); }); - test("inherits the routed provider's registry backfills; only the key changes", () => { - // The persisted config predates the registry scalar flags and merged metadata — - // routedProviderConfig backfilled them at request time. Rotation must not fall back - // to the bare persisted snapshot and silently drop them for the retried request. + test("drops stale routed configuration while persisted fields stay authoritative", () => { + // Request-time configuration cannot revive fields absent from the committed snapshot. + // Registry providers still receive their canonical backfills from routedProviderConfig. const config = makeConfig({ apiKey: "key-alpha-000111222333", apiKeyPool: pool3() }); const routedProvider = { ...config.providers.p, @@ -220,16 +274,38 @@ describe("rotateProviderTransportOn429", () => { }); expect(rotated?.apiKey).toBe("key-beta-444555666777"); - expect(rotated?.baseUrl).toBe("https://registry-pinned.example/v1"); - expect(rotated?.promptCacheKey).toBe(true); - expect(rotated?.parallelToolCalls).toBe(false); - expect(rotated?.modelContextWindows).toEqual({ "some-model": 262_144 }); - expect(rotated?.noTemperatureModels).toEqual(["some-model"]); + expect(rotated?.baseUrl).toBe("https://api.example.com/v1"); + expect(rotated?.promptCacheKey).toBeUndefined(); + expect(rotated?.parallelToolCalls).toBeUndefined(); + expect(rotated?.modelContextWindows).toBeUndefined(); + expect(rotated?.noTemperatureModels).toBeUndefined(); // The pool swap still lands in the persisted config. expect(config.providers.p.apiKey).toBe("key-beta-444555666777"); expect(config.providers.p.promptCacheKey).toBeUndefined(); }); + test("does not resurrect an optional provider field removed before rotation", () => { + const config = makeConfig({ + apiKey: "key-alpha-000111222333", + apiKeyPool: pool3(), + headers: { "x-user-header": "stale" }, + }); + const routedProvider = { ...config.providers.p }; + const edit = mutatePersistedConfig(fresh => { + delete fresh.providers.p.headers; + return { changed: true, value: undefined }; + }); + expect(edit.status).toBe("committed"); + + const rotated = rotateProviderTransportOn429(config, "p", routedProvider, { + now: 1_000_000, + attemptedKey: "key-alpha-000111222333", + }); + + expect(rotated?.apiKey).toBe("key-beta-444555666777"); + expect(rotated?.headers).toBeUndefined(); + }); + test("re-applies xAI cache affinity without OAuth CLI headers after key rotation", () => { const promptCacheKey = "stable-conversation-429"; const config = makeConfig({ @@ -239,6 +315,8 @@ describe("rotateProviderTransportOn429", () => { }); config.providers.xai = config.providers.p; delete config.providers.p; + config.defaultProvider = "xai"; + writeFileSync(getConfigPath(), `${JSON.stringify(config, null, 2)}\n`); const rotated = rotateProviderTransportOn429(config, "xai", { ...config.providers.xai }, { now: 1_000_000, diff --git a/tests/adapters/openai/openai-chat-native-policy.test.ts b/tests/adapters/openai/openai-chat-native-policy.test.ts index 95b39c065f..92fc1bf371 100644 --- a/tests/adapters/openai/openai-chat-native-policy.test.ts +++ b/tests/adapters/openai/openai-chat-native-policy.test.ts @@ -2,6 +2,7 @@ import { afterEach, describe, expect, test } from "bun:test"; import { mkdtempSync} from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { saveConfig } from "../../../src/config"; import { buildOpenAIChatPassthroughRequest, createOpenAIChatAdapter, @@ -285,6 +286,7 @@ describe("native Chat passthrough service-tier policy", () => { } as OcxConfig; try { + saveConfig(config); const response = await handleChatCompletions( new Request("http://localhost/v1/chat/completions", { method: "POST", diff --git a/tests/providers/openrouter-provider-routing.test.ts b/tests/providers/openrouter-provider-routing.test.ts index 8211c8f008..18633733e1 100644 --- a/tests/providers/openrouter-provider-routing.test.ts +++ b/tests/providers/openrouter-provider-routing.test.ts @@ -2,6 +2,7 @@ import { describe, expect, test } from "bun:test"; import { mkdtempSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; +import { saveConfig } from "../../src/config"; import { buildOpenAIChatPassthroughRequest, createOpenAIChatAdapter } from "../../src/adapters/openai-chat"; import { openRouterRoutingConfigError, @@ -187,6 +188,7 @@ describe("OpenRouter configurable provider routing", () => { providers: { openrouter }, }; try { + saveConfig(config); const rotated = rotateProviderTransportOn429(config, "openrouter", openrouter, { attemptedKey: "key-one", now: 1_000_000, diff --git a/tests/server/server-combo-failover-e2e.test.ts b/tests/server/server-combo-failover-e2e.test.ts index 6c84ccab30..f01878c36a 100644 --- a/tests/server/server-combo-failover-e2e.test.ts +++ b/tests/server/server-combo-failover-e2e.test.ts @@ -37,6 +37,7 @@ import { } from "../../src/responses/state"; import { clearCursorThreadContinuityForTests } from "../../src/adapters/cursor/thread-continuity"; import { COMPACT_PROMPT, encodeCompactionSummary } from "../../src/responses/compaction"; +import { clearKeyCooldowns } from "../../src/providers/key-failover"; // Full-suite Windows load: startServer + combo rename/delete management flows exceed the // default 5s per-test budget (same flake class as 810fa115 / claude-management-api). @@ -138,6 +139,7 @@ beforeEach(() => { process.env.OPENCODEX_HOME = testDir; clearComboSelectionState(); clearComboTargetCooldowns(); + clearKeyCooldowns(); clearCodexUpstreamHealth(); customRunTurn = undefined; customFetchResponse = undefined; @@ -169,6 +171,7 @@ afterEach(async () => { if (testDir) removeTreeWithRetry(testDir); clearComboSelectionState(); clearComboTargetCooldowns(); + clearKeyCooldowns(); clearCodexUpstreamHealth(); clearRequestLogsForTests(); } @@ -1355,6 +1358,7 @@ describe("server combo failover 030 activation matrix", () => { const config = comboConfig({ a: provider("test-response", "https://test.invalid/v1", pool[0]!.key, { apiKeyPool: pool }), }); + saveConfig(config); const response = await postLogged(config); expect(response.status).toBe(200); await response.text(); diff --git a/tests/server/terminal-guard-server.test.ts b/tests/server/terminal-guard-server.test.ts index 96a82b05d8..c2c889c3d8 100644 --- a/tests/server/terminal-guard-server.test.ts +++ b/tests/server/terminal-guard-server.test.ts @@ -1,6 +1,12 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { saveConfig } from "../../src/config"; +import { clearKeyCooldowns } from "../../src/providers/key-failover"; import { handleResponses } from "../../src/server/responses"; import type { OcxConfig } from "../../src/types"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; const config = { port: 0, @@ -166,6 +172,10 @@ describe("server terminal guard integration", () => { }); test("terminal-guard continuation retry budget stays per request across key failover", async () => { + const previousHome = process.env.OPENCODEX_HOME; + const home = mkdtempSync(join(tmpdir(), "ocx-terminal-guard-failover-")); + process.env.OPENCODEX_HOME = home; + clearKeyCooldowns("claude-se"); const budgetConfig = { ...config, providers: { @@ -193,21 +203,29 @@ describe("server terminal guard integration", () => { }); }) as typeof fetch; - const response = await handleResponses(new Request("http://localhost/v1/responses", { - method: "POST", - headers: { "content-type": "application/json" }, - body: JSON.stringify({ - model: "se-claude-opus-4.8", - input: "请检查这个问题并修复代码", - stream: true, - tools: [{ type: "function", name: "exec_command", description: "run a command", parameters: { type: "object" } }], - }), - }), budgetConfig, { model: "", provider: "" }); - - await response.text(); - // initial turn + continuation retry (same key) + failover continuation (second key) = 4. - // A per-iteration budget would replay on the second key too (5+ sends). - expect(sends).toBe(4); + try { + saveConfig(budgetConfig); + const response = await handleResponses(new Request("http://localhost/v1/responses", { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + model: "se-claude-opus-4.8", + input: "请检查这个问题并修复代码", + stream: true, + tools: [{ type: "function", name: "exec_command", description: "run a command", parameters: { type: "object" } }], + }), + }), budgetConfig, { model: "", provider: "" }); + + await response.text(); + // initial turn + continuation retry (same key) + failover continuation (second key) = 4. + // A per-iteration budget would replay on the second key too (5+ sends). + expect(sends).toBe(4); + } finally { + clearKeyCooldowns("claude-se"); + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + removeTreeWithRetry(home); + } }); test("terminal-guard continuation shares the request-wide 429 budget with the main loop", async () => {