fix(agent): harden undici stream timeouts for long openai-completions runs

src/infra/net/undici-global-dispatcher.test.ts

@@ -0,0 +1,138 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const {
+  Agent,
+  EnvHttpProxyAgent,
+  ProxyAgent,
+  getGlobalDispatcher,
+  setGlobalDispatcher,
+  setCurrentDispatcher,
+  getCurrentDispatcher,
+  getDefaultAutoSelectFamily,
+} = vi.hoisted(() => {
+  class Agent {
+    constructor(public readonly options?: Record<string, unknown>) {}
+  }
+
+  class EnvHttpProxyAgent {
+    constructor(public readonly options?: Record<string, unknown>) {}
+  }
+
+  class ProxyAgent {
+    constructor(public readonly url: string) {}
+  }
+
+  let currentDispatcher: unknown = new Agent();
+
+  const getGlobalDispatcher = vi.fn(() => currentDispatcher);
+  const setGlobalDispatcher = vi.fn((next: unknown) => {
+    currentDispatcher = next;
+  });
+  const setCurrentDispatcher = (next: unknown) => {
+    currentDispatcher = next;
+  };
+  const getCurrentDispatcher = () => currentDispatcher;
+  const getDefaultAutoSelectFamily = vi.fn(() => undefined as boolean | undefined);
+
+  return {
+    Agent,
+    EnvHttpProxyAgent,
+    ProxyAgent,
+    getGlobalDispatcher,
+    setGlobalDispatcher,
+    setCurrentDispatcher,
+    getCurrentDispatcher,
+    getDefaultAutoSelectFamily,
+  };
+});
+
+vi.mock("undici", () => ({
+  Agent,
+  EnvHttpProxyAgent,
+  getGlobalDispatcher,
+  setGlobalDispatcher,
+}));
+
+vi.mock("node:net", () => ({
+  getDefaultAutoSelectFamily,
+}));
+
+import {
+  DEFAULT_UNDICI_STREAM_TIMEOUT_MS,
+  ensureGlobalUndiciStreamTimeouts,
+  resetGlobalUndiciStreamTimeoutsForTests,
+} from "./undici-global-dispatcher.js";
+
+describe("ensureGlobalUndiciStreamTimeouts", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    resetGlobalUndiciStreamTimeoutsForTests();
+    setCurrentDispatcher(new Agent());
+    getDefaultAutoSelectFamily.mockReturnValue(undefined);
+  });
+
+  it("replaces default Agent dispatcher with extended stream timeouts", () => {
+    getDefaultAutoSelectFamily.mockReturnValue(true);
+
+    ensureGlobalUndiciStreamTimeouts();
+
+    expect(setGlobalDispatcher).toHaveBeenCalledTimes(1);
+    const next = getCurrentDispatcher() as { options?: Record<string, unknown> };
+    expect(next).toBeInstanceOf(Agent);
+    expect(next.options?.bodyTimeout).toBe(DEFAULT_UNDICI_STREAM_TIMEOUT_MS);
+    expect(next.options?.headersTimeout).toBe(DEFAULT_UNDICI_STREAM_TIMEOUT_MS);
+    expect(next.options?.connect).toEqual({
+      autoSelectFamily: true,
+      autoSelectFamilyAttemptTimeout: 300,
+    });
+  });
+
+  it("replaces EnvHttpProxyAgent dispatcher while preserving env-proxy mode", () => {
+    getDefaultAutoSelectFamily.mockReturnValue(false);
+    setCurrentDispatcher(new EnvHttpProxyAgent());
+
+    ensureGlobalUndiciStreamTimeouts();
+
+    expect(setGlobalDispatcher).toHaveBeenCalledTimes(1);
+    const next = getCurrentDispatcher() as { options?: Record<string, unknown> };
+    expect(next).toBeInstanceOf(EnvHttpProxyAgent);
+    expect(next.options?.bodyTimeout).toBe(DEFAULT_UNDICI_STREAM_TIMEOUT_MS);
+    expect(next.options?.headersTimeout).toBe(DEFAULT_UNDICI_STREAM_TIMEOUT_MS);
+    expect(next.options?.connect).toEqual({
+      autoSelectFamily: false,
+      autoSelectFamilyAttemptTimeout: 300,
+    });
+  });
+
+  it("does not override unsupported custom proxy dispatcher types", () => {
+    setCurrentDispatcher(new ProxyAgent("http://proxy.test:8080"));
+
+    ensureGlobalUndiciStreamTimeouts();
+
+    expect(setGlobalDispatcher).not.toHaveBeenCalled();
+  });
+
+  it("is idempotent for unchanged dispatcher kind and network policy", () => {
+    getDefaultAutoSelectFamily.mockReturnValue(true);
+
+    ensureGlobalUndiciStreamTimeouts();
+    ensureGlobalUndiciStreamTimeouts();
+
+    expect(setGlobalDispatcher).toHaveBeenCalledTimes(1);
+  });
+
+  it("re-applies when autoSelectFamily decision changes", () => {
+    getDefaultAutoSelectFamily.mockReturnValue(true);
+    ensureGlobalUndiciStreamTimeouts();
+
+    getDefaultAutoSelectFamily.mockReturnValue(false);
+    ensureGlobalUndiciStreamTimeouts();
+
+    expect(setGlobalDispatcher).toHaveBeenCalledTimes(2);
+    const next = getCurrentDispatcher() as { options?: Record<string, unknown> };
+    expect(next.options?.connect).toEqual({
+      autoSelectFamily: false,
+      autoSelectFamilyAttemptTimeout: 300,
+    });
+  });
+});

