diff --git a/src/server/relay.ts b/src/server/relay.ts index 3b5aae236..d3d99902d 100644 --- a/src/server/relay.ts +++ b/src/server/relay.ts @@ -517,11 +517,14 @@ export function relaySseWithHeartbeat( if (!body) return null; const reader = body.getReader(); const decoder = new TextDecoder(); - const heartbeat = new TextEncoder().encode(": opencodex keepalive\n\n"); + const encoder = new TextEncoder(); + const heartbeat = encoder.encode(": opencodex keepalive\n\n"); + const doneSentinel = encoder.encode("data: [DONE]\n\n"); let timer: ReturnType | undefined; let closed = false; let clientCancelled = false; let terminalReported = false; + let doneSeen = false; let buffer = ""; const reportTerminal = (status: ResponsesTerminalStatus) => { @@ -532,6 +535,10 @@ export function relaySseWithHeartbeat( const inspectPayload = (payload: string | null) => { if (!payload) return; + if (payload === "[DONE]") { + doneSeen = true; + return; + } const status = terminalStatusFromSsePayload(payload); if (status) reportTerminal(status); }; @@ -571,6 +578,10 @@ export function relaySseWithHeartbeat( if (done) { buffer += decoder.decode(); if (buffer.trim()) inspectPayload(sseDataPayload(buffer)); + if (terminalReported && !doneSeen) { + controller.enqueue(doneSentinel); + doneSeen = true; + } if (!terminalReported && !clientCancelled) reportTerminal("incomplete"); cleanup(); controller.close(); @@ -578,6 +589,13 @@ export function relaySseWithHeartbeat( } inspectChunk(value); controller.enqueue(value); + if (terminalReported && !doneSeen) { + controller.enqueue(doneSentinel); + doneSeen = true; + cleanup(); + controller.close(); + reader.cancel("Responses terminal event received").catch(() => {}); + } } catch (err) { if (!clientCancelled) reportTerminal("incomplete"); cleanup(); diff --git a/tests/passthrough-abort.test.ts b/tests/passthrough-abort.test.ts index 5c782b607..fb8b951c5 100644 --- a/tests/passthrough-abort.test.ts +++ b/tests/passthrough-abort.test.ts @@ -221,6 +221,21 @@ describe("passthrough relayWithAbort (RC2, passthrough path)", () => { expect(terminals).toEqual(["completed"]); }); + test("SSE passthrough appends DONE after a terminal payload without upstream DONE", async () => { + const enc = new TextEncoder(); + const ac = new AbortController(); + const terminals: string[] = []; + const relayed = relaySseWithHeartbeat(streamFromChunks([ + enc.encode('event: response.completed\ndata: {"type":"response.completed","response":{"id":"r1","status":"completed"}}\n\n'), + ]), ac, 15_000, status => terminals.push(status))!; + + const raw = await readAll(relayed); + + expect(raw).toContain('event: response.completed\ndata: {"type":"response.completed","response":{"id":"r1","status":"completed"}}\n\n'); + expect(raw).toEndWith("data: [DONE]\n\n"); + expect(terminals).toEqual(["completed"]); + }); + test("SSE passthrough treats DONE without a terminal as incomplete", async () => { const enc = new TextEncoder(); const ac = new AbortController();