fix(cron): prevent local model preflight socket leaks (#114540)

Co-authored-by: hugenshen <16300669+hugenshen@users.noreply.github.com>

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<void>(() => {});
+    });
+    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"));
 

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<void>[];
+};
+
+const activeServers = new Set<Server>();
+
+async function startStreamingProviderServer(status = 200): Promise<StreamingProviderServer> {
+  const requestedPaths: string[] = [];
+  const socketClosures: Promise<void>[] = [];
+  const server = createServer((request, response) => {
+    requestedPaths.push(request.url ?? "");
+    socketClosures.push(
+      new Promise<void>((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<void>((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<void>((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",
+    );
+  });
+});

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();
   }
 }