src/infra/net/undici-global-dispatcher.ts

@@ -0,0 +1,113 @@
+import * as net from "node:net";
+import { Agent, EnvHttpProxyAgent, getGlobalDispatcher, setGlobalDispatcher } from "undici";
+
+export const DEFAULT_UNDICI_STREAM_TIMEOUT_MS = 30 * 60 * 1000;
+
+const AUTO_SELECT_FAMILY_ATTEMPT_TIMEOUT_MS = 300;
+
+let lastAppliedDispatcherKey: string | null = null;
+
+type DispatcherKind = "agent" | "env-proxy" | "unsupported";
+
+function resolveDispatcherKind(dispatcher: unknown): DispatcherKind {
+  const ctorName = (dispatcher as { constructor?: { name?: string } })?.constructor?.name;
+  if (typeof ctorName !== "string" || ctorName.length === 0) {
+    return "unsupported";
+  }
+  if (ctorName.includes("EnvHttpProxyAgent")) {
+    return "env-proxy";
+  }
+  if (ctorName.includes("ProxyAgent")) {
+    return "unsupported";
+  }
+  if (ctorName.includes("Agent")) {
+    return "agent";
+  }
+  return "unsupported";
+}
+
+function resolveAutoSelectFamily(): boolean | undefined {
+  if (typeof net.getDefaultAutoSelectFamily !== "function") {
+    return undefined;
+  }
+  try {
+    return net.getDefaultAutoSelectFamily();
+  } catch {
+    return undefined;
+  }
+}
+
+function resolveConnectOptions(
+  autoSelectFamily: boolean | undefined,
+): { autoSelectFamily: boolean; autoSelectFamilyAttemptTimeout: number } | undefined {
+  if (autoSelectFamily === undefined) {
+    return undefined;
+  }
+  return {
+    autoSelectFamily,
+    autoSelectFamilyAttemptTimeout: AUTO_SELECT_FAMILY_ATTEMPT_TIMEOUT_MS,
+  };
+}
+
+function resolveDispatcherKey(params: {
+  kind: DispatcherKind;
+  timeoutMs: number;
+  autoSelectFamily: boolean | undefined;
+}): string {
+  const autoSelectToken =
+    params.autoSelectFamily === undefined ? "na" : params.autoSelectFamily ? "on" : "off";
+  return `${params.kind}:${params.timeoutMs}:${autoSelectToken}`;
+}
+
+export function ensureGlobalUndiciStreamTimeouts(opts?: { timeoutMs?: number }): void {
+  const timeoutMsRaw = opts?.timeoutMs ?? DEFAULT_UNDICI_STREAM_TIMEOUT_MS;
+  const timeoutMs = Math.max(1, Math.floor(timeoutMsRaw));
+  if (!Number.isFinite(timeoutMsRaw)) {
+    return;
+  }
+
+  let dispatcher: unknown;
+  try {
+    dispatcher = getGlobalDispatcher();
+  } catch {
+    return;
+  }
+
+  const kind = resolveDispatcherKind(dispatcher);
+  if (kind === "unsupported") {
+    return;
+  }
+
+  const autoSelectFamily = resolveAutoSelectFamily();
+  const nextKey = resolveDispatcherKey({ kind, timeoutMs, autoSelectFamily });
+  if (lastAppliedDispatcherKey === nextKey) {
+    return;
+  }
+
+  const connect = resolveConnectOptions(autoSelectFamily);
+  try {
+    if (kind === "env-proxy") {
+      const proxyOptions = {
+        bodyTimeout: timeoutMs,
+        headersTimeout: timeoutMs,
+        ...(connect ? { connect } : {}),
+      } as ConstructorParameters<typeof EnvHttpProxyAgent>[0];
+      setGlobalDispatcher(new EnvHttpProxyAgent(proxyOptions));
+    } else {
+      setGlobalDispatcher(
+        new Agent({
+          bodyTimeout: timeoutMs,
+          headersTimeout: timeoutMs,
+          ...(connect ? { connect } : {}),
+        }),
+      );
+    }
+    lastAppliedDispatcherKey = nextKey;
+  } catch {
+    // Best-effort hardening only.
+  }
+}
+
+export function resetGlobalUndiciStreamTimeoutsForTests(): void {
+  lastAppliedDispatcherKey = null;
+}