fix: cancel agent runs when the HTTP client disconnects

src/http.ts

@@ -30,6 +30,15 @@ function sendJson(response: ServerResponse, status: number, body: unknown): void
   response.end(JSON.stringify(body));
 }
 
+// The response closes before it ends only when the client goes away mid-run.
+function abortOnDisconnect(response: ServerResponse): AbortController {
+  const controller = new AbortController();
+  response.once("close", () => {
+    if (!response.writableEnded) controller.abort();
+  });
+  return controller;
+}
+
 export async function startHttp(
   config: AgentConfig,
   agent: ImpossibleAgent,
@@ -52,10 +61,7 @@ export async function startHttp(
     if (request.method === "POST" && url.pathname === "/v1/chat/stream") {
       try {
         const body = requestSchema.parse(await readJson(request));
-        const controller = new AbortController();
-        request.once("aborted", () => {
-          controller.abort();
-        });
+        const controller = abortOnDisconnect(response);
         response.writeHead(200, {
           "content-type": "text/event-stream; charset=utf-8",
           "cache-control": "no-cache",
@@ -87,10 +93,7 @@ export async function startHttp(
 
     try {
       const body = requestSchema.parse(await readJson(request));
-      const controller = new AbortController();
-      request.once("aborted", () => {
-        controller.abort();
-      });
+      const controller = abortOnDisconnect(response);
       const result = await runAgent(agent, config, body.input, {
         ...(body.threadId ? { threadId: body.threadId } : {}),
         signal: controller.signal,

tests/http.test.ts

@@ -1,4 +1,5 @@
 import { AIMessage, AIMessageChunk } from "@langchain/core/messages";
+import { request } from "node:http";
 import pino from "pino";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import type { ImpossibleAgent } from "../src/agent.js";
@@ -73,4 +74,33 @@ describe("agent HTTP API", () => {
     expect(body).toContain('"text":"hello"');
     expect(body).toContain("event: done");
   });
+
+  it("cancels the agent run when the client disconnects", async () => {
+    const config = loadConfig({ AGENT_PORT: "41421", AGENT_LOG_LEVEL: "silent" });
+    let signal: AbortSignal | undefined;
+    let markStarted: () => void = () => undefined;
+    const started = new Promise<void>((resolve) => {
+      markStarted = resolve;
+    });
+    const agent = {
+      invoke: vi.fn().mockImplementation((_input: unknown, options: { signal: AbortSignal }) => {
+        signal = options.signal;
+        markStarted();
+        return new Promise((_resolve, reject) => {
+          options.signal.addEventListener("abort", () => {
+            reject(new Error("aborted"));
+          });
+        });
+      }),
+    } as unknown as ImpossibleAgent;
+    runtime = await startHttp(config, agent, pino({ level: "silent" }));
+    const client = request({ host: "127.0.0.1", port: 41421, path: "/v1/chat", method: "POST" });
+    client.on("error", () => undefined);
+    client.end(JSON.stringify({ input: "wait" }));
+    await started;
+    client.destroy();
+    await vi.waitFor(() => {
+      expect(signal?.aborted).toBe(true);
+    });
+  });
 });