1
import { describe, expect, it } from "vitest";
2
import type { AiUsage } from "@/shared/ai";
3
import { parseProviderSseFrames, parseProviderTokenUsage } from "@/worker/lib/ai/sse-transform";
4
import type { AiFrame } from "@/worker/lib/ai/types";
5
 
6
function providerStream(frames: string[]): ReadableStream<Uint8Array> {
7
  const encoder = new TextEncoder();
8
  return new ReadableStream<Uint8Array>({
9
    start(controller) {
10
      for (const f of frames) controller.enqueue(encoder.encode(f));
11
      controller.close();
12
    },
13
  });
14
}
15
 
16
async function collect(iter: AsyncIterable<AiFrame>): Promise<AiFrame[]> {
17
  const out: AiFrame[] = [];
18
  for await (const frame of iter) out.push(frame);
19
  return out;
20
}
21
 
22
describe("parseProviderSseFrames — workers-ai shape", () => {
23
  const extractChunkText = (p: unknown): string | null => {
24
    if (typeof p !== "object" || p === null) return null;
25
    const r = (p as { response?: unknown }).response;
26
    return typeof r === "string" && r.length > 0 ? r : null;
27
  };
28
  const extractUsage = (p: unknown): AiUsage | null => {
29
    if (typeof p !== "object" || p === null) return null;
30
    return parseProviderTokenUsage((p as { usage?: unknown }).usage);
31
  };
32
 
33
  it("emits chunks and a usage frame when [DONE] sentinel arrives", async () => {
34
    const source = providerStream([
35
      `data: ${JSON.stringify({ response: "Hello" })}\n\n`,
36
      `data: ${JSON.stringify({ response: " world" })}\n\n`,
37
      `data: ${JSON.stringify({ response: "", usage: { prompt_tokens: 12, completion_tokens: 3, total_tokens: 15 } })}\n\n`,
38
      `data: [DONE]\n\n`,
39
    ]);
40
    const out = await collect(
41
      parseProviderSseFrames(source, { extractChunkText, extractUsage, errorLabel: "workers-ai" }),
42
    );
43
    expect(out).toEqual([
44
      { type: "chunk", text: "Hello" },
45
      { type: "chunk", text: " world" },
46
      { type: "usage", usage: { promptTokens: 12, completionTokens: 3, totalTokens: 15 } },
47
    ]);
48
  });
49
 
50
  it("omits a usage frame when the provider doesn't report any", async () => {
51
    const source = providerStream([`data: ${JSON.stringify({ response: "hi" })}\n\n`, `data: [DONE]\n\n`]);
52
    const out = await collect(
53
      parseProviderSseFrames(source, { extractChunkText, extractUsage, errorLabel: "workers-ai" }),
54
    );
55
    expect(out).toEqual([{ type: "chunk", text: "hi" }]);
56
  });
57
 
58
  it("ends cleanly when the source closes without [DONE]", async () => {
59
    const source = providerStream([`data: ${JSON.stringify({ response: "abc" })}\n\n`]);
60
    const out = await collect(parseProviderSseFrames(source, { extractChunkText, errorLabel: "workers-ai" }));
61
    expect(out).toEqual([{ type: "chunk", text: "abc" }]);
62
  });
63
 
64
  it("ignores malformed JSON payloads", async () => {
65
    const source = providerStream([
66
      `data: not json\n\n`,
67
      `data: ${JSON.stringify({ response: "ok" })}\n\n`,
68
      `data: [DONE]\n\n`,
69
    ]);
70
    const out = await collect(parseProviderSseFrames(source, { extractChunkText, errorLabel: "workers-ai" }));
71
    expect(out).toEqual([{ type: "chunk", text: "ok" }]);
72
  });
73
});
74
 
75
describe("parseProviderSseFrames — openai-compat shape", () => {
76
  const extractChunkText = (p: unknown): string | null => {
77
    if (typeof p !== "object" || p === null) return null;
78
    const choices = (p as { choices?: unknown }).choices;
79
    if (!Array.isArray(choices) || choices.length === 0) return null;
80
    const delta = (choices[0] as { delta?: { content?: unknown } } | null)?.delta;
81
    return typeof delta?.content === "string" && delta.content.length > 0 ? delta.content : null;
82
  };
83
  const extractUsage = (p: unknown): AiUsage | null => {
84
    if (typeof p !== "object" || p === null) return null;
85
    return parseProviderTokenUsage((p as { usage?: unknown }).usage);
86
  };
87
 
88
  it("handles include_usage frame (empty choices + usage)", async () => {
89
    const source = providerStream([
90
      `data: ${JSON.stringify({ choices: [{ delta: { content: "he" } }] })}\n\n`,
91
      `data: ${JSON.stringify({ choices: [{ delta: { content: "llo" } }] })}\n\n`,
92
      `data: ${JSON.stringify({ choices: [], usage: { prompt_tokens: 5, completion_tokens: 2, total_tokens: 7 } })}\n\n`,
93
      `data: [DONE]\n\n`,
94
    ]);
95
    const out = await collect(
96
      parseProviderSseFrames(source, { extractChunkText, extractUsage, errorLabel: "openai-compat" }),
97
    );
98
    expect(out).toEqual([
99
      { type: "chunk", text: "he" },
100
      { type: "chunk", text: "llo" },
101
      { type: "usage", usage: { promptTokens: 5, completionTokens: 2, totalTokens: 7 } },
102
    ]);
103
  });
104
});
105
 
106
describe("parseProviderTokenUsage", () => {
107
  it("maps snake_case provider tokens to camelCase", () => {
108
    expect(parseProviderTokenUsage({ prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 })).toEqual({
109
      promptTokens: 1,
110
      completionTokens: 2,
111
      totalTokens: 3,
112
    });
113
  });
114
 
115
  it("returns null when no recognized token fields are present", () => {
116
    expect(parseProviderTokenUsage({})).toBeNull();
117
    expect(parseProviderTokenUsage(null)).toBeNull();
118
    expect(parseProviderTokenUsage({ prompt_tokens: "many" })).toBeNull();
119
  });
120
 
121
  it("partial usage is preserved", () => {
122
    expect(parseProviderTokenUsage({ total_tokens: 42 })).toEqual({ totalTokens: 42 });
123
  });
124
});