Merge pull request #1505 from LZong-tw/fix/client-abort-stream-cleanup

fix: handle downstream aborts in gateway streams

packages/core/src/gateway/service.ts

@@ -706,28 +706,28 @@ class GatewayService {
     let responseCompleted = false;
     let onClientDisconnect: (() => void) | undefined;
     let onResponseFinish: (() => void) | undefined;
+    const handleClientDisconnect = () => {
+      if (responseCompleted || response.writableEnded) {
+        return;
+      }
+      if (!clientDisconnected) {
+        clientDisconnected = true;
+        upstreamAbortController.abort(new Error(clientDisconnectMessage));
+      }
+      onClientDisconnect?.();
+    };
 
     response.once("finish", () => {
       responseCompleted = true;
       onResponseFinish?.();
     });
-    response.once("close", () => {
-      if (responseCompleted || response.writableEnded) {
-        return;
-      }
-      clientDisconnected = true;
-      upstreamAbortController.abort(new Error(clientDisconnectMessage));
-      onClientDisconnect?.();
-    });
-    response.on("error", (error) => {
+    response.once("close", handleClientDisconnect);
+    response.on("error", () => {
       // Client-side write failures (EPIPE / ECONNRESET when the client closes
       // mid-stream, common during tool execution) must not crash the main
       // process as an Uncaught Exception. Swallow them here; the close handler
       // above already records the disconnect via writeStreamLog.
-      if (!clientDisconnected) {
-        clientDisconnected = true;
-        upstreamAbortController.abort(new Error(clientDisconnectMessage));
-      }
+      handleClientDisconnect();
     });
 
     const writeRequestLog = (
@@ -976,6 +976,11 @@ class GatewayService {
       this.config
     );
     const upstreamResponse = upstreamResult.response;
+    if (clientDisconnected || upstreamAbortController.signal.aborted) {
+      await cancelResponseBody(upstreamResponse);
+      writeRequestLog(clientClosedRequestStatusCode, responseHeaders, "", false, clientDisconnectMessage);
+      return;
+    }
     if (codexApplyPatchBridgeActive) {
       responseHeaders.delete("content-length");
     }
@@ -989,6 +994,11 @@ class GatewayService {
       responseHeaders.delete("content-length");
     }
     recordProviderCredentialOutcome(this.config, method, upstreamResult.attempt, upstreamResponse.status, responseHeaders);
+    if (clientDisconnected || response.destroyed) {
+      await cancelResponseBody(upstreamResponse);
+      writeRequestLog(clientClosedRequestStatusCode, responseHeaders, "", false, clientDisconnectMessage);
+      return;
+    }
     response.writeHead(upstreamResponse.status, Object.fromEntries(filteredResponseHeaders(responseHeaders)));
     if (!upstreamResponse.body) {
       if (shouldCaptureUsage) {
@@ -1023,6 +1033,7 @@ class GatewayService {
           this.browserWebSearchMcpIntegration
         )
       : patchedResponseBody;
+    const responseStreams = uniqueStreams([upstreamBody, patchedResponseBody, responseBody]);
     const sampler = createBodySampler();
     const sseErrorDetector = createSseErrorDetector(responseHeaders.get("content-type") ?? undefined);
     let streamDetectedError: string | undefined;
@@ -1034,7 +1045,7 @@ class GatewayService {
       }
       logRecorded = true;
       writeRequestLog(
-        upstreamResponse.status,
+        clientDisconnected ? clientClosedRequestStatusCode : upstreamResponse.status,
         responseHeaders,
         sampler.read(),
         sampler.isTruncated(),
@@ -1043,13 +1054,21 @@ class GatewayService {
     };
     onClientDisconnect = () => {
       writeStreamLog(clientDisconnectMessage);
-      responseBody.destroy(new Error(clientDisconnectMessage));
+      responseBody.unpipe(response);
+      destroyResponseStreams(responseStreams);
     };
     onResponseFinish = () => {
       if (upstreamStreamEnded) {
         writeStreamLog();
       }
     };
+    const onResponseStreamError = (error: Error) => {
+      streamDetectedError ??= sseErrorDetector.finish();
+      writeStreamLog(clientDisconnected ? clientDisconnectMessage : formatError(error));
+    };
+    for (const stream of responseStreams) {
+      stream.on("error", onResponseStreamError);
+    }
     responseBody.on("data", (chunk) => {
       sampler.append(chunk);
       streamDetectedError ??= sseErrorDetector.append(chunk);
@@ -1061,10 +1080,6 @@ class GatewayService {
         writeStreamLog();
       }
     });
-    responseBody.on("error", (error) => {
-      streamDetectedError ??= sseErrorDetector.finish();
-      writeStreamLog(clientDisconnected ? clientDisconnectMessage : formatError(error));
-    });
     if (shouldCaptureUsage) {
       responseBody.once("end", () => {
         void recordGatewayUsageCapture({
@@ -1082,6 +1097,10 @@ class GatewayService {
         });
       });
     }
+    if (clientDisconnected || response.destroyed) {
+      onClientDisconnect();
+      return;
+    }
     responseBody.pipe(response);
   }
 
@@ -6205,6 +6224,29 @@ async function drainResponseBody(response: Response): Promise<void> {
   }
 }
 
+async function cancelResponseBody(response: Response): Promise<void> {
+  try {
+    await response.body?.cancel();
+  } catch {
+    // The client already disconnected; best-effort upstream cleanup must not mask that expected path.
+  }
+}
+
+function uniqueStreams(streams: Readable[]): Readable[] {
+  return [...new Set(streams)];
+}
+
+function destroyResponseStreams(streams: Readable[]): void {
+  for (const stream of streams) {
+    if (!stream.destroyed) {
+      // A downstream client close is an expected abort path. Destroying with
+      // an Error would emit another error event on Readable/Transform stages,
+      // and intermediate stages may not be the final responseBody listener.
+      stream.destroy();
+    }
+  }
+}
+
 function parseJsonObjectSafe(buffer: Buffer | undefined): Record<string, unknown> | undefined {
   if (!buffer || buffer.byteLength === 0) {
     return undefined;

tests/main/gateway-client-disconnect.test.mjs

@@ -0,0 +1,155 @@
+import assert from "node:assert/strict";
+import { createServer } from "node:http";
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import path from "node:path";
+import test from "node:test";
+import { createDefaultAppConfig } from "../../packages/core/src/config/default-config.ts";
+import { gatewayService } from "../../packages/core/src/gateway/service.ts";
+
+test("gateway treats downstream client aborts as expected stream cleanup", async () => {
+  const dir = mkdtempSync(path.join(tmpdir(), "ccr-gateway-client-abort-test-"));
+  const uncaughtErrors = [];
+  const unhandledRejections = [];
+  const onUncaughtException = (error) => uncaughtErrors.push(error);
+  const onUnhandledRejection = (error) => unhandledRejections.push(error);
+  process.prependListener("uncaughtException", onUncaughtException);
+  process.prependListener("unhandledRejection", onUnhandledRejection);
+  let upstreamResponseClosed = false;
+  const patch = "*** Begin Patch\n*** Add File: foo.txt\n+hi\n*** End Patch\n";
+
+  const upstream = createServer((request, response) => {
+    request.resume();
+    response.writeHead(200, { "content-type": "text/event-stream" });
+    response.write("event: response.output_text.delta\n");
+    response.write(`data: ${JSON.stringify({
+      item: {
+        arguments: JSON.stringify({ patch }),
+        call_id: "call_patch",
+        name: "virtual_apply_patch",
+        type: "function_call"
+      },
+      type: "response.output_item.done"
+    })}\n\n`);
+    const interval = setInterval(() => {
+      response.write("event: response.output_text.delta\n");
+      response.write(`data: ${JSON.stringify({ delta: "tick", type: "response.output_text.delta" })}\n\n`);
+    }, 20);
+    response.on("close", () => {
+      upstreamResponseClosed = true;
+      clearInterval(interval);
+    });
+  });
+  const gateway = createServer((request, response) => {
+    const requestPath = new URL(request.url ?? "/", "http://127.0.0.1").pathname;
+    void gatewayService.proxyRequest(request, response, requestPath).catch((error) => {
+      if (!response.headersSent) {
+        response.writeHead(502, { "content-type": "application/json" });
+      }
+      if (!response.writableEnded) {
+        response.end(`${JSON.stringify({ error: { message: String(error?.message ?? error) } })}\n`);
+      }
+    });
+  });
+
+  try {
+    await listen(upstream);
+    const upstreamPort = serverPort(upstream);
+    const config = createDefaultAppConfig({ generatedConfigFile: path.join(dir, "gateway.config.json") });
+    config.APIKEY = "test-api-key";
+    config.gateway.coreHost = "127.0.0.1";
+    config.gateway.corePort = upstreamPort;
+    config.gateway.host = "127.0.0.1";
+    config.gateway.port = 0;
+    gatewayService.updateConfig(config);
+    gatewayService.coreAuthToken = "test-core-auth-token";
+
+    await listen(gateway);
+    const gatewayUrl = `http://127.0.0.1:${serverPort(gateway)}/v1/responses`;
+    const controller = new AbortController();
+    const response = await fetch(gatewayUrl, {
+      body: JSON.stringify({
+        input: "hello",
+        model: "provider-deepseek::openai_chat_completions/deepseek-v4-flash",
+        stream: true,
+        tools: [{ type: "custom", name: "apply_patch", format: { type: "grammar", syntax: "lark", definition: "start: begin_patch" } }]
+      }),
+      headers: {
+        authorization: "Bearer test-api-key",
+        "content-type": "application/json",
+        "user-agent": "codex-test"
+      },
+      method: "POST",
+      signal: controller.signal
+    });
+
+    const errorText = response.status === 200 ? "" : await response.text();
+    assert.equal(response.status, 200, errorText);
+    const reader = response.body.getReader();
+    const firstChunk = await reader.read();
+    assert.equal(firstChunk.done, false);
+    const firstChunkText = new TextDecoder().decode(firstChunk.value);
+    assert.match(firstChunkText, /"type":"custom_tool_call"/, "expected the response to use the Codex apply_patch response transform");
+    assert.match(firstChunkText, /"name":"apply_patch"/, "expected virtual_apply_patch to be rewritten back to apply_patch");
+
+    await reader.cancel("test downstream disconnect");
+    controller.abort();
+    await waitFor(() => upstreamResponseClosed, 1000);
+
+    assert.deepEqual(uncaughtErrors.map((error) => error?.message ?? String(error)), []);
+    assert.deepEqual(unhandledRejections.map((error) => error?.message ?? String(error)), []);
+  } finally {
+    process.off("uncaughtException", onUncaughtException);
+    process.off("unhandledRejection", onUnhandledRejection);
+    await closeServer(gateway);
+    await closeServer(upstream);
+    await gatewayService.stop();
+    rmSync(dir, { force: true, recursive: true });
+  }
+});
+
+function listen(server) {
+  return new Promise((resolve, reject) => {
+    server.once("error", reject);
+    server.listen(0, "127.0.0.1", () => {
+      server.off("error", reject);
+      resolve();
+    });
+  });
+}
+
+function serverPort(server) {
+  const address = server.address();
+  assert.equal(typeof address, "object");
+  assert.ok(address);
+  return address.port;
+}
+
+function closeServer(server) {
+  return new Promise((resolve, reject) => {
+    if (!server.listening) {
+      resolve();
+      return;
+    }
+    const timeout = setTimeout(() => server.closeAllConnections?.(), 1000);
+    server.close((error) => {
+      clearTimeout(timeout);
+      error ? reject(error) : resolve();
+    });
+  });
+}
+
+function sleep(ms) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}
+
+async function waitFor(predicate, timeoutMs) {
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline) {
+    if (predicate()) {
+      return;
+    }
+    await sleep(10);
+  }
+  assert.equal(predicate(), true);
+}