@dbx-tools/cli-model-proxy 0.3.32 → 0.3.34

package/README.md

@@ -145,6 +145,35 @@ 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.
+
 ## Unsupported Request Fields
 
 Databricks Model Serving validates the chat body strictly, so a single

package/package.json

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

package/src/server.ts

@@ -34,6 +34,7 @@
 import { createServer, type IncomingMessage, type Server, type ServerResponse } from "node:http";
 import { 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";
@@ -43,6 +44,24 @@ const { chatToResponsesRequest, responseToChatCompletion, sanitizeOpenResponsesR
 
 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",
@@ -95,9 +114,74 @@ 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 by the client hanging up.
+ *
+ * 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: unknown,
+  res: ServerResponse,
+): Promise<Response> {
+  const init: RequestInit = {
+    method: "POST",
+    headers,
+    body: JSON.stringify(body),
+    signal: clientAbortSignal(res),
+  };
+  // `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);
 }
 
 /**
@@ -256,11 +340,12 @@ async function handleProxy(
     ...(dropped.length > 0 ? { dropped } : {}),
   });
 
-  const upstream = await fetch(backend.invocationsUrl(resolved.modelId), {
-    method: "POST",
+  const upstream = await upstreamFetch(
+    backend.invocationsUrl(resolved.modelId),
     headers,
-    body: JSON.stringify(body),
-  });
+    body,
+    res,
+  );
 
   res.writeHead(upstream.status, {
     "content-type": upstream.headers.get("content-type") ?? "application/json",
@@ -314,11 +399,7 @@ async function proxyChatViaResponses(
     stream: false,
   });
 
-  const upstream = await fetch(upstreamUrl, {
-    method: "POST",
-    headers,
-    body: JSON.stringify(forward),
-  });
+  const upstream = await upstreamFetch(upstreamUrl, headers, forward, res);
 
   if (!upstream.ok) {
     const text = await upstream.text();
@@ -383,11 +464,7 @@ async function handleResponses(
     ...(stripped ? { strippedNonFunctionTools: true } : {}),
   });
 
-  const upstream = await fetch(upstreamUrl, {
-    method: "POST",
-    headers,
-    body: JSON.stringify(forward),
-  });
+  const upstream = await upstreamFetch(upstreamUrl, headers, forward, res);
 
   res.writeHead(upstream.status, {
     "content-type": upstream.headers.get("content-type") ?? "application/json",
@@ -396,7 +473,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 +497,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,73 @@
+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 { createProxyServer } 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");
+  });
+});