From 844044d58c8a30f21d24462354ebc9c23072f67a Mon Sep 17 00:00:00 2001 From: Roomote Date: Sun, 2 Aug 2026 17:15:24 +0000 Subject: [PATCH] refactor: roll out stream helpers to qwen lmstudio fireworks specs --- src/api/providers/__tests__/fireworks.spec.ts | 106 +++--------- .../__tests__/lmstudio-native-tools.spec.ts | 108 +++++------- .../__tests__/qwen-code-native-tools.spec.ts | 159 ++++++++---------- src/eslint-suppressions.json | 2 +- 4 files changed, 138 insertions(+), 237 deletions(-) diff --git a/src/api/providers/__tests__/fireworks.spec.ts b/src/api/providers/__tests__/fireworks.spec.ts index 33d50ab7b2..a066b139cf 100644 --- a/src/api/providers/__tests__/fireworks.spec.ts +++ b/src/api/providers/__tests__/fireworks.spec.ts @@ -6,6 +6,7 @@ import OpenAI from "openai" import { type FireworksModelId, fireworksDefaultModelId, fireworksModels } from "@roo-code/types" import { FireworksHandler } from "../fireworks" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" // Create mock functions const mockCreate = vi.fn() @@ -29,9 +30,9 @@ describe("FireworksHandler", () => { beforeEach(() => { vi.clearAllMocks() // Set up default mock implementation - mockCreate.mockImplementation(async () => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementation(async () => + asyncStreamFrom([ + { choices: [ { delta: { content: "Test response" }, @@ -39,8 +40,8 @@ describe("FireworksHandler", () => { }, ], usage: null, - } - yield { + }, + { choices: [ { delta: {}, @@ -52,9 +53,9 @@ describe("FireworksHandler", () => { completion_tokens: 5, total_tokens: 15, }, - } - }, - })) + }, + ]), + ) handler = new FireworksHandler({ fireworksApiKey: "test-key" }) }) @@ -436,19 +437,7 @@ describe("FireworksHandler", () => { it("createMessage should yield text content from stream", async () => { const testContent = "This is test content from Fireworks stream" - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - next: vi - .fn() - .mockResolvedValueOnce({ - done: false, - value: { choices: [{ delta: { content: testContent } }] }, - }) - .mockResolvedValueOnce({ done: true }), - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([{ choices: [{ delta: { content: testContent } }] }])) const stream = handler.createMessage("system prompt", []) const firstChunk = await stream.next() @@ -458,19 +447,9 @@ describe("FireworksHandler", () => { }) it("createMessage should yield usage data from stream", async () => { - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - next: vi - .fn() - .mockResolvedValueOnce({ - done: false, - value: { choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 20 } }, - }) - .mockResolvedValueOnce({ done: true }), - }), - } - }) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: {} }], usage: { prompt_tokens: 10, completion_tokens: 20 } }]), + ) const stream = handler.createMessage("system prompt", []) const firstChunk = await stream.next() @@ -487,15 +466,7 @@ describe("FireworksHandler", () => { fireworksApiKey: "test-fireworks-api-key", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const systemPrompt = "Test system prompt for Fireworks" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Fireworks" }] @@ -523,13 +494,7 @@ describe("FireworksHandler", () => { fireworksApiKey: "test-fireworks-api-key", }) - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - })) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system", []) await messageGenerator.next() @@ -549,13 +514,7 @@ describe("FireworksHandler", () => { fireworksApiKey: "test-fireworks-api-key", }) - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - })) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system", []) await messageGenerator.next() @@ -577,13 +536,7 @@ describe("FireworksHandler", () => { modelTemperature: 0.7, }) - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - })) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system", []) await messageGenerator.next() @@ -610,9 +563,9 @@ describe("FireworksHandler", () => { }) it("createMessage should handle stream with multiple chunks", async () => { - mockCreate.mockImplementationOnce(async () => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { choices: [ { delta: { content: "Hello" }, @@ -620,8 +573,8 @@ describe("FireworksHandler", () => { }, ], usage: null, - } - yield { + }, + { choices: [ { delta: { content: " world" }, @@ -629,8 +582,8 @@ describe("FireworksHandler", () => { }, ], usage: null, - } - yield { + }, + { choices: [ { delta: {}, @@ -642,18 +595,15 @@ describe("FireworksHandler", () => { completion_tokens: 10, total_tokens: 15, }, - } - }, - })) + }, + ]), + ) const systemPrompt = "You are a helpful assistant." const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hi" }] const stream = handler.createMessage(systemPrompt, messages) - const chunks = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks[0]).toEqual({ type: "text", text: "Hello" }) expect(chunks[1]).toEqual({ type: "text", text: " world" }) diff --git a/src/api/providers/__tests__/lmstudio-native-tools.spec.ts b/src/api/providers/__tests__/lmstudio-native-tools.spec.ts index c6870e80ef..2e04399f98 100644 --- a/src/api/providers/__tests__/lmstudio-native-tools.spec.ts +++ b/src/api/providers/__tests__/lmstudio-native-tools.spec.ts @@ -2,6 +2,7 @@ // Mock OpenAI client - must come before other imports const mockCreate = vi.fn() +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" vi.mock("openai", () => { return { __esModule: true, @@ -58,13 +59,9 @@ describe("LmStudioHandler Native Tools", () => { describe("Native Tool Calling Support", () => { it("should include tools in request when model supports native tools and tools are provided", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -90,13 +87,9 @@ describe("LmStudioHandler Native Tools", () => { }) it("should include tool_choice when provided", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -113,13 +106,9 @@ describe("LmStudioHandler Native Tools", () => { }) it("should always include tools and tool_choice in request (tools are always present after PR #10841)", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -135,9 +124,9 @@ describe("LmStudioHandler Native Tools", () => { }) it("should yield tool_call_partial chunks during streaming", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [ { delta: { @@ -154,8 +143,8 @@ describe("LmStudioHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: { @@ -170,19 +159,16 @@ describe("LmStudioHandler Native Tools", () => { }, }, ], - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", tools: testTools, }) - const chunks = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks).toContainEqual({ type: "tool_call_partial", @@ -202,13 +188,9 @@ describe("LmStudioHandler Native Tools", () => { }) it("should set parallel_tool_calls based on metadata", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -225,9 +207,9 @@ describe("LmStudioHandler Native Tools", () => { }) it("should yield tool_call_end events when finish_reason is tool_calls", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [ { delta: { @@ -244,17 +226,17 @@ describe("LmStudioHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: {}, finish_reason: "tool_calls", }, ], - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -286,13 +268,9 @@ describe("LmStudioHandler Native Tools", () => { }) it("should work with parallel tool calls disabled (sends false)", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -307,9 +285,9 @@ describe("LmStudioHandler Native Tools", () => { }) it("should handle reasoning content alongside tool calls", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [ { delta: { @@ -317,8 +295,8 @@ describe("LmStudioHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: { @@ -335,17 +313,17 @@ describe("LmStudioHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: {}, finish_reason: "tool_calls", }, ], - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", diff --git a/src/api/providers/__tests__/qwen-code-native-tools.spec.ts b/src/api/providers/__tests__/qwen-code-native-tools.spec.ts index 3615c0f92d..6c7caba260 100644 --- a/src/api/providers/__tests__/qwen-code-native-tools.spec.ts +++ b/src/api/providers/__tests__/qwen-code-native-tools.spec.ts @@ -9,6 +9,7 @@ vi.mock("node:fs", () => ({ })) const mockCreate = vi.fn() +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" vi.mock("openai", () => { return { __esModule: true, @@ -77,13 +78,9 @@ describe("QwenCodeHandler Native Tools", () => { describe("Native Tool Calling Support", () => { it("should include tools in request when model supports native tools and tools are provided", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -107,13 +104,9 @@ describe("QwenCodeHandler Native Tools", () => { }) it("should include tool_choice when provided", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -130,13 +123,9 @@ describe("QwenCodeHandler Native Tools", () => { }) it("should always include tools and tool_choice (tools are guaranteed to be present after ALWAYS_AVAILABLE_TOOLS)", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -151,9 +140,9 @@ describe("QwenCodeHandler Native Tools", () => { }) it("should yield tool_call_partial chunks during streaming", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [ { delta: { @@ -170,8 +159,8 @@ describe("QwenCodeHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: { @@ -186,19 +175,16 @@ describe("QwenCodeHandler Native Tools", () => { }, }, ], - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", tools: testTools, }) - const chunks = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks).toContainEqual({ type: "tool_call_partial", @@ -218,13 +204,9 @@ describe("QwenCodeHandler Native Tools", () => { }) it("should set parallel_tool_calls based on metadata", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { - choices: [{ delta: { content: "Test response" } }], - } - }, - })) + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([{ choices: [{ delta: { content: "Test response" } }] }]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -241,9 +223,9 @@ describe("QwenCodeHandler Native Tools", () => { }) it("should yield tool_call_end events when finish_reason is tool_calls", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [ { delta: { @@ -260,8 +242,8 @@ describe("QwenCodeHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: {}, @@ -269,9 +251,9 @@ describe("QwenCodeHandler Native Tools", () => { }, ], usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", @@ -303,50 +285,44 @@ describe("QwenCodeHandler Native Tools", () => { }) it("streams reasoning chunks from delta.reasoning_content", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { choices: [{ delta: { reasoning_content: "thinking..." }, index: 0 }] } - yield { choices: [{ delta: { content: "answer" }, index: 0 }] } - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [{ delta: { reasoning_content: "thinking..." }, index: 0 }] }, + { choices: [{ delta: { content: "answer" }, index: 0 }] }, + { choices: [{ delta: {}, index: 0 }], usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", []) - const chunks: any[] = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks).toContainEqual({ type: "reasoning", text: "thinking..." }) }) it("falls back to delta.reasoning when reasoning_content is absent", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { choices: [{ delta: { reasoning: "router-style thought" }, index: 0 }] } - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [{ delta: { reasoning: "router-style thought" }, index: 0 }] }, + { choices: [{ delta: {}, index: 0 }], usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", []) - const chunks: any[] = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks).toContainEqual({ type: "reasoning", text: "router-style thought" }) }) it("prefers delta.reasoning_content over delta.reasoning when both are present", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [ { delta: { @@ -356,28 +332,25 @@ describe("QwenCodeHandler Native Tools", () => { index: 0, }, ], - } - yield { + }, + { choices: [{ delta: {}, index: 0 }], usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", []) - const chunks: any[] = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning") expect(reasoningChunks).toEqual([{ type: "reasoning", text: "primary thought" }]) }) it("should preserve thinking block handling alongside tool calls", async () => { - mockCreate.mockImplementationOnce(() => ({ - [Symbol.asyncIterator]: async function* () { - yield { + mockCreate.mockImplementationOnce(() => + asyncStreamFrom([ + { choices: [ { delta: { @@ -385,8 +358,8 @@ describe("QwenCodeHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: { @@ -403,17 +376,17 @@ describe("QwenCodeHandler Native Tools", () => { }, }, ], - } - yield { + }, + { choices: [ { delta: {}, finish_reason: "tool_calls", }, ], - } - }, - })) + }, + ]), + ) const stream = handler.createMessage("test prompt", [], { taskId: "test-task-id", diff --git a/src/eslint-suppressions.json b/src/eslint-suppressions.json index 7558fb6d57..a10b98b445 100644 --- a/src/eslint-suppressions.json +++ b/src/eslint-suppressions.json @@ -256,7 +256,7 @@ }, "api/providers/__tests__/qwen-code-native-tools.spec.ts": { "@typescript-eslint/no-explicit-any": { - "count": 5 + "count": 2 } }, "api/providers/__tests__/sambanova.spec.ts": {