fix: charge key cooldown to the credential sent

src/providers/api-key-response-ownership.ts

@@ -15,15 +15,19 @@ registerProviderResponseProvenance(owners);
 /** Capture the credential after connection-policy rewrites, immediately before fetch. */
 export function distributedKeyPhysicalFetch(
   physicalFetch: typeof globalThis.fetch, providerName: string, provider: OcxProviderConfig, target: string,
+  authHeader: "authorization" | "x-api-key" = "authorization",
 ): typeof globalThis.fetch {
   if (!provider._apiKeyAttempt?.distributed) return physicalFetch;
   const wrapped = (async (input, init) => {
     const headers = new Headers(init?.headers ?? (input instanceof Request ? input.headers : undefined));
     const destination = input instanceof Request ? input.url : String(input);
-    const ownsCredential = destination === target && headers.get("authorization") === `Bearer ${provider.apiKey}`
-      && !headers.has("x-api-key") && !headers.has("api-key");
+    const ownsCredential = destination === target && !headers.has("api-key")
+      && (authHeader === "authorization"
+        ? headers.get("authorization") === `Bearer ${provider.apiKey}` && !headers.has("x-api-key")
+        : headers.get("x-api-key") === provider.apiKey && !headers.has("authorization"));
     const owner: KeyResponseOwner | null = provider._apiKeyAttempt?.distributed && ownsCredential
-      ? { providerName, provider: { ...provider, _apiKeyAttempt: { ...provider._apiKeyAttempt } } } : null;
+      ? Object.freeze({ providerName, provider: Object.freeze({ ...provider,
+        _apiKeyAttempt: Object.freeze({ ...provider._apiKeyAttempt }) }) }) : null;
     const response = await physicalFetch(input, init);
     // Reused response identities cannot safely describe two physical sends.
     owners.set(response, owners.has(response) ? null : owner);
@@ -32,14 +36,22 @@ export function distributedKeyPhysicalFetch(
   return isEgressTransparentExecutor(physicalFetch) ? markEgressTransparentExecutor(wrapped) : wrapped;
 }
 
+/** Return only a still-configured credential captured at the physical send. */
+export function distributedKeyResponseOwner(
+  response: Response, config: OcxConfig, providerName: string,
+): OcxProviderConfig | null {
+  const owner = owners.get(response);
+  return owner?.providerName === providerName
+    && providerApiKeySelectionOwnsFailure(config, providerName, owner.provider) ? owner.provider : null;
+}
+
 /** Settle only the terminal refusal; configured same-key 429 retries retain their key. */
 export function recordDistributedResponseFailure(
   response: Response, config: OcxConfig, providerName: string, signal?: AbortSignal,
 ): void {
   if (signal?.aborted || isNonReplayableResponse(response) || (response.status !== 401 && response.status !== 429)) return;
-  const owner = owners.get(response);
-  if (!owner || owner.providerName !== providerName
-    || !providerApiKeySelectionOwnsFailure(config, providerName, owner.provider)) return;
-  recordDistributedKeyFailure(config, providerName, owner.provider._apiKeyAttempt,
+  const owner = distributedKeyResponseOwner(response, config, providerName);
+  if (!owner) return;
+  recordDistributedKeyFailure(config, providerName, owner._apiKeyAttempt,
     response.status, Date.now(), response.headers.get("retry-after"));
 }

src/server/chat-native.ts

@@ -62,6 +62,7 @@ import {
 import { fastPolicyForModel } from "../providers/service-tier";
 import { RequestPacingQueueOverloadError } from "../providers/request-pacing";
 import { stampApiKeyAccountLabel } from "../providers/label";
+import { distributedKeyPhysicalFetch, distributedKeyResponseOwner } from "../providers/api-key-response-ownership";
 import { providerApiKeySelectionIsCurrent, resolveCurrentProviderApiKeyTransport } from "../providers/api-key-selection";
 import { enrichOpenCodeZenFreeTierMessage } from "../providers/opencode-zen-rate-limit";
 import type { OcxProviderTransport } from "../providers/xai-transport";
@@ -515,8 +516,6 @@ export async function runNativeChatAttempt(
               const spendReport = createPhysicalSendReporter(physicalBudget, () => ({
                 poolId: logCtx.spendPoolId ?? route.providerName, identityId: logCtx.accountLogLabel,
               }), execution.comboDispatchPermit);
-              const ownsDistributedKey = !activeProvider._apiKeyAttempt?.distributed
-                || headers.get("authorization") === `Bearer ${activeProvider.apiKey}`;
               let dispatched: Response;
               try {
                 // Capture this request's spend policy at its first physical boundary, including
@@ -534,7 +533,8 @@ export async function runNativeChatAttempt(
                 // A reselected provider transport is still a physical send: the connection policy
                 // and manual-redirect ownership wrap the selected implementation (#4992).
                 dispatched = await sendWithConnectionPolicy(
-                  (activeProvider as OcxProviderTransport).fetch ?? execute,
+                  distributedKeyPhysicalFetch((activeProvider as OcxProviderTransport).fetch ?? execute,
+                    route.providerName, activeProvider, request.url),
                   request.url,
                   applyUpstreamRecoveryInit({
                     ...init, method: request.method, headers, body: request.body,
@@ -549,7 +549,8 @@ export async function runNativeChatAttempt(
                 try { if (physicalBudget.spendEnforced) spendReport(1); }
                 finally { spendReport.close?.(); }
               }
-              responseOwnsDistributedKey = ownsDistributedKey;
+              responseOwnsDistributedKey = !activeProvider._apiKeyAttempt?.distributed
+                || distributedKeyResponseOwner(dispatched, config, route.providerName) !== null;
               if (!dispatched.ok) await recordKeyAttemptFailure(logCtx, dispatched, init.signal ?? upstream.signal);
               return dispatched;
             },

src/server/messages-native.ts

@@ -61,13 +61,13 @@ import {
   type UpstreamSendRecovery,
   UPSTREAM_RESET_REPLAY_REFUSED_CODE,
 } from "../lib/upstream-retry";
+import { distributedKeyPhysicalFetch, distributedKeyResponseOwner, recordDistributedResponseFailure } from "../providers/api-key-response-ownership";
 import { providerApiKeySelectionIsCurrent, resolveCurrentProviderApiKeyTransport } from "../providers/api-key-selection";
 import { RequestPacingQueueOverloadError } from "../providers/request-pacing";
 import {
   hasKeyPoolFailover,
   rateLimitRetryDelayMs,
   rateLimitRetryPolicyFor,
-  recordDistributedKeyFailure,
   rotateProviderTransportOn401,
   rotateProviderTransportOn429,
   selectProactiveApiKeyTransport,
@@ -616,10 +616,6 @@ export async function handleNativeMessages(options: HandleNativeMessagesOptions)
             const releaseFamily = snapshot && oauthBinding ? anthropicModelQuotaFor(oauthBinding.instance).claimAnthropicFamilyRevalidation(snapshot.accountId, route.modelId) : () => {};
             if (!releaseFamily) throw new AnthropicAccountCooldownError(1);
             const selection = activeProvider._apiKeyAttempt;
-            const ownsDistributedKey = !selection?.distributed
-              || (activeProvider.apiKeyTransport === "bearer"
-                ? headers.get("authorization") === `Bearer ${activeProvider.apiKey}` && !headers.has("x-api-key")
-                : headers.get("x-api-key") === activeProvider.apiKey && !headers.has("authorization"));
             let dispatched: Response;
             const spendReport = createPhysicalSendReporter(physicalBudget, () => ({
               poolId: logCtx.spendPoolId ?? route.providerName, identityId: logCtx.accountLogLabel,
@@ -640,7 +636,9 @@ export async function handleNativeMessages(options: HandleNativeMessagesOptions)
               if (snapshot && !sendOwner) throw new NativeOAuthSelectionChangedError();
               physicalSendsUsed++;
               dispatched = await sendWithConnectionPolicy(
-                anthropicBillingPhysicalFetch((activeProvider as OcxProviderTransport).fetch ?? execute, route.providerName, activeProvider),
+                distributedKeyPhysicalFetch(
+                  anthropicBillingPhysicalFetch((activeProvider as OcxProviderTransport).fetch ?? execute, route.providerName, activeProvider),
+                  route.providerName, activeProvider, wire.url, activeProvider.apiKeyTransport === "bearer" ? "authorization" : "x-api-key"),
                 wire.url,
                 applyUpstreamRecoveryInit({ ...init, method: "POST", headers, body: wire.body }, transportRecovery),
                 { providerName: route.providerName, provider: activeProvider },
@@ -664,12 +662,13 @@ export async function handleNativeMessages(options: HandleNativeMessagesOptions)
                 }
               } catch { /* Passive observation must not fail the response. */ }
             }
+            const ownsDistributedKey = !selection?.distributed
+              || distributedKeyResponseOwner(dispatched, config, route.providerName) !== null;
             distributedResponseOwnsKey = ownsDistributedKey;
             const failureStatus = dispatched.status;
             pendingDistributedFailure = selection?.distributed && ownsDistributedKey && (failureStatus === 401 || failureStatus === 429)
               && !isNonReplayableResponse(dispatched)
-              ? () => { recordDistributedKeyFailure(config, route.providerName, selection, failureStatus,
-                Date.now(), dispatched.headers.get("retry-after")); }
+              ? () => { recordDistributedResponseFailure(dispatched, config, route.providerName, req.signal); }
               : undefined;
             if (!dispatched.ok) await recordKeyAttemptFailure(logCtx, dispatched, init.signal ?? upstream.signal);
             return dispatched;

src/server/responses/compact.ts

@@ -164,8 +164,9 @@ import {
   UnsupportedContentEncodingError,
 } from "../request-decompress";
 import { resolveAdapter, resolveWireProtocolOverride } from "../adapter-resolve";
-import { hasKeyPoolFailover, recordDistributedKeyFailure, rotateProviderTransportOn429, selectProactiveApiKeyTransport, transientRetryPolicyFor } from "../../providers/key-failover";
+import { hasKeyPoolFailover, rotateProviderTransportOn429, selectProactiveApiKeyTransport, transientRetryPolicyFor } from "../../providers/key-failover";
 import { transientSendCapFor } from "./request-send-budget";
+import { distributedKeyPhysicalFetch, recordDistributedResponseFailure } from "../../providers/api-key-response-ownership";
 import { providerApiKeySelectionIsCurrent, resolveCurrentProviderApiKeyTransport } from "../../providers/api-key-selection";
 import { stampApiKeyAccountLabel } from "../../providers/label";
 import { shouldAttemptImageTierRetry } from "../image-retry";
@@ -1203,7 +1204,8 @@ export async function handleResponsesCompact(
               if (!rebindPhysicalSend(sendBudget, { poolId: logCtx.spendPoolId ?? route.providerName,
                 identityId: logCtx.accountLogLabel })) throw new SendBudgetExhaustedError(safeHostLabel(compactUrl));
               noteProviderAttemptSend(logCtx, route.providerName, sendProvider, logCtx.usageLogInputTokens, upstreamRecovery);
-              return sendWithConnectionPolicy((sendProvider as { fetch?: typeof fetch }).fetch ?? execute, compactUrl,
+              return sendWithConnectionPolicy(
+                distributedKeyPhysicalFetch((sendProvider as { fetch?: typeof fetch }).fetch ?? execute, route.providerName, sendProvider, compactUrl), compactUrl,
                 { ...init, headers: sendHeaders, body: JSON.stringify({ ...buildCompactSendBody() as typeof compactBody, model: route.modelId }) },
                 { providerName: route.providerName, provider: sendProvider });
             } : undefined,
@@ -1283,9 +1285,7 @@ export async function handleResponsesCompact(
       return formatErrorResponse(502, "upstream_error", "Failed to connect to compact upstream");
     }
 
-    if ((upstream.status === 401 || upstream.status === 429) && !isNonReplayableResponse(upstream)) {
-      recordDistributedKeyFailure(config, route.providerName, compactProvider._apiKeyAttempt, upstream.status, Date.now(), upstream.headers.get("retry-after"));
-    }
+    recordDistributedResponseFailure(upstream, config, route.providerName, req.signal);
     if (
       upstream.status === 401
       && (authCtx.kind === "main-pool" || authCtx.kind === "pool")

structure/providers-and-adapters.md

@@ -281,7 +281,7 @@ Cursor, one lease spans the whole turn: RunSSE and BidiAppend may overlap inside
 turns wait at the cap. Follow-up sends still obey start intervals. The turn owner returns its lease
 after `runTurn` settles, so RunSSE body completion cannot admit another turn early. A capped
 canonical Codex WebSocket turn uses HTTP/SSE because the socket has no response-body lifecycle.
-Capacity waits use the same bounded queue and retryable queue-overload errors as interval waits. Optional `requestPacing.keys` adds a shared cross-model limit for each credential reference; duplicate references share capacity, but distinct references resolving to one upstream project do not imply separate quota. `keySelection: round-robin` assigns request-local healthy keys without mutating persisted selection, skips known cooldowns, and refuses exhausted pools with retryable 429. Provider, model and key gates reserve atomically; key counters and the allocation cursor are process-local, not token or billing balances. OAuth and forwarded credentials retain their existing selection. Cooldown state is shared through the dependency-light `src/providers/api-key-cooldowns.ts` (retaining the longest outstanding expiry until explicit manual clearing) owner. Distributed selections are checked again at physical dispatch; a cooled queued key releases its old permit, selects a healthy sibling, and reacquires pacing before sending. Native Messages follows the same readmission rule. Native Responses settles terminal distributed refusals from immutable physical-send metadata; successful same-key retry and unowned refusals do not cool keys. Failure ownership remains valid for a still-configured cooling key, so late in-flight rejections retain the longest cooldown without admitting another send. CLI-backed adapters check `IncomingMeta.beforePhysicalSend` after asynchronous setup and before spawn. An established distributed turn refuses later sends after its key cools instead of changing credentials mid-turn.
+Capacity waits use the same bounded queue and retryable queue-overload errors as interval waits. Optional `requestPacing.keys` adds a shared cross-model limit for each credential reference; duplicate references share capacity, but distinct references resolving to one upstream project do not imply separate quota. `keySelection: round-robin` assigns request-local healthy keys without mutating persisted selection, skips known cooldowns, and refuses exhausted pools with retryable 429. Provider, model and key gates reserve atomically; key counters and the allocation cursor are process-local, not token or billing balances. OAuth and forwarded credentials retain their existing selection. Cooldown state is shared through the dependency-light `src/providers/api-key-cooldowns.ts` (retaining the longest outstanding expiry until explicit manual clearing) owner. Distributed selections are checked again at physical dispatch; a cooled queued key releases its old permit, selects a healthy sibling, and reacquires pacing before sending. Native Messages follows the same readmission rule. Native Responses, compact, Chat and Messages settle distributed refusals from immutable credential and destination metadata captured by `src/providers/api-key-response-ownership.ts` at physical fetch; successful same-key retry and unowned refusals do not cool keys. Failure ownership remains valid for a still-configured cooling key, so late in-flight rejections retain the longest cooldown without admitting another send. CLI-backed adapters check `IncomingMeta.beforePhysicalSend` after asynchronous setup and before spawn. An established distributed turn refuses later sends after its key cools instead of changing credentials mid-turn.
 
 Command Code effort defaults in `src/providers/command-code-efforts.ts` combine public-profile
 facts with the live API measurements from #5096. Both presets share the exact per-model rows;

tests/claude-integration/messages-native-key-selection.test.ts

@@ -3,6 +3,7 @@ import { mkdtempSync, readFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
 import { saveConfig } from "../../src/config";
+import { registerUpstreamRewriter } from "../../src/plugins/upstream-hooks";
 import { apiKeyAccountLogLabel } from "../../src/codex/account-label";
 import { configureSharedSpendLedger, spendPolicyFromConfig } from "../../src/lib/spend-reservation-ledger";
 import { captureProviderApiKeySelection } from "../../src/providers/api-key-selection";
@@ -321,6 +322,28 @@ for (const statusCode of [401, 429]) for (const overridden of [false, true]) tes
   unchanged(config); noLeaks(config);
 });
 
+for (const statusCode of [401, 429]) for (const transport of ["bearer", "x-api-key"] as const) {
+  for (const rewrite of ["credential", "destination"] as const) test(`Messages physical ${transport} ${statusCode} ownership after ${rewrite} rewrite`, async () => {
+    const config = fixture(2, false, true, [statusCode, 200]);
+    provider(config).apiKeyTransport = transport;
+    saveConfig(config);
+    const off = registerUpstreamRewriter("fixture-key-owner", target => {
+      if (rewrite === "credential") target.headers.set(transport === "bearer" ? "authorization" : "x-api-key",
+        transport === "bearer" ? "Bearer fixture-plugin-key" : "fixture-plugin-key");
+      if (rewrite === "destination") target.url += "?physical-owner=other";
+    });
+    try {
+      const result = await start(config, new AbortController());
+      expect(result.response.status).toBe(statusCode);
+      expect(seen).toHaveLength(1);
+      expect(seen[0]!.key).toBe(rewrite === "credential" ? "fixture-plugin-key" : KEYS[0]);
+      expect(getKeyCooldownUntil(NAME, "k1")).toBeNull();
+      expect(getKeyCooldownUntil(NAME, "k2")).toBeNull();
+      unchanged(config); noLeaks(config);
+    } finally { off(); }
+  });
+}
+
 for (const cancel of [false, true]) test(`active manual key switch readmits the selected key (cancel=${cancel})`, async () => {
   const config = fixture();
   config.providers[NAME]!.requestPacing = { enabled: true, keys: { maxConcurrentRequests: 1 } };

tests/responses/compact-key-selection.test.ts

@@ -247,6 +247,23 @@ test("a synthetic non-replayable compact 429 never cools a distributed key", asy
   unchanged(config); noLeaks(config);
 });
 
+for (const statusCode of [401, 429]) test(`compact static Authorization ${statusCode} retains pool availability`, async () => {
+  const config = fixture();
+  config.providers[NAME]!.supportsNativeCompactEndpoint = true;
+  config.providers[NAME]!.headers = { Authorization: "Bearer fixture-override" };
+  globalThis.fetch = (async (input, init) => {
+    seen.push({ key: new Headers(init?.headers).get("authorization"), url: String(input), body: JSON.parse(String(init?.body)), headers: new Headers(init?.headers) });
+    return Response.json({ error: { message: "fixture refusal" } }, { status: statusCode });
+  }) as typeof fetch;
+  for (let request = 0; request < 4; request++) {
+    expect((await start(config, new AbortController())).response.status).toBe(statusCode);
+    expect(seen).toHaveLength(request + 1);
+    for (const id of ["k1", "k2", "k3"]) expect(getKeyCooldownUntil(NAME, id)).toBeNull();
+  }
+  expect(seen.map(row => row.key)).toEqual(Array(4).fill("Bearer fixture-override"));
+  unchanged(config); noLeaks(config);
+});
+
 for (const cancel of [false, true]) test(`active manual key switch readmits the selected key (cancel=${cancel})`, async () => {
   const config = fixture();
   config.providers[NAME]!.requestPacing = { enabled: true, keys: { maxConcurrentRequests: 1 } };