diff --git a/apps/server/src/serve-node.test.ts b/apps/server/src/serve-node.test.ts index 250d0b75e..3cafa78e5 100644 --- a/apps/server/src/serve-node.test.ts +++ b/apps/server/src/serve-node.test.ts @@ -120,4 +120,117 @@ describe("serve", () => { await server.stop(); } }); + + // A long-lived response has no other way to learn its reader is gone. Without these two + // signals an SSE producer keeps streaming into a dead socket for the life of the process, + // holding its upstream engine subscription open with it. + test("aborts request.signal when the client hangs up mid-stream", async () => { + const encoder = new TextEncoder(); + let aborted = false; + let sawRequest = (): void => {}; + const requestSeen = new Promise((resolve) => { + sawRequest = resolve; + }); + + const server = await serve({ + hostname: "127.0.0.1", + port: 0, + fetch: (request) => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode("open\n")); + sawRequest(); + request.signal.addEventListener("abort", () => { + aborted = true; + try { + controller.close(); + } catch { + // Already closed by the transport. + } + }, { once: true }); + }, + }), + { headers: { "content-type": "text/event-stream" } }, + ), + }); + + try { + const controller = new AbortController(); + const response = await fetch(`http://127.0.0.1:${server.port}/stream`, { signal: controller.signal }); + await response.body!.getReader().read(); + await requestSeen; + + controller.abort(); + for (let attempt = 0; attempt < 50 && !aborted; attempt += 1) await delay(10); + + expect(aborted).toBe(true); + } finally { + await server.stop(); + } + }); + + test("cancels the response stream when the client hangs up mid-stream", async () => { + const encoder = new TextEncoder(); + let cancelled = false; + let sawRequest = (): void => {}; + const requestSeen = new Promise((resolve) => { + sawRequest = resolve; + }); + + const server = await serve({ + hostname: "127.0.0.1", + port: 0, + fetch: () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode("open\n")); + sawRequest(); + }, + cancel() { + cancelled = true; + }, + }), + { headers: { "content-type": "text/event-stream" } }, + ), + }); + + try { + const controller = new AbortController(); + const response = await fetch(`http://127.0.0.1:${server.port}/stream`, { signal: controller.signal }); + await response.body!.getReader().read(); + await requestSeen; + + controller.abort(); + for (let attempt = 0; attempt < 50 && !cancelled; attempt += 1) await delay(10); + + expect(cancelled).toBe(true); + } finally { + await server.stop(); + } + }); + + test("a normal response still completes and does not report a disconnect", async () => { + let aborted = false; + const server = await serve({ + hostname: "127.0.0.1", + port: 0, + fetch: (request) => { + request.signal.addEventListener("abort", () => { + aborted = true; + }, { once: true }); + return Response.json({ ok: true }); + }, + }); + + try { + const response = await fetch(`http://127.0.0.1:${server.port}/health`); + expect(await response.json()).toEqual({ ok: true }); + await delay(50); + expect(aborted).toBe(false); + } finally { + await server.stop(); + } + }); }); diff --git a/apps/server/src/serve-node.ts b/apps/server/src/serve-node.ts index 2b1440f9e..8ca1a6264 100644 --- a/apps/server/src/serve-node.ts +++ b/apps/server/src/serve-node.ts @@ -79,7 +79,31 @@ async function waitForDrainOrClose(nodeRes: ServerResponse): Promise { /** * Convert a Node.js IncomingMessage into a Web API Request. */ -function toWebRequest(nodeReq: IncomingMessage, hostname: string, port: number): Request { +/** + * Signals that the client went away. + * + * `close` also fires on a perfectly normal response, so the finished flag is what + * separates "the exchange completed" from "the socket died under us". Without this a + * long-lived response — an SSE stream — has no way to learn its reader is gone, and its + * producer keeps running for the life of the process. + */ +function createDisconnectSignal(nodeReq: IncomingMessage, nodeRes: ServerResponse): AbortSignal { + const controller = new AbortController(); + const abort = (): void => { + if (nodeRes.writableFinished) return; + if (!controller.signal.aborted) controller.abort(); + }; + nodeRes.once("close", abort); + nodeReq.once("aborted", abort); + return controller.signal; +} + +function toWebRequest( + nodeReq: IncomingMessage, + hostname: string, + port: number, + signal?: AbortSignal, +): Request { const url = `http://${hostname}:${port}${nodeReq.url ?? "/"}`; const method = nodeReq.method ?? "GET"; const headers = new Headers(); @@ -106,6 +130,7 @@ function toWebRequest(nodeReq: IncomingMessage, hostname: string, port: number): method, headers, body, + ...(signal ? { signal } : {}), // @ts-expect-error duplex is required for streaming request bodies in Node duplex: hasBody ? "half" : undefined, }); @@ -114,7 +139,11 @@ function toWebRequest(nodeReq: IncomingMessage, hostname: string, port: number): /** * Write a Web API Response to a Node.js ServerResponse. */ -async function writeWebResponse(webRes: Response, nodeRes: ServerResponse): Promise { +async function writeWebResponse( + webRes: Response, + nodeRes: ServerResponse, + disconnected?: AbortSignal, +): Promise { const headersObj: Record = {}; webRes.headers.forEach((value, key) => { const existing = headersObj[key]; @@ -135,17 +164,36 @@ async function writeWebResponse(webRes: Response, nodeRes: ServerResponse): Prom } const reader = webRes.body.getReader(); + let drained = false; + + // An idle stream parks in `reader.read()` with nothing to wake it, so noticing the + // disconnect between chunks is not enough — the read itself has to be cut short. + // Cancelling does both: it resolves the pending read and reaches the stream's + // `cancel()` callback, which is how the producer learns to release its upstream. + const cancelOnDisconnect = (): void => { + void reader.cancel().catch(() => undefined); + }; + disconnected?.addEventListener("abort", cancelOnDisconnect, { once: true }); + try { while (true) { const { done, value } = await reader.read(); - if (done) break; + if (done) { + drained = true; + break; + } if (!isResponseWritable(nodeRes)) break; if (!nodeRes.write(value)) { await waitForDrainOrClose(nodeRes); } } } finally { - reader.releaseLock(); + disconnected?.removeEventListener("abort", cancelOnDisconnect); + if (!drained) { + await reader.cancel().catch(() => undefined); + } else { + reader.releaseLock(); + } endResponse(nodeRes); } } @@ -168,9 +216,10 @@ export function serve(options: ServeOptions): Promise { }); try { - const webReq = toWebRequest(nodeReq, hostname, boundPort); + const disconnected = createDisconnectSignal(nodeReq, nodeRes); + const webReq = toWebRequest(nodeReq, hostname, boundPort, disconnected); const webRes = await fetchHandler(webReq); - await writeWebResponse(webRes, nodeRes); + await writeWebResponse(webRes, nodeRes, disconnected); } catch (error) { if (isExpectedConnectionAbort(error)) { if (isResponseWritable(nodeRes) && !nodeRes.headersSent) {