fix(chat): share transient retry budget across native recovery legs

src/server/chat-native.ts

@@ -204,12 +204,24 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio
     return fail(400, error instanceof Error ? error.message : String(error), "invalid_request_error");
   }
 
+  // One inbound request owns one transient send allowance. Capture the policy before any
+  // key rotation so recovery cannot replace the ceiling along with the active credential.
+  const requestTransientPolicy = transientRetryPolicyFor(activeProvider);
+  let transientSendsUsed = 0;
+  const remainingTransientSends = (): number => requestTransientPolicy
+    ? Math.max(0, requestTransientPolicy.attempts - transientSendsUsed)
+    : Number.POSITIVE_INFINITY;
+  const transientSendAvailable = (): boolean => remainingTransientSends() > 0;
+
   const send = async (request: AdapterRequest, recovery?: "rate-limit-429" | "key-429"): Promise<Response> => {
     try {
       // #2643: opted-in key-auth openai-chat providers retry pre-stream transient statuses on
       // the native chat lane too; everyone else keeps reset-only semantics.
-      const transientPolicy = transientRetryPolicyFor(activeProvider);
-      const fetchWithPolicy = transientPolicy ? fetchWithTransientRetry : fetchWithResetRetry;
+      const remaining = remainingTransientSends();
+      if (requestTransientPolicy && remaining <= 0) {
+        throw new Error("native Chat transient send budget exhausted before recovery dispatch");
+      }
+      const fetchWithPolicy = requestTransientPolicy ? fetchWithTransientRetry : fetchWithResetRetry;
       return await fetchWithPolicy(
         (transportRecovery?: UpstreamSendRecovery) => {
           noteAttemptSend(attempt, logCtx.usageLogInputTokens, transportRecovery ?? recovery);
@@ -232,7 +244,12 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio
         {
           abortSignal: upstream.signal,
           label: safeHostLabel(request.url),
-          ...(transientPolicy ? { attempts: transientPolicy.attempts } : {}),
+          ...(requestTransientPolicy
+            ? {
+              attempts: remaining,
+              onSendsConsumed: (sends: number) => { transientSendsUsed += Math.max(0, sends); },
+            }
+            : {}),
         },
       );
     } finally {
@@ -245,7 +262,12 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio
     response = await send(activeRequest);
     const retryPolicy = rateLimitRetryPolicyFor(activeProvider);
     let retries = 0;
-    while (response.status === 429 && retryPolicy && retries < retryPolicy.attempts) {
+    while (
+      response.status === 429
+      && retryPolicy
+      && retries < retryPolicy.attempts
+      && transientSendAvailable()
+    ) {
       retries += 1;
       for await (const _ of prepareSameTarget429Wait({
         body: response.body,
@@ -263,6 +285,10 @@ export async function handleNativeChatCompletions(options: HandleNativeChatOptio
         promptCacheKey: typeof options.chatBody.prompt_cache_key === "string" ? options.chatBody.prompt_cache_key : undefined,
       });
       if (!rotated) break;
+      // Rotation also records the failed key's cooldown and persists the next healthy key.
+      // Keep that bookkeeping when this request has spent its final send, but preserve the
+      // terminal 429 body and do not dispatch with the replacement credential.
+      if (!transientSendAvailable()) break;
       try { void response.body?.cancel().catch(() => {}); } catch { /* already closed */ }
       activeProvider = rotated;
       activeAdapter = createOpenAIChatAdapter(activeProvider);

tests/chat-completions-endpoint.test.ts

@@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test";
 import { mkdtempSync, rmSync, writeFileSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { saveConfig } from "../src/config";
+import { loadConfig, saveConfig } from "../src/config";
 import { startServer } from "../src/server";
 import { ownedServiceHomeInspection } from "./helpers/owned-service-home-inspection";
 import type { OcxConfig, OcxProviderConfig } from "../src/types";
@@ -1382,6 +1382,149 @@ test("chat-native preserves a structured cyber_policy type on JSON and SSE failu
   }
 });
 
+test("chat-native shares the transient send budget across same-target 429 recovery", async () => {
+  let upstreamSends = 0;
+  const upstream = Bun.serve({
+    port: 0,
+    fetch() {
+      upstreamSends += 1;
+      if (upstreamSends === 1) {
+        return Response.json({ error: { message: "rate limited", type: "rate_limit_error" } }, {
+          status: 429,
+          headers: { "retry-after": "0" },
+        });
+      }
+      return Response.json({ error: { message: "temporarily unavailable", type: "server_error" } }, {
+        status: 503,
+        headers: { "retry-after": "0" },
+      });
+    },
+  });
+  saveConfig(mockConfig(`${upstream.url.toString().replace(/\/$/, "")}/v1`, {
+    authMode: "key",
+    transientRetryOn5xx: { attempts: 3 },
+    retryOn429: { attempts: 1, intervalMs: 100, maxIntervalMs: 100, respectRetryAfter: false },
+  }));
+  const server = startServer(0);
+  try {
+    const response = await fetch(new URL("/v1/chat/completions", server.url), {
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      body: JSON.stringify({ model: "mock/test-model", stream: false, messages: [{ role: "user", content: "hi" }] }),
+    });
+    expect(response.status).toBe(503);
+    await response.text();
+    expect(upstreamSends).toBe(3);
+  } finally {
+    await server.stop(true);
+    upstream.stop(true);
+  }
+});
+
+test("chat-native shares the transient send budget across key rotation", async () => {
+  const { clearKeyCooldowns } = await import("../src/providers/key-failover");
+  clearKeyCooldowns("mock");
+  const authorizations: Array<string | null> = [];
+  const upstream = Bun.serve({
+    port: 0,
+    fetch(req) {
+      authorizations.push(req.headers.get("authorization"));
+      if (authorizations.length === 1) {
+        return Response.json({ error: { message: "temporarily unavailable", type: "server_error" } }, {
+          status: 503,
+          headers: { "retry-after": "0" },
+        });
+      }
+      if (authorizations.length <= 3) {
+        return Response.json({ error: { message: "rate limited", type: "rate_limit_error" } }, {
+          status: 429,
+          headers: { "retry-after": "60" },
+        });
+      }
+      return Response.json({
+        id: "chatcmpl_budget_escape",
+        object: "chat.completion",
+        choices: [{
+          index: 0,
+          message: { role: "assistant", content: "escaped budget" },
+          finish_reason: "stop",
+        }],
+      });
+    },
+  });
+  saveConfig(mockConfig(`${upstream.url.toString().replace(/\/$/, "")}/v1`, {
+    authMode: "key",
+    apiKey: "key-one",
+    apiKeyPool: [
+      { id: "one", key: "key-one" },
+      { id: "two", key: "key-two" },
+      { id: "three", key: "key-three" },
+    ],
+    transientRetryOn5xx: { attempts: 3 },
+  }));
+  const server = startServer(0);
+  try {
+    const response = await fetch(new URL("/v1/chat/completions", server.url), {
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      body: JSON.stringify({ model: "mock/test-model", stream: false, messages: [{ role: "user", content: "hi" }] }),
+    });
+    expect(response.status).toBe(429);
+    await response.text();
+    expect(authorizations).toEqual([
+      "Bearer key-one",
+      "Bearer key-one",
+      "Bearer key-two",
+    ]);
+  } finally {
+    await server.stop(true);
+    upstream.stop(true);
+    clearKeyCooldowns("mock");
+  }
+});
+
+test("chat-native records terminal key cooldown after the send budget is exhausted", async () => {
+  const { clearKeyCooldowns, getKeyCooldownUntil } = await import("../src/providers/key-failover");
+  clearKeyCooldowns("mock");
+  let upstreamSends = 0;
+  const upstream = Bun.serve({
+    port: 0,
+    fetch() {
+      upstreamSends += 1;
+      return Response.json({ error: { message: "rate limited", type: "rate_limit_error" } }, {
+        status: 429,
+        headers: { "retry-after": "60" },
+      });
+    },
+  });
+  saveConfig(mockConfig(`${upstream.url.toString().replace(/\/$/, "")}/v1`, {
+    authMode: "key",
+    apiKey: "key-one",
+    apiKeyPool: [
+      { id: "one", key: "key-one" },
+      { id: "two", key: "key-two" },
+    ],
+    transientRetryOn5xx: { attempts: 1 },
+  }));
+  const server = startServer(0);
+  try {
+    const response = await fetch(new URL("/v1/chat/completions", server.url), {
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      body: JSON.stringify({ model: "mock/test-model", stream: false, messages: [{ role: "user", content: "hi" }] }),
+    });
+    expect(response.status).toBe(429);
+    await response.text();
+    expect(upstreamSends).toBe(1);
+    expect(getKeyCooldownUntil("mock", "one")).not.toBeNull();
+    expect(loadConfig().providers.mock?.apiKey).toBe("key-two");
+  } finally {
+    await server.stop(true);
+    upstream.stop(true);
+    clearKeyCooldowns("mock");
+  }
+});
+
 test("chat-native preserves same-key retry, key rotation, usage, and request logging", async () => {
   const { clearRequestLogsForTests, getRequestLogEntries } = await import("../src/server/request-log");
   const { clearKeyCooldowns } = await import("../src/providers/key-failover");