1
import { describe, expect, it } from "vitest";
2
import { encodeAiSseChunk, encodeAiSseDone, encodeAiSseError, parseAiSseStream, type AiSseFrame } from "@/shared/ai";
3
 
4
function streamFromBytes(bytes: Uint8Array[]): ReadableStream<Uint8Array> {
5
  return new ReadableStream<Uint8Array>({
6
    start(controller) {
7
      for (const b of bytes) controller.enqueue(b);
8
      controller.close();
9
    },
10
  });
11
}
12
 
13
async function collect(stream: ReadableStream<Uint8Array>): Promise<AiSseFrame[]> {
14
  const frames: AiSseFrame[] = [];
15
  for await (const frame of parseAiSseStream(stream)) {
16
    frames.push(frame);
17
  }
18
  return frames;
19
}
20
 
21
describe("ai sse envelope", () => {
22
  it("round-trips chunk + done frames", async () => {
23
    const stream = streamFromBytes([encodeAiSseChunk("Hello"), encodeAiSseChunk(" world"), encodeAiSseDone()]);
24
    const frames = await collect(stream);
25
    expect(frames).toEqual([
26
      { event: "chunk", data: { text: "Hello" } },
27
      { event: "chunk", data: { text: " world" } },
28
      { event: "done", data: {} },
29
    ]);
30
  });
31
 
32
  it("parses error frames with and without code", async () => {
33
    const stream = streamFromBytes([encodeAiSseError("something broke", "ai_chat_failed")]);
34
    const [frame] = await collect(stream);
35
    expect(frame).toEqual({ event: "error", data: { message: "something broke", code: "ai_chat_failed" } });
36
  });
37
 
38
  it("round-trips done frame with usage payload", async () => {
39
    const stream = streamFromBytes([
40
      encodeAiSseChunk("hi"),
41
      encodeAiSseDone({ promptTokens: 42, completionTokens: 17, totalTokens: 59 }),
42
    ]);
43
    const frames = await collect(stream);
44
    expect(frames).toEqual([
45
      { event: "chunk", data: { text: "hi" } },
46
      { event: "done", data: { usage: { promptTokens: 42, completionTokens: 17, totalTokens: 59 } } },
47
    ]);
48
  });
49
 
50
  it("ignores malformed usage in done frames", async () => {
51
    const malformed = new TextEncoder().encode(
52
      `event: done\ndata: ${JSON.stringify({ usage: { promptTokens: "lots" } })}\n\n`,
53
    );
54
    const stream = streamFromBytes([malformed]);
55
    const [frame] = await collect(stream);
56
    expect(frame).toEqual({ event: "done", data: {} });
57
  });
58
 
59
  it("handles split chunks that cross frame boundaries", async () => {
60
    const full =
61
      new TextDecoder().decode(encodeAiSseChunk("Alpha")) +
62
      new TextDecoder().decode(encodeAiSseChunk("Beta")) +
63
      new TextDecoder().decode(encodeAiSseDone());
64
    const encoder = new TextEncoder();
65
    const mid = Math.floor(full.length / 2);
66
    const stream = streamFromBytes([encoder.encode(full.slice(0, mid)), encoder.encode(full.slice(mid))]);
67
    const frames = await collect(stream);
68
    expect(frames.map((f) => f.event)).toEqual(["chunk", "chunk", "done"]);
69
    expect(frames[0]).toEqual({ event: "chunk", data: { text: "Alpha" } });
70
  });
71
});