1
import { afterEach, describe, expect, it, vi } from "vitest";
2
import { createOpenAiCompatClient } from "@/worker/lib/ai/openai-compat";
3
import { AiMisconfiguredError, type AiFrame } from "@/worker/lib/ai/types";
4
 
5
function sseBody(frames: string[]): ReadableStream<Uint8Array> {
6
  const encoder = new TextEncoder();
7
  return new ReadableStream<Uint8Array>({
8
    start(controller) {
9
      for (const frame of frames) {
10
        controller.enqueue(encoder.encode(frame));
11
      }
12
      controller.close();
13
    },
14
  });
15
}
16
 
17
async function collect(iter: AsyncIterable<AiFrame>): Promise<string> {
18
  let buffer = "";
19
  for await (const frame of iter) {
20
    if (frame.type === "chunk") buffer += frame.text;
21
  }
22
  return buffer;
23
}
24
 
25
afterEach(() => {
26
  vi.unstubAllGlobals();
27
});
28
 
29
describe("openai-compat AI client", () => {
30
  it("throws when endpoint is missing", () => {
31
    expect(() =>
32
      createOpenAiCompatClient({ endpoint: "", apiKey: "", chatModel: "llama3", summarizeModel: "llama3" }),
33
    ).toThrow(AiMisconfiguredError);
34
  });
35
 
36
  it("translates OpenAI SSE delta frames into the normalized envelope", async () => {
37
    vi.stubGlobal(
38
      "fetch",
39
      vi.fn(async () => {
40
        return new Response(
41
          sseBody([
42
            `data: ${JSON.stringify({ choices: [{ delta: { content: "Hello" } }] })}\n\n`,
43
            `data: ${JSON.stringify({ choices: [{ delta: { content: " world" } }] })}\n\n`,
44
            "data: [DONE]\n\n",
45
          ]),
46
          { status: 200, headers: { "content-type": "text/event-stream" } },
47
        );
48
      }),
49
    );
50
 
51
    const client = createOpenAiCompatClient({
52
      endpoint: "http://127.0.0.1:11434/v1",
53
      apiKey: "",
54
      chatModel: "llama3",
55
      summarizeModel: "llama3",
56
    });
57
    const stream = await client.chat([{ role: "user", content: "hi" }]);
58
    await expect(collect(stream)).resolves.toBe("Hello world");
59
  });
60
 
61
  it("returns the summary text from a non-streamed chat completion", async () => {
62
    vi.stubGlobal(
63
      "fetch",
64
      vi.fn(async () => {
65
        return new Response(JSON.stringify({ choices: [{ message: { content: "Summary body." } }] }), {
66
          status: 200,
67
          headers: { "content-type": "application/json" },
68
        });
69
      }),
70
    );
71
 
72
    const client = createOpenAiCompatClient({
73
      endpoint: "http://127.0.0.1:11434/v1",
74
      apiKey: "",
75
      chatModel: "llama3",
76
      summarizeModel: "llama3",
77
    });
78
    await expect(client.summarize("long text body")).resolves.toEqual({ summary: "Summary body." });
79
  });
80
 
81
  it("ignores delta.reasoning_content frames from reasoning-enabled servers", async () => {
82
    vi.stubGlobal(
83
      "fetch",
84
      vi.fn(async () => {
85
        return new Response(
86
          sseBody([
87
            `data: ${JSON.stringify({ choices: [{ delta: { reasoning_content: "Let me think." } }] })}\n\n`,
88
            `data: ${JSON.stringify({ choices: [{ delta: { reasoning_content: " Drafting…" } }] })}\n\n`,
89
            `data: ${JSON.stringify({ choices: [{ delta: { content: "Answer." } }] })}\n\n`,
90
            "data: [DONE]\n\n",
91
          ]),
92
          { status: 200, headers: { "content-type": "text/event-stream" } },
93
        );
94
      }),
95
    );
96
 
97
    const client = createOpenAiCompatClient({
98
      endpoint: "http://127.0.0.1:11434/v1",
99
      apiKey: "",
100
      chatModel: "gemma-reasoning",
101
      summarizeModel: "gemma-reasoning",
102
    });
103
    const stream = await client.chat([{ role: "user", content: "hi" }]);
104
    await expect(collect(stream)).resolves.toBe("Answer.");
105
  });
106
 
107
  it("returns only message.content when summarize response also carries reasoning_content", async () => {
108
    vi.stubGlobal(
109
      "fetch",
110
      vi.fn(async () => {
111
        return new Response(
112
          JSON.stringify({
113
            choices: [
114
              {
115
                message: {
116
                  content: "Three-sentence summary.",
117
                  reasoning_content: "Thinking Process: 1. Analyze… 2. Draft…",
118
                },
119
              },
120
            ],
121
          }),
122
          { status: 200, headers: { "content-type": "application/json" } },
123
        );
124
      }),
125
    );
126
 
127
    const client = createOpenAiCompatClient({
128
      endpoint: "http://127.0.0.1:11434/v1",
129
      apiKey: "",
130
      chatModel: "gemma-reasoning",
131
      summarizeModel: "gemma-reasoning",
132
    });
133
    await expect(client.summarize("doc")).resolves.toEqual({ summary: "Three-sentence summary." });
134
  });
135
 
136
  it("does not include upstream body content in chat error messages", async () => {
137
    const sensitiveBody = "Bearer sk-leakyToken-shouldNotEscape";
138
    vi.stubGlobal(
139
      "fetch",
140
      vi.fn(async () => new Response(sensitiveBody, { status: 500 })),
141
    );
142
 
143
    const client = createOpenAiCompatClient({
144
      endpoint: "http://127.0.0.1:11434/v1",
145
      apiKey: "",
146
      chatModel: "llama3",
147
      summarizeModel: "llama3",
148
    });
149
 
150
    let captured: unknown = null;
151
    try {
152
      await client.chat([{ role: "user", content: "hi" }]);
153
    } catch (err) {
154
      captured = err;
155
    }
156
 
157
    expect(captured).toBeInstanceOf(Error);
158
    const message = (captured as Error).message;
159
    expect(message).not.toContain("sk-leakyToken");
160
    expect(message).not.toContain("Bearer");
161
    expect(message).toMatch(/openai-compat chat failed: 500/);
162
  });
163
 
164
  it("does not include upstream body content in summarize error messages", async () => {
165
    const sensitiveBody = "Authorization: Bearer sk-secret\nrequest_id=abc";
166
    vi.stubGlobal(
167
      "fetch",
168
      vi.fn(async () => new Response(sensitiveBody, { status: 503 })),
169
    );
170
 
171
    const client = createOpenAiCompatClient({
172
      endpoint: "http://127.0.0.1:11434/v1",
173
      apiKey: "",
174
      chatModel: "llama3",
175
      summarizeModel: "llama3",
176
    });
177
 
178
    let captured: unknown = null;
179
    try {
180
      await client.summarize("doc body");
181
    } catch (err) {
182
      captured = err;
183
    }
184
 
185
    expect(captured).toBeInstanceOf(Error);
186
    const message = (captured as Error).message;
187
    expect(message).not.toContain("sk-secret");
188
    expect(message).not.toContain("Authorization");
189
    expect(message).not.toContain("request_id");
190
    expect(message).toMatch(/openai-compat summarize failed: 503/);
191
  });
192
 
193
  it("forwards the abort signal to fetch so upstream is cancelled", async () => {
194
    let received: AbortSignal | undefined;
195
    vi.stubGlobal(
196
      "fetch",
197
      vi.fn(async (_url: string, init: RequestInit) => {
198
        received = init.signal ?? undefined;
199
        // Never resolve until aborted.
200
        return new Promise<Response>((_, reject) => {
201
          init.signal?.addEventListener("abort", () => reject(new Error("aborted")));
202
        });
203
      }),
204
    );
205
 
206
    const client = createOpenAiCompatClient({
207
      endpoint: "http://127.0.0.1:11434/v1",
208
      apiKey: "",
209
      chatModel: "llama3",
210
      summarizeModel: "llama3",
211
    });
212
 
213
    const controller = new AbortController();
214
    const promise = client.chat([{ role: "user", content: "hi" }], { signal: controller.signal });
215
    controller.abort();
216
    await expect(promise).rejects.toThrow();
217
    expect(received).toBe(controller.signal);
218
  });
219
 
220
  it("defaults max_tokens high enough to fit reasoning plus answer", async () => {
221
    let capturedBody: Record<string, unknown> | null = null;
222
    vi.stubGlobal(
223
      "fetch",
224
      vi.fn(async (_url: string, init: RequestInit) => {
225
        capturedBody = JSON.parse(String(init.body)) as Record<string, unknown>;
226
        return new Response(JSON.stringify({ choices: [{ message: { content: "ok." } }] }), {
227
          status: 200,
228
          headers: { "content-type": "application/json" },
229
        });
230
      }),
231
    );
232
 
233
    const client = createOpenAiCompatClient({
234
      endpoint: "http://127.0.0.1:11434/v1",
235
      apiKey: "",
236
      chatModel: "gemma-reasoning",
237
      summarizeModel: "gemma-reasoning",
238
    });
239
    await client.summarize("doc");
240
    expect(capturedBody).not.toBeNull();
241
    expect(capturedBody!.max_tokens).toBeGreaterThanOrEqual(512);
242
  });
243
});