refactor(cliproxy): extract SSE lifecycle tracker and add backpressure handling

Consolidate 5 duplicate boolean tracking variables into a lifecycle
object with an update() method, reducing event detection from 4
locations to centralized calls. Add pause/drain/resume backpressure
to the passthrough path to prevent unbounded memory buffering with
slow clients.
This commit is contained in:
kaitranntt committed 2026-02-07 05:40:24 -05:00
1 parent 68f37efda9
commit 1b3440b939
1 file changed
+33 -53
+33 -53
View File
@@ -439,12 +439,21 @@ export class ToolSanitizationProxy {
(upstreamRes) => { (upstreamRes) => {
clientRes.writeHead(upstreamRes.statusCode || 200, upstreamRes.headers); clientRes.writeHead(upstreamRes.statusCode || 200, upstreamRes.headers);
// Track whether upstream emitted any content (guards against empty proxy responses) // Track upstream SSE lifecycle events (guards against empty proxy responses)
let hasReceivedContent = false; const lifecycle = {
let hasReceivedData = false; hasContent: false,
let hasReceivedMessageStart = false; hasData: false,
let hasReceivedMessageDelta = false; hasMessageStart: false,
let hasReceivedMessageStop = false; hasMessageDelta: false,
hasMessageStop: false,
/** Scan text for SSE lifecycle events and update tracking flags */
update(text: string) {
if (text.includes('"content_block_start"')) this.hasContent = true;
if (text.includes('"message_start"')) this.hasMessageStart = true;
if (text.includes('"message_delta"')) this.hasMessageDelta = true;
if (text.includes('"message_stop"')) this.hasMessageStop = true;
},
};
const isSuccessResponse = const isSuccessResponse =
(upstreamRes.statusCode || 200) >= 200 && (upstreamRes.statusCode || 200) < 300; (upstreamRes.statusCode || 200) >= 200 && (upstreamRes.statusCode || 200) < 300;
@@ -455,34 +464,27 @@ export class ToolSanitizationProxy {
// synthetic block as additional content in the same conversation turn. // synthetic block as additional content in the same conversation turn.
if (!mapper.hasChanges()) { if (!mapper.hasChanges()) {
upstreamRes.on('data', (chunk: Buffer) => { upstreamRes.on('data', (chunk: Buffer) => {
hasReceivedData = true; lifecycle.hasData = true;
const chunkStr = chunk.toString('utf8'); lifecycle.update(chunk.toString('utf8'));
if (chunkStr.includes('"content_block_start"')) { // Respect backpressure: pause upstream if client can't keep up
hasReceivedContent = true; const canContinue = clientRes.write(chunk);
if (!canContinue) {
upstreamRes.pause();
clientRes.once('drain', () => upstreamRes.resume());
} }
if (chunkStr.includes('"message_start"')) {
hasReceivedMessageStart = true;
}
if (chunkStr.includes('"message_delta"')) {
hasReceivedMessageDelta = true;
}
if (chunkStr.includes('"message_stop"')) {
hasReceivedMessageStop = true;
}
clientRes.write(chunk);
}); });
upstreamRes.on('end', () => { upstreamRes.on('end', () => {
try { try {
if (!hasReceivedContent && isSuccessResponse && hasReceivedData) { if (!lifecycle.hasContent && isSuccessResponse && lifecycle.hasData) {
this.writeLog( this.writeLog(
'warn', 'warn',
'[tool-sanitization-proxy] Empty response detected from upstream (no content blocks). Injecting synthetic response to prevent client crash.' '[tool-sanitization-proxy] Empty response detected from upstream (no content blocks). Injecting synthetic response to prevent client crash.'
); );
clientRes.write( clientRes.write(
this.buildSyntheticErrorResponse( this.buildSyntheticErrorResponse(
hasReceivedMessageStart, lifecycle.hasMessageStart,
hasReceivedMessageDelta, lifecycle.hasMessageDelta,
hasReceivedMessageStop lifecycle.hasMessageStop
) )
); );
} }
@@ -500,7 +502,7 @@ export class ToolSanitizationProxy {
let buffer = ''; let buffer = '';
upstreamRes.on('data', (chunk: Buffer) => { upstreamRes.on('data', (chunk: Buffer) => {
hasReceivedData = true; lifecycle.hasData = true;
buffer += chunk.toString('utf8'); buffer += chunk.toString('utf8');
// Process complete SSE events // Process complete SSE events
@@ -510,18 +512,7 @@ export class ToolSanitizationProxy {
for (const event of events) { for (const event of events) {
if (!event.trim()) continue; if (!event.trim()) continue;
if (event.includes('"content_block_start"')) { lifecycle.update(event);
hasReceivedContent = true;
}
if (event.includes('"message_start"')) {
hasReceivedMessageStart = true;
}
if (event.includes('"message_delta"')) {
hasReceivedMessageDelta = true;
}
if (event.includes('"message_stop"')) {
hasReceivedMessageStop = true;
}
const processedEvent = this.processSSEEvent(event, mapper); const processedEvent = this.processSSEEvent(event, mapper);
clientRes.write(processedEvent + '\n\n'); clientRes.write(processedEvent + '\n\n');
@@ -532,33 +523,22 @@ export class ToolSanitizationProxy {
try { try {
// Process any remaining buffer // Process any remaining buffer
if (buffer.trim()) { if (buffer.trim()) {
if (buffer.includes('"content_block_start"')) { lifecycle.update(buffer);
hasReceivedContent = true;
}
if (buffer.includes('"message_start"')) {
hasReceivedMessageStart = true;
}
if (buffer.includes('"message_delta"')) {
hasReceivedMessageDelta = true;
}
if (buffer.includes('"message_stop"')) {
hasReceivedMessageStop = true;
}
const processedEvent = this.processSSEEvent(buffer, mapper); const processedEvent = this.processSSEEvent(buffer, mapper);
clientRes.write(processedEvent + '\n\n'); clientRes.write(processedEvent + '\n\n');
} }
// Safety net: if upstream sent data but no content blocks, inject synthetic response // Safety net: if upstream sent data but no content blocks, inject synthetic response
if (!hasReceivedContent && isSuccessResponse && hasReceivedData) { if (!lifecycle.hasContent && isSuccessResponse && lifecycle.hasData) {
this.writeLog( this.writeLog(
'warn', 'warn',
'[tool-sanitization-proxy] Empty response detected from upstream (no content blocks). Injecting synthetic response to prevent client crash.' '[tool-sanitization-proxy] Empty response detected from upstream (no content blocks). Injecting synthetic response to prevent client crash.'
); );
clientRes.write( clientRes.write(
this.buildSyntheticErrorResponse( this.buildSyntheticErrorResponse(
hasReceivedMessageStart, lifecycle.hasMessageStart,
hasReceivedMessageDelta, lifecycle.hasMessageDelta,
hasReceivedMessageStop lifecycle.hasMessageStop
) )
); );
} }