1
9import { describe, it, expect } from "vitest";
10import { parseSseFrames, type AiSseFrame } from "@/lib/ai";
11
12const wire = (frames: AiSseFrame[], keepaliveEvery: number) =>
13 frames
14 .map((f, i) => (keepaliveEvery > 0 && i > 0 && i % keepaliveEvery === 0 ? `: keepalive\n\ndata: ${JSON.stringify(f)}\n\n` : `data: ${JSON.stringify(f)}\n\n`))
15 .join("");
16
17describe("parseSseFrames", () => {
18 it("parses every delta type in the documented contract", () => {
19 const frames: AiSseFrame[] = [
20 { type: "delta", text: "Hel" },
21 { type: "delta", text: "lo" },
22 { type: "tool_call", id: "mock_call_0", name: "createPage", arguments: { title: "X" } },
23 { type: "usage", tokensIn: 12, tokensOut: 4 },
24 { type: "error", code: "TIMEOUT", message: "The AI request took too long.", retryable: true, retryAfterSeconds: 5 },
25 { type: "done", assistantMessageId: "a1", userMessageId: "u1", reason: "ok" }
26 ];
27 const { frames: out, rest } = parseSseFrames(wire(frames, 2));
28 expect(rest).toBe("");
29 expect(out).toEqual(frames);
30 });
31
32 it("skips `: keepalive` comment lines", () => {
33 const { frames } = parseSseFrames(": keepalive\n\ndata: {\"type\":\"delta\",\"text\":\"a\"}\n\n: keepalive\n\ndata: {\"type\":\"done\",\"reason\":\"ok\"}\n\n");
34 expect(frames).toEqual([
35 { type: "delta", text: "a" },
36 { type: "done", reason: "ok" }
37 ]);
38 });
39
40 it("keeps a partial trailing frame for the next chunk (chunk boundary)", () => {
41
42 const full = `data: {"type":"usage","tokensIn":7,"tokensOut":3}\n\n`;
43 const { frames: f1, rest: r1 } = parseSseFrames(full.slice(0, 14));
44 expect(f1).toEqual([]);
45 expect(r1).toBe(full.slice(0, 14));
46 const { frames: f2, rest: r2 } = parseSseFrames(r1 + full.slice(14));
47 expect(f2).toEqual([{ type: "usage", tokensIn: 7, tokensOut: 3 }]);
48 expect(r2).toBe("");
49 });
50
51 it("handles multiple frames in one chunk and frames split across chunks", () => {
52 const full = `data: {"type":"delta","text":"a"}\n\ndata: {"type":"delta","text":"b"}\n\ndata: {"type":"done","reason":"ok"}\n\n`;
53
54 const split = full.indexOf("b") + 3;
55 const { frames: f1, rest: r1 } = parseSseFrames(full.slice(0, split));
56 expect(f1).toEqual([{ type: "delta", text: "a" }]);
57 const { frames: f2, rest: r2 } = parseSseFrames(r1 + full.slice(split));
58 expect(f2).toEqual([{ type: "delta", text: "b" }, { type: "done", reason: "ok" }]);
59 expect(r2).toBe("");
60 });
61
62 it("skips a malformed data line without dropping the stream", () => {
63 const { frames } = parseSseFrames('data: {not-json}\n\ndata: {"type":"delta","text":"ok"}\n\n');
64 expect(frames).toEqual([{ type: "delta", text: "ok" }]);
65 });
66
67 it("round-trips the exact mock-provider stream (word deltas + sys tail)", () => {
68
69 const words = "what is the launch plan".split(" ");
70 const stream = words.map((w) => `data: ${JSON.stringify({ type: "delta", text: w === "what" ? w : ` ${w}` })}\n\n`).join("") + `data: ${JSON.stringify({ type: "usage", tokensIn: 3, tokensOut: 3 })}\n\ndata: ${JSON.stringify({ type: "done", assistantMessageId: "a", userMessageId: "u", reason: "ok" })}\n\n`;
71 const { frames, rest } = parseSseFrames(stream);
72 expect(rest).toBe("");
73 expect(frames.filter((f) => f.type === "delta")).toHaveLength(words.length);
74 const text = frames.filter((f) => f.type === "delta").map((f) => (f as { text: string }).text).join("");
75 expect(text).toBe("what is the launch plan");
76 expect(frames[frames.length - 1]).toEqual({ type: "done", assistantMessageId: "a", userMessageId: "u", reason: "ok" });
77 });
78});
79
