From: Stefan Gasser Date: Tue, 23 Jun 2026 09:42:43 +0000 (+0200) Subject: Fix OpenAI stream line buffering (#112) X-Git-Tag: v0.7.0~5 X-Git-Url: http://git.99rst.org/?a=commitdiff_plain;h=1f7ac64682c1b91a0a56ecc2162b32f931f7836b;p=sgasser-llm-shield.git Fix OpenAI stream line buffering (#112) --- diff --git a/src/providers/openai/stream-transformer.test.ts b/src/providers/openai/stream-transformer.test.ts index 2566fcc..d85acc5 100644 --- a/src/providers/openai/stream-transformer.test.ts +++ b/src/providers/openai/stream-transformer.test.ts @@ -10,9 +10,6 @@ const defaultConfig: MaskingConfig = { denylist: [], }; -/** - * Helper to create a ReadableStream from SSE data - */ function createSSEStream(chunks: string[]): ReadableStream { const encoder = new TextEncoder(); let index = 0; @@ -29,9 +26,6 @@ function createSSEStream(chunks: string[]): ReadableStream { }); } -/** - * Helper to consume a stream and return all chunks as string - */ async function consumeStream(stream: ReadableStream): Promise { const reader = stream.getReader(); const decoder = new TextDecoder(); @@ -60,6 +54,21 @@ describe("createUnmaskingStream", () => { expect(result).toContain("Hello test@test.com!"); }); + test("buffers a data line split mid-json across chunks", async () => { + const context = createMaskingContext(); + context.mapping["[[EMAIL_ADDRESS_1]]"] = "a@b.com"; + + const event = `data: {"choices":[{"delta":{"content":"Hello [[EMAIL_ADDRESS_1]]"}}]}\n\n`; + const splitAt = event.indexOf("ADDRESS_1"); + const source = createSSEStream([event.slice(0, splitAt), event.slice(splitAt)]); + + const unmaskedStream = createUnmaskingStream(source, context, defaultConfig); + const result = await consumeStream(unmaskedStream); + + expect(result).toContain("Hello a@b.com"); + expect(result).not.toContain("[[EMAIL_ADDRESS_1]]"); + }); + test("handles [DONE] message", async () => { const context = createMaskingContext(); diff --git a/src/providers/openai/stream-transformer.ts b/src/providers/openai/stream-transformer.ts index c0c6a81..1051c41 100644 --- a/src/providers/openai/stream-transformer.ts +++ b/src/providers/openai/stream-transformer.ts @@ -40,14 +40,6 @@ function unmaskTextContent( return { text: processedText, piiBuffer: nextPiiBuffer, secretsBuffer: nextSecretsBuffer }; } -/** - * Creates a transform stream that unmasks SSE content - * - * Processes Server-Sent Events (SSE) chunks, buffering partial placeholders - * and unmasking complete ones before forwarding to the client. - * - * Supports both PII unmasking and secrets unmasking, or either alone. - */ export function createUnmaskingStream( source: ReadableStream, piiContext: PlaceholderContext | undefined, @@ -58,27 +50,100 @@ export function createUnmaskingStream( const encoder = new TextEncoder(); let piiBuffer = ""; let secretsBuffer = ""; + let lineBuffer = ""; return new ReadableStream({ async start(controller) { const reader = source.getReader(); + function processLine(line: string) { + if (line.startsWith("data: ")) { + const data = line.slice(6); + + if (data === "[DONE]") { + controller.enqueue(encoder.encode("data: [DONE]\n\n")); + return; + } + + try { + const parsed = JSON.parse(data); + const content = parsed.choices?.[0]?.delta?.content; + + if (typeof content === "string" && content !== "") { + const unmasked = unmaskTextContent( + content, + piiBuffer, + piiContext, + config, + secretsBuffer, + secretsContext, + ); + piiBuffer = unmasked.piiBuffer; + secretsBuffer = unmasked.secretsBuffer; + + if (unmasked.text) { + parsed.choices[0].delta.content = unmasked.text; + controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); + } + } else if (Array.isArray(content)) { + const processedContent = content.flatMap((part: OpenAIContentPart) => { + if (part.type !== "text" || typeof part.text !== "string") { + return [part]; + } + + const unmasked = unmaskTextContent( + part.text, + piiBuffer, + piiContext, + config, + secretsBuffer, + secretsContext, + ); + piiBuffer = unmasked.piiBuffer; + secretsBuffer = unmasked.secretsBuffer; + + if (!unmasked.text) { + return []; + } + + return [{ ...part, text: unmasked.text }]; + }); + + if (processedContent.length > 0) { + parsed.choices[0].delta.content = processedContent; + controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); + } + } else { + controller.enqueue(encoder.encode(`data: ${data}\n\n`)); + } + } catch { + controller.enqueue(encoder.encode(`${line}\n`)); + } + } else if (line.trim()) { + controller.enqueue(encoder.encode(`${line}\n`)); + } + } + try { while (true) { const { done, value } = await reader.read(); if (done) { - // Flush remaining buffer content before closing + lineBuffer += decoder.decode(); + + if (lineBuffer) { + processLine(lineBuffer); + lineBuffer = ""; + } + let flushed = ""; - // Flush PII buffer first if (piiBuffer && piiContext) { flushed = flushMaskingBuffer(piiBuffer, piiContext, config); } else if (piiBuffer) { flushed = piiBuffer; } - // Then flush secrets buffer if (secretsBuffer && secretsContext) { flushed += flushSecretsMaskingBuffer(secretsBuffer, secretsContext); } else if (secretsBuffer) { @@ -104,77 +169,12 @@ export function createUnmaskingStream( break; } - const chunk = decoder.decode(value, { stream: true }); - const lines = chunk.split("\n"); + lineBuffer += decoder.decode(value, { stream: true }); + const lines = lineBuffer.split("\n"); + lineBuffer = lines.pop() || ""; for (const line of lines) { - if (line.startsWith("data: ")) { - const data = line.slice(6); - - if (data === "[DONE]") { - controller.enqueue(encoder.encode("data: [DONE]\n\n")); - continue; - } - - try { - const parsed = JSON.parse(data); - const content = parsed.choices?.[0]?.delta?.content; - - if (typeof content === "string" && content !== "") { - const unmasked = unmaskTextContent( - content, - piiBuffer, - piiContext, - config, - secretsBuffer, - secretsContext, - ); - piiBuffer = unmasked.piiBuffer; - secretsBuffer = unmasked.secretsBuffer; - - if (unmasked.text) { - parsed.choices[0].delta.content = unmasked.text; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); - } - } else if (Array.isArray(content)) { - const processedContent = content.flatMap((part: OpenAIContentPart) => { - if (part.type !== "text" || typeof part.text !== "string") { - return [part]; - } - - const unmasked = unmaskTextContent( - part.text, - piiBuffer, - piiContext, - config, - secretsBuffer, - secretsContext, - ); - piiBuffer = unmasked.piiBuffer; - secretsBuffer = unmasked.secretsBuffer; - - if (!unmasked.text) { - return []; - } - - return [{ ...part, text: unmasked.text }]; - }); - - if (processedContent.length > 0) { - parsed.choices[0].delta.content = processedContent; - controller.enqueue(encoder.encode(`data: ${JSON.stringify(parsed)}\n\n`)); - } - } else { - // Pass through non-content events - controller.enqueue(encoder.encode(`data: ${data}\n\n`)); - } - } catch { - // Pass through unparseable data - controller.enqueue(encoder.encode(`${line}\n`)); - } - } else if (line.trim()) { - controller.enqueue(encoder.encode(`${line}\n`)); - } + processLine(line); } } } catch (error) {