From 5e4b0d3beaad09dd82a9b0bb4bf2cc0b0da3cf3f Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 27 Jul 2026 07:42:01 -0400 Subject: [PATCH] fix(cron): prevent local model preflight socket leaks (#114540) Co-authored-by: hugenshen <16300669+hugenshen@users.noreply.github.com> --- .../model-preflight.runtime.test.ts | 143 ++++++++++++++++ .../model-preflight.runtime.transport.test.ts | 162 ++++++++++++++++++ .../isolated-agent/model-preflight.runtime.ts | 5 + 3 files changed, 310 insertions(+) create mode 100644 src/cron/isolated-agent/model-preflight.runtime.transport.test.ts diff --git a/src/cron/isolated-agent/model-preflight.runtime.test.ts b/src/cron/isolated-agent/model-preflight.runtime.test.ts index 3e31890f6643..49ffdcab2340 100644 --- a/src/cron/isolated-agent/model-preflight.runtime.test.ts +++ b/src/cron/isolated-agent/model-preflight.runtime.test.ts @@ -1,5 +1,6 @@ // Runtime model preflight tests cover provider/model checks before cron execution. import { beforeEach, describe, expect, it, vi } from "vitest"; +import { withTestTimeout } from "../../../test/helpers/promise.js"; const { fetchWithSsrFGuardMock } = vi.hoisted(() => ({ fetchWithSsrFGuardMock: vi.fn(), @@ -87,6 +88,148 @@ describe("preflightCronModelProvider", () => { expect(request.timeoutMs).toBe(2500); }); + it("starts unread-body cancellation before release without waiting for a split stream", async () => { + const cleanupOrder: string[] = []; + const cancel = vi.fn(() => { + cleanupOrder.push("cancel"); + return new Promise(() => {}); + }); + const release = vi.fn(async () => { + cleanupOrder.push("release"); + }); + fetchWithSsrFGuardMock.mockResolvedValueOnce({ + response: { status: 200, bodyUsed: false, body: { cancel } }, + release, + }); + const cfg = { + models: { + providers: { + vllm: { + api: "openai-completions" as const, + baseUrl: "http://127.0.0.1:8000/v1", + models: [], + }, + }, + }, + }; + + const result = await withTestTimeout( + preflightCronModelProvider({ cfg, provider: "vllm", model: "llama" }), + 1_000, + "cron provider preflight waited for unread response-body cancellation", + ); + const cached = await preflightCronModelProvider({ + cfg, + provider: "vllm", + model: "llama-cached", + }); + + expect(result).toEqual({ status: "available" }); + expect(cached).toEqual({ status: "available" }); + expect(cancel).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + expect(fetchWithSsrFGuardMock).toHaveBeenCalledOnce(); + expect(cleanupOrder).toEqual(["cancel", "release"]); + }); + + it("keeps a reachable provider available when response cancellation rejects", async () => { + const cancel = vi.fn(async () => { + throw new Error("provider response was already closed"); + }); + const release = vi.fn(async () => {}); + fetchWithSsrFGuardMock.mockResolvedValueOnce({ + response: { status: 401, bodyUsed: false, body: { cancel } }, + release, + }); + + const result = await preflightCronModelProvider({ + cfg: { + models: { + providers: { + vllm: { + api: "openai-completions", + baseUrl: "http://127.0.0.1:8000/v1", + models: [], + }, + }, + }, + }, + provider: "vllm", + model: "llama", + }); + + expect(result).toEqual({ status: "available" }); + expect(cancel).toHaveBeenCalledOnce(); + expect(release).toHaveBeenCalledOnce(); + }); + + it("does not cancel a response body that has already been consumed", async () => { + const cancel = vi.fn(async () => {}); + const release = vi.fn(async () => {}); + fetchWithSsrFGuardMock.mockResolvedValueOnce({ + response: { status: 200, bodyUsed: true, body: { cancel } }, + release, + }); + + const result = await preflightCronModelProvider({ + cfg: { + models: { + providers: { + vllm: { + api: "openai-completions", + baseUrl: "http://127.0.0.1:8000/v1", + models: [], + }, + }, + }, + }, + provider: "vllm", + model: "llama", + }); + + expect(result).toEqual({ status: "available" }); + expect(cancel).not.toHaveBeenCalled(); + expect(release).toHaveBeenCalledOnce(); + }); + + it("cancels and releases every response during concurrent local-provider probes", async () => { + const cancel = vi.fn(async () => {}); + const release = vi.fn(async () => {}); + fetchWithSsrFGuardMock.mockImplementation(async () => ({ + response: { status: 200, bodyUsed: false, body: { cancel } }, + release, + })); + + const results = await withTestTimeout( + Promise.all( + Array.from({ length: 32 }, (_, index) => + preflightCronModelProvider({ + cfg: { + models: { + providers: { + vllm: { + api: "openai-completions", + baseUrl: `http://127.0.0.1:${18_000 + index}/v1`, + models: [], + }, + }, + }, + }, + provider: "vllm", + model: `model-${index}`, + }), + ), + ), + 1_000, + "concurrent cron provider preflights did not release their response bodies", + ); + + expect(results).toEqual(Array.from({ length: 32 }, () => ({ status: "available" }))); + expect(fetchWithSsrFGuardMock).toHaveBeenCalledTimes(32); + expect(cancel).toHaveBeenCalledTimes(32); + expect(release).toHaveBeenCalledTimes(32); + }); + it("marks unreachable local Ollama endpoints unavailable and caches the result", async () => { fetchWithSsrFGuardMock.mockRejectedValueOnce(new Error("ECONNREFUSED")); diff --git a/src/cron/isolated-agent/model-preflight.runtime.transport.test.ts b/src/cron/isolated-agent/model-preflight.runtime.transport.test.ts new file mode 100644 index 000000000000..052cdc0291a1 --- /dev/null +++ b/src/cron/isolated-agent/model-preflight.runtime.transport.test.ts @@ -0,0 +1,162 @@ +import { createServer, type Server } from "node:http"; +import type { AddressInfo } from "node:net"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { withTestTimeout } from "../../../test/helpers/promise.js"; +import { + preflightCronModelProvider, + resetCronModelProviderPreflightCacheForTest, +} from "./model-preflight.runtime.js"; + +type StreamingProviderServer = { + baseUrl: string; + requestedPaths: string[]; + socketClosures: Promise[]; +}; + +const activeServers = new Set(); + +async function startStreamingProviderServer(status = 200): Promise { + const requestedPaths: string[] = []; + const socketClosures: Promise[] = []; + const server = createServer((request, response) => { + requestedPaths.push(request.url ?? ""); + socketClosures.push( + new Promise((resolve) => { + request.socket.once("close", resolve); + }), + ); + response.writeHead(status, { "content-type": "application/json" }); + // Keep the real HTTP response open: the probe needs only its status. + response.write('{"models":['); + }); + + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(0, "127.0.0.1", () => { + server.removeListener("error", reject); + resolve(); + }); + }); + activeServers.add(server); + const address = server.address() as AddressInfo; + + return { + baseUrl: `http://127.0.0.1:${address.port}`, + requestedPaths, + socketClosures, + }; +} + +describe("local cron provider preflight HTTP transport", () => { + beforeEach(() => { + resetCronModelProviderPreflightCacheForTest(); + }); + + afterEach(async () => { + resetCronModelProviderPreflightCacheForTest(); + const servers = [...activeServers]; + activeServers.clear(); + await Promise.all( + servers.map(async (server) => { + server.closeAllConnections(); + await new Promise((resolve, reject) => { + server.close((error) => { + if (error) { + reject(error); + return; + } + resolve(); + }); + }); + }), + ); + }); + + it.each([ + { + label: "OpenAI-compatible /models", + provider: "vllm", + api: "openai-completions" as const, + basePath: "/v1", + expectedPath: "/v1/models", + status: 401, + }, + { + label: "Ollama /api/tags", + provider: "ollama", + api: "ollama" as const, + basePath: "", + expectedPath: "/api/tags", + status: 200, + }, + ])("closes a never-ending $label response", async (scenario) => { + const server = await startStreamingProviderServer(scenario.status); + + const result = await withTestTimeout( + preflightCronModelProvider({ + cfg: { + models: { + providers: { + [scenario.provider]: { + api: scenario.api, + baseUrl: `${server.baseUrl}${scenario.basePath}`, + models: [], + }, + }, + }, + }, + provider: scenario.provider, + model: "streaming-model", + }), + 2_500, + `cron ${scenario.label} preflight stalled on a never-ending response body`, + ); + + expect(result).toEqual({ status: "available" }); + expect(server.requestedPaths).toEqual([scenario.expectedPath]); + expect(server.socketClosures).toHaveLength(1); + await withTestTimeout( + Promise.all(server.socketClosures), + 2_500, + `cron ${scenario.label} preflight left its real HTTP socket open`, + ); + }); + + it("closes every socket during a concurrent burst of streaming probes", async () => { + const server = await startStreamingProviderServer(); + const probeCount = 24; + + const results = await withTestTimeout( + Promise.all( + Array.from({ length: probeCount }, (_, index) => + preflightCronModelProvider({ + cfg: { + models: { + providers: { + vllm: { + api: "openai-completions", + baseUrl: `${server.baseUrl}/v1/provider-${index}`, + models: [], + }, + }, + }, + }, + provider: "vllm", + model: `streaming-model-${index}`, + }), + ), + ), + 5_000, + "concurrent cron provider preflights stalled on streaming HTTP responses", + ); + + expect(results).toEqual(Array.from({ length: probeCount }, () => ({ status: "available" }))); + expect(server.requestedPaths).toHaveLength(probeCount); + expect(server.socketClosures).toHaveLength(probeCount); + await withTestTimeout( + Promise.all(server.socketClosures), + 5_000, + "concurrent cron provider preflights left streaming HTTP sockets open", + ); + }); +}); diff --git a/src/cron/isolated-agent/model-preflight.runtime.ts b/src/cron/isolated-agent/model-preflight.runtime.ts index 73b170c88974..adc21af1445f 100644 --- a/src/cron/isolated-agent/model-preflight.runtime.ts +++ b/src/cron/isolated-agent/model-preflight.runtime.ts @@ -254,6 +254,11 @@ async function probeLocalProviderEndpoint(params: { // have the full provider context. void response.status; } finally { + // Captured responses can tee their body, so awaiting branch cancellation + // would hang the cron probe; start cancellation before closing the agent. + if (!response.bodyUsed) { + void response.body?.cancel().catch(() => undefined); + } await release(); } }