From bb018a9141d0d7d631e66692e93b71b5513c43e1 Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Thu, 30 Jul 2026 16:14:53 +0200 Subject: [PATCH] fix(ai-gateway): log partial response body when client disconnects mid-stream --- apps/web/src/lib/rewriteModelResponse.test.ts | 20 ++++++- apps/web/src/lib/rewriteModelResponse.ts | 54 ++++++++++++++----- 2 files changed, 59 insertions(+), 15 deletions(-) diff --git a/apps/web/src/lib/rewriteModelResponse.test.ts b/apps/web/src/lib/rewriteModelResponse.test.ts index d6ce3f3feb..ecc14ceb4a 100644 --- a/apps/web/src/lib/rewriteModelResponse.test.ts +++ b/apps/web/src/lib/rewriteModelResponse.test.ts @@ -914,7 +914,25 @@ describe('request log capture', () => { const reader = result.body?.getReader(); await reader?.cancel(); - expect(capture.setReadError).toHaveBeenCalled(); + expect(capture.setReadError).toHaveBeenCalledWith(expect.any(Error), undefined); expect(capture.setBody).not.toHaveBeenCalled(); }); + + test.each(rewriters)( + '%s: records the chunks received before the response stream is cancelled', + async (_name, rewrite) => { + const capture = makeCapture(); + const receivedChunks = 'data: {"id":"gen-1","choices":[]}\n\n'; + const { response: upstream } = hangingSseResponse(receivedChunks); + + const result = await rewrite(upstream, true, capture, null); + const reader = result.body?.getReader(); + await reader?.read(); + await reader?.cancel(); + + expect(capture.setReadError).toHaveBeenCalledTimes(1); + expect(capture.setReadError).toHaveBeenCalledWith(expect.any(Error), receivedChunks); + expect(capture.setBody).not.toHaveBeenCalled(); + } + ); }); diff --git a/apps/web/src/lib/rewriteModelResponse.ts b/apps/web/src/lib/rewriteModelResponse.ts index e65740d40d..b6a870c9f9 100644 --- a/apps/web/src/lib/rewriteModelResponse.ts +++ b/apps/web/src/lib/rewriteModelResponse.ts @@ -263,6 +263,10 @@ async function readResponseText( } } +function partialCapturedBody(capturedChunks: string[] | null): string | undefined { + return capturedChunks && capturedChunks.length > 0 ? capturedChunks.join('') : undefined; +} + async function rewriteSseStream( reader: ReadableStreamDefaultReader, parser: ReturnType, @@ -272,17 +276,12 @@ async function rewriteSseStream( serializeError: (error: ResponseReadError) => string, onFinally: () => void, vercelRequestId: string | null | undefined, - capture: RequestLogCapture | null + capture: RequestLogCapture | null, + capturedChunks: string[] | null ) { const decoder = new TextDecoder(); - // Accumulate the raw upstream text for request logging while the stream is - // being processed anyway, so it doesn't have to be processed a second time. - const capturedChunks: string[] | null = capture ? [] : null; const settleReadError = (error: unknown) => - capture?.setReadError( - error, - capturedChunks && capturedChunks.length > 0 ? capturedChunks.join('') : undefined - ); + capture?.setReadError(error, partialCapturedBody(capturedChunks)); const settleBody = () => { if (capturedChunks) { capturedChunks.push(decoder.decode()); @@ -398,6 +397,11 @@ export async function rewriteModelResponse_ChatCompletions( }); } + // Accumulate the raw upstream text for request logging while the stream is + // being processed anyway, so it doesn't have to be processed a second time. + // Shared with the stream's cancel() callback so a client disconnect still + // logs the partially received response body. + const capturedChunks: string[] | null = capture ? [] : null; const stream = new ReadableStream({ async start(controller) { const reader = response.body?.getReader(); @@ -480,11 +484,15 @@ export async function rewriteModelResponse_ChatCompletions( '\n\n', progress.stop, vercelRequestId, - capture + capture, + capturedChunks ); }, cancel() { - capture?.setReadError(new Error('response stream was cancelled')); + capture?.setReadError( + new Error('response stream was cancelled'), + partialCapturedBody(capturedChunks) + ); }, }); @@ -559,6 +567,11 @@ export async function rewriteModelResponse_Messages( }); } + // Accumulate the raw upstream text for request logging while the stream is + // being processed anyway, so it doesn't have to be processed a second time. + // Shared with the stream's cancel() callback so a client disconnect still + // logs the partially received response body. + const capturedChunks: string[] | null = capture ? [] : null; const stream = new ReadableStream({ async start(controller) { const reader = response.body?.getReader(); @@ -648,11 +661,15 @@ export async function rewriteModelResponse_Messages( '\n\n', progress.stop, vercelRequestId, - capture + capture, + capturedChunks ); }, cancel() { - capture?.setReadError(new Error('response stream was cancelled')); + capture?.setReadError( + new Error('response stream was cancelled'), + partialCapturedBody(capturedChunks) + ); }, }); @@ -708,6 +725,11 @@ export async function rewriteModelResponse_Responses( }); } + // Accumulate the raw upstream text for request logging while the stream is + // being processed anyway, so it doesn't have to be processed a second time. + // Shared with the stream's cancel() callback so a client disconnect still + // logs the partially received response body. + const capturedChunks: string[] | null = capture ? [] : null; const stream = new ReadableStream({ async start(controller) { const reader = response.body?.getReader(); @@ -792,11 +814,15 @@ export async function rewriteModelResponse_Responses( '\n\n', progress.stop, vercelRequestId, - capture + capture, + capturedChunks ); }, cancel() { - capture?.setReadError(new Error('response stream was cancelled')); + capture?.setReadError( + new Error('response stream was cancelled'), + partialCapturedBody(capturedChunks) + ); }, });