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

package/README.md

@@ -174,6 +174,43 @@ 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
@@ -205,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

@@ -19,16 +19,16 @@
     "commander": "^15.0.0",
     "tsx": "^4.23.0",
     "undici": "^7.17.0",
-    "@dbx-tools/model": "0.3.34",
-    "@dbx-tools/shared-core": "0.3.34",
-    "@dbx-tools/shared-model": "0.3.34"
+    "@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.34",
+  "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,12 +32,12 @@
  */
 
 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;
@@ -101,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. */
@@ -155,7 +160,7 @@ function clientAbortSignal(res: ServerResponse): AbortSignal {
 
 /**
  * POST a JSON body to a serving endpoint on the no-timeout
- * {@link upstreamDispatcher}, cancelled by the client hanging up.
+ * {@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
@@ -165,15 +170,10 @@ function clientAbortSignal(res: ServerResponse): AbortSignal {
 async function upstreamFetch(
   url: string,
   headers: Record<string, string>,
-  body: unknown,
-  res: ServerResponse,
+  body: string,
+  signal: AbortSignal,
 ): Promise<Response> {
-  const init: RequestInit = {
-    method: "POST",
-    headers,
-    body: JSON.stringify(body),
-    signal: clientAbortSignal(res),
-  };
+  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
@@ -184,6 +184,98 @@ async function upstreamFetch(
   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();
+}
+
 /**
  * Build and start the proxy, resolving once it is accepting connections.
  * Returns the server and its base URL (with the actually-bound port, so a
@@ -233,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(
@@ -305,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;
@@ -317,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;
   }
 
@@ -340,11 +442,13 @@ async function handleProxy(
     ...(dropped.length > 0 ? { dropped } : {}),
   });
 
-  const upstream = await upstreamFetch(
+  const upstream = await upstreamFetchRetrying(
+    backend,
     backend.invocationsUrl(resolved.modelId),
     headers,
     body,
     res,
+    retry,
   );
 
   res.writeHead(upstream.status, {
@@ -366,6 +470,7 @@ async function proxyChatViaResponses(
   modelId: string,
   matched: boolean,
   res: ServerResponse,
+  retry: RetryConfig,
 ): Promise<void> {
   const { responses, stream } = chatToResponsesRequest(body);
   responses.model = modelId;
@@ -399,7 +504,7 @@ async function proxyChatViaResponses(
     stream: false,
   });
 
-  const upstream = await upstreamFetch(upstreamUrl, headers, forward, res);
+  const upstream = await upstreamFetchRetrying(backend, upstreamUrl, headers, forward, res, retry);
 
   if (!upstream.ok) {
     const text = await upstream.text();
@@ -428,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;
@@ -464,7 +570,7 @@ async function handleResponses(
     ...(stripped ? { strippedNonFunctionTools: true } : {}),
   });
 
-  const upstream = await upstreamFetch(upstreamUrl, headers, forward, res);
+  const upstream = await upstreamFetchRetrying(backend, upstreamUrl, headers, forward, res, retry);
 
   res.writeHead(upstream.status, {
     "content-type": upstream.headers.get("content-type") ?? "application/json",

package/test/server.test.ts

@@ -3,7 +3,8 @@ import { createServer, type Server } from "node:http";
 import { describe, it } from "node:test";
 
 import type { DatabricksBackend } from "../src/backend";
-import { createProxyServer } from "../src/server";
+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> {
@@ -71,3 +72,102 @@ describe("upstream dispatcher", () => {
     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),
+    );
+  });
+});