From d7d5ad2ad2170dc95b35984968dfe2362f4c7454 Mon Sep 17 00:00:00 2001 From: "Kai (Tam Nhu) Tran" <61256810+kaitranntt@users.noreply.github.com> Date: Wed, 27 May 2026 11:14:40 -0400 Subject: [PATCH] fix: time out stalled CLIProxy response streams --- ...x-reasoning-proxy-extended-context.test.ts | 125 ++++++++++++++++++ .../ai-providers/codex-reasoning-proxy.ts | 83 ++++++++++-- ...ool-sanitization-proxy-integration.test.ts | 103 +++++++++++++++ src/cliproxy/proxy/tool-sanitization-proxy.ts | 74 +++++++++-- .../proxy/upstream-response-timeout.ts | 87 ++++++++++++ 5 files changed, 454 insertions(+), 18 deletions(-) create mode 100644 src/cliproxy/proxy/upstream-response-timeout.ts diff --git a/src/cliproxy/ai-providers/__tests__/codex-reasoning-proxy-extended-context.test.ts b/src/cliproxy/ai-providers/__tests__/codex-reasoning-proxy-extended-context.test.ts index 78e9056a..0a3abf3b 100644 --- a/src/cliproxy/ai-providers/__tests__/codex-reasoning-proxy-extended-context.test.ts +++ b/src/cliproxy/ai-providers/__tests__/codex-reasoning-proxy-extended-context.test.ts @@ -76,6 +76,53 @@ function postJson( }); } +function postJsonText( + url: string, + body: JsonRecord, + timeoutMs = 2_000 +): Promise<{ statusCode: number; body: string }> { + return new Promise((resolve, reject) => { + const parsed = new URL(url); + const payload = JSON.stringify(body); + let timeout: ReturnType; + + const req = http.request( + { + hostname: parsed.hostname, + port: parsed.port, + path: parsed.pathname + parsed.search, + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'Content-Length': Buffer.byteLength(payload), + }, + }, + (res) => { + let responseBody = ''; + res.setEncoding('utf8'); + res.on('data', (chunk) => { + responseBody += chunk; + }); + res.on('end', () => { + clearTimeout(timeout); + resolve({ statusCode: res.statusCode ?? 0, body: responseBody }); + }); + } + ); + + timeout = setTimeout(() => { + req.destroy(new Error(`Timed out waiting for proxy response after ${timeoutMs}ms`)); + }, timeoutMs); + + req.on('error', (error) => { + clearTimeout(timeout); + reject(error); + }); + req.write(payload); + req.end(); + }); +} + describe('CodexReasoningProxy extended-context compatibility', () => { const cleanupServers: http.Server[] = []; @@ -147,6 +194,84 @@ describe('CodexReasoningProxy extended-context compatibility', () => { expect((capturedBody?.reasoning as JsonRecord | undefined)?.effort).toBe('high'); }); + it('ends streaming responses when upstream stalls after headers', async () => { + const upstream = http.createServer((req, res) => { + req.resume(); + req.on('end', () => { + res.writeHead(200, { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }); + res.write( + 'event: message_start\n' + + 'data: {"type":"message_start","message":{"id":"msg_stall","type":"message","role":"assistant","content":[],"model":"test","stop_reason":null}}\n\n' + ); + }); + }); + cleanupServers.push(upstream); + + const upstreamPort = await listenOnRandomPort(upstream); + const proxy = new CodexReasoningProxy({ + upstreamBaseUrl: `http://127.0.0.1:${upstreamPort}`, + modelMap: { defaultModel: 'gpt-5.3-codex' }, + timeoutMs: 100, + }); + + const proxyPort = await proxy.start(); + try { + const response = await postJsonText( + `http://127.0.0.1:${proxyPort}/api/provider/codex/v1/messages`, + { + model: 'gpt-5.3-codex', + messages: [], + stream: true, + } + ); + + expect(response.statusCode).toBe(200); + expect(response.body).toContain('message_start'); + expect(response.body).toContain('timeout_error'); + expect(response.body).toContain('Upstream response timed out'); + } finally { + proxy.stop(); + } + }); + + it('returns 504 when non-stream upstream stalls after headers', async () => { + const upstream = http.createServer((req, res) => { + req.resume(); + req.on('end', () => { + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.flushHeaders(); + }); + }); + cleanupServers.push(upstream); + + const upstreamPort = await listenOnRandomPort(upstream); + const proxy = new CodexReasoningProxy({ + upstreamBaseUrl: `http://127.0.0.1:${upstreamPort}`, + modelMap: { defaultModel: 'gpt-5.3-codex' }, + timeoutMs: 100, + }); + + const proxyPort = await proxy.start(); + try { + const response = await postJsonText( + `http://127.0.0.1:${proxyPort}/api/provider/codex/v1/messages`, + { + model: 'gpt-5.3-codex', + messages: [], + } + ); + + expect(response.statusCode).toBe(504); + expect(response.body).toContain('Upstream response timed out'); + } finally { + proxy.stop(); + } + }); + it('translates codex fast model suffixes into service_tier', async () => { const capturedBodies: JsonRecord[] = []; diff --git a/src/cliproxy/ai-providers/codex-reasoning-proxy.ts b/src/cliproxy/ai-providers/codex-reasoning-proxy.ts index fe91e40b..a9d7bd6a 100644 --- a/src/cliproxy/ai-providers/codex-reasoning-proxy.ts +++ b/src/cliproxy/ai-providers/codex-reasoning-proxy.ts @@ -7,6 +7,10 @@ import { resolveRuntimeCodexFallbackModel, } from './codex-plan-compatibility'; import { getModelMaxLevel } from '../model-catalog'; +import { + attachUpstreamResponseTimeout, + writeForwardResponseHead, +} from '../proxy/upstream-response-timeout'; export type CodexReasoningEffort = 'minimal' | 'low' | 'medium' | 'high' | 'xhigh'; export type CodexServiceTier = 'fast'; @@ -301,7 +305,7 @@ export class CodexReasoningProxy { headers: http.IncomingHttpHeaders, responseBody: string ): void { - clientRes.writeHead(statusCode, headers); + writeForwardResponseHead(clientRes, statusCode, headers); clientRes.end(responseBody); } @@ -635,10 +639,38 @@ export class CodexReasoningProxy { ), (upstreamRes) => { clearResponseTimeout(); - clientRes.writeHead(upstreamRes.statusCode || 200, upstreamRes.headers); - upstreamRes.pipe(clientRes); - upstreamRes.on('end', () => resolve()); - upstreamRes.on('error', reject); + const statusCode = upstreamRes.statusCode || 200; + let responseStarted = false; + const writeResponseHead = () => { + if (responseStarted) return; + responseStarted = true; + writeForwardResponseHead(clientRes, statusCode, upstreamRes.headers); + }; + const clearUpstreamResponseTimeout = attachUpstreamResponseTimeout({ + upstreamReq, + upstreamRes, + clientRes, + timeoutMs: this.config.timeoutMs, + onTimeout: () => resolve(), + }); + upstreamRes.on('data', (chunk: Buffer) => { + writeResponseHead(); + const canContinue = clientRes.write(chunk); + if (!canContinue) { + upstreamRes.pause(); + clientRes.once('drain', () => upstreamRes.resume()); + } + }); + upstreamRes.on('end', () => { + clearUpstreamResponseTimeout(); + writeResponseHead(); + clientRes.end(); + resolve(); + }); + upstreamRes.on('error', (error) => { + clearUpstreamResponseTimeout(); + reject(error); + }); } ); @@ -672,11 +704,38 @@ export class CodexReasoningProxy { (upstreamRes) => { clearResponseTimeout(); const statusCode = upstreamRes.statusCode || 200; + const clearUpstreamResponseTimeout = attachUpstreamResponseTimeout({ + upstreamReq, + upstreamRes, + clientRes, + timeoutMs: this.config.timeoutMs, + onTimeout: () => resolve(504), + }); if (statusCode >= 200 && statusCode < 300) { - clientRes.writeHead(statusCode, upstreamRes.headers); - upstreamRes.pipe(clientRes); - upstreamRes.on('end', () => resolve(statusCode)); - upstreamRes.on('error', reject); + let responseStarted = false; + const writeResponseHead = () => { + if (responseStarted) return; + responseStarted = true; + writeForwardResponseHead(clientRes, statusCode, upstreamRes.headers); + }; + upstreamRes.on('data', (chunk: Buffer) => { + writeResponseHead(); + const canContinue = clientRes.write(chunk); + if (!canContinue) { + upstreamRes.pause(); + clientRes.once('drain', () => upstreamRes.resume()); + } + }); + upstreamRes.on('end', () => { + clearUpstreamResponseTimeout(); + writeResponseHead(); + clientRes.end(); + resolve(statusCode); + }); + upstreamRes.on('error', (error) => { + clearUpstreamResponseTimeout(); + reject(error); + }); return; } @@ -694,6 +753,7 @@ export class CodexReasoningProxy { chunks.push(chunk); }); upstreamRes.on('end', async () => { + clearUpstreamResponseTimeout(); if (responseTooLarge) { reject(new Error('Upstream error response exceeded 10MB limit')); return; @@ -756,7 +816,10 @@ export class CodexReasoningProxy { reject(error); } }); - upstreamRes.on('error', reject); + upstreamRes.on('error', (error) => { + clearUpstreamResponseTimeout(); + reject(error); + }); } ); diff --git a/src/cliproxy/proxy/__tests__/tool-sanitization-proxy-integration.test.ts b/src/cliproxy/proxy/__tests__/tool-sanitization-proxy-integration.test.ts index 6da226e2..335019cf 100644 --- a/src/cliproxy/proxy/__tests__/tool-sanitization-proxy-integration.test.ts +++ b/src/cliproxy/proxy/__tests__/tool-sanitization-proxy-integration.test.ts @@ -825,6 +825,49 @@ describe('ToolSanitizationProxy Integration', () => { proxy.stop(); } }); + + it('returns 504 when raw upstream passthrough stalls after headers', async () => { + const upstream = http.createServer((req, res) => { + req.resume(); + req.on('end', () => { + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.flushHeaders(); + }); + }); + const upstreamPort = await new Promise((resolve, reject) => { + upstream.once('error', reject); + upstream.listen(0, '127.0.0.1', () => { + const address = upstream.address(); + if (typeof address !== 'object' || !address) { + reject(new Error('Failed to resolve upstream port')); + return; + } + resolve(address.port); + }); + }); + + const proxy = new ToolSanitizationProxy({ + upstreamBaseUrl: `http://127.0.0.1:${upstreamPort}`, + timeoutMs: 100, + }); + const port = await proxy.start(); + + try { + const response = await fetch(`http://127.0.0.1:${port}/health`); + const text = await Promise.race([ + response.text(), + new Promise((_, reject) => + setTimeout(() => reject(new Error('Timed out waiting for proxy response')), 2_000) + ), + ]); + + expect(response.status).toBe(504); + expect(text).toContain('Upstream response timed out'); + } finally { + proxy.stop(); + await new Promise((resolve) => upstream.close(() => resolve())); + } + }); }); describe('Multiple Tools Sanitization', () => { @@ -1086,5 +1129,65 @@ describe('ToolSanitizationProxy Integration', () => { proxy.stop(); } }); + + it('ends stalled streaming responses after upstream headers', async () => { + const upstream = http.createServer((req, res) => { + req.resume(); + req.on('end', () => { + res.writeHead(200, { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }); + res.write( + 'event: message_start\n' + + 'data: {"type":"message_start","message":{"id":"msg_stall","type":"message","role":"assistant","content":[],"model":"test","stop_reason":null}}\n\n' + ); + }); + }); + const upstreamPort = await new Promise((resolve, reject) => { + upstream.once('error', reject); + upstream.listen(0, '127.0.0.1', () => { + const address = upstream.address(); + if (typeof address !== 'object' || !address) { + reject(new Error('Failed to resolve upstream port')); + return; + } + resolve(address.port); + }); + }); + + const proxy = new ToolSanitizationProxy({ + upstreamBaseUrl: `http://127.0.0.1:${upstreamPort}`, + timeoutMs: 100, + }); + const port = await proxy.start(); + + try { + const response = await fetch(`http://127.0.0.1:${port}/v1/messages`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + stream: true, + tools: [{ name: 'valid_tool' }], + }), + }); + + const text = await Promise.race([ + response.text(), + new Promise((_, reject) => + setTimeout(() => reject(new Error('Timed out waiting for proxy response')), 2_000) + ), + ]); + + expect(response.status).toBe(200); + expect(text).toContain('message_start'); + expect(text).toContain('timeout_error'); + expect(text).toContain('Upstream response timed out'); + } finally { + proxy.stop(); + await new Promise((resolve) => upstream.close(() => resolve())); + } + }); }); }); diff --git a/src/cliproxy/proxy/tool-sanitization-proxy.ts b/src/cliproxy/proxy/tool-sanitization-proxy.ts index 3fed56e4..a49ca71a 100644 --- a/src/cliproxy/proxy/tool-sanitization-proxy.ts +++ b/src/cliproxy/proxy/tool-sanitization-proxy.ts @@ -28,6 +28,10 @@ import { getModelMaxLevel } from '../model-catalog'; import { createLogger } from '../../services/logging'; import { getCcsDir } from '../../config/config-loader-facade'; +import { + attachUpstreamResponseTimeout, + writeForwardResponseHead, +} from './upstream-response-timeout'; export interface ToolSanitizationProxyConfig { /** Upstream CLIProxy URL */ @@ -521,10 +525,38 @@ export class ToolSanitizationProxy { ), (upstreamRes) => { clearResponseTimeout(); - clientRes.writeHead(upstreamRes.statusCode || 200, upstreamRes.headers); - upstreamRes.pipe(clientRes); - upstreamRes.on('end', () => resolve()); - upstreamRes.on('error', reject); + const statusCode = upstreamRes.statusCode || 200; + let responseStarted = false; + const writeResponseHead = () => { + if (responseStarted) return; + responseStarted = true; + writeForwardResponseHead(clientRes, statusCode, upstreamRes.headers); + }; + const clearUpstreamResponseTimeout = attachUpstreamResponseTimeout({ + upstreamReq, + upstreamRes, + clientRes, + timeoutMs: this.config.timeoutMs, + onTimeout: () => resolve(), + }); + upstreamRes.on('data', (chunk: Buffer) => { + writeResponseHead(); + const canContinue = clientRes.write(chunk); + if (!canContinue) { + upstreamRes.pause(); + clientRes.once('drain', () => upstreamRes.resume()); + } + }); + upstreamRes.on('end', () => { + clearUpstreamResponseTimeout(); + writeResponseHead(); + clientRes.end(); + resolve(); + }); + upstreamRes.on('error', (error) => { + clearUpstreamResponseTimeout(); + reject(error); + }); } ); @@ -560,10 +592,18 @@ export class ToolSanitizationProxy { ), (upstreamRes) => { clearResponseTimeout(); + const clearUpstreamResponseTimeout = attachUpstreamResponseTimeout({ + upstreamReq, + upstreamRes, + clientRes, + timeoutMs: this.config.timeoutMs, + onTimeout: () => resolve(), + }); const chunks: Buffer[] = []; upstreamRes.on('data', (chunk: Buffer) => chunks.push(chunk)); upstreamRes.on('end', () => { + clearUpstreamResponseTimeout(); try { const responseBody = Buffer.concat(chunks).toString('utf8'); const contentType = upstreamRes.headers['content-type'] || ''; @@ -598,7 +638,10 @@ export class ToolSanitizationProxy { reject(err); } }); - upstreamRes.on('error', reject); + upstreamRes.on('error', (error) => { + clearUpstreamResponseTimeout(); + reject(error); + }); } ); @@ -636,7 +679,14 @@ export class ToolSanitizationProxy { ), (upstreamRes) => { clearResponseTimeout(); - clientRes.writeHead(upstreamRes.statusCode || 200, upstreamRes.headers); + const clearUpstreamResponseTimeout = attachUpstreamResponseTimeout({ + upstreamReq, + upstreamRes, + clientRes, + timeoutMs: this.config.timeoutMs, + onTimeout: () => resolve(), + }); + writeForwardResponseHead(clientRes, upstreamRes.statusCode || 200, upstreamRes.headers); // Track upstream SSE lifecycle events (guards against empty proxy responses) const lifecycle = { @@ -673,6 +723,7 @@ export class ToolSanitizationProxy { } }); upstreamRes.on('end', () => { + clearUpstreamResponseTimeout(); try { if (!lifecycle.hasContent && isSuccessResponse && lifecycle.hasData) { this.writeLog( @@ -693,7 +744,10 @@ export class ToolSanitizationProxy { } resolve(); }); - upstreamRes.on('error', reject); + upstreamRes.on('error', (error) => { + clearUpstreamResponseTimeout(); + reject(error); + }); return; } @@ -719,6 +773,7 @@ export class ToolSanitizationProxy { }); upstreamRes.on('end', () => { + clearUpstreamResponseTimeout(); try { // Process any remaining buffer if (buffer.trim()) { @@ -749,7 +804,10 @@ export class ToolSanitizationProxy { resolve(); }); - upstreamRes.on('error', reject); + upstreamRes.on('error', (error) => { + clearUpstreamResponseTimeout(); + reject(error); + }); } ); diff --git a/src/cliproxy/proxy/upstream-response-timeout.ts b/src/cliproxy/proxy/upstream-response-timeout.ts new file mode 100644 index 00000000..31d06f3c --- /dev/null +++ b/src/cliproxy/proxy/upstream-response-timeout.ts @@ -0,0 +1,87 @@ +import * as http from 'http'; + +export const UPSTREAM_RESPONSE_TIMEOUT_MESSAGE = + 'Upstream response timed out while streaming response body'; + +export function buildTimeoutSafeResponseHeaders( + headers: http.IncomingHttpHeaders +): http.OutgoingHttpHeaders { + const safeHeaders: http.OutgoingHttpHeaders = {}; + for (const [name, value] of Object.entries(headers)) { + const normalized = name.toLowerCase(); + if (normalized === 'content-length' || normalized === 'transfer-encoding') { + continue; + } + safeHeaders[name] = value; + } + return safeHeaders; +} + +export function writeForwardResponseHead( + clientRes: http.ServerResponse, + statusCode: number, + headers: http.IncomingHttpHeaders +): void { + if (clientRes.headersSent) return; + clientRes.writeHead(statusCode, buildTimeoutSafeResponseHeaders(headers)); +} + +export function writeTimeoutResponse( + clientRes: http.ServerResponse, + headers: http.IncomingHttpHeaders, + message = UPSTREAM_RESPONSE_TIMEOUT_MESSAGE +): void { + if (clientRes.destroyed || clientRes.writableEnded) return; + + const contentType = String(headers['content-type'] ?? '').toLowerCase(); + + try { + if (contentType.includes('text/event-stream')) { + const payload = { + type: 'error', + error: { + type: 'timeout_error', + message, + }, + }; + clientRes.write(`event: error\ndata: ${JSON.stringify(payload)}\n\n`); + } else if (!clientRes.headersSent) { + clientRes.writeHead(504, { 'Content-Type': 'application/json' }); + clientRes.write(JSON.stringify({ error: message })); + } else { + clientRes.write(`\n${JSON.stringify({ error: message })}`); + } + + clientRes.end(); + } catch { + // Client may have disconnected while the upstream response was stalled. + } +} + +export function attachUpstreamResponseTimeout(options: { + upstreamReq: http.ClientRequest; + upstreamRes: http.IncomingMessage; + clientRes: http.ServerResponse; + timeoutMs: number; + onTimeout?: (error: Error) => void; +}): () => void { + const { upstreamReq, upstreamRes, clientRes, timeoutMs, onTimeout } = options; + let settled = false; + + const clear = () => { + settled = true; + upstreamRes.setTimeout(0); + }; + + upstreamRes.setTimeout(timeoutMs, () => { + if (settled) return; + settled = true; + const error = new Error(UPSTREAM_RESPONSE_TIMEOUT_MESSAGE); + writeTimeoutResponse(clientRes, upstreamRes.headers, error.message); + onTimeout?.(error); + upstreamRes.destroy(error); + upstreamReq.destroy(error); + }); + + return clear; +}