feat(ai-provider): messages.stream — the adapter half that was missing #5627
2 changed files+351−1
Addedsrc/__tests__/ai-provider-streaming.test.ts+126−0View fileUnifiedSplit
@@ -0,0 +1,126 @@
1/**
2 * ai-provider streaming — the half of the adapter that was missing.
3 *
4 * The 2026-09-02 audit found `AiClientLike` implemented only
5 * `messages.create`, so setting AI_PROVIDER=openai crashed the four
6 * streaming features (personal-chat, repo-chat, spec-ai, streaming-review)
7 * with `messages.stream is not a function`. These tests pin the stream
8 * contract in BOTH consumption styles those features use:
9 * - async iteration over Anthropic-shaped events, and
10 * - `.on("text")` + `.finalMessage()`.
11 */
12
13import { describe, expect, it } from "bun:test";
14import { createOpenAiCompatibleClient } from "../lib/ai-provider";
15
16function sseResponse(frames: string[]): Response {
17 const body = new ReadableStream<Uint8Array>({
18 start(controller) {
19 const enc = new TextEncoder();
20 for (const f of frames) controller.enqueue(enc.encode(f));
21 controller.close();
22 },
23 });
24 return new Response(body, {
25 status: 200,
26 headers: { "Content-Type": "text/event-stream" },
27 });
28}
29
30const FRAMES = [
31 `data: {"id":"c1","choices":[{"delta":{"content":"Hel"}}]}\n\n`,
32 `data: {"id":"c1","choices":[{"delta":{"content":"lo"}}]}\n\n`,
33 // finish + usage arrive on the tail chunks (stream_options.include_usage)
34 `data: {"id":"c1","choices":[{"delta":{},"finish_reason":"stop"}]}\n\n`,
35 `data: {"id":"c1","choices":[],"usage":{"prompt_tokens":7,"completion_tokens":2}}\n\n`,
36 `data: [DONE]\n\n`,
37];
38
39function client(frames: string[] = FRAMES, status = 200) {
40 return createOpenAiCompatibleClient({
41 baseUrl: "http://localhost:9999",
42 fetchImpl: (async () =>
43 status === 200
44 ? sseResponse(frames)
45 : new Response("nope", { status, statusText: "Bad" })) as unknown as typeof fetch,
46 });
47}
48
49const REQ = {
50 model: "test-model",
51 max_tokens: 32,
52 messages: [{ role: "user", content: "hi" }],
53};
54
55describe("messages.stream — async iteration (streaming-review style)", () => {
56 it("yields Anthropic-shaped text deltas and a usage-bearing message_delta", async () => {
57 const stream = client().messages.stream(REQ);
58 const events: any[] = [];
59 for await (const ev of stream) events.push(ev);
60
61 const deltas = events.filter((e) => e.type === "content_block_delta");
62 expect(deltas.map((d) => d.delta.text).join("")).toBe("Hello");
63 expect(deltas.every((d) => d.delta.type === "text_delta")).toBe(true);
64
65 const md = events.find((e) => e.type === "message_delta");
66 expect(md.usage).toEqual({ input_tokens: 7, output_tokens: 2 });
67 expect(events.at(-1)?.type).toBe("message_stop");
68 });
69});
70
71describe("messages.stream — on('text') + finalMessage (spec-ai style)", () => {
72 it("fires text listeners and resolves finalMessage with the full text + usage", async () => {
73 const stream = client().messages.stream(REQ);
74 let heard = "";
75 stream.on("text", (d) => {
76 heard += d;
77 });
78 const final = await stream.finalMessage();
79 expect(heard).toBe("Hello");
80 expect(final.content[0].text).toBe("Hello");
81 expect(final.usage).toEqual({ input_tokens: 7, output_tokens: 2 });
82 expect(final.stop_reason).toBe("stop");
83 });
84
85 it("finalMessage alone drains the wire — no iteration required", async () => {
86 const final = await client().messages.stream(REQ).finalMessage();
87 expect(final.content[0].text).toBe("Hello");
88 });
89});
90
91describe("messages.stream — failure paths", () => {
92 it("a non-2xx response rejects finalMessage and throws from iteration", async () => {
93 const s1 = client(FRAMES, 500).messages.stream(REQ);
94 await expect(s1.finalMessage()).rejects.toThrow(/stream failed \(500/);
95
96 const s2 = client(FRAMES, 500).messages.stream(REQ);
97 let threw = false;
98 try {
99 for await (const _ of s2) {
100 /* drain */
101 }
102 } catch {
103 threw = true;
104 }
105 expect(threw).toBe(true);
106 });
107
108 it("a torn SSE frame is skipped, not fatal", async () => {
109 const frames = [
110 `data: {"choices":[{"delta":{"content":"ok"}}]}\n\n`,
111 `data: {not json\n\n`,
112 `data: [DONE]\n\n`,
113 ];
114 const final = await client(frames).messages.stream(REQ).finalMessage();
115 expect(final.content[0].text).toBe("ok");
116 });
117
118 it("frames split across reads reassemble", async () => {
119 const frames = [
120 `data: {"choices":[{"delta":{"con`,
121 `tent":"ab"}}]}\n\ndata: [DONE]\n\n`,
122 ];
123 const final = await client(frames).messages.stream(REQ).finalMessage();
124 expect(final.content[0].text).toBe("ab");
125 });
126});
Modifiedsrc/lib/ai-provider.ts+225−1View fileUnifiedSplit
@@ -51,9 +51,30 @@ export interface AiMessageResponse {
5151 usage: { input_tokens: number; output_tokens: number };
5252}
5353
54/**
55 * The streaming slice of the Anthropic surface the platform consumes.
56 *
57 * Four features stream (personal-chat, repo-chat, spec-ai,
58 * streaming-review), and between them they use BOTH consumption styles the
59 * Anthropic SDK offers:
60 * - `for await (const event of stream)` over Anthropic-shaped events
61 * (content_block_delta / message_delta with usage), and
62 * - `stream.on("text", cb)` + `await stream.finalMessage()`.
63 * An adapter stream must serve both or those features crash the moment
64 * AI_PROVIDER=openai — which is exactly what the 2026-09-02 audit found:
65 * `messages.stream is not a function` on all four.
66 */
67export interface AiStreamLike extends AsyncIterable<unknown> {
68 on(event: "text", cb: (delta: string) => void): AiStreamLike;
69 finalMessage(): Promise<AiMessageResponse>;
70}
71
5472/** The slice of the Anthropic client the platform actually consumes. */
5573export interface AiClientLike {
56 messages: { create(req: AiMessageRequest): Promise<AiMessageResponse> };
74 messages: {
75 create(req: AiMessageRequest): Promise<AiMessageResponse>;
76 stream(req: AiMessageRequest): AiStreamLike;
77 };
5778}
5879
5980/**
@@ -204,6 +225,209 @@ export function createOpenAiCompatibleClient(
204225
205226 return fromOpenAiResponse(await res.json(), req.model);
206227 },
228
229 stream(req: AiMessageRequest): AiStreamLike {
230 return createOpenAiSseStream({
231 url,
232 apiKey: opts.apiKey,
233 fetchImpl: doFetch,
234 timeoutMs,
235 req,
236 });
237 },
238 },
239 };
240}
241
242// ---------------------------------------------------------------------------
243// SSE streaming — OpenAI chat-completions chunks translated live into the
244// Anthropic event shapes the four streaming features already consume.
245// ---------------------------------------------------------------------------
246
247/**
248 * One pump drives the connection regardless of how the caller consumes:
249 * events are teed to a queue (for `for await`) AND to text listeners
250 * (`on("text")`), while `finalMessage()` resolves when the wire closes.
251 * That mirrors the Anthropic SDK's MessageStream contract — a caller that
252 * only awaits finalMessage without iterating must still drain the socket.
253 */
254function createOpenAiSseStream(args: {
255 url: string;
256 apiKey?: string;
257 fetchImpl: typeof fetch;
258 timeoutMs: number;
259 req: AiMessageRequest;
260}): AiStreamLike {
261 const textListeners: Array<(delta: string) => void> = [];
262 const queue: unknown[] = [];
263 let queueResolve: (() => void) | null = null;
264 let done = false;
265 let failure: unknown = null;
266
267 let accumulated = "";
268 let usageIn = 0;
269 let usageOut = 0;
270 let stopReason: string | null = null;
271
272 const notify = () => {
273 queueResolve?.();
274 queueResolve = null;
275 };
276 const push = (event: unknown) => {
277 queue.push(event);
278 notify();
279 };
280
281 const finalPromise = (async (): Promise<AiMessageResponse> => {
282 const headers: Record<string, string> = {
283 "content-type": "application/json",
284 };
285 if (args.apiKey) headers.authorization = `Bearer ${args.apiKey}`;
286
287 const controller = new AbortController();
288 const timer = setTimeout(() => controller.abort(), args.timeoutMs);
289 try {
290 const res = await args.fetchImpl(args.url, {
291 method: "POST",
292 headers,
293 body: JSON.stringify({
294 ...toOpenAiBody(args.req),
295 stream: true,
296 // Ask for usage on the final chunk (OpenAI, vLLM honour this;
297 // servers that don't simply omit it and cost falls back to 0 —
298 // visible as "unknown usage", never NaN in the ledger).
299 stream_options: { include_usage: true },
300 }),
301 signal: controller.signal,
302 });
303
304 if (!res.ok || !res.body) {
305 let detail = "";
306 try {
307 detail = (await res.text()).slice(0, 500);
308 } catch {
309 detail = "";
310 }
311 throw new Error(
312 `AI provider stream failed (${res.status} ${res.statusText})` +
313 (detail ? `: ${detail}` : "")
314 );
315 }
316
317 push({
318 type: "message_start",
319 message: {
320 role: "assistant",
321 model: args.req.model,
322 usage: { input_tokens: 0, output_tokens: 0 },
323 },
324 });
325
326 const reader = res.body.getReader();
327 const decoder = new TextDecoder();
328 let buffer = "";
329 for (;;) {
330 const { value, done: eof } = await reader.read();
331 if (eof) break;
332 buffer += decoder.decode(value, { stream: true });
333 // SSE frames are newline-delimited; keep the trailing partial line.
334 const lines = buffer.split("\n");
335 buffer = lines.pop() ?? "";
336 for (const line of lines) {
337 const payload = line.startsWith("data:") ? line.slice(5).trim() : "";
338 if (!payload) continue;
339 if (payload === "[DONE]") continue;
340 let chunk: Record<string, any>;
341 try {
342 chunk = JSON.parse(payload);
343 } catch {
344 continue; // a torn frame is the server's bug, not a crash of ours
345 }
346 const choice = Array.isArray(chunk.choices)
347 ? chunk.choices[0]
348 : undefined;
349 const delta: string =
350 typeof choice?.delta?.content === "string"
351 ? choice.delta.content
352 : "";
353 if (delta) {
354 accumulated += delta;
355 for (const cb of textListeners) {
356 try {
357 cb(delta);
358 } catch {
359 /* a listener error must not kill the stream */
360 }
361 }
362 push({
363 type: "content_block_delta",
364 index: 0,
365 delta: { type: "text_delta", text: delta },
366 });
367 }
368 if (choice?.finish_reason) stopReason = choice.finish_reason;
369 const u = chunk.usage as Record<string, unknown> | undefined;
370 if (u) {
371 if (typeof u.prompt_tokens === "number") usageIn = u.prompt_tokens;
372 if (typeof u.completion_tokens === "number")
373 usageOut = u.completion_tokens;
374 }
375 }
376 }
377
378 push({
379 type: "message_delta",
380 delta: { stop_reason: stopReason },
381 usage: { input_tokens: usageIn, output_tokens: usageOut },
382 });
383 push({ type: "message_stop" });
384
385 return {
386 id: "aiprov_stream",
387 model: args.req.model,
388 role: "assistant",
389 type: "message",
390 stop_reason: stopReason,
391 content: [{ type: "text", text: accumulated }],
392 usage: { input_tokens: usageIn, output_tokens: usageOut },
393 };
394 } catch (err) {
395 // Recorded HERE, not in a .catch on the promise: the finally below
396 // wakes the iterator, and a microtask-later .catch would lose the
397 // race — the iterator would see done=true, failure=null, and end
398 // cleanly on an error. Silent truncation is the worst failure mode a
399 // stream has.
400 failure = err;
401 throw err;
402 } finally {
403 clearTimeout(timer);
404 done = true;
405 notify();
406 }
407 })();
408
409 // A finalMessage nobody awaits must not surface as an unhandled rejection.
410 finalPromise.catch(() => {});
411
412 const stream: AiStreamLike = {
413 on(event, cb) {
414 if (event === "text") textListeners.push(cb);
415 return stream;
416 },
417 finalMessage() {
418 return finalPromise;
419 },
420 async *[Symbol.asyncIterator]() {
421 let cursor = 0;
422 for (;;) {
423 while (cursor < queue.length) yield queue[cursor++];
424 if (failure) throw failure;
425 if (done) return;
426 await new Promise<void>((resolve) => {
427 queueResolve = resolve;
428 });
429 }
207430 },
208431 };
432 return stream;
209433}
210434
c comment · e edit title · m merge · a approve · r request changes · ? shortcuts