Skip to content

fix(responses-ws): handle abort during upstream handshake #1334

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ding113 merged 1 commit into from
Jul 22, 2026
Merged

fix(responses-ws): handle abort during upstream handshake #1334

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .env.example
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,10 @@ STORE_SESSION_RESPONSE_BODY=true # 是否在 Redis 中存储会话响应
# Dashboard 配置
DASHBOARD_LOGS_POLL_INTERVAL_MS=5000 # 日志页自动刷新轮询间隔（毫秒，默认 5000，范围 250-60000）

# OpenAI Responses WebSocket（可选三态覆盖）
# 未设置时读取数据库系统设置；仅接受 true/false/1/0，设置后强制覆盖数据库值。
# ENABLE_OPENAI_RESPONSES_WEBSOCKET=false

# 熔断器配置
# 功能说明：控制网络错误是否计入熔断器失败计数
# - false (默认)：网络错误（DNS 解析失败、连接超时、代理连接失败等）不计入熔断器，仅供应商错误（4xx/5xx HTTP 响应）计入
Expand Down
331 changes: 329 additions & 2 deletions src/app/v1/_lib/responses-ws/__tests__/upstream-adapter.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import type { AddressInfo } from "node:net";
import { afterEach, describe, expect, it } from "vitest";
import { WebSocketServer } from "ws";
import { afterEach, describe, expect, it, vi } from "vitest";
import { WebSocket, WebSocketServer } from "ws";
import type { Provider } from "@/types/provider";
import {
clearResponsesWsSessionsForTests,
Expand All @@ -21,6 +21,12 @@ type ServerHandle = {
close: () => Promise<void>;
};

type StalledUpgradeServerHandle = {
port: number;
upgradeReceived: Promise<void>;
close: () => Promise<void>;
};

function startMockServer(
handler: (socket: import("ws").WebSocket, req: import("http").IncomingMessage) => void
): Promise<ServerHandle> {
Expand All @@ -42,6 +48,67 @@ function startMockServer(
});
}

async function startStalledUpgradeServer(): Promise<StalledUpgradeServerHandle> {
const http = await import("node:http");
const sockets = new Set<import("node:net").Socket>();
let resolveUpgradeReceived: (() => void) | null = null;
const upgradeReceived = new Promise<void>((resolve) => {
resolveUpgradeReceived = resolve;
});
const httpServer = http.createServer();

httpServer.on("connection", (socket) => {
sockets.add(socket);
socket.once("close", () => sockets.delete(socket));
});
httpServer.on("upgrade", () => {
resolveUpgradeReceived?.();
});

await new Promise<void>((resolve, reject) => {
const onError = (error: Error) => reject(error);
httpServer.once("error", onError);
httpServer.listen(0, "127.0.0.1", () => {
httpServer.off("error", onError);
resolve();
});
});

const address = httpServer.address() as AddressInfo;
return {
port: address.port,
upgradeReceived,
close: async () => {
for (const socket of sockets) socket.destroy();
await new Promise<void>((resolve) => httpServer.close(() => resolve()));
},
};
}

function captureProcessErrors() {
const uncaughtExceptions: unknown[] = [];
const unhandledRejections: unknown[] = [];
const onUncaughtException = (error: unknown) => uncaughtExceptions.push(error);
const onUnhandledRejection = (reason: unknown) => unhandledRejections.push(reason);

process.on("uncaughtException", onUncaughtException);
process.on("unhandledRejection", onUnhandledRejection);

return {
uncaughtExceptions,
unhandledRejections,
stop() {
process.off("uncaughtException", onUncaughtException);
process.off("unhandledRejection", onUnhandledRejection);
},
};
}

async function flushProcessEvents(): Promise<void> {
await new Promise<void>((resolve) => setImmediate(resolve));
await new Promise<void>((resolve) => setImmediate(resolve));
}

