From 5a7f681088109e646b8244eb801d607ac395360a Mon Sep 17 00:00:00 2001 From: Roomote Date: Sat, 1 Aug 2026 17:31:35 +0000 Subject: [PATCH 1/2] refactor: roll out stream helpers to anthropic vertex spec --- .../__tests__/anthropic-vertex.spec.ts | 328 ++++++------------ 1 file changed, 105 insertions(+), 223 deletions(-) diff --git a/src/api/providers/__tests__/anthropic-vertex.spec.ts b/src/api/providers/__tests__/anthropic-vertex.spec.ts index 706bce58e7..307db6ba37 100644 --- a/src/api/providers/__tests__/anthropic-vertex.spec.ts +++ b/src/api/providers/__tests__/anthropic-vertex.spec.ts @@ -9,6 +9,7 @@ import { VERTEX_1M_CONTEXT_MODEL_IDS } from "@roo-code/types" import { ApiStreamChunk } from "../../transform/stream" import { AnthropicVertexHandler } from "../anthropic-vertex" +import { asyncStreamFrom, collectStream } from "../../../test-utils/stream" vitest.mock("../utils/timeout-config", () => ({ getApiRequestTimeout: vitest.fn().mockReturnValue(300_000), @@ -39,26 +40,24 @@ vitest.mock("@anthropic-ai/vertex-sdk", () => ({ }, } } - return { - async *[Symbol.asyncIterator]() { - yield { - type: "message_start", - message: { - usage: { - input_tokens: 10, - output_tokens: 5, - }, - }, - } - yield { - type: "content_block_start", - content_block: { - type: "text", - text: "Test response", + return asyncStreamFrom([ + { + type: "message_start", + message: { + usage: { + input_tokens: 10, + output_tokens: 5, }, - } + }, }, - } + { + type: "content_block_start", + content_block: { + type: "text", + text: "Test response", + }, + }, + ]) }), }, } @@ -194,23 +193,13 @@ describe("VertexHandler", () => { ] // Setup async iterator for mock stream - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate const stream = handler.createMessage(systemPrompt, mockMessages) - const chunks: ApiStreamChunk[] = [] - - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks.length).toBe(4) expect(chunks[0]).toEqual({ @@ -296,23 +285,13 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate const stream = handler.createMessage(systemPrompt, mockMessages) - const chunks: ApiStreamChunk[] = [] - - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks.length).toBe(3) expect(chunks[0]).toEqual({ @@ -343,9 +322,7 @@ describe("VertexHandler", () => { const stream = handler.createMessage(systemPrompt, mockMessages) await expect(async () => { - for await (const _chunk of stream) { - // Should throw before yielding any chunks - } + await collectStream(stream) }).rejects.toThrow("Vertex API error") }) @@ -391,13 +368,7 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate @@ -417,10 +388,7 @@ describe("VertexHandler", () => { }, ]) - const chunks: ApiStreamChunk[] = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) // Verify usage information const usageChunks = chunks.filter((chunk) => chunk.type === "usage") @@ -514,23 +482,13 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate const stream = handler.createMessage(systemPrompt, mockMessages) - const chunks: ApiStreamChunk[] = [] - - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) // Check for cache-related metrics in usage chunk const usageChunks = chunks.filter((chunk) => chunk.type === "usage") @@ -593,23 +551,13 @@ describe("VertexHandler", () => { ] // Setup async iterator for mock stream - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate const stream = handler.createMessage(systemPrompt, mockMessages) - const chunks: ApiStreamChunk[] = [] - - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) // Verify thinking content is processed correctly const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning") @@ -650,23 +598,13 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate const stream = handler.createMessage(systemPrompt, mockMessages) - const chunks: ApiStreamChunk[] = [] - - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) expect(chunks.length).toBe(3) expect(chunks[0]).toEqual({ @@ -690,29 +628,27 @@ describe("VertexHandler", () => { vertexRegion: "us-central1", }) - const mockCreate = vitest.fn().mockImplementation(async (options) => { - return { - async *[Symbol.asyncIterator]() { - yield { - type: "message_start", - message: { - usage: { - input_tokens: 10, - output_tokens: 0, - }, - }, - } - yield { - type: "content_block_start", - index: 0, - content_block: { - type: "text", - text: "Response", + const mockCreate = vitest.fn().mockImplementation(async (options) => + asyncStreamFrom([ + { + type: "message_start", + message: { + usage: { + input_tokens: 10, + output_tokens: 0, }, - } + }, }, - } - }) + { + type: "content_block_start", + index: 0, + content_block: { + type: "text", + text: "Response", + }, + }, + ]), + ) ;(handler["client"].messages as any).create = mockCreate // Messages with internal reasoning blocks (from stored conversation history) @@ -741,11 +677,7 @@ describe("VertexHandler", () => { ] const stream = handler.createMessage(systemPrompt, messagesWithReasoning) - const chunks: ApiStreamChunk[] = [] - - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) // Verify the API was called with filtered messages (no reasoning blocks) const calledMessages = mockCreate.mock.calls[0][0].messages @@ -772,21 +704,19 @@ describe("VertexHandler", () => { vertexRegion: "us-central1", }) - const mockCreate = vitest.fn().mockImplementation(async (options) => { - return { - async *[Symbol.asyncIterator]() { - yield { - type: "message_start", - message: { - usage: { - input_tokens: 10, - output_tokens: 0, - }, + const mockCreate = vitest.fn().mockImplementation(async (options) => + asyncStreamFrom([ + { + type: "message_start", + message: { + usage: { + input_tokens: 10, + output_tokens: 0, }, - } + }, }, - } - }) + ]), + ) ;(handler["client"].messages as any).create = mockCreate // Message with only reasoning content (should be completely filtered) @@ -811,11 +741,7 @@ describe("VertexHandler", () => { ] const stream = handler.createMessage(systemPrompt, messagesWithOnlyReasoning) - const chunks: ApiStreamChunk[] = [] - - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) // Verify empty message was filtered out const calledMessages = mockCreate.mock.calls[0][0].messages @@ -1161,22 +1087,14 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate const stream = handler.createMessage(systemPrompt, mockMessages) - for await (const _chunk of stream) { - // Just consume - } + await collectStream(stream) // Verify the API was called with the beta header expect(mockCreate).toHaveBeenCalledWith( @@ -1207,22 +1125,14 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate const stream = handler.createMessage(systemPrompt, mockMessages) - for await (const _chunk of stream) { - // Just consume - } + await collectStream(stream) // Verify the API was called without the beta header expect(mockCreate).toHaveBeenCalledWith(expect.anything(), undefined) @@ -1298,11 +1208,9 @@ describe("VertexHandler", () => { usage: { input_tokens: 10, output_tokens: 5 }, } } - return { - async *[Symbol.asyncIterator]() { - yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } } - }, - } + return asyncStreamFrom([ + { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }, + ]) }) ;(thinkingHandler["client"].messages as any).create = mockCreate @@ -1327,11 +1235,13 @@ describe("VertexHandler", () => { enableReasoningEffort: true, }) - const mockCreate = vitest.fn().mockImplementation(async () => ({ - async *[Symbol.asyncIterator]() { - yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } } - }, - })) + const mockCreate = vitest + .fn() + .mockImplementation(async () => + asyncStreamFrom([ + { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }, + ]), + ) ;(opus48Handler["client"].messages as any).create = mockCreate await opus48Handler @@ -1358,11 +1268,13 @@ describe("VertexHandler", () => { enableReasoningEffort: true, }) - const mockCreate = vitest.fn().mockImplementation(async () => ({ - async *[Symbol.asyncIterator]() { - yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } } - }, - })) + const mockCreate = vitest + .fn() + .mockImplementation(async () => + asyncStreamFrom([ + { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }, + ]), + ) ;(fableHandler["client"].messages as any).create = mockCreate await fableHandler.createMessage("You are a helpful assistant", [{ role: "user", content: "Hello" }]).next() @@ -1387,11 +1299,13 @@ describe("VertexHandler", () => { enableReasoningEffort: true, }) - const mockCreate = vitest.fn().mockImplementation(async () => ({ - async *[Symbol.asyncIterator]() { - yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } } - }, - })) + const mockCreate = vitest + .fn() + .mockImplementation(async () => + asyncStreamFrom([ + { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }, + ]), + ) ;(sonnetHandler["client"].messages as any).create = mockCreate await sonnetHandler @@ -1418,11 +1332,13 @@ describe("VertexHandler", () => { enableReasoningEffort: true, }) - const mockCreate = vitest.fn().mockImplementation(async () => ({ - async *[Symbol.asyncIterator]() { - yield { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } } - }, - })) + const mockCreate = vitest + .fn() + .mockImplementation(async () => + asyncStreamFrom([ + { type: "message_start", message: { usage: { input_tokens: 10, output_tokens: 5 } } }, + ]), + ) ;(opusHandler["client"].messages as any).create = mockCreate await opusHandler.createMessage("You are a helpful assistant", [{ role: "user", content: "Hello" }]).next() @@ -1485,13 +1401,7 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate @@ -1502,9 +1412,7 @@ describe("VertexHandler", () => { }) // Consume the stream to trigger the API call - for await (const _chunk of stream) { - // Just consume - } + await collectStream(stream) expect(mockCreate).toHaveBeenCalledWith( expect.objectContaining({ @@ -1545,13 +1453,7 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate @@ -1562,9 +1464,7 @@ describe("VertexHandler", () => { }) // Consume the stream to trigger the API call - for await (const _chunk of stream) { - // Just consume - } + await collectStream(stream) // Tool calling is request-driven: if tools are provided, we should include them. expect(mockCreate).toHaveBeenCalledWith( @@ -1607,13 +1507,7 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate @@ -1623,10 +1517,7 @@ describe("VertexHandler", () => { tools: mockTools, }) - const chunks: ApiStreamChunk[] = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) // Find the tool_call_partial chunk const toolCallChunk = chunks.find((chunk) => chunk.type === "tool_call_partial") @@ -1688,13 +1579,7 @@ describe("VertexHandler", () => { }, ] - const asyncIterator = { - async *[Symbol.asyncIterator]() { - for (const chunk of mockStream) { - yield chunk - } - }, - } + const asyncIterator = asyncStreamFrom(mockStream) const mockCreate = vitest.fn().mockResolvedValue(asyncIterator) ;(handler["client"].messages as any).create = mockCreate @@ -1704,10 +1589,7 @@ describe("VertexHandler", () => { tools: mockTools, }) - const chunks: ApiStreamChunk[] = [] - for await (const chunk of stream) { - chunks.push(chunk) - } + const chunks = await collectStream(stream) // Find the tool_call_partial chunks const toolCallChunks = chunks.filter((chunk) => chunk.type === "tool_call_partial") From 0cbe77fbda650d578593b819db90a7a520431814 Mon Sep 17 00:00:00 2001 From: Roomote Date: Sat, 1 Aug 2026 17:50:33 +0000 Subject: [PATCH 2/2] refactor: remove unused anthropic vertex test import --- src/api/providers/__tests__/anthropic-vertex.spec.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/api/providers/__tests__/anthropic-vertex.spec.ts b/src/api/providers/__tests__/anthropic-vertex.spec.ts index 307db6ba37..3e98f3ec5b 100644 --- a/src/api/providers/__tests__/anthropic-vertex.spec.ts +++ b/src/api/providers/__tests__/anthropic-vertex.spec.ts @@ -6,8 +6,6 @@ import { GoogleAuth } from "google-auth-library" import { VERTEX_1M_CONTEXT_MODEL_IDS } from "@roo-code/types" -import { ApiStreamChunk } from "../../transform/stream" - import { AnthropicVertexHandler } from "../anthropic-vertex" import { asyncStreamFrom, collectStream } from "../../../test-utils/stream"