diff --git a/src/lib/sseTextTransform.ts b/src/lib/sseTextTransform.ts index cc692ae6c3..daa03638cb 100644 --- a/src/lib/sseTextTransform.ts +++ b/src/lib/sseTextTransform.ts @@ -54,6 +54,7 @@ export function createSseTextTransform( let errored = false; let currentEventLine = ""; let lastEventLine = ""; + let pendingEventLine = ""; const handleLine = (line: string, controller: TransformStreamDefaultController) => { const trimmed = line.trim(); @@ -62,6 +63,10 @@ export function createSseTextTransform( if (trimmed === "") { currentEventLine = ""; } + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + pendingEventLine = ""; + } controller.enqueue(encoder.encode(line + "\n")); return; } @@ -83,6 +88,10 @@ export function createSseTextTransform( } flushed = true; } + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + pendingEventLine = ""; + } controller.enqueue(encoder.encode(line + "\n")); return; } @@ -98,10 +107,6 @@ export function createSseTextTransform( const isStopSignal = checkIfStopSignal(json); const isSnapshot = checkIfSnapshot(json); - if (!isStopSignal && !isSnapshot) { - lastEventLine = currentEventLine; - } - const METADATA_KEYS = [ "id", "model", "object", "created", "finish_reason", "finishReason", "role", "type", "index", "stop_reason", "stop_sequence", @@ -166,7 +171,15 @@ export function createSseTextTransform( flushed = true; } + if (!isStopSignal && !isSnapshot) { + lastEventLine = currentEventLine; + } + lastJson = json; + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + pendingEventLine = ""; + } controller.enqueue(encoder.encode(prefix + JSON.stringify(json) + "\n")); } catch (err: any) { if (err?.message?.startsWith("[PII]")) { @@ -176,7 +189,12 @@ export function createSseTextTransform( // JSON parsing failed. Check if it looks like JSON that failed to parse. if (trimmedSegment.startsWith("{") || trimmedSegment.startsWith("[")) { console.warn("[SSE-TRANSFORM] Dropping malformed JSON chunk to prevent syntax injection:", trimmedSegment.slice(0, 100)); + pendingEventLine = ""; } else { + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + pendingEventLine = ""; + } // Treat segment as raw text delta (fail-open) const processed = processor(segment, "content"); controller.enqueue(encoder.encode(prefix + processed + "\n")); @@ -189,14 +207,27 @@ export function createSseTextTransform( // Starts with data: but not JSON, process as raw text lastEventLine = currentEventLine; const processed = processor(segment, "content"); + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + pendingEventLine = ""; + } controller.enqueue(encoder.encode(prefix + processed + "\n")); } } else { // Non-data line, pass through (e.g. event: content_block_delta) if (line.startsWith("event:")) { + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + } currentEventLine = line; + pendingEventLine = line; + } else { + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + pendingEventLine = ""; + } + controller.enqueue(encoder.encode(line + "\n")); } - controller.enqueue(encoder.encode(line + "\n")); } }; @@ -235,6 +266,10 @@ export function createSseTextTransform( if (remaining) { handleLine(remaining, controller); } + if (pendingEventLine) { + controller.enqueue(encoder.encode(pendingEventLine + "\n")); + pendingEventLine = ""; + } if (onFlush && !flushed) { const flushedValue = onFlush(lastJson, isJsonStream, lastContentJson); if (flushedValue) { diff --git a/src/lib/streamingPiiTransform.ts b/src/lib/streamingPiiTransform.ts index 4d4592115b..b1d2d310d7 100644 --- a/src/lib/streamingPiiTransform.ts +++ b/src/lib/streamingPiiTransform.ts @@ -24,7 +24,7 @@ export function createPiiSseTransform(options?: PiiTransformOptions): TransformS }; let windowSize = Math.max(200, options?.windowSize ?? (parseInt(process.env.PII_WINDOW_SIZE || "", 10) || 200)); - if (options?.windowSize && process.env.PII_TEST_BYPASS_MIN_WINDOW === "true") { + if (options?.windowSize !== undefined && process.env.PII_TEST_BYPASS_MIN_WINDOW === "true") { windowSize = options.windowSize; } const W = windowSize; diff --git a/tests/unit/streamingPiiTransform.test.ts b/tests/unit/streamingPiiTransform.test.ts index a8484f2dd7..5fe28e9f2c 100644 --- a/tests/unit/streamingPiiTransform.test.ts +++ b/tests/unit/streamingPiiTransform.test.ts @@ -360,6 +360,70 @@ test("two consecutive events with different names each get their own event name assert.ok(output.includes("event: event.type.alpha"), "event-A name should appear in output"); assert.ok(output.includes("event: event.type.beta"), "event-B name should appear in output"); }); + +test("stop signal event name is enqueued correctly without misattribution or loss", async () => { + const transform = (createPiiSseTransform as any)({ windowSize: 10 }); + + const contentEventLine = "event: response.output_text.delta\n"; + const contentData = `data: {"choices":[{"delta":{"content":"abcdefghijklmno"}}]}\n\n`; + + const stopEventLine = "event: response.done\n"; + const stopData = `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n\n`; + const doneLine = `data: [DONE]\n\n`; + + const output = await testTransform(transform, [ + contentEventLine + contentData, + stopEventLine + stopData + doneLine, + ]); + + // The stop signal itself should be preceded by its own event name "response.done" + const stopSignalIndex = output.indexOf('"finish_reason":"stop"'); + assert.ok(stopSignalIndex !== -1, "stop signal should be present in output"); + const sectionBeforeStop = output.slice(0, stopSignalIndex); + const lastEventBeforeStop = sectionBeforeStop.slice(sectionBeforeStop.lastIndexOf("event:")); + assert.ok( + lastEventBeforeStop.includes("event: response.done"), + "stop signal payload must be immediately preceded by event: response.done" + ); +}); + +test("verify keep-alive event preservation (no-data event)", async () => { + const transform = (createPiiSseTransform as any)({ windowSize: 10 }); + + const eventLine = "event: keep-alive\n\n"; + const output = await testTransform(transform, [eventLine]); + + assert.ok(output.includes("event: keep-alive"), "keep-alive event should be preserved"); +}); + +test("verify event line flushed before other non-data lines (e.g. id, retry)", async () => { + const transform = (createPiiSseTransform as any)({ windowSize: 0 }); + + const inputLines = "event: foo\nid: 123\ndata: bar\n\n"; + const output = await testTransform(transform, [inputLines]); + + assert.ok(output.includes("event: foo\nid: 123\ndata: bar"), "event line must be flushed before non-data lines like id"); +}); + +test("verify trailing event line is flushed on stream close", async () => { + const transform = (createPiiSseTransform as any)({ windowSize: 10 }); + + const inputLines = "event: some-trailing-event\n"; + const output = await testTransform(transform, [inputLines]); + + assert.ok(output.includes("event: some-trailing-event"), "trailing event line should be flushed on stream close"); +}); + +test("verify consecutive event lines without intervening data are both preserved", async () => { + const transform = (createPiiSseTransform as any)({ windowSize: 10 }); + + const inputLines = "event: first-event\nevent: second-event\ndata: some-data\n\n"; + const output = await testTransform(transform, [inputLines]); + + assert.ok(output.includes("event: first-event"), "first event should be preserved"); + assert.ok(output.includes("event: second-event"), "second event should be preserved"); +}); + test.after(async () => { if (originalEnv !== undefined) { process.env.PII_RESPONSE_SANITIZATION = originalEnv;