function codexProvider(): Provider {
return {
id: 1,
Expand Down Expand Up @@ -773,6 +840,105 @@ describe("tryResponsesWebsocketUpstream", () => {
}
});

it("safely terminates a CONNECTING socket when aborted after construction", async () => {
const stalledServer = await startStalledUpgradeServer();
const abortController = new AbortController();
const processErrors = captureProcessErrors();
const closeSpy = vi.spyOn(WebSocket.prototype, "close");
const terminateSpy = vi.spyOn(WebSocket.prototype, "terminate");

try {
const resultPromise = tryResponsesWebsocketUpstream({
provider: codexProvider(),
upstreamUrl: `http://127.0.0.1:${stalledServer.port}/v1/responses`,
upstreamHeaders: new Headers({ authorization: "Bearer sk-mock" }),
abortSignal: abortController.signal,
body: { model: "gpt-5.5", input: "hi" },
});

await withTimeout(
stalledServer.upgradeReceived,
1_000,
"upstream WebSocket did not enter the stalled handshake"
);
abortController.abort();

const result = await withTimeout(
resultPromise,
1_000,
"CONNECTING upstream WebSocket attempt hung after abort"
);
await flushProcessEvents();

expect("failed" in result).toBe(true);
if (!("failed" in result)) return;
expect(result.message).toContain("aborted before first upstream WebSocket event");
expect(terminateSpy).toHaveBeenCalledTimes(1);
expect(closeSpy).not.toHaveBeenCalled();
expect(processErrors.uncaughtExceptions).toEqual([]);
expect(processErrors.unhandledRejections).toEqual([]);

const clientSocket = terminateSpy.mock.instances[0] as WebSocket;
expect(clientSocket.listenerCount("open")).toBe(0);
expect(clientSocket.listenerCount("message")).toBe(0);
expect(clientSocket.listenerCount("error")).toBe(0);
expect(clientSocket.listenerCount("close")).toBe(0);
expect(clientSocket.listenerCount("unexpected-response")).toBe(0);
} finally {
processErrors.stop();
closeSpy.mockRestore();
terminateSpy.mockRestore();
await stalledServer.close();
}
});

it("safely terminates a CONNECTING socket when the signal is already aborted", async () => {
const stalledServer = await startStalledUpgradeServer();
const abortController = new AbortController();
abortController.abort();
const processErrors = captureProcessErrors();
const timeoutSpy = vi.spyOn(globalThis, "setTimeout");
const closeSpy = vi.spyOn(WebSocket.prototype, "close");
const terminateSpy = vi.spyOn(WebSocket.prototype, "terminate");

try {
const result = await withTimeout(
tryResponsesWebsocketUpstream({
provider: codexProvider(),
upstreamUrl: `http://127.0.0.1:${stalledServer.port}/v1/responses`,
upstreamHeaders: new Headers({ authorization: "Bearer sk-mock" }),
abortSignal: abortController.signal,
body: { model: "gpt-5.5", input: "hi" },
}),
1_000,
"already-aborted upstream WebSocket attempt hung"
);
await flushProcessEvents();

expect("failed" in result).toBe(true);
if (!("failed" in result)) return;
expect(result.message).toContain("aborted before first upstream WebSocket event");
Comment on lines +914 to +920

@greptile-apps greptile-apps Bot Jul 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Direct mutation of private _readyState field is fragile

The CLOSING-race test sets mutableClientSocket._readyState = WebSocket.CLOSING by casting through an intersection type to reach the ws library's internal private field. This works today because ws stores readyState in _readyState, but the library makes no API promise about that name. If a future ws patch renames or wraps the field, the cast silently succeeds while the socket stays OPEN, and the test becomes a false-positive. Consider whether there is a higher-level way to reach CLOSING state (e.g. by having the mock server close its end and waiting for the half-close) or, if the direct field mutation is necessary, adding a runtime assertion right after (expect(clientSocket.readyState).toBe(WebSocket.CLOSING)) to fail loudly when the internal shape changes.

Prompt To Fix With AI 
This is a comment left during a code review.
Path: src/app/v1/_lib/responses-ws/__tests__/upstream-adapter.test.ts
Line: 914-920

Comment:
**Direct mutation of private `_readyState` field is fragile**

The CLOSING-race test sets `mutableClientSocket._readyState = WebSocket.CLOSING` by casting through an intersection type to reach the `ws` library's internal private field. This works today because `ws` stores `readyState` in `_readyState`, but the library makes no API promise about that name. If a future `ws` patch renames or wraps the field, the cast silently succeeds while the socket stays `OPEN`, and the test becomes a false-positive. Consider whether there is a higher-level way to reach `CLOSING` state (e.g. by having the mock server close its end and waiting for the half-close) or, if the direct field mutation is necessary, adding a runtime assertion right after (`expect(clientSocket.readyState).toBe(WebSocket.CLOSING)`) to fail loudly when the internal shape changes.

How can I resolve this? If you propose a fix, please make it concise.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

expect(terminateSpy).toHaveBeenCalledTimes(1);
expect(closeSpy).not.toHaveBeenCalled();
expect(timeoutSpy.mock.calls.some(([, delay]) => delay === 20_000)).toBe(false);
expect(processErrors.uncaughtExceptions).toEqual([]);
expect(processErrors.unhandledRejections).toEqual([]);

const clientSocket = terminateSpy.mock.instances[0] as WebSocket;
expect(clientSocket.listenerCount("open")).toBe(0);
expect(clientSocket.listenerCount("message")).toBe(0);
expect(clientSocket.listenerCount("error")).toBe(0);
expect(clientSocket.listenerCount("close")).toBe(0);
expect(clientSocket.listenerCount("unexpected-response")).toBe(0);
} finally {
processErrors.stop();
timeoutSpy.mockRestore();
closeSpy.mockRestore();
terminateSpy.mockRestore();
await stalledServer.close();
}
});

it("resolves and closes upstream when aborted before the first WS event", async () => {
let resolveMessageReceived: (() => void) | null = null;
const messageReceived = new Promise<void>((resolve) => {
Expand Down Expand Up @@ -821,6 +987,167 @@ describe("tryResponsesWebsocketUpstream", () => {
expect(upstreamCloseCode).toBe(1000);
});

it("closes only once when abort races the first-event timeout", async () => {
let resolveMessageReceived: (() => void) | null = null;
const messageReceived = new Promise<void>((resolve) => {
resolveMessageReceived = resolve;
});
let resolveUpstreamClosed: (() => void) | null = null;
const upstreamClosed = new Promise<void>((resolve) => {
resolveUpstreamClosed = resolve;
});
server = await startMockServer((socket) => {
socket.on("message", () => {
resolveMessageReceived?.();
});
socket.on("close", () => {
resolveUpstreamClosed?.();
});
});

const realSetTimeout = globalThis.setTimeout.bind(globalThis);
let firstEventTimeoutCallback: (() => void) | null = null;
const placeholderTimer = realSetTimeout(() => {}, 60_000);
const timeoutSpy = vi.spyOn(globalThis, "setTimeout").mockImplementation(((
callback: (...args: unknown[]) => void,
delay?: number,
...args: unknown[]
) => {
if (delay === 20_000) {
firstEventTimeoutCallback = () => callback(...args);
return placeholderTimer;
}
return realSetTimeout(callback, delay, ...args);
}) as typeof setTimeout);
const abortController = new AbortController();
const processErrors = captureProcessErrors();
const closeSpy = vi.spyOn(WebSocket.prototype, "close");
const terminateSpy = vi.spyOn(WebSocket.prototype, "terminate");

try {
const resultPromise = tryResponsesWebsocketUpstream({
provider: codexProvider(),
upstreamUrl: `http://127.0.0.1:${server.port}/v1/responses`,
upstreamHeaders: new Headers({ authorization: "Bearer sk-mock" }),
abortSignal: abortController.signal,
body: { model: "gpt-5.5", input: "hi" },
});

await withTimeout(messageReceived, 1_000, "upstream did not receive the WS request frame");
expect(firstEventTimeoutCallback).not.toBeNull();
const triggerFirstEventTimeout = firstEventTimeoutCallback as unknown as () => void;

abortController.abort();
triggerFirstEventTimeout();
const closeInstances = closeSpy.mock.instances as unknown as WebSocket[];
const clientCloseCallsAfterRace = closeInstances.filter(
(socket) => socket.url === `ws://127.0.0.1:${server?.port}/v1/responses`
);
expect(clientCloseCallsAfterRace).toHaveLength(1);
const clientSocket = clientCloseCallsAfterRace[0] as WebSocket;

const result = await withTimeout(
resultPromise,
1_000,
"upstream WS attempt hung during abort and timeout race"
);
await withTimeout(upstreamClosed, 1_000, "upstream WS did not close during abort race");
await flushProcessEvents();

expect("failed" in result).toBe(true);
if (!("failed" in result)) return;
expect(result.message).toContain("aborted before first upstream WebSocket event");
expect(terminateSpy).not.toHaveBeenCalled();
expect(processErrors.uncaughtExceptions).toEqual([]);
expect(processErrors.unhandledRejections).toEqual([]);

expect(clientSocket.listenerCount("open")).toBe(0);
expect(clientSocket.listenerCount("message")).toBe(0);
expect(clientSocket.listenerCount("error")).toBe(0);
expect(clientSocket.listenerCount("close")).toBe(0);
expect(clientSocket.listenerCount("unexpected-response")).toBe(0);
} finally {
clearTimeout(placeholderTimer);
processErrors.stop();
timeoutSpy.mockRestore();
closeSpy.mockRestore();
terminateSpy.mockRestore();
}
});

it("keeps an error sink when abort races a socket already in CLOSING", async () => {
let resolveMessageReceived: (() => void) | null = null;
const messageReceived = new Promise<void>((resolve) => {
resolveMessageReceived = resolve;
});
server = await startMockServer((socket) => {
socket.on("message", () => {
resolveMessageReceived?.();
});
});

const abortController = new AbortController();
const processErrors = captureProcessErrors();
const sendSpy = vi.spyOn(WebSocket.prototype, "send");
const closeSpy = vi.spyOn(WebSocket.prototype, "close");
const terminateSpy = vi.spyOn(WebSocket.prototype, "terminate");
let clientSocket: WebSocket | null = null;

try {
const resultPromise = tryResponsesWebsocketUpstream({
provider: codexProvider(),
upstreamUrl: `http://127.0.0.1:${server.port}/v1/responses`,
upstreamHeaders: new Headers({ authorization: "Bearer sk-mock" }),
abortSignal: abortController.signal,
body: { model: "gpt-5.5", input: "hi" },
});

await withTimeout(messageReceived, 1_000, "upstream did not receive the WS request frame");
const sendInstances = sendSpy.mock.instances as unknown as WebSocket[];
clientSocket =
sendInstances.find(
(socket) => socket.url === `ws://127.0.0.1:${server?.port}/v1/responses`
) ?? null;
expect(clientSocket).not.toBeNull();
if (!clientSocket) return;

const mutableClientSocket = clientSocket as WebSocket & { _readyState: number };
mutableClientSocket._readyState = WebSocket.CLOSING;
abortController.abort();

const result = await withTimeout(
resultPromise,
1_000,
"CLOSING upstream WebSocket attempt hung after abort"
);

expect("failed" in result).toBe(true);
if (!("failed" in result)) return;
expect(closeSpy).not.toHaveBeenCalled();
expect(terminateSpy).not.toHaveBeenCalled();
expect(clientSocket.listenerCount("error")).toBe(1);

clientSocket.emit("error", new Error("synthetic closing-race error"));
clientSocket.emit("close", 1000, Buffer.alloc(0));
await flushProcessEvents();

expect(clientSocket.listenerCount("error")).toBe(0);
expect(clientSocket.listenerCount("close")).toBe(0);
expect(processErrors.uncaughtExceptions).toEqual([]);
expect(processErrors.unhandledRejections).toEqual([]);
} finally {
sendSpy.mockRestore();
closeSpy.mockRestore();
terminateSpy.mockRestore();
if (clientSocket && clientSocket.readyState !== WebSocket.CLOSED) {
clientSocket.once("error", () => {});
clientSocket.terminate();
await flushProcessEvents();
}
processErrors.stop();
}
});

it("keeps the persistent session map bounded when every retained session is active", async () => {
setResponsesWsSessionMaxEntriesForTests(1);
let connectionCount = 0;
Expand Down
Loading