diff --git a/apps/web/src/lib/rewriteModelResponse.test.ts b/apps/web/src/lib/rewriteModelResponse.test.ts index 935f161e95..817ed77ab6 100644 --- a/apps/web/src/lib/rewriteModelResponse.test.ts +++ b/apps/web/src/lib/rewriteModelResponse.test.ts @@ -39,6 +39,22 @@ function sseResponse(body: string, status = 200): Response { }); } +function hangingSseResponse(body: string): { response: Response; cancel: jest.Mock } { + const encoder = new TextEncoder(); + const cancel = jest.fn(); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(body)); + }, + cancel, + }); + + return { + response: new Response(stream, { headers: { 'content-type': 'text/event-stream' } }), + cancel, + }; +} + function failingResponse(contentType: string, errorName: string, initialBody?: string): Response { const encoder = new TextEncoder(); let pullCount = 0; @@ -302,6 +318,24 @@ describe('rewriteModelResponse_ChatCompletions', () => { expect(dataPayloads(sse)).toContain('[DONE]'); }); + test('cancels upstream and closes immediately after [DONE]', async () => { + const capture = makeCapture(); + const body = + 'data: {"id":"gen-chat","model":"upstream-model","choices":[]}\n\n' + + 'data: [DONE]\n\n' + + 'data: {"id":"ignored","model":"upstream-model","choices":[]}\n\n'; + const { response: upstream, cancel } = hangingSseResponse(body); + + const result = await rewriteModelResponse_ChatCompletions(upstream, true, capture, null); + const sse = await readOutputStream(result); + + expect(dataObjects(sse)).toEqual([{ id: 'gen-chat', model: 'upstream-model', choices: [] }]); + expect(dataPayloads(sse)).toContain('[DONE]'); + expect(cancel).toHaveBeenCalledTimes(1); + expect(capture.setBody).toHaveBeenCalledWith(body); + expect(capture.setReadError).not.toHaveBeenCalled(); + }); + test('adds an empty choices array and strips cost on usage-only chunks', async () => { const upstream = sseResponse( 'data: {"model":"upstream-model","usage":{"cost":1,"is_byok":true,"prompt_tokens":4,"completion_tokens":2,"total_tokens":6,"prompt_tokens_details":{}}}\n\n' @@ -464,6 +498,40 @@ describe('rewriteModelResponse_Messages', () => { expect(dataPayloads(sse)).not.toContain('[DONE]'); }); + + test('cancels upstream and closes immediately after message_stop', async () => { + const capture = makeCapture(); + const body = + 'event: message_stop\ndata: {"type":"message_stop"}\n\n' + + 'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":10},"delta":{}}\n\n'; + const { response: upstream, cancel } = hangingSseResponse(body); + + const result = await rewriteModelResponse_Messages(upstream, true, capture, null); + const sse = await readOutputStream(result); + + expect(dataObjects(sse)).toEqual([{ type: 'message_stop' }]); + expect(dataPayloads(sse)).not.toContain('[DONE]'); + expect(cancel).toHaveBeenCalledTimes(1); + expect(capture.setBody).toHaveBeenCalledWith(body); + expect(capture.setReadError).not.toHaveBeenCalled(); + }); + + test('cancels upstream and closes immediately after a compatible [DONE] sentinel', async () => { + const capture = makeCapture(); + const body = + 'data: [DONE]\n\n' + + 'event: message_delta\ndata: {"type":"message_delta","usage":{"output_tokens":10},"delta":{}}\n\n'; + const { response: upstream, cancel } = hangingSseResponse(body); + + const result = await rewriteModelResponse_Messages(upstream, true, capture, null); + const sse = await readOutputStream(result); + + expect(dataObjects(sse)).toEqual([]); + expect(dataPayloads(sse)).toEqual(['[DONE]']); + expect(cancel).toHaveBeenCalledTimes(1); + expect(capture.setBody).toHaveBeenCalledWith(body); + expect(capture.setReadError).not.toHaveBeenCalled(); + }); }); describe('rewriteModelResponse_Responses', () => { @@ -526,7 +594,7 @@ describe('rewriteModelResponse_Responses', () => { expect(json.usage.prompt_tokens_details.cached_tokens).toBe(0); }); - test('strips the nested response usage in stream events and emits [DONE]', async () => { + test('strips nested response usage and closes after the completed event', async () => { const upstream = sseResponse( 'event: response.completed\n' + 'data: {"type":"response.completed","response":{"model":"upstream-model","usage":{"cost":0.5,"is_byok":true,"prompt_tokens":3,"completion_tokens":1,"total_tokens":4,"prompt_tokens_details":{"cached_tokens":1}}}}\n\n' + @@ -552,8 +620,27 @@ describe('rewriteModelResponse_Responses', () => { expect(event.response.usage.is_byok).toBeUndefined(); expect(event.response.usage.prompt_tokens_details.cached_tokens).toBe(1); expect(sse).toContain('event: response.completed'); - expect(dataPayloads(sse)).toContain('[DONE]'); + expect(dataPayloads(sse)).not.toContain('[DONE]'); }); + + test.each(['response.completed', 'response.incomplete', 'response.failed'])( + 'forwards %s, cancels upstream, and closes without waiting for EOF', + async type => { + const capture = makeCapture(); + const body = + `event: ${type}\ndata: ${JSON.stringify({ type })}\n\n` + + 'event: response.output_text.delta\ndata: {"type":"response.output_text.delta","delta":"ignored"}\n\n'; + const { response: upstream, cancel } = hangingSseResponse(body); + + const result = await rewriteModelResponse_Responses(upstream, true, capture, null); + const sse = await readOutputStream(result); + + expect(dataObjects(sse)).toEqual([{ type }]); + expect(cancel).toHaveBeenCalledTimes(1); + expect(capture.setBody).toHaveBeenCalledWith(body); + expect(capture.setReadError).not.toHaveBeenCalled(); + } + ); }); function makeLogging(overrides?: Partial): RequestLoggingParams { diff --git a/apps/web/src/lib/rewriteModelResponse.ts b/apps/web/src/lib/rewriteModelResponse.ts index 724dd2ff60..c884142a56 100644 --- a/apps/web/src/lib/rewriteModelResponse.ts +++ b/apps/web/src/lib/rewriteModelResponse.ts @@ -254,6 +254,7 @@ async function rewriteSseStream( parser: ReturnType, controller: ReadableStreamDefaultController, doneReceived: () => boolean, + terminalEventReceived: () => boolean, serializeError: (error: ResponseReadError) => string, onFinally: () => void, vercelRequestId: string | null | undefined, @@ -268,6 +269,12 @@ async function rewriteSseStream( error, capturedChunks && capturedChunks.length > 0 ? capturedChunks.join('') : undefined ); + const settleBody = () => { + if (capturedChunks) { + capturedChunks.push(decoder.decode()); + capture?.setBody(capturedChunks.join('')); + } + }; try { while (true) { const { done, value } = await reader.read(); @@ -279,15 +286,29 @@ async function rewriteSseStream( controller.enqueue('data: [DONE]\n\n'); } controller.close(); - if (capturedChunks) { - capturedChunks.push(decoder.decode()); - capture?.setBody(capturedChunks.join('')); - } + settleBody(); return; } const chunk = decoder.decode(value, { stream: true }); capturedChunks?.push(chunk); parser.feed(chunk); + if (terminalEventReceived()) { + if (doneReceived()) { + controller.enqueue('data: [DONE]\n\n'); + } + const cancellation = reader.cancel(); + controller.close(); + settleBody(); + try { + await cancellation; + } catch (error) { + errorExceptInTest( + '[rewriteModelResponse] failed to cancel terminal upstream stream', + error + ); + } + return; + } } } catch (error) { const responseReadError = getResponseReadError(error, vercelRequestId); @@ -377,6 +398,9 @@ export async function rewriteModelResponse_ChatCompletions( const progress = createStreamProgressLogger(); const parser = createParser({ onEvent(event: EventSourceMessage) { + if (doneReceived) { + return; + } progress.eventProcessed(); if (event.data === '[DONE]') { doneReceived = true; @@ -412,6 +436,9 @@ export async function rewriteModelResponse_ChatCompletions( controller.enqueue(eventLine + 'data: ' + JSON.stringify(json) + '\n\n'); }, onComment() { + if (doneReceived) { + return; + } controller.enqueue(': KILO PROCESSING\n\n'); }, }); @@ -421,6 +448,7 @@ export async function rewriteModelResponse_ChatCompletions( parser, controller, () => doneReceived, + () => doneReceived, responseReadError => 'data: ' + JSON.stringify({ @@ -526,10 +554,14 @@ export async function rewriteModelResponse_Messages( } let doneReceived = false; + let terminalEventReceived = false; let generationId: string | undefined; const progress = createStreamProgressLogger(); const parser = createParser({ onEvent(event: EventSourceMessage) { + if (doneReceived || terminalEventReceived) { + return; + } progress.eventProcessed(); if (event.data === '[DONE]') { doneReceived = true; @@ -563,8 +595,12 @@ export async function rewriteModelResponse_Messages( const eventLine = event.event ? 'event: ' + event.event + '\n' : ''; controller.enqueue(eventLine + 'data: ' + JSON.stringify(json) + '\n\n'); + terminalEventReceived = json.type === 'message_stop'; }, onComment() { + if (doneReceived || terminalEventReceived) { + return; + } controller.enqueue(': KILO PROCESSING\n\n'); }, }); @@ -574,6 +610,7 @@ export async function rewriteModelResponse_Messages( parser, controller, () => doneReceived, + () => doneReceived || terminalEventReceived, responseReadError => 'event: error\n' + 'data: ' + @@ -662,11 +699,15 @@ export async function rewriteModelResponse_Responses( } let doneReceived = false; + let terminalEventReceived = false; let generationId: string | undefined; let nextSequenceNumber = 0; const progress = createStreamProgressLogger(); const parser = createParser({ onEvent(event: EventSourceMessage) { + if (doneReceived || terminalEventReceived) { + return; + } progress.eventProcessed(); if (event.data === '[DONE]') { doneReceived = true; @@ -690,8 +731,15 @@ export async function rewriteModelResponse_Responses( } const eventLine = event.event ? 'event: ' + event.event + '\n' : ''; controller.enqueue(eventLine + 'data: ' + JSON.stringify(json) + '\n\n'); + terminalEventReceived = + json.type === 'response.completed' || + json.type === 'response.incomplete' || + json.type === 'response.failed'; }, onComment() { + if (doneReceived || terminalEventReceived) { + return; + } controller.enqueue(': KILO PROCESSING\n\n'); }, }); @@ -701,6 +749,7 @@ export async function rewriteModelResponse_Responses( parser, controller, () => doneReceived, + () => doneReceived || terminalEventReceived, responseReadError => 'event: error\n' + 'data: ' +