fix: handle gateway client aborts during streaming

This commit is contained in:
Lim, Un-tiong 2026-07-09 18:15:22 +08:00
parent 990bc0d958
commit 768baa8575
2 changed files with 216 additions and 19 deletions

View file

@ -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;

View file

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