writer-work-flow/apps/web/lib/stream/sse.test.ts

import { describe, expect, it } from "vitest";

import {
  SseFrameBuffer,
  initialStreamState,
  parseSseBlock,
  reduceStream,
  type SseEvent,
} from "./sse";

describe("parseSseBlock", () => {
  it("parses a token frame", () => {
    const evt = parseSseBlock('event:token\ndata:{"text":"夜"}');
    expect(evt).toEqual({ event: "token", data: { text: "夜" } });
  });

  it("parses a done frame", () => {
    const evt = parseSseBlock('event:done\ndata:{"length":12}');
    expect(evt).toEqual({ event: "done", data: { length: 12 } });
  });

  it("parses an error frame with request_id", () => {
    const evt = parseSseBlock(
      'event:error\ndata:{"code":"INTERNAL","message":"boom","request_id":"r1"}',
    );
    expect(evt).toEqual({
      event: "error",
      data: { code: "INTERNAL", message: "boom", request_id: "r1" },
    });
  });

  it("ignores comment/heartbeat lines and tolerates leading space after colon", () => {
    const evt = parseSseBlock(': keep-alive\nevent: token\ndata: {"text":"x"}');
    expect(evt).toEqual({ event: "token", data: { text: "x" } });
  });

  it("returns null for unknown event types", () => {
    expect(parseSseBlock('event:section\ndata:{"x":1}')).toBeNull();
  });

  it("returns null for malformed json", () => {
    expect(parseSseBlock("event:token\ndata:{not json")).toBeNull();
  });
});

describe("SseFrameBuffer", () => {
  it("emits complete blocks split by blank line and buffers the tail", () => {
    const buf = new SseFrameBuffer();
    const first = buf.push('event:token\ndata:{"text":"a"}\n\nevent:tok');
    expect(first).toEqual([{ event: "token", data: { text: "a" } }]);
    const second = buf.push('en\ndata:{"text":"b"}\n\n');
    expect(second).toEqual([{ event: "token", data: { text: "b" } }]);
  });

  it("handles CRLF boundaries", () => {
    const buf = new SseFrameBuffer();
    const out = buf.push('event:token\r\ndata:{"text":"z"}\r\n\r\n');
    expect(out).toEqual([{ event: "token", data: { text: "z" } }]);
  });
});

describe("reduceStream", () => {
  it("accumulates token text and marks streaming", () => {
    const tokens: SseEvent[] = [
      { event: "token", data: { text: "夜色" } },
      { event: "token", data: { text: "如墨" } },
    ];
    const state = tokens.reduce(reduceStream, initialStreamState);
    expect(state.text).toBe("夜色如墨");
    expect(state.phase).toBe("streaming");
  });

  it("marks done on done event preserving text", () => {
    let s = reduceStream(initialStreamState, {
      event: "token",
      data: { text: "abc" },
    });
    s = reduceStream(s, { event: "done", data: { length: 3 } });
    expect(s.phase).toBe("done");
    expect(s.text).toBe("abc");
  });

  it("captures error and retains partial text", () => {
    let s = reduceStream(initialStreamState, {
      event: "token",
      data: { text: "partial" },
    });
    s = reduceStream(s, {
      event: "error",
      data: { code: "LLM_UNAVAILABLE", message: "no key" },
    });
    expect(s.phase).toBe("error");
    expect(s.text).toBe("partial");
    expect(s.error?.code).toBe("LLM_UNAVAILABLE");
  });
});