fix(providers): rebase key failover on persisted state

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<Rotation | null>(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). */

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.
 

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<OcxProviderConfig>): OcxConfig {
-  return {
+  const config = {
     port: 10199,
     defaultProvider: "p",
     providers: {
@@ -30,6 +36,8 @@ function makeConfig(provider: Partial<OcxProviderConfig>): 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,

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",