From a426876ad2fb54a78cb6edbd95327c7e48b5e27e Mon Sep 17 00:00:00 2001 From: Roomote Date: Sat, 1 Aug 2026 20:34:39 +0000 Subject: [PATCH 1/2] refactor: roll out stream helpers to lite llm and zai specs --- src/api/providers/__tests__/lite-llm.spec.ts | 431 ++++++++----------- src/api/providers/__tests__/zai.spec.ts | 171 ++------ 2 files changed, 197 insertions(+), 405 deletions(-) diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index a5898513c5..b8b98bf287 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -4,6 +4,7 @@ import { Anthropic } from "@anthropic-ai/sdk" import { LiteLLMHandler } from "../lite-llm" import { ApiHandlerOptions } from "../../../shared/api" import { litellmDefaultModelId, litellmDefaultModelInfo } from "@roo-code/types" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" // Mock vscode first to avoid import errors vi.mock("vscode", () => ({ @@ -89,29 +90,24 @@ describe("LiteLLMHandler", () => { ] // Mock the stream response - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "I'm doing well!" } }], - usage: { - prompt_tokens: 100, - completion_tokens: 50, - cache_creation_input_tokens: 20, - cache_read_input_tokens: 30, - }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "I'm doing well!" } }], + usage: { + prompt_tokens: 100, + completion_tokens: 50, + cache_creation_input_tokens: 20, + cache_read_input_tokens: 30, + }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage(systemPrompt, messages) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) // Verify that create was called with cache control headers const createCall = mockCreate.mock.calls[0][0] @@ -186,27 +182,22 @@ describe("LiteLLMHandler", () => { const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Hello" }] // Mock the stream response - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Hello!" } }], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Hello!" } }], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage(systemPrompt, messages) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) // Verify that create was called with max_completion_tokens instead of max_tokens const createCall = mockCreate.mock.calls[0][0] @@ -241,26 +232,22 @@ describe("LiteLLMHandler", () => { const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test" }] // Mock the stream response - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Response" } }], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Response" } }], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage(systemPrompt, messages) - for await (const chunk of generator) { - // Consume the generator - } + await collectStream(generator) // Verify that create was called with max_completion_tokens for this model variation const createCall = mockCreate.mock.calls[0][0] @@ -286,26 +273,22 @@ describe("LiteLLMHandler", () => { const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test" }] // Mock the stream response - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Response" } }], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Response" } }], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage(systemPrompt, messages) - for await (const chunk of generator) { - // Consume the generator - } + await collectStream(generator) // Verify that create was called with max_tokens for non-GPT-5 models const createCall = mockCreate.mock.calls[0][0] @@ -349,17 +332,15 @@ describe("LiteLLMHandler", () => { }) // Mock the stream response - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Hello!" } }], - usage: { - prompt_tokens: 10, - completion_tokens: 5, - }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Hello!" } }], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), @@ -368,9 +349,7 @@ describe("LiteLLMHandler", () => { const generator = handler.createMessage("You are a helpful assistant", [ { role: "user", content: "Hello" } as unknown as Anthropic.Messages.MessageParam, ]) - for await (const _chunk of generator) { - // consume - } + await collectStream(generator) // Should not include either token field const createCall = mockCreate.mock.calls[0][0] @@ -616,17 +595,15 @@ describe("LiteLLMHandler", () => { ] // Mock the stream response - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "You're welcome!" } }], - usage: { - prompt_tokens: 100, - completion_tokens: 20, - }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "You're welcome!" } }], + usage: { + prompt_tokens: 100, + completion_tokens: 20, + }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), @@ -643,9 +620,7 @@ describe("LiteLLMHandler", () => { } const generator = handler.createMessage(systemPrompt, messages, metadata as any) - for await (const _chunk of generator) { - // Consume the generator - } + await collectStream(generator) // Verify that the assistant message with tool_calls has thought_signature injected const createCall = mockCreate.mock.calls[0][0] @@ -687,14 +662,12 @@ describe("LiteLLMHandler", () => { }, ] - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Response" } }], - usage: { prompt_tokens: 100, completion_tokens: 20 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Response" } }], + usage: { prompt_tokens: 100, completion_tokens: 20 }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), @@ -710,9 +683,7 @@ describe("LiteLLMHandler", () => { } const generator = handler.createMessage(systemPrompt, messages, metadata as any) - for await (const _chunk of generator) { - // Consume - } + await collectStream(generator) // Verify that thought_signature was NOT injected for non-Gemini model const createCall = mockCreate.mock.calls[0][0] @@ -729,28 +700,23 @@ describe("LiteLLMHandler", () => { describe("reasoning field handling", () => { it("should yield reasoning chunks from reasoning_content delta", async () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { reasoning_content: "Let me think..." } }], - usage: null, - } - yield { - choices: [{ delta: { content: "The answer is 42." } }], - usage: { prompt_tokens: 20, completion_tokens: 10 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { reasoning_content: "Let me think..." } }], + usage: null, }, - } + { + choices: [{ delta: { content: "The answer is 42." } }], + usage: { prompt_tokens: 20, completion_tokens: 10 }, + }, + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage("system", [{ role: "user", content: "What is the answer?" }]) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) const reasoningChunk = results.find((c) => c.type === "reasoning") expect(reasoningChunk).toBeDefined() @@ -761,28 +727,23 @@ describe("LiteLLMHandler", () => { }) it("should yield reasoning chunks from reasoning delta field", async () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { reasoning: "Analyzing the problem..." } }], - usage: null, - } - yield { - choices: [{ delta: { content: "Done." } }], - usage: { prompt_tokens: 10, completion_tokens: 5 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { reasoning: "Analyzing the problem..." } }], + usage: null, }, - } + { + choices: [{ delta: { content: "Done." } }], + usage: { prompt_tokens: 10, completion_tokens: 5 }, + }, + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage("system", [{ role: "user", content: "Solve this." }]) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) const reasoningChunk = results.find((c) => c.type === "reasoning") expect(reasoningChunk).toBeDefined() @@ -790,26 +751,19 @@ describe("LiteLLMHandler", () => { }) it("should prefer reasoning_content over reasoning when both are present", async () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [ - { delta: { reasoning_content: "from_reasoning_content", reasoning: "from_reasoning" } }, - ], - usage: { prompt_tokens: 5, completion_tokens: 5 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { reasoning_content: "from_reasoning_content", reasoning: "from_reasoning" } }], + usage: { prompt_tokens: 5, completion_tokens: 5 }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage("system", [{ role: "user", content: "Test." }]) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) const reasoningChunks = results.filter((c) => c.type === "reasoning") expect(reasoningChunks).toHaveLength(1) @@ -817,104 +771,89 @@ describe("LiteLLMHandler", () => { }) it("should not yield reasoning chunk when reasoning field is present but falsy", async () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { reasoning_content: undefined } }], - usage: null, - } - yield { - choices: [{ delta: { reasoning: "" } }], - usage: null, - } - yield { - choices: [{ delta: { content: "Hello" } }], - usage: { prompt_tokens: 5, completion_tokens: 5 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { reasoning_content: undefined } }], + usage: null, }, - } + { + choices: [{ delta: { reasoning: "" } }], + usage: null, + }, + { + choices: [{ delta: { content: "Hello" } }], + usage: { prompt_tokens: 5, completion_tokens: 5 }, + }, + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage("system", [{ role: "user", content: "Hi" }]) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) const reasoningChunks = results.filter((c) => c.type === "reasoning") expect(reasoningChunks).toHaveLength(0) }) it("should preserve whitespace-only reasoning chunks so streamed boundaries survive concatenation", async () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { reasoning_content: "Let's" } }], - usage: null, - } - yield { - choices: [{ delta: { reasoning_content: " " } }], - usage: null, - } - yield { - choices: [{ delta: { reasoning_content: "think" } }], - usage: null, - } - yield { - choices: [{ delta: { reasoning_content: "\n\n" } }], - usage: null, - } - yield { - choices: [{ delta: { reasoning_content: "next" } }], - usage: null, - } - yield { - choices: [{ delta: { content: "Hello" } }], - usage: { prompt_tokens: 5, completion_tokens: 5 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { reasoning_content: "Let's" } }], + usage: null, }, - } + { + choices: [{ delta: { reasoning_content: " " } }], + usage: null, + }, + { + choices: [{ delta: { reasoning_content: "think" } }], + usage: null, + }, + { + choices: [{ delta: { reasoning_content: "\n\n" } }], + usage: null, + }, + { + choices: [{ delta: { reasoning_content: "next" } }], + usage: null, + }, + { + choices: [{ delta: { content: "Hello" } }], + usage: { prompt_tokens: 5, completion_tokens: 5 }, + }, + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage("system", [{ role: "user", content: "Hi" }]) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) const reasoningChunks = results.filter((c) => c.type === "reasoning") expect(reasoningChunks.map((c) => (c as { text: string }).text).join("")).toBe("Let's think\n\nnext") }) it("should fall back to reasoning when reasoning_content is null on the same delta", async () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { reasoning_content: null, reasoning: "fallback thinking" } }], - usage: null, - } - yield { - choices: [{ delta: { content: "Answer." } }], - usage: { prompt_tokens: 5, completion_tokens: 5 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { reasoning_content: null, reasoning: "fallback thinking" } }], + usage: null, }, - } + { + choices: [{ delta: { content: "Answer." } }], + usage: { prompt_tokens: 5, completion_tokens: 5 }, + }, + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage("system", [{ role: "user", content: "Test." }]) - const results = [] - for await (const chunk of generator) { - results.push(chunk) - } + const results = await collectStream(generator) const reasoningChunks = results.filter((c) => c.type === "reasoning") expect(reasoningChunks).toHaveLength(1) @@ -954,23 +893,19 @@ describe("LiteLLMHandler", () => { }, ] - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Response" } }], - usage: { prompt_tokens: 100, completion_tokens: 20 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Response" } }], + usage: { prompt_tokens: 100, completion_tokens: 20 }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage(systemPrompt, messages) - for await (const _chunk of generator) { - // Consume - } + await collectStream(generator) // Verify that tool IDs are truncated to 64 characters or less const createCall = mockCreate.mock.calls[0][0] @@ -1017,23 +952,19 @@ describe("LiteLLMHandler", () => { }, ] - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Response" } }], - usage: { prompt_tokens: 100, completion_tokens: 20 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Response" } }], + usage: { prompt_tokens: 100, completion_tokens: 20 }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage(systemPrompt, messages) - for await (const _chunk of generator) { - // Consume - } + await collectStream(generator) // Verify that tool IDs are unchanged const createCall = mockCreate.mock.calls[0][0] @@ -1085,23 +1016,19 @@ describe("LiteLLMHandler", () => { }, ] - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "Response" } }], - usage: { prompt_tokens: 100, completion_tokens: 20 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "Response" } }], + usage: { prompt_tokens: 100, completion_tokens: 20 }, }, - } + ]) mockCreate.mockReturnValue({ withResponse: vi.fn().mockResolvedValue({ data: mockStream }), }) const generator = handler.createMessage(systemPrompt, messages) - for await (const _chunk of generator) { - // Consume - } + await collectStream(generator) // Verify that truncated tool IDs are unique (hash suffix ensures this) const createCall = mockCreate.mock.calls[0][0] @@ -1125,14 +1052,12 @@ describe("LiteLLMHandler", () => { }) describe("preserveReasoning message conversion", () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "ok" } }], - usage: { prompt_tokens: 1, completion_tokens: 1 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "ok" } }], + usage: { prompt_tokens: 1, completion_tokens: 1 }, }, - } + ]) it("uses convertToR1Format (merging tool-result text) when the model info sets preserveReasoning", async () => { const optionsWithReasoning: ApiHandlerOptions = { @@ -1172,9 +1097,7 @@ describe("LiteLLMHandler", () => { }) const generator = handler.createMessage(systemPrompt, messages) - for await (const _chunk of generator) { - // Consume - } + await collectStream(generator) const createCall = mockCreate.mock.calls[0][0] @@ -1234,9 +1157,7 @@ describe("LiteLLMHandler", () => { }) const generator = handler.createMessage(systemPrompt, messages) - for await (const _chunk of generator) { - // Consume - } + await collectStream(generator) const createCall = mockCreate.mock.calls[0][0] @@ -1264,14 +1185,12 @@ describe("LiteLLMHandler", () => { }) describe("session ID header", () => { - const mockStream = { - async *[Symbol.asyncIterator]() { - yield { - choices: [{ delta: { content: "ok" } }], - usage: { prompt_tokens: 1, completion_tokens: 1 }, - } + const mockStream = asyncStreamFrom([ + { + choices: [{ delta: { content: "ok" } }], + usage: { prompt_tokens: 1, completion_tokens: 1 }, }, - } + ]) it("should send the X-Zoo-Session-ID header when a taskId is provided", async () => { mockCreate.mockReturnValue({ @@ -1281,9 +1200,7 @@ describe("LiteLLMHandler", () => { const generator = handler.createMessage("system", [{ role: "user", content: "hi" }], { taskId: "task-123", }) - for await (const _chunk of generator) { - // drain the stream - } + await collectStream(generator) const requestHeaders = mockCreate.mock.calls[0][1]?.headers expect(requestHeaders).toMatchObject({ "X-Zoo-Session-ID": "task-123" }) @@ -1295,9 +1212,7 @@ describe("LiteLLMHandler", () => { }) const generator = handler.createMessage("system", [{ role: "user", content: "hi" }]) - for await (const _chunk of generator) { - // drain the stream - } + await collectStream(generator) const requestHeaders = mockCreate.mock.calls[0][1]?.headers expect(requestHeaders).not.toHaveProperty("X-Zoo-Session-ID") @@ -1311,9 +1226,7 @@ describe("LiteLLMHandler", () => { const generator = handler.createMessage("system", [{ role: "user", content: "hi" }], { taskId: "", }) - for await (const _chunk of generator) { - // drain the stream - } + await collectStream(generator) const requestHeaders = mockCreate.mock.calls[0][1]?.headers expect(requestHeaders).not.toHaveProperty("X-Zoo-Session-ID") diff --git a/src/api/providers/__tests__/zai.spec.ts b/src/api/providers/__tests__/zai.spec.ts index 66266a2fee..db13bedccf 100644 --- a/src/api/providers/__tests__/zai.spec.ts +++ b/src/api/providers/__tests__/zai.spec.ts @@ -14,6 +14,7 @@ import { } from "@roo-code/types" import { ZAiHandler } from "../zai" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" vitest.mock("openai", () => { const createMock = vitest.fn() @@ -430,19 +431,9 @@ describe("ZAiHandler", () => { it("createMessage should yield text content from stream", async () => { const testContent = "This is test content from Z AI stream" - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - next: vitest - .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() @@ -452,22 +443,14 @@ describe("ZAiHandler", () => { }) it("createMessage should yield usage data from stream", async () => { - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - next: vitest - .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() @@ -485,15 +468,7 @@ describe("ZAiHandler", () => { zaiApiLine: "international_coding", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const systemPrompt = "Test system prompt for Z AI" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message for Z AI" }] @@ -526,15 +501,7 @@ describe("ZAiHandler", () => { zaiApiLine: "international_coding", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -566,15 +533,7 @@ describe("ZAiHandler", () => { modelMaxTokens: 100_000, }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -595,15 +554,7 @@ describe("ZAiHandler", () => { // No reasoningEffort setting - should use model default (medium) }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -625,15 +576,7 @@ describe("ZAiHandler", () => { // No reasoningEffort setting - should use model default (high) }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -655,15 +598,7 @@ describe("ZAiHandler", () => { reasoningEffort: "max", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -685,15 +620,7 @@ describe("ZAiHandler", () => { reasoningEffort: "disable", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -711,15 +638,7 @@ describe("ZAiHandler", () => { reasoningEffort: "medium", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -742,15 +661,7 @@ describe("ZAiHandler", () => { reasoningEffort: "disable", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -773,15 +684,7 @@ describe("ZAiHandler", () => { reasoningEffort: "medium", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -802,15 +705,7 @@ describe("ZAiHandler", () => { zaiApiLine: "international_coding", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -827,15 +722,7 @@ describe("ZAiHandler", () => { zaiApiLine: "international_coding", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() @@ -857,15 +744,7 @@ describe("ZAiHandler", () => { reasoningEffort: "disable", }) - mockCreate.mockImplementationOnce(() => { - return { - [Symbol.asyncIterator]: () => ({ - async next() { - return { done: true } - }, - }), - } - }) + mockCreate.mockImplementationOnce(() => asyncStreamFrom([])) const messageGenerator = handlerWithModel.createMessage("system prompt", []) await messageGenerator.next() From 838d3fba585ed84909df39794da2f4f5519b24a8 Mon Sep 17 00:00:00 2001 From: Roomote Date: Sat, 1 Aug 2026 20:51:43 +0000 Subject: [PATCH 2/2] refactor: address lite zai review findings --- src/api/providers/__tests__/lite-llm.spec.ts | 36 +++++++++++--------- src/api/providers/__tests__/zai.spec.ts | 2 +- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/src/api/providers/__tests__/lite-llm.spec.ts b/src/api/providers/__tests__/lite-llm.spec.ts index b8b98bf287..03399b0256 100644 --- a/src/api/providers/__tests__/lite-llm.spec.ts +++ b/src/api/providers/__tests__/lite-llm.spec.ts @@ -1052,12 +1052,13 @@ describe("LiteLLMHandler", () => { }) describe("preserveReasoning message conversion", () => { - const mockStream = asyncStreamFrom([ - { - choices: [{ delta: { content: "ok" } }], - usage: { prompt_tokens: 1, completion_tokens: 1 }, - }, - ]) + const makeMockStream = () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "ok" } }], + usage: { prompt_tokens: 1, completion_tokens: 1 }, + }, + ]) it("uses convertToR1Format (merging tool-result text) when the model info sets preserveReasoning", async () => { const optionsWithReasoning: ApiHandlerOptions = { @@ -1093,7 +1094,7 @@ describe("LiteLLMHandler", () => { ] mockCreate.mockReturnValue({ - withResponse: vi.fn().mockResolvedValue({ data: mockStream }), + withResponse: vi.fn().mockResolvedValue({ data: makeMockStream() }), }) const generator = handler.createMessage(systemPrompt, messages) @@ -1153,7 +1154,7 @@ describe("LiteLLMHandler", () => { ] mockCreate.mockReturnValue({ - withResponse: vi.fn().mockResolvedValue({ data: mockStream }), + withResponse: vi.fn().mockResolvedValue({ data: makeMockStream() }), }) const generator = handler.createMessage(systemPrompt, messages) @@ -1185,16 +1186,17 @@ describe("LiteLLMHandler", () => { }) describe("session ID header", () => { - const mockStream = asyncStreamFrom([ - { - choices: [{ delta: { content: "ok" } }], - usage: { prompt_tokens: 1, completion_tokens: 1 }, - }, - ]) + const makeMockStream = () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "ok" } }], + usage: { prompt_tokens: 1, completion_tokens: 1 }, + }, + ]) it("should send the X-Zoo-Session-ID header when a taskId is provided", async () => { mockCreate.mockReturnValue({ - withResponse: vi.fn().mockResolvedValue({ data: mockStream }), + withResponse: vi.fn().mockResolvedValue({ data: makeMockStream() }), }) const generator = handler.createMessage("system", [{ role: "user", content: "hi" }], { @@ -1208,7 +1210,7 @@ describe("LiteLLMHandler", () => { it("should not send the X-Zoo-Session-ID header when no taskId is provided", async () => { mockCreate.mockReturnValue({ - withResponse: vi.fn().mockResolvedValue({ data: mockStream }), + withResponse: vi.fn().mockResolvedValue({ data: makeMockStream() }), }) const generator = handler.createMessage("system", [{ role: "user", content: "hi" }]) @@ -1220,7 +1222,7 @@ describe("LiteLLMHandler", () => { it("should not send the X-Zoo-Session-ID header when taskId is an empty string", async () => { mockCreate.mockReturnValue({ - withResponse: vi.fn().mockResolvedValue({ data: mockStream }), + withResponse: vi.fn().mockResolvedValue({ data: makeMockStream() }), }) const generator = handler.createMessage("system", [{ role: "user", content: "hi" }], { diff --git a/src/api/providers/__tests__/zai.spec.ts b/src/api/providers/__tests__/zai.spec.ts index db13bedccf..4e7aa1ca46 100644 --- a/src/api/providers/__tests__/zai.spec.ts +++ b/src/api/providers/__tests__/zai.spec.ts @@ -14,7 +14,7 @@ import { } from "@roo-code/types" import { ZAiHandler } from "../zai" -import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" +import { asyncStreamFrom } from "../../../test-utils/stream" vitest.mock("openai", () => { const createMock = vitest.fn()