fix: time out stalled CLIProxy response streams

2026-07-16 06:16:37 +00:00 · 2026-05-27 11:14:40 -04:00
parent b4a81327b2
commit d7d5ad2ad2
5 changed files with 454 additions and 18 deletions
@@ -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<typeof setTimeout>;
+
+    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[] = [];

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

@@ -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<number>((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<string>((_, 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<void>((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<number>((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<string>((_, 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<void>((resolve) => upstream.close(() => resolve()));
+      }
+    });
  });
 });
@@ -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);
+          });
        }
      );

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