diff --git a/src/ai/providers/openai-compat.ts b/src/ai/providers/openai-compat.ts index d434b8b..53ce267 100644 --- a/src/ai/providers/openai-compat.ts +++ b/src/ai/providers/openai-compat.ts @@ -115,6 +115,56 @@ export async function parseOpenAIStream( let usageData: { promptTokens: number; completionTokens: number; totalTokens: number } | undefined; let finishReason = ""; + // Handle one already-stripped SSE data payload. Extracted so it can run + // both inside the read loop and once more on the trailing buffer. + const handleData = (data: string): void => { + if (data === "[DONE]") return; + try { + const parsed = JSON.parse(data); + const delta = parsed.choices?.[0]?.delta; + if (delta?.content) { + fullContent.push(delta.content); + onChunk?.({ content: delta.content, done: false }); + } + if (delta?.tool_calls) { + for (const tc of delta.tool_calls) { + const idx = tc.index ?? 0; + if (!toolCallMap.has(idx)) { + toolCallMap.set(idx, { + id: tc.id || `call_${idx}`, + name: tc.function?.name || "", + args: "", + }); + } + const existing = toolCallMap.get(idx)!; + if (tc.id) existing.id = tc.id; + if (tc.function?.name) existing.name = tc.function.name; + if (tc.function?.arguments) existing.args += tc.function.arguments; + } + } + if (parsed.model) model = parsed.model; + if (parsed.choices?.[0]?.finish_reason) finishReason = parsed.choices[0].finish_reason; + if (parsed.usage) { + usageData = { + promptTokens: parsed.usage.prompt_tokens || 0, + completionTokens: parsed.usage.completion_tokens || 0, + totalTokens: parsed.usage.total_tokens || 0, + }; + } + } catch { + // skip malformed event + } + }; + + // Accept `data:` with OR without the spec-optional space. Requiring the + // space silently dropped every event from providers (e.g. Gemini) that emit + // `data:{...}`. + const handleLine = (line: string): void => { + if (line.startsWith("data:")) { + handleData(line.slice(5).trim()); + } + }; + if (response.body) { const reader = response.body.getReader(); const decoder = new TextDecoder(); @@ -129,48 +179,18 @@ export async function parseOpenAIStream( buffer = lines.pop() || ""; for (const line of lines) { - if (line.startsWith("data: ")) { - const data = line.slice(6).trim(); - if (data === "[DONE]") continue; - - try { - const parsed = JSON.parse(data); - const delta = parsed.choices?.[0]?.delta; - if (delta?.content) { - fullContent.push(delta.content); - onChunk?.({ content: delta.content, done: false }); - } - if (delta?.tool_calls) { - for (const tc of delta.tool_calls) { - const idx = tc.index ?? 0; - if (!toolCallMap.has(idx)) { - toolCallMap.set(idx, { - id: tc.id || `call_${idx}`, - name: tc.function?.name || "", - args: "", - }); - } - const existing = toolCallMap.get(idx)!; - if (tc.id) existing.id = tc.id; - if (tc.function?.name) existing.name = tc.function.name; - if (tc.function?.arguments) existing.args += tc.function.arguments; - } - } - if (parsed.model) model = parsed.model; - if (parsed.choices?.[0]?.finish_reason) finishReason = parsed.choices[0].finish_reason; - if (parsed.usage) { - usageData = { - promptTokens: parsed.usage.prompt_tokens || 0, - completionTokens: parsed.usage.completion_tokens || 0, - totalTokens: parsed.usage.total_tokens || 0, - }; - } - } catch { - // skip - } - } + handleLine(line); } } + + // Flush any bytes the decoder is still holding, then process the final + // line. A stream that ends without a trailing newline leaves its last + // event in `buffer` — which carries the `usage` record on many providers, + // so dropping it broke cost/token accounting. + buffer += decoder.decode(); + if (buffer.length > 0) { + handleLine(buffer); + } } onChunk?.({ content: "", done: true }); diff --git a/tests/sse-edge-cases.test.ts b/tests/sse-edge-cases.test.ts new file mode 100644 index 0000000..e5d6fa6 --- /dev/null +++ b/tests/sse-edge-cases.test.ts @@ -0,0 +1,56 @@ +import { describe, it, expect } from "vitest"; +import { parseOpenAIStream } from "../src/ai/providers/openai-compat.js"; + +/** + * Regression tests for two SSE parsing edge cases. Ports the fixes from the + * `agentic-hardening` lineage onto the current v3.x stream parser: + * 1. `data:` with no space after the colon (e.g. Gemini's OpenAI-compat mode) + * 2. a final event with no trailing newline (carries the `usage` record) + */ +function streamFrom(chunks: string[]): Response { + const body = new ReadableStream({ + start(controller) { + const enc = new TextEncoder(); + for (const c of chunks) controller.enqueue(enc.encode(c)); + controller.close(); + }, + }); + return new Response(body); +} + +describe("parseOpenAIStream SSE edge cases", () => { + it("accepts `data:` without a space after the colon", async () => { + // No space after `data:` — the old parser required `data: ` and dropped these. + const res = await parseOpenAIStream( + streamFrom([ + 'data:{"choices":[{"delta":{"content":"Hel"}}]}\n', + 'data:{"choices":[{"delta":{"content":"lo"}}]}\n', + "data:[DONE]\n", + ]) + ); + expect(res.content).toBe("Hello"); + }); + + it("processes a final event that has no trailing newline", async () => { + // The usage record arrives as the last event with no closing newline; + // the old loop left it in `buffer` and never parsed it. + const res = await parseOpenAIStream( + streamFrom([ + 'data: {"choices":[{"delta":{"content":"hi"}}]}\n', + 'data: {"choices":[{"delta":{}}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}', + ]) + ); + expect(res.content).toBe("hi"); + expect(res.usage?.totalTokens).toBe(4); + }); + + it("still handles the well-formed `data: ` case", async () => { + const res = await parseOpenAIStream( + streamFrom([ + 'data: {"choices":[{"delta":{"content":"ok"}}]}\n', + "data: [DONE]\n", + ]) + ); + expect(res.content).toBe("ok"); + }); +});