@dbx-tools/cli-model-proxy 0.3.33 → 0.3.35

package/README.md

@@ -145,6 +145,72 @@ Use this when tests or local developer tools need a managed proxy lifecycle.
 This keeps the package small: Databricks already speaks the OpenAI schema, so
 the useful work is auth, endpoint resolution, and routing to the right surface.
 
+## Timeouts And Cancellation
+
+A proxied turn takes as long as the model behind it, so the proxy imposes **no
+deadline of its own** in either direction. The stream ends when Databricks ends
+it, or when the client hangs up.
+
+That is a deliberate override of two sets of defaults that otherwise truncate
+long turns:
+
+- **Inbound** (client → proxy). Node's `requestTimeout` (300s) and
+  `headersTimeout` (60s) are sized for ordinary web traffic, not for holding a
+  streamed model response open. Both are set to `0`.
+- **Upstream** (proxy → Databricks). Every upstream call goes out on an
+  `undici` `Agent` with `headersTimeout: 0` and `bodyTimeout: 0`. `bodyTimeout`
+  is the important one: it measures the gap **between** chunks, so its 300s
+  default fires on a model that pauses mid-stream - extended reasoning, a long
+  tool round trip, a slow Genie or SQL step - and tears the socket down with
+  `UND_ERR_BODY_TIMEOUT`. The client sees a stream that simply stops.
+
+Because nothing times out, client disconnect is the backstop: each upstream
+request carries an `AbortSignal` tied to the response, so a cancelled turn
+(Ctrl-C, a closed tab, a killed CLI) releases the Databricks-side stream instead
+of leaking it. If a turn should have a deadline, send one from the client - the
+same as you would to OpenAI.
+
+An upstream failure that happens _after_ the response headers are sent is
+logged (`stream ended early`) rather than thrown, since the status line is
+already committed and cannot be turned into an HTTP error.
+
+## Rate Limits And 429 Backoff
+
+Databricks Foundation Model endpoints are pay-per-token with an account-level
+throughput ceiling. That ceiling is **not** exposed as a readable number, so the
+proxy cannot pace against it in advance - it can only react. An agentic client
+that bursts (Codex, for one) trips the limit and, left alone, retries a few
+times and gives up with "exceeded retry limit".
+
+So **by default the proxy absorbs the `429` itself**: instead of relaying it, it
+retries the upstream call in-proxy with exponential backoff and up to +50%
+jitter, honoring a server-sent `Retry-After` when present. Because the retry
+happens before any status line is written, it is transparent to both streaming
+and non-streaming callers - the client sees a slower success, not a 429. A fresh
+auth header is minted per attempt so a long backoff can't outlive the token, and
+a client disconnect during a backoff ends the wait immediately.
+
+Retries are exhausted after `maxRetries` attempts, at which point the final 429
+is relayed unchanged.
+
+Disable it (relay 429s straight through) with the flag or the env var:
+
+```sh
+dbx-tools-model-proxy --no-retry-429
+PROXY_RETRY_ON_429=false dbx-tools-model-proxy
+```
+
+Tune the policy with environment variables (all optional):
+
+| Variable | Default | Meaning |
+| --- | --- | --- |
+| `PROXY_RETRY_ON_429` | `true` | Master switch (loose boolean: `false`/`off`/`0`/`no`). `--no-retry-429` overrides it. |
+| `PROXY_RETRY_MAX` | `5` | Max retry attempts after the initial try. |
+| `PROXY_RETRY_BASE_MS` | `500` | First backoff, doubled each attempt. |
+| `PROXY_RETRY_MAX_MS` | `30000` | Ceiling for any single backoff, including a `Retry-After`. |
+
+Precedence is CLI flag → env → built-in default.
+
 ## Unsupported Request Fields
 
 Databricks Model Serving validates the chat body strictly, so a single
@@ -176,7 +242,7 @@ PROXY_DROP_FIELDS=some_new_field,another dbx-tools-model-proxy
 - `backend` - `DatabricksBackend`, auth, model resolution, and upstream request
   forwarding.
 - `server` - Express proxy app and `startProxyServer()`.
-- `defaults` - bind host, port, and invocation path constants.
+- `defaults` - bind host, port, and the 429-retry policy (`resolveRetryConfig`).
 
 Endpoint ranking and fuzzy matching come from
 [`@dbx-tools/model`](../../node/model).

package/index.ts

@@ -7,4 +7,5 @@ export * as cli from "./src/cli";
 export * as defaults from "./src/defaults";
 export * as server from "./src/server";
 export type { BackendOptions } from "./src/backend";
+export type { RetryConfig } from "./src/defaults";
 export type { ProxyServerOptions, StartProxyOptions } from "./src/server";

package/package.json

@@ -18,16 +18,17 @@
     "@databricks/sdk-experimental": "^0.17.0",
     "commander": "^15.0.0",
     "tsx": "^4.23.0",
-    "@dbx-tools/model": "0.3.33",
-    "@dbx-tools/shared-core": "0.3.33",
-    "@dbx-tools/shared-model": "0.3.33"
+    "undici": "^7.17.0",
+    "@dbx-tools/model": "0.3.35",
+    "@dbx-tools/shared-model": "0.3.35",
+    "@dbx-tools/shared-core": "0.3.35"
   },
   "main": "index.ts",
   "license": "UNLICENSED",
   "publishConfig": {
     "access": "public"
   },
-  "version": "0.3.33",
+  "version": "0.3.35",
   "types": "index.ts",
   "type": "module",
   "exports": {

package/src/cli.ts

@@ -24,7 +24,7 @@ import { classify, type ServingEndpointSummary } from "@dbx-tools/shared-model";
 import { Command, CommanderError } from "commander";
 
 import { DatabricksBackend, type BackendOptions } from "./backend";
-import { DEFAULT_BIND_HOST, DEFAULT_PORT } from "./defaults";
+import { DEFAULT_BIND_HOST, DEFAULT_PORT, resolveRetryConfig } from "./defaults";
 import { startProxyServer } from "./server";
 
 /**
@@ -47,6 +47,11 @@ interface ServeOpts extends CommonOpts {
   port: string;
   host: string;
   apiKey?: string;
+  /**
+   * Commander's negatable-flag value for `--no-retry-429`: `true` by default,
+   * `false` once the flag is passed. Off means relay upstream 429s unchanged.
+   */
+  retry429: boolean;
 }
 
 /** Map shared CLI flags onto {@link BackendOptions}. */
@@ -73,9 +78,13 @@ async function startProxy(
 ): Promise<{ backend: DatabricksBackend; server: Server; url: string }> {
   const backend = await DatabricksBackend.create(backendOptions(opts));
   const apiKey = opts.apiKey ?? process.env.PROXY_API_KEY;
+  // `--no-retry-429` (opts.retry429 === false) is the only explicit override;
+  // otherwise resolveRetryConfig layers PROXY_RETRY_* env then the on-default.
+  const retry = resolveRetryConfig(opts.retry429 === false ? { enabled: false } : {});
   const { server, url } = await startProxyServer(backend, {
     host: opts.host,
     port: Number(opts.port),
+    retry,
     ...(apiKey ? { apiKey } : {}),
   });
   return { backend, server, url };
@@ -99,7 +108,11 @@ export function buildProgram(): Command {
       .description("Local OpenAI-compatible proxy to Databricks Model Serving.")
       .option("-p, --port <port>", "port to listen on", String(DEFAULT_PORT))
       .option("-H, --host <host>", "address to bind", DEFAULT_BIND_HOST)
-      .option("-k, --api-key <key>", "require this bearer token from local clients"),
+      .option("-k, --api-key <key>", "require this bearer token from local clients")
+      .option(
+        "--no-retry-429",
+        "relay upstream 429s instead of retrying with backoff (default: retry)",
+      ),
   ).action(async (opts: ServeOpts) => {
     const { backend, url } = await startProxy(opts);
     process.stderr.write(`model-proxy -> ${backend.host}\n`);
@@ -113,6 +126,10 @@ export function buildProgram(): Command {
       .option("-p, --port <port>", "proxy port", String(DEFAULT_PORT))
       .option("-H, --host <host>", "proxy bind host", DEFAULT_BIND_HOST)
       .option("-m, --model <name>", "default model (fuzzy name ok)")
+      .option(
+        "--no-retry-429",
+        "relay upstream 429s instead of retrying with backoff (default: retry)",
+      )
       .option(
         "--client <cmd>",
         "terminal chat CLI to launch (run via your shell)",

package/src/defaults.ts

@@ -1,5 +1,7 @@
 // Default values for the local Databricks model proxy.
 
+import { object } from "@dbx-tools/shared-core";
+
 /**
  * Loopback address the proxy binds to by default. Keeps the OpenAI-compatible
  * endpoint private to the machine unless the operator explicitly opts into a
@@ -9,3 +11,63 @@ export const DEFAULT_BIND_HOST = "127.0.0.1";
 
 /** Default TCP port for the local proxy. */
 export const DEFAULT_PORT = 4000;
+
+/**
+ * Resolved policy for absorbing upstream `429 Too Many Requests` responses.
+ *
+ * Databricks Foundation Model endpoints are pay-per-token with an account-level
+ * throughput ceiling that is not exposed as a readable number, so the proxy
+ * cannot pace against it proactively - it can only react. When `enabled`, a 429
+ * is retried in-proxy with exponential backoff (honoring any `Retry-After`)
+ * instead of being relayed to the client, so an agentic caller like Codex sees
+ * a slow success rather than "exceeded retry limit".
+ */
+export interface RetryConfig {
+  /** Retry 429s in-proxy rather than relaying them. */
+  enabled: boolean;
+  /** Maximum retry attempts after the initial try (so N+1 total requests). */
+  maxRetries: number;
+  /** First backoff, doubled each attempt (before jitter and the `Retry-After` override). */
+  baseDelayMs: number;
+  /** Ceiling for any single backoff, including a server-sent `Retry-After`. */
+  maxDelayMs: number;
+}
+
+/** Built-in retry policy: on, with a bounded exponential backoff. */
+export const DEFAULT_RETRY: RetryConfig = {
+  enabled: true,
+  maxRetries: 5,
+  baseDelayMs: 500,
+  maxDelayMs: 30_000,
+};
+
+/** Positive integer from an env var, or `undefined` when unset/unparseable. */
+function envInt(name: string): number | undefined {
+  const raw = process.env[name];
+  if (raw === undefined || raw.trim() === "") return undefined;
+  const value = Number(raw);
+  return Number.isFinite(value) && value >= 0 ? Math.floor(value) : undefined;
+}
+
+/**
+ * Resolve the 429-retry policy from, in decreasing precedence: an explicit
+ * override (the CLI flag), environment variables, then {@link DEFAULT_RETRY}.
+ *
+ * `enabled` layers as: a CLI `--no-retry-429` (`override.enabled === false`)
+ * always wins; otherwise `PROXY_RETRY_ON_429` (loose boolean via
+ * {@link object.toBoolean}) may switch it off; otherwise it stays on. There is
+ * deliberately no enable flag - the default is on, so the only meaningful
+ * action is disabling.
+ *
+ * Tunables read `PROXY_RETRY_MAX`, `PROXY_RETRY_BASE_MS`, `PROXY_RETRY_MAX_MS`.
+ */
+export function resolveRetryConfig(override: Partial<RetryConfig> = {}): RetryConfig {
+  const envEnabled = object.toBoolean(process.env.PROXY_RETRY_ON_429);
+  const enabled = override.enabled ?? envEnabled ?? DEFAULT_RETRY.enabled;
+  return {
+    enabled,
+    maxRetries: override.maxRetries ?? envInt("PROXY_RETRY_MAX") ?? DEFAULT_RETRY.maxRetries,
+    baseDelayMs: override.baseDelayMs ?? envInt("PROXY_RETRY_BASE_MS") ?? DEFAULT_RETRY.baseDelayMs,
+    maxDelayMs: override.maxDelayMs ?? envInt("PROXY_RETRY_MAX_MS") ?? DEFAULT_RETRY.maxDelayMs,
+  };
+}

package/src/server.ts

@@ -32,17 +32,36 @@
  */
 
 import { createServer, type IncomingMessage, type Server, type ServerResponse } from "node:http";
-import { error, json, log } from "@dbx-tools/shared-core";
+import { async as asyncUtil, error, json, log } from "@dbx-tools/shared-core";
 import { classify, openaiChat, openaiResponses } from "@dbx-tools/shared-model";
+import { Agent } from "undici";
 
 import type { DatabricksBackend } from "./backend";
-import { DEFAULT_BIND_HOST, DEFAULT_PORT } from "./defaults";
+import { DEFAULT_BIND_HOST, DEFAULT_PORT, DEFAULT_RETRY, type RetryConfig } from "./defaults";
 
 const { chatToResponsesRequest, responseToChatCompletion, sanitizeOpenResponsesRequest } =
   openaiResponses;
 
 const logger = log.logger("model-proxy/server");
 
+/**
+ * Dispatcher for every upstream call, with undici's inactivity timeouts
+ * DISABLED (`0`).
+ *
+ * A proxied turn is only as fast as the model behind it, and undici's
+ * defaults (300s `headersTimeout`, 300s `bodyTimeout`) are both wrong here.
+ * `bodyTimeout` is the killer: it measures the gap BETWEEN chunks, so a model
+ * that thinks for a while mid-stream - extended reasoning, a long tool round
+ * trip, a slow Genie/SQL step - trips it and undici tears the socket down
+ * with `UND_ERR_BODY_TIMEOUT`, which the client sees as a truncated stream.
+ *
+ * The proxy has no opinion on how long a model may take, so it imposes no
+ * deadline: the stream ends when the upstream ends it, or when the client
+ * hangs up (see {@link requestAbortSignal}). Callers that DO want a deadline
+ * should send one, the same as they would to OpenAI.
+ */
+const upstreamDispatcher = new Agent({ headersTimeout: 0, bodyTimeout: 0 });
+
 /** POST routes forwarded to a serving endpoint (chat or Responses-only). */
 const PROXY_PATHS = new Set([
   "/v1/chat/completions",
@@ -82,6 +101,11 @@ export interface ProxyServerOptions {
    * for a loopback bind but should be paired with a key on a wider one.
    */
   apiKey?: string;
+  /**
+   * Policy for absorbing upstream 429s. Omitted defaults to {@link
+   * DEFAULT_RETRY} (retry on). See {@link RetryConfig}.
+   */
+  retry?: RetryConfig;
 }
 
 /** Options for {@link startProxyServer}, adding the listen address. */
@@ -95,9 +119,161 @@ export function createProxyServer(
   backend: DatabricksBackend,
   options: ProxyServerOptions = {},
 ): Server {
-  return createServer((req, res) => {
+  const server = createServer((req, res) => {
     void handleRequest(backend, options, req, res);
   });
+  // Match the upstream policy on the INBOUND side: never cut a client off
+  // mid-turn. Node's defaults (300s `requestTimeout`, 60s `headersTimeout`)
+  // are sized for ordinary web traffic, not for holding a streamed model
+  // response open, and a client that hits one gets a dead socket rather than
+  // an error it can act on. `0` disables each; the turn now ends when the
+  // model is done or the client disconnects.
+  server.requestTimeout = 0;
+  server.headersTimeout = 0;
+  server.timeout = 0;
+  return server;
+}
+
+/**
+ * Signal that aborts when the client hangs up.
+ *
+ * Removing the timeouts means nothing else will ever end a forgotten request,
+ * so the client's own disconnect becomes the only backstop against an upstream
+ * call running on with no one to receive it. Passing this to `fetch` is what
+ * makes a cancelled turn (Ctrl-C, a closed tab, a killed CLI) release the
+ * Databricks-side stream instead of leaking it for the life of the process.
+ *
+ * Keyed off the RESPONSE, not the request: `IncomingMessage` emits `close` as
+ * soon as its body has been fully consumed, which on every POST here happens
+ * before the upstream call is even made - watching that would abort each turn
+ * instantly. `ServerResponse` emits `close` when the response finishes or the
+ * connection drops, so the `writableFinished` check separates "client left"
+ * from "we replied".
+ */
+function clientAbortSignal(res: ServerResponse): AbortSignal {
+  const controller = new AbortController();
+  res.once("close", () => {
+    if (!res.writableFinished) controller.abort();
+  });
+  return controller.signal;
+}
+
+/**
+ * POST a JSON body to a serving endpoint on the no-timeout
+ * {@link upstreamDispatcher}, cancelled via `signal`.
+ *
+ * Every upstream call goes through here so the timeout and cancellation policy
+ * is stated once. `dispatcher` is an undici extension that the DOM
+ * `RequestInit` Node compiles against does not declare, hence the local
+ * intersection type rather than a cast at each call site.
+ */
+async function upstreamFetch(
+  url: string,
+  headers: Record<string, string>,
+  body: string,
+  signal: AbortSignal,
+): Promise<Response> {
+  const init: RequestInit = { method: "POST", headers, body, signal };
+  // `dispatcher` is an undici extension the DOM `RequestInit` doesn't declare,
+  // and it can't simply be typed on: `@types/node` pins its own `undici-types`
+  // copy, so the `Agent` from our `undici` dependency is a structurally
+  // identical but nominally different type than the one the global `fetch`
+  // signature expects. Attaching it structurally sidesteps a version skew that
+  // no cast expresses cleanly, and `fetch` reads it at runtime all the same.
+  Reflect.set(init, "dispatcher", upstreamDispatcher);
+  return fetch(url, init);
+}
+
+/**
+ * Delay (ms) a `429` response asks us to wait, from its `Retry-After` header,
+ * or `undefined` when absent/unparseable. Handles both header forms: an integer
+ * count of seconds, and an HTTP date to wait until (clamped to `>= 0`).
+ */
+export function parseRetryAfterMs(header: string | null, nowMs: number): number | undefined {
+  if (!header) return undefined;
+  const trimmed = header.trim();
+  if (/^\d+$/.test(trimmed)) return Number(trimmed) * 1000;
+  const when = Date.parse(trimmed);
+  if (Number.isNaN(when)) return undefined;
+  return Math.max(0, when - nowMs);
+}
+
+/**
+ * Backoff (ms) before retry `attempt` (0-based): exponential from
+ * `baseDelayMs`, capped at `maxDelayMs`, with up to +50% jitter so concurrent
+ * retries from one agent don't resynchronize into the next burst. A server-sent
+ * `Retry-After` wins outright (still capped), since the server knows better
+ * than our schedule when it will accept traffic again.
+ */
+export function backoffDelayMs(
+  attempt: number,
+  retry: RetryConfig,
+  retryAfterMs: number | undefined,
+  jitter: number,
+): number {
+  if (retryAfterMs !== undefined) return Math.min(retryAfterMs, retry.maxDelayMs);
+  const exponential = retry.baseDelayMs * 2 ** attempt;
+  const capped = Math.min(exponential, retry.maxDelayMs);
+  return Math.round(capped * (1 + jitter * 0.5));
+}
+
+/**
+ * {@link upstreamFetch} with in-proxy `429` handling. When `retry.enabled`, a
+ * 429 is retried up to `retry.maxRetries` times with {@link backoffDelayMs}
+ * backoff instead of being surfaced; the throttled response body is drained
+ * each time so the connection is released. The retry happens BEFORE the caller
+ * writes any status line, so it is transparent to both streaming and
+ * non-streaming paths - the client only ever sees the final response.
+ *
+ * Bails early (returning the 429) when retries are exhausted, when the client
+ * has hung up (`signal.aborted`), or when the response carries no `Retry-After`
+ * and we would otherwise loop blind past the cap.
+ */
+async function upstreamFetchRetrying(
+  backend: DatabricksBackend,
+  url: string,
+  headers: Record<string, string>,
+  body: unknown,
+  res: ServerResponse,
+  retry: RetryConfig,
+): Promise<Response> {
+  const signal = clientAbortSignal(res);
+  const payload = JSON.stringify(body);
+  let response = await upstreamFetch(url, headers, payload, signal);
+  if (!retry.enabled) return response;
+
+  for (let attempt = 0; response.status === 429 && attempt < retry.maxRetries; attempt++) {
+    const retryAfterMs = parseRetryAfterMs(response.headers.get("retry-after"), Date.now());
+    const delayMs = backoffDelayMs(attempt, retry, retryAfterMs, jitterFraction());
+    // Release the throttled response's socket before waiting; its body is an
+    // error payload we are choosing not to relay.
+    await response.body?.cancel().catch(() => {});
+    logger.warn("upstream 429; backing off", {
+      attempt: attempt + 1,
+      maxRetries: retry.maxRetries,
+      delayMs,
+      ...(retryAfterMs !== undefined ? { retryAfterMs } : {}),
+    });
+    try {
+      await asyncUtil.sleep(delayMs, signal);
+    } catch {
+      // Client hung up mid-wait: return the last 429 rather than issue a
+      // request no one is listening for.
+      return response;
+    }
+    // Fresh auth header per attempt so a long backoff can't outlive the token.
+    const retryHeaders = { ...(await backend.authHeaders()), ...headers };
+    response = await upstreamFetch(url, retryHeaders, payload, signal);
+  }
+  return response;
+}
+
+/**
+ * Jitter fraction in `[0, 1)`. Split out so tests can stub determinism;
+ * `Math.random` is fine for spreading retry timing.
+ */
+function jitterFraction(): number {
+  return Math.random();
 }
 
 /**
@@ -149,12 +325,13 @@ async function handleRequest(
       await handleModels(backend, res, wantsCodexShape);
       return;
     }
+    const retry = options.retry ?? DEFAULT_RETRY;
     if (req.method === "POST" && RESPONSES_PATHS.has(routePath)) {
-      await handleResponses(backend, req, res);
+      await handleResponses(backend, req, res, retry);
       return;
     }
     if (req.method === "POST" && PROXY_PATHS.has(routePath)) {
-      await handleProxy(backend, req, res);
+      await handleProxy(backend, req, res, retry);
       return;
     }
     sendJson(
@@ -221,6 +398,7 @@ async function handleProxy(
   backend: DatabricksBackend,
   req: IncomingMessage,
   res: ServerResponse,
+  retry: RetryConfig,
 ): Promise<void> {
   const body = await readJsonBody(req);
   const requested = typeof body.model === "string" ? body.model : undefined;
@@ -233,7 +411,15 @@ async function handleProxy(
   body.model = resolved.modelId;
 
   if (backend.isResponsesOnly(resolved.modelId)) {
-    await proxyChatViaResponses(backend, body, requested, resolved.modelId, resolved.matched, res);
+    await proxyChatViaResponses(
+      backend,
+      body,
+      requested,
+      resolved.modelId,
+      resolved.matched,
+      res,
+      retry,
+    );
     return;
   }
 
@@ -256,11 +442,14 @@ async function handleProxy(
     ...(dropped.length > 0 ? { dropped } : {}),
   });
 
-  const upstream = await fetch(backend.invocationsUrl(resolved.modelId), {
-    method: "POST",
+  const upstream = await upstreamFetchRetrying(
+    backend,
+    backend.invocationsUrl(resolved.modelId),
     headers,
-    body: JSON.stringify(body),
-  });
+    body,
+    res,
+    retry,
+  );
 
   res.writeHead(upstream.status, {
     "content-type": upstream.headers.get("content-type") ?? "application/json",
@@ -281,6 +470,7 @@ async function proxyChatViaResponses(
   modelId: string,
   matched: boolean,
   res: ServerResponse,
+  retry: RetryConfig,
 ): Promise<void> {
   const { responses, stream } = chatToResponsesRequest(body);
   responses.model = modelId;
@@ -314,11 +504,7 @@ async function proxyChatViaResponses(
     stream: false,
   });
 
-  const upstream = await fetch(upstreamUrl, {
-    method: "POST",
-    headers,
-    body: JSON.stringify(forward),
-  });
+  const upstream = await upstreamFetchRetrying(backend, upstreamUrl, headers, forward, res, retry);
 
   if (!upstream.ok) {
     const text = await upstream.text();
@@ -347,6 +533,7 @@ async function handleResponses(
   backend: DatabricksBackend,
   req: IncomingMessage,
   res: ServerResponse,
+  retry: RetryConfig,
 ): Promise<void> {
   const body = await readJsonBody(req);
   const requested = typeof body.model === "string" ? body.model : undefined;
@@ -383,11 +570,7 @@ async function handleResponses(
     ...(stripped ? { strippedNonFunctionTools: true } : {}),
   });
 
-  const upstream = await fetch(upstreamUrl, {
-    method: "POST",
-    headers,
-    body: JSON.stringify(forward),
-  });
+  const upstream = await upstreamFetchRetrying(backend, upstreamUrl, headers, forward, res, retry);
 
   res.writeHead(upstream.status, {
     "content-type": upstream.headers.get("content-type") ?? "application/json",
@@ -396,7 +579,16 @@ async function handleResponses(
   await streamBody(upstream, res);
 }
 
-/** Pump an upstream `fetch` Response body to the Node response, chunk by chunk. */
+/**
+ * Pump an upstream `fetch` Response body to the Node response, chunk by chunk.
+ *
+ * A mid-stream upstream failure is logged rather than thrown: the status line
+ * and some number of chunks have already gone out, so there is no way left to
+ * turn it into an HTTP error and the only honest move is to end the response
+ * and say why in the log. Without this the read loop rejected into the caller's
+ * catch, which found `headersSent` and ended the response silently - making an
+ * upstream teardown indistinguishable from a clean finish.
+ */
 async function streamBody(upstream: Response, res: ServerResponse): Promise<void> {
   const body = upstream.body;
   if (!body) {
@@ -411,6 +603,13 @@ async function streamBody(upstream: Response, res: ServerResponse): Promise<void
       if (value && !res.writableEnded) res.write(Buffer.from(value));
       if (res.writableEnded) break;
     }
+  } catch (err) {
+    // An aborted request is the expected shape of a client hanging up, not a
+    // fault worth reporting at error level.
+    const aborted = res.writableEnded || !res.writable;
+    logger[aborted ? "info" : "error"]("stream ended early", {
+      error: error.errorMessage(err),
+    });
   } finally {
     await reader.cancel().catch(() => {});
     res.end();

package/test/server.test.ts

@@ -0,0 +1,173 @@
+import assert from "node:assert/strict";
+import { createServer, type Server } from "node:http";
+import { describe, it } from "node:test";
+
+import type { DatabricksBackend } from "../src/backend";
+import { DEFAULT_RETRY, resolveRetryConfig } from "../src/defaults";
+import { backoffDelayMs, createProxyServer, parseRetryAfterMs } from "../src/server";
+
+/** Listen on an ephemeral port and resolve the bound port. */
+async function listen(server: Server): Promise<number> {
+  await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", () => resolve()));
+  const address = server.address();
+  assert.ok(typeof address === "object" && address);
+  return address.port;
+}
+
+describe("createProxyServer timeouts", () => {
+  it("disables every inbound timeout that could cut a stream short", () => {
+    const server = createProxyServer({} as DatabricksBackend);
+    // Node's defaults (300s request, 60s headers) are sized for ordinary web
+    // traffic and would kill a long model turn mid-stream.
+    assert.equal(server.requestTimeout, 0);
+    assert.equal(server.headersTimeout, 0);
+    assert.equal(server.timeout, 0);
+    server.close();
+  });
+
+  it("differs from a stock node server, so the override is doing the work", () => {
+    const stock = createServer();
+    assert.notEqual(stock.requestTimeout, 0);
+    assert.notEqual(stock.headersTimeout, 0);
+    stock.close();
+  });
+});
+
+describe("upstream dispatcher", () => {
+  /**
+   * Serve one chunk, stall past the supplied body timeout, then finish. Proves
+   * `bodyTimeout` measures the gap BETWEEN chunks - the failure mode that made
+   * long-thinking turns look like dropped streams.
+   */
+  async function readThroughStall(bodyTimeout: number): Promise<string> {
+    const upstream = createServer((_req, res) => {
+      res.writeHead(200, { "content-type": "text/event-stream" });
+      res.write("data: first\n\n");
+      setTimeout(() => {
+        res.write("data: second\n\n");
+        res.end();
+      }, 1_000);
+    });
+    const port = await listen(upstream);
+    try {
+      const { Agent } = await import("undici");
+      const init: RequestInit = { method: "GET" };
+      Reflect.set(init, "dispatcher", new Agent({ bodyTimeout, headersTimeout: 0 }));
+      const response = await fetch(`http://127.0.0.1:${port}/`, init);
+      let body = "";
+      for await (const chunk of response.body as AsyncIterable<Uint8Array>) {
+        body += Buffer.from(chunk).toString();
+      }
+      return body;
+    } finally {
+      upstream.close();
+    }
+  }
+
+  it("aborts a stalled stream when bodyTimeout is set", async () => {
+    await assert.rejects(() => readThroughStall(100));
+  });
+
+  it("survives the same stall with bodyTimeout disabled, as the proxy configures it", async () => {
+    assert.equal(await readThroughStall(0), "data: first\n\ndata: second\n\n");
+  });
+});
+
+describe("parseRetryAfterMs", () => {
+  const now = Date.UTC(2026, 0, 1, 0, 0, 0);
+
+  it("reads an integer count of seconds", () => {
+    assert.equal(parseRetryAfterMs("2", now), 2000);
+    assert.equal(parseRetryAfterMs("  30 ", now), 30_000);
+  });
+
+  it("reads an HTTP date as ms-until, clamped at zero", () => {
+    const future = new Date(now + 5000).toUTCString();
+    assert.equal(parseRetryAfterMs(future, now), 5000);
+    const past = new Date(now - 5000).toUTCString();
+    assert.equal(parseRetryAfterMs(past, now), 0);
+  });
+
+  it("returns undefined for a missing or unparseable header", () => {
+    assert.equal(parseRetryAfterMs(null, now), undefined);
+    assert.equal(parseRetryAfterMs("soon", now), undefined);
+  });
+});
+
+describe("backoffDelayMs", () => {
+  const retry = { enabled: true, maxRetries: 5, baseDelayMs: 500, maxDelayMs: 30_000 };
+
+  it("grows exponentially from the base with jitter=0", () => {
+    assert.equal(backoffDelayMs(0, retry, undefined, 0), 500);
+    assert.equal(backoffDelayMs(1, retry, undefined, 0), 1000);
+    assert.equal(backoffDelayMs(2, retry, undefined, 0), 2000);
+  });
+
+  it("caps the exponential at maxDelayMs", () => {
+    assert.equal(backoffDelayMs(20, retry, undefined, 0), retry.maxDelayMs);
+  });
+
+  it("adds up to +50% jitter", () => {
+    assert.equal(backoffDelayMs(0, retry, undefined, 1), 750);
+  });
+
+  it("lets a Retry-After win outright, still capped", () => {
+    assert.equal(backoffDelayMs(0, retry, 3000, 1), 3000);
+    assert.equal(backoffDelayMs(0, retry, 99_000, 0), retry.maxDelayMs);
+  });
+});
+
+describe("resolveRetryConfig", () => {
+  /** Run `fn` with `env` applied to `process.env`, restoring after. */
+  function withEnv(env: Record<string, string | undefined>, fn: () => void): void {
+    const saved = new Map<string, string | undefined>();
+    for (const key of Object.keys(env)) {
+      saved.set(key, process.env[key]);
+      if (env[key] === undefined) delete process.env[key];
+      else process.env[key] = env[key];
+    }
+    try {
+      fn();
+    } finally {
+      for (const [key, value] of saved) {
+        if (value === undefined) delete process.env[key];
+        else process.env[key] = value;
+      }
+    }
+  }
+
+  it("defaults to on with the built-in policy", () => {
+    withEnv(
+      {
+        PROXY_RETRY_ON_429: undefined,
+        PROXY_RETRY_MAX: undefined,
+        PROXY_RETRY_BASE_MS: undefined,
+        PROXY_RETRY_MAX_MS: undefined,
+      },
+      () => assert.deepEqual(resolveRetryConfig(), DEFAULT_RETRY),
+    );
+  });
+
+  it("honors a loose PROXY_RETRY_ON_429 to switch it off", () => {
+    withEnv({ PROXY_RETRY_ON_429: "no" }, () => assert.equal(resolveRetryConfig().enabled, false));
+    withEnv({ PROXY_RETRY_ON_429: "off" }, () => assert.equal(resolveRetryConfig().enabled, false));
+  });
+
+  it("reads numeric tunables from env", () => {
+    withEnv(
+      { PROXY_RETRY_MAX: "9", PROXY_RETRY_BASE_MS: "250", PROXY_RETRY_MAX_MS: "60000" },
+      () => {
+        const config = resolveRetryConfig();
+        assert.equal(config.maxRetries, 9);
+        assert.equal(config.baseDelayMs, 250);
+        assert.equal(config.maxDelayMs, 60_000);
+      },
+    );
+  });
+
+  it("lets an explicit override beat env (the CLI --no-retry-429 path)", () => {
+    withEnv({ PROXY_RETRY_ON_429: "true" }, () =>
+      assert.equal(resolveRetryConfig({ enabled: false }).enabled, false),
+    );
+  });
+});