diff --git a/src/api/providers/__tests__/openai-compatible.spec.ts b/src/api/providers/__tests__/openai-compatible.spec.ts new file mode 100644 index 0000000000..ef0a6999a8 --- /dev/null +++ b/src/api/providers/__tests__/openai-compatible.spec.ts @@ -0,0 +1,313 @@ +// Use vi.hoisted to define mock functions that can be referenced in hoisted vi.mock() calls +const { mockStreamText, mockGenerateText } = vi.hoisted(() => ({ + mockStreamText: vi.fn(), + mockGenerateText: vi.fn(), +})) + +vi.mock("ai", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + streamText: mockStreamText, + generateText: mockGenerateText, + } +}) + +vi.mock("@ai-sdk/openai-compatible", () => ({ + createOpenAICompatible: vi.fn(function () { + // Return a function that returns a mock language model + return vi.fn(() => ({ + modelId: "test-model", + provider: "openai-compatible", + })) + }), +})) + +import type { Anthropic } from "@anthropic-ai/sdk" + +import { OpenAICompatibleHandler, OpenAICompatibleConfig } from "../openai-compatible" +import type { ApiHandlerOptions } from "../../../shared/api" + +// Concrete implementation for testing +class TestOpenAICompatibleHandler extends OpenAICompatibleHandler { + constructor(options: ApiHandlerOptions, config: OpenAICompatibleConfig) { + super(options, config) + } + + override getModel() { + return { + id: this.config.modelId, + info: this.config.modelInfo, + maxTokens: this.config.modelMaxTokens, + temperature: this.config.temperature, + } + } +} + +describe("OpenAICompatibleHandler", () => { + let handler: TestOpenAICompatibleHandler + let mockOptions: ApiHandlerOptions + let mockConfig: OpenAICompatibleConfig + + const systemPrompt = "You are a helpful assistant." + const messages: Anthropic.Messages.MessageParam[] = [ + { + role: "user", + content: [{ type: "text", text: "Hello!" }], + }, + ] + + beforeEach(() => { + mockOptions = { + apiModelId: "test-model", + apiKey: "test-api-key", + } + mockConfig = { + providerName: "TestProvider", + baseURL: "https://api.test.com/v1", + apiKey: "test-api-key", + modelId: "test-model", + modelInfo: { + maxTokens: 8192, + contextWindow: 128000, + supportsImages: false, + supportsPromptCache: true, + }, + } + handler = new TestOpenAICompatibleHandler(mockOptions, mockConfig) + vi.clearAllMocks() + }) + + describe("constructor", () => { + it("should initialize with provided options and config", () => { + expect(handler).toBeInstanceOf(TestOpenAICompatibleHandler) + expect(handler.getModel().id).toBe(mockConfig.modelId) + }) + }) + + describe("createMessage", () => { + it("should handle streaming responses", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test response" } + } + + const mockUsage = Promise.resolve({ + inputTokens: 10, + outputTokens: 5, + details: {}, + raw: {}, + }) + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: mockUsage, + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(0) + const textChunks = chunks.filter((chunk) => chunk.type === "text") + expect(textChunks).toHaveLength(1) + expect(textChunks[0].text).toBe("Test response") + }) + + it("should handle multiple stream parts and yield usage metrics", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "First part" } + yield { type: "text-delta", text: "Second part" } + yield { type: "tool-call", toolCallId: "123", name: "test_tool", args: "{}" } + } + + const mockUsage = Promise.resolve({ + inputTokens: 20, + outputTokens: 10, + details: {}, + raw: {}, + }) + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: mockUsage, + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(2) + const textChunks = chunks.filter((chunk) => chunk.type === "text") + expect(textChunks).toHaveLength(2) + expect(textChunks[0].text).toBe("First part") + expect(textChunks[1].text).toBe("Second part") + + const usageChunks = chunks.filter((chunk) => chunk.type === "usage") + expect(usageChunks).toHaveLength(1) + }) + + it("should handle stream without usage metrics", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Test response" } + } + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: Promise.resolve(undefined), + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(0) + const usageChunks = chunks.filter((chunk) => chunk.type === "usage") + expect(usageChunks).toHaveLength(0) + }) + + it("should handle tool-call events in stream", async () => { + async function* mockFullStream() { + yield { type: "text-delta", text: "Calling tool" } + yield { type: "tool-call-start", toolCallId: "tc_1", name: "read_file" } + yield { type: "tool-call-delta", toolCallId: "tc_1", delta: '{"path":"test.ts"}' } + yield { type: "tool-call-end", toolCallId: "tc_1" } + } + + const mockUsage = Promise.resolve({ + inputTokens: 15, + outputTokens: 8, + details: {}, + raw: {}, + }) + + mockStreamText.mockReturnValue({ + fullStream: mockFullStream(), + usage: mockUsage, + }) + + const stream = handler.createMessage(systemPrompt, messages) + const chunks: any[] = [] + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(0) + }) + + // Test 1: createMessage() with mock 429 response → verify thrown error has .status === 429 and provider name in message + it("should throw error with .status 429 when API returns 429", async () => { + const rateLimitError = new Error("Rate limited") as Error & { status: number } + rateLimitError.status = 429 + + mockStreamText.mockReturnValue({ + // eslint-disable-next-line require-yield + fullStream: (async function* () { + throw rateLimitError + })(), + usage: Promise.resolve({ inputTokens: 0, outputTokens: 0, details: {}, raw: {} }), + }) + + let thrownError: any + try { + for await (const chunk of handler.createMessage(systemPrompt, messages)) { + void chunk // Use void to satisfy no-unused-expressions rule + } + } catch (e: any) { + thrownError = e + } + + expect(thrownError).toBeInstanceOf(Error) + expect(thrownError.status).toBe(429) + expect(thrownError.message).toContain("TestProvider") + }) + + // Test 2: createMessage() with mock 500 response → verify error is properly tagged + it("should throw error with .status 500 and provider name when API returns 500", async () => { + const serverError = new Error("Internal Server Error") as Error & { status: number } + serverError.status = 500 + + mockStreamText.mockReturnValue({ + // eslint-disable-next-line require-yield + fullStream: (async function* () { + throw serverError + })(), + usage: Promise.resolve({ inputTokens: 0, outputTokens: 0, details: {}, raw: {} }), + }) + + let thrownError: any + try { + for await (const chunk of handler.createMessage(systemPrompt, messages)) { + void chunk + } + } catch (e: any) { + thrownError = e + } + + expect(thrownError).toBeInstanceOf(Error) + expect(thrownError.status).toBe(500) + expect(thrownError.message).toContain("TestProvider") + }) + }) + + describe("completePrompt", () => { + it("should complete a prompt using generateText", async () => { + mockGenerateText.mockResolvedValue({ + text: "Test completion", + }) + + const result = await handler.completePrompt("Test prompt") + + expect(result).toBe("Test completion") + expect(mockGenerateText).toHaveBeenCalledWith( + expect.objectContaining({ + prompt: "Test prompt", + }), + ) + }) + + // Test 3: completePrompt() with mock 4xx/5xx → verify error carries .status and provider name + it("should throw error with .status and provider name when generateText throws 400", async () => { + const badRequestError = new Error("Bad Request") as Error & { status: number } + badRequestError.status = 400 + + mockGenerateText.mockRejectedValue(badRequestError) + + await expect(handler.completePrompt("Test prompt")).rejects.toThrow("TestProvider") + + let thrownError: any + try { + await handler.completePrompt("Test prompt") + } catch (e: any) { + thrownError = e + } + + expect(thrownError).toBeInstanceOf(Error) + expect((thrownError as any).status).toBe(400) + expect(thrownError.message).toContain("TestProvider") + }) + + it("should throw error with .status and provider name when generateText throws 500", async () => { + const serverError = new Error("Internal Server Error") as Error & { status: number } + serverError.status = 500 + + mockGenerateText.mockRejectedValue(serverError) + + let thrownError: any + try { + await handler.completePrompt("Test prompt") + } catch (e: any) { + thrownError = e + } + + expect(thrownError).toBeInstanceOf(Error) + expect((thrownError as any).status).toBe(500) + expect(thrownError.message).toContain("TestProvider") + }) + }) +}) diff --git a/src/api/providers/openai-compatible.ts b/src/api/providers/openai-compatible.ts index d129e72452..8693e0ffff 100644 --- a/src/api/providers/openai-compatible.ts +++ b/src/api/providers/openai-compatible.ts @@ -16,6 +16,7 @@ import { convertToAiSdkMessages, convertToolsForAiSdk, processAiSdkStreamPart } import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { DEFAULT_HEADERS } from "./constants" +import { handleOpenAIError } from "./utils/error-handler" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" @@ -177,20 +178,24 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si } // Use streamText for streaming responses - const result = streamText(requestOptions) - - // Process the full stream to get all events - for await (const part of result.fullStream) { - // Use the processAiSdkStreamPart utility to convert stream parts - for (const chunk of processAiSdkStreamPart(part)) { - yield chunk + try { + const result = streamText(requestOptions) + + // Process the full stream to get all events + for await (const part of result.fullStream) { + // Use the processAiSdkStreamPart utility to convert stream parts + for (const chunk of processAiSdkStreamPart(part)) { + yield chunk + } } - } - // Yield usage metrics at the end - const usage = await result.usage - if (usage) { - yield this.processUsageMetrics(usage) + // Yield usage metrics at the end + const usage = await result.usage + if (usage) { + yield this.processUsageMetrics(usage) + } + } catch (error) { + throw handleOpenAIError(error, this.config.providerName) } } @@ -198,15 +203,19 @@ export abstract class OpenAICompatibleHandler extends BaseProvider implements Si * Complete a prompt using the AI SDK generateText. */ async completePrompt(prompt: string): Promise { - const languageModel = this.getLanguageModel() - - const { text } = await generateText({ - model: languageModel, - prompt, - maxOutputTokens: this.getMaxOutputTokens(), - temperature: this.config.temperature ?? 0, - }) - - return text + try { + const languageModel = this.getLanguageModel() + + const { text } = await generateText({ + model: languageModel, + prompt, + maxOutputTokens: this.getMaxOutputTokens(), + temperature: this.config.temperature ?? 0, + }) + + return text + } catch (error) { + throw handleOpenAIError(error, this.config.providerName) + } } } diff --git a/src/api/transform/stream.ts b/src/api/transform/stream.ts index 960ebbe770..cecf517281 100644 --- a/src/api/transform/stream.ts +++ b/src/api/transform/stream.ts @@ -92,6 +92,8 @@ export interface ApiStreamToolCallDeltaChunk { export interface ApiStreamToolCallEndChunk { type: "tool_call_end" id: string + /** Tool name for compound-key deduplication (present when emitted by finalizeRawChunks) */ + name?: string } /** diff --git a/src/core/assistant-message/NativeToolCallParser.ts b/src/core/assistant-message/NativeToolCallParser.ts index 9639ae1baa..65e5509c1f 100644 --- a/src/core/assistant-message/NativeToolCallParser.ts +++ b/src/core/assistant-message/NativeToolCallParser.ts @@ -51,7 +51,8 @@ export type ToolCallStreamEvent = ApiStreamToolCallStartChunk | ApiStreamToolCal * provider-level raw chunks into start/delta/end events. */ export class NativeToolCallParser { - // Streaming state management for argument accumulation (keyed by tool call id) + // Streaming state management for argument accumulation (keyed by compound id+name) + // Using compound key to distinguish different tools that may share the same backend id // Note: name is string to accommodate dynamic MCP tools (mcp--serverName--toolName) private static streamingToolCalls = new Map< string, @@ -62,6 +63,13 @@ export class NativeToolCallParser { } >() + /** + * Generate a compound key from id and name for streaming tool call tracking. + */ + public static makeStreamingKey(id: string, name: string): string { + return `${id}::${name}` + } + // Raw chunk tracking state (keyed by index from API stream) private static rawChunkTracker = new Map< number, @@ -175,6 +183,7 @@ export class NativeToolCallParser { events.push({ type: "tool_call_end", id: tracked.id, + name: tracked.name, }) } } @@ -195,6 +204,7 @@ export class NativeToolCallParser { events.push({ type: "tool_call_end", id: tracked.id, + name: tracked.name, }) } } @@ -218,7 +228,8 @@ export class NativeToolCallParser { * Accepts string to support both ToolName and dynamic MCP tools (mcp--serverName--toolName). */ public static startStreamingToolCall(id: string, name: string): void { - this.streamingToolCalls.set(id, { + const key = this.makeStreamingKey(id, name) + this.streamingToolCalls.set(key, { id, name, argumentsAccumulator: "", @@ -242,13 +253,37 @@ export class NativeToolCallParser { return this.streamingToolCalls.size > 0 } + /** + * Get the name of a streaming tool call by its compound key (id + name). + * Returns undefined if the tool call is not found. + */ + public static getStreamingToolName(key: string): string | undefined { + return this.streamingToolCalls.get(key)?.name + } + + /** + * Get a streaming tool call entry by its id (legacy lookup for backward compatibility). + * Returns the first matching entry if multiple tools share the same id. + * @deprecated Use getStreamingToolName with compound key instead. + */ + public static getStreamingToolCallById( + id: string, + ): { id: string; name: string; argumentsAccumulator: string } | null { + for (const entry of this.streamingToolCalls.values()) { + if (entry.id === id) { + return entry + } + } + return null + } + /** * Process a chunk of JSON arguments for a streaming tool call. * Uses partial-json-parser to extract values from incomplete JSON immediately. * Returns a partial ToolUse with currently parsed parameters. */ - public static processStreamingChunk(id: string, chunk: string): ToolUse | null { - const toolCall = this.streamingToolCalls.get(id) + public static processStreamingChunk(key: string, chunk: string): ToolUse | null { + const toolCall = this.streamingToolCalls.get(key) if (!toolCall) { return null } @@ -291,8 +326,8 @@ export class NativeToolCallParser { * Finalize a streaming tool call. * Parses the complete JSON and returns the final ToolUse or McpToolUse. */ - public static finalizeStreamingToolCall(id: string): ToolUse | McpToolUse | null { - const toolCall = this.streamingToolCalls.get(id) + public static finalizeStreamingToolCall(key: string): ToolUse | McpToolUse | null { + const toolCall = this.streamingToolCalls.get(key) if (!toolCall) { return null } @@ -306,7 +341,7 @@ export class NativeToolCallParser { }) // Clean up streaming state - this.streamingToolCalls.delete(id) + this.streamingToolCalls.delete(key) return finalToolUse } diff --git a/src/core/assistant-message/__tests__/NativeToolCallParser-additional.spec.ts b/src/core/assistant-message/__tests__/NativeToolCallParser-additional.spec.ts new file mode 100644 index 0000000000..47d7e07448 --- /dev/null +++ b/src/core/assistant-message/__tests__/NativeToolCallParser-additional.spec.ts @@ -0,0 +1,156 @@ +import { NativeToolCallParser } from "../NativeToolCallParser" + +describe("NativeToolCallParser - Additional Coverage", () => { + beforeEach(() => { + NativeToolCallParser.clearAllStreamingToolCalls() + NativeToolCallParser.clearRawChunkState() + }) + + describe("hasActiveStreamingToolCalls", () => { + it("should return false when no streaming tool calls are active", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + }) + + it("should return true when there is an active streaming tool call", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + NativeToolCallParser.startStreamingToolCall("toolu_123", "read_file") + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(true) + + // Finalize to clear state (needs compound key) + NativeToolCallParser.finalizeStreamingToolCall( + NativeToolCallParser.makeStreamingKey("toolu_123", "read_file"), + ) + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + }) + }) + + describe("getStreamingToolName", () => { + it("should return undefined for a non-existent tool call id", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + expect(NativeToolCallParser.getStreamingToolName("toolu_nonexistent")).toBeUndefined() + }) + + it("should return the name of an active streaming tool call", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + const testId = "toolu_stream_name_123" + const testName = "write_to_file" + NativeToolCallParser.startStreamingToolCall(testId, testName) + + expect( + NativeToolCallParser.getStreamingToolName(NativeToolCallParser.makeStreamingKey(testId, testName)), + ).toBe("write_to_file") + + // Cleanup + NativeToolCallParser.finalizeStreamingToolCall(NativeToolCallParser.makeStreamingKey(testId, testName)) + }) + + it("should return undefined after the tool call is finalized", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + const testId = "toolu_stream_name_456" + NativeToolCallParser.startStreamingToolCall(testId, "codebase_search") + NativeToolCallParser.finalizeStreamingToolCall(testId) + + expect(NativeToolCallParser.getStreamingToolName(testId)).toBeUndefined() + }) + + it("should return the name for MCP tools", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + const testId = "toolu_mcp_name_789" + const mcpToolName = "mcp--my-server--my-tool" + NativeToolCallParser.startStreamingToolCall(testId, mcpToolName) + + expect( + NativeToolCallParser.getStreamingToolName(NativeToolCallParser.makeStreamingKey(testId, mcpToolName)), + ).toBe(mcpToolName) + + // Cleanup + NativeToolCallParser.finalizeStreamingToolCall(NativeToolCallParser.makeStreamingKey(testId, mcpToolName)) + }) + }) + + describe("processStreamingChunk with MCP tools", () => { + it("should return null for MCP tools during streaming (wait for final)", () => { + const id = "toolu_mcp_stream_123" + NativeToolCallParser.startStreamingToolCall(id, "mcp--my-server--get_config") + + // For MCP tools, processStreamingChunk should return null + const result = NativeToolCallParser.processStreamingChunk(id, '{"key":"value"}') + expect(result).toBeNull() + }) + + it("should return final McpToolUse for MCP tools on finalize", () => { + const id = "toolu_mcp_finalize_123" + const name = "mcp--my-server--get_config" + NativeToolCallParser.startStreamingToolCall(id, name) + + // Add arguments via processStreamingChunk (doesn't return partial for MCP) + NativeToolCallParser.processStreamingChunk( + NativeToolCallParser.makeStreamingKey(id, name), + '{"key":"value"}', + ) + + const result = NativeToolCallParser.finalizeStreamingToolCall( + NativeToolCallParser.makeStreamingKey(id, name), + ) + expect(result).not.toBeNull() + // MCP tools return "mcp_tool_use" type, not "tool_use" + expect(result?.type).toMatch(/^(tool_use|mcp_tool_use)$/) + }) + }) + + describe("finalizeRawChunks edge cases", () => { + it("should return empty array when no raw chunks are tracked", () => { + NativeToolCallParser.clearRawChunkState() + const events = NativeToolCallParser.finalizeRawChunks() + expect(events).toHaveLength(0) + }) + + it("should emit an end event for a tool call that already started (name provided upfront)", () => { + NativeToolCallParser.clearRawChunkState() + NativeToolCallParser.processRawChunk({ + index: 1, + id: "toolu_partial", + name: "read_file", + arguments: undefined, + }) + const events = NativeToolCallParser.finalizeRawChunks() + expect(events).toEqual([{ type: "tool_call_end", id: "toolu_partial", name: "read_file" }]) + }) + }) + + describe("clearRawChunkState", () => { + it("should clear all raw chunk tracking state", () => { + // Add some raw chunks first + NativeToolCallParser.processRawChunk({ + index: 1, + id: "toolu_clear_test", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + // Clear state + NativeToolCallParser.clearRawChunkState() + + // Finalize should return empty array + const events = NativeToolCallParser.finalizeRawChunks() + expect(events).toHaveLength(0) + }) + }) + + describe("clearAllStreamingToolCalls", () => { + it("should clear all streaming tool call state", () => { + NativeToolCallParser.startStreamingToolCall("toolu_clear_stream_1", "read_file") + NativeToolCallParser.startStreamingToolCall("toolu_clear_stream_2", "write_to_file") + + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(true) + + // Clear all + NativeToolCallParser.clearAllStreamingToolCalls() + + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + expect(NativeToolCallParser.getStreamingToolName("toolu_clear_stream_1")).toBeUndefined() + expect(NativeToolCallParser.getStreamingToolName("toolu_clear_stream_2")).toBeUndefined() + }) + }) +}) diff --git a/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts b/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts index 2c15e12069..5f2931f17a 100644 --- a/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts +++ b/src/core/assistant-message/__tests__/NativeToolCallParser.spec.ts @@ -302,8 +302,9 @@ describe("NativeToolCallParser", () => { // Simulate streaming chunks const fullArgs = JSON.stringify({ path: "src/test.ts" }) - // Process the complete args as a single chunk for simplicity - const result = NativeToolCallParser.processStreamingChunk(id, fullArgs) + // Process the complete args as a single chunk for simplicity using compound key + const key = NativeToolCallParser.makeStreamingKey(id, "read_file") + const result = NativeToolCallParser.processStreamingChunk(key, fullArgs) expect(result).not.toBeNull() expect(result?.nativeArgs).toBeDefined() @@ -319,9 +320,10 @@ describe("NativeToolCallParser", () => { const id = "toolu_finalize_123" NativeToolCallParser.startStreamingToolCall(id, "read_file") - // Add the complete arguments + // Add the complete arguments using compound key + const key = NativeToolCallParser.makeStreamingKey(id, "read_file") NativeToolCallParser.processStreamingChunk( - id, + key, JSON.stringify({ path: "finalized.ts", mode: "slice", @@ -330,7 +332,7 @@ describe("NativeToolCallParser", () => { }), ) - const result = NativeToolCallParser.finalizeStreamingToolCall(id) + const result = NativeToolCallParser.finalizeStreamingToolCall(key) expect(result).not.toBeNull() expect(result?.type).toBe("tool_use") @@ -343,4 +345,110 @@ describe("NativeToolCallParser", () => { }) }) }) + + describe("streamingToolCalls collision", () => { + it("should keep separate accumulated arguments for tools with same id but different names", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + NativeToolCallParser.clearRawChunkState() + + const id = "toolu_collision_123" + + // Start two tool calls with the same id but different names + NativeToolCallParser.startStreamingToolCall(id, "read_file") + NativeToolCallParser.startStreamingToolCall(id, "write_to_file") + + // Accumulate arguments for each using compound keys + const key1 = NativeToolCallParser.makeStreamingKey(id, "read_file") + const key2 = NativeToolCallParser.makeStreamingKey(id, "write_to_file") + + NativeToolCallParser.processStreamingChunk(key1, JSON.stringify({ path: "file_a.ts" })) + NativeToolCallParser.processStreamingChunk(key2, JSON.stringify({ path: "file_b.ts", content: "hello" })) + + // Finalize both and verify they have distinct arguments + const result1 = NativeToolCallParser.finalizeStreamingToolCall(key1) + const result2 = NativeToolCallParser.finalizeStreamingToolCall(key2) + + expect(result1).not.toBeNull() + expect(result2).not.toBeNull() + expect(result1?.type).toBe("tool_use") + expect(result2?.type).toBe("tool_use") + + if (result1?.type === "tool_use" && result2?.type === "tool_use") { + const nativeArgs1 = result1.nativeArgs as { path: string } + const nativeArgs2 = result2.nativeArgs as { path: string; content: string } + expect(nativeArgs1.path).toBe("file_a.ts") + expect(nativeArgs2.path).toBe("file_b.ts") + expect(nativeArgs2.content).toBe("hello") + } + + // Verify streaming state is cleaned up for both + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + }) + }) + + describe("finalizeRawChunks", () => { + it("should include name field in events for compound-key deduplication", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + NativeToolCallParser.clearRawChunkState() + + // Simulate two different tools with the same toolCallId via raw chunk tracking. + // processRawChunk returns [start, delta] when arguments are provided alongside name. + const result1 = NativeToolCallParser.processRawChunk({ + index: 1, + id: "toolu_sameid", + name: "read_file", + arguments: '{"path":"a.ts"}', + }) + // First call for index 1 returns start + delta events + expect(result1.some((e) => e.type === "tool_call_start")).toBe(true) + + const result2 = NativeToolCallParser.processRawChunk({ + index: 2, + id: "toolu_sameid", + name: "write_to_file", + arguments: '{"path":"b.ts","content":"hello"}', + }) + // Second call for index 2 also returns start + delta events + expect(result2.some((e) => e.type === "tool_call_start")).toBe(true) + + // Finalize raw chunks — both should be included with their names + const finalizeEvents = NativeToolCallParser.finalizeRawChunks() + + expect(finalizeEvents).toHaveLength(2) + expect(finalizeEvents.every((e) => e.type === "tool_call_end")).toBe(true) + + // Each event must carry its name for compound-key deduplication + const names = finalizeEvents.map((e) => (e as any).name) + expect(names).toContain("read_file") + expect(names).toContain("write_to_file") + }) + + it("should emit separate end events when two tools share the same toolCallId", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + NativeToolCallParser.clearRawChunkState() + + // Both tools use the exact same call ID but different names + NativeToolCallParser.processRawChunk({ + index: 10, + id: "toolu_dup", + name: "codebase_search", + arguments: '{"query":"foo"}', + }) + NativeToolCallParser.processRawChunk({ + index: 11, + id: "toolu_dup", + name: "search_files", + arguments: '{"path":".","regex":"bar"}', + }) + + const events = NativeToolCallParser.finalizeRawChunks() + + // Should produce two distinct end events + expect(events).toHaveLength(2) + + // Verify compound keys would be unique + const dedupKeys = new Set(events.map((e) => `${e.id}::${(e as any).name}`)) + expect(dedupKeys.size).toBe(2) + }) + }) }) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 0bc3130169..2672ac78bc 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -60,7 +60,7 @@ import { CloudService } from "@roo-code/cloud" // api import { ApiHandler, ApiHandlerCreateMessageMetadata, buildApiHandler } from "../../api" -import { ApiStream, GroundingSource } from "../../api/transform/stream" +import { ApiStream, ApiStreamToolCallEndChunk, GroundingSource } from "../../api/transform/stream" import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning" // shared @@ -2828,7 +2828,9 @@ export class Task extends EventEmitter implements TaskLike { // Without this check, duplicate tool_use blocks with the same ID would // be added to assistantMessageContent, causing API 400 errors: // "tool_use ids must be unique" - if (this.streamingToolCallIndices.has(event.id)) { + // Use compound key (id, name) to distinguish different tools with the same call ID. + const dedupKey = `${event.id}::${event.name}` + if (this.streamingToolCallIndices.has(dedupKey)) { console.warn( `[Task#${this.taskId}] Ignoring duplicate tool_call_start for ID: ${event.id} (tool: ${event.name})`, ) @@ -2848,7 +2850,7 @@ export class Task extends EventEmitter implements TaskLike { // Track the index where this tool will be stored const toolUseIndex = this.assistantMessageContent.length - this.streamingToolCallIndices.set(event.id, toolUseIndex) + this.streamingToolCallIndices.set(dedupKey, toolUseIndex) // Create initial partial tool use const partialToolUse: ToolUse = { @@ -2867,15 +2869,28 @@ export class Task extends EventEmitter implements TaskLike { /* v8 ignore next -- streaming presenter; .catch lives in presentAssistantMessageSafe (covered) */ this.presentAssistantMessageSafe() } else if (event.type === "tool_call_delta") { - // Process chunk using streaming JSON parser + // Look up the streaming entry by id to get its compound key + // (delta events don't carry name, but we stored it during startStreamingToolCall) + const existingEntry = NativeToolCallParser.getStreamingToolCallById(event.id) + const streamingKey = existingEntry + ? NativeToolCallParser.makeStreamingKey( + existingEntry.id, + existingEntry.name, + ) + : undefined + const partialToolUse = NativeToolCallParser.processStreamingChunk( - event.id, + streamingKey ?? event.id, event.delta, ) if (partialToolUse) { - // Get the index for this tool call - const toolUseIndex = this.streamingToolCallIndices.get(event.id) + // Retrieve name from NativeToolCallParser's streaming state + const name = NativeToolCallParser.getStreamingToolName( + streamingKey ?? event.id, + ) + const dedupKey = `${event.id}::${name}` + const toolUseIndex = this.streamingToolCallIndices.get(dedupKey) if (toolUseIndex !== undefined) { // Store the ID for native protocol ;(partialToolUse as any).id = event.id @@ -2890,10 +2905,19 @@ export class Task extends EventEmitter implements TaskLike { } } else if (event.type === "tool_call_end") { // Finalize the streaming tool call - const finalToolUse = NativeToolCallParser.finalizeStreamingToolCall(event.id) + // Build compound key from event to match the key used in startStreamingToolCall + // Use event.name or fallback to id for events that don't carry name + const eventName = (event as ApiStreamToolCallEndChunk).name ?? event.id + const streamingKey = NativeToolCallParser.makeStreamingKey(event.id, eventName) + + // Finalize the streaming tool call + const finalToolUse = + NativeToolCallParser.finalizeStreamingToolCall(streamingKey) - // Get the index for this tool call - const toolUseIndex = this.streamingToolCallIndices.get(event.id) + // Retrieve name from NativeToolCallParser's streaming state + const name = NativeToolCallParser.getStreamingToolName(streamingKey) + const dedupKey = `${event.id}::${name}` + const toolUseIndex = this.streamingToolCallIndices.get(dedupKey) if (finalToolUse) { // Store the tool call ID @@ -2905,7 +2929,7 @@ export class Task extends EventEmitter implements TaskLike { } // Clean up tracking - this.streamingToolCallIndices.delete(event.id) + this.streamingToolCallIndices.delete(dedupKey) // Mark that we have new content to process this.userMessageContentReady = false @@ -2924,8 +2948,8 @@ export class Task extends EventEmitter implements TaskLike { ;(existingToolUse as any).id = event.id } - // Clean up tracking - this.streamingToolCallIndices.delete(event.id) + // Clean up tracking (use compound key to match the if branch above) + this.streamingToolCallIndices.delete(dedupKey) // Mark that we have new content to process this.userMessageContentReady = false @@ -3291,11 +3315,14 @@ export class Task extends EventEmitter implements TaskLike { const finalizeEvents = NativeToolCallParser.finalizeRawChunks() for (const event of finalizeEvents) { if (event.type === "tool_call_end") { + // Create compound key for deduplication (same pattern as streaming handler) + const dedupKey = `${event.id}::${event.name}` + // Finalize the streaming tool call const finalToolUse = NativeToolCallParser.finalizeStreamingToolCall(event.id) - // Get the index for this tool call - const toolUseIndex = this.streamingToolCallIndices.get(event.id) + // Get the index for this tool call using compound key + const toolUseIndex = this.streamingToolCallIndices.get(dedupKey) if (finalToolUse) { // Store the tool call ID @@ -3306,8 +3333,8 @@ export class Task extends EventEmitter implements TaskLike { this.assistantMessageContent[toolUseIndex] = finalToolUse } - // Clean up tracking - this.streamingToolCallIndices.delete(event.id) + // Clean up tracking using compound key + this.streamingToolCallIndices.delete(dedupKey) // Mark that we have new content to process this.userMessageContentReady = false @@ -3326,8 +3353,8 @@ export class Task extends EventEmitter implements TaskLike { ;(existingToolUse as any).id = event.id } - // Clean up tracking - this.streamingToolCallIndices.delete(event.id) + // Clean up tracking using compound key + this.streamingToolCallIndices.delete(dedupKey) // Mark that we have new content to process this.userMessageContentReady = false diff --git a/src/core/task/__tests__/Task.streaming-tool-calls.spec.ts b/src/core/task/__tests__/Task.streaming-tool-calls.spec.ts new file mode 100644 index 0000000000..b8dacd6774 --- /dev/null +++ b/src/core/task/__tests__/Task.streaming-tool-calls.spec.ts @@ -0,0 +1,850 @@ +import * as os from "os" +import * as path from "path" +import * as vscode from "vscode" +import { Task } from "../Task" +import { NativeToolCallParser } from "../../assistant-message/NativeToolCallParser" +import { ClineProvider } from "../../webview/ClineProvider" +import { ApiStreamChunk, type ApiStreamToolCallPartialChunk } from "../../../api/transform/stream" +import { ContextProxy } from "../../config/ContextProxy" +import { TelemetryService } from "@roo-code/telemetry" + +// Mock delay before any imports that might use it +vi.mock("delay", () => ({ + __esModule: true, + default: vi.fn().mockResolvedValue(undefined), +})) + +import delay from "delay" + +vi.mock("uuid", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + v7: vi.fn(() => "00000000-0000-7000-8000-000000000000"), + } +}) + +vi.mock("execa", () => ({ + execa: vi.fn(), +})) + +vi.mock("fs/promises", async (importOriginal) => { + const actual = (await importOriginal()) as Record + const mockFunctions = { + mkdir: vi.fn().mockResolvedValue(undefined), + writeFile: vi.fn().mockResolvedValue(undefined), + readFile: vi.fn().mockImplementation((filePath) => { + if (filePath.includes("ui_messages.json")) { + return Promise.resolve(JSON.stringify([])) + } + if (filePath.includes("api_conversation_history.json")) { + return Promise.resolve(JSON.stringify([])) + } + return Promise.resolve("[]") + }), + unlink: vi.fn().mockResolvedValue(undefined), + rmdir: vi.fn().mockResolvedValue(undefined), + stat: vi.fn().mockRejectedValue({ code: "ENOENT" }), + readdir: vi.fn().mockResolvedValue([]), + } + + return { + ...actual, + ...mockFunctions, + default: mockFunctions, + } +}) + +vi.mock("p-wait-for", () => ({ + default: vi.fn().mockImplementation(async () => Promise.resolve()), +})) + +vi.mock("vscode", () => { + const mockDisposable = { dispose: vi.fn() } + const mockEventEmitter = { event: vi.fn(), fire: vi.fn() } + const mockTextDocument = { uri: { fsPath: "/mock/workspace/path/file.ts" } } + const mockTextEditor = { document: mockTextDocument } + const mockTab = { input: { uri: { fsPath: "/mock/workspace/path/file.ts" } } } + const mockTabGroup = { tabs: [mockTab] } + + return { + TabInputTextDiff: vi.fn(), + CodeActionKind: { + QuickFix: { value: "quickfix" }, + RefactorRewrite: { value: "refactor.rewrite" }, + }, + window: { + createTextEditorDecorationType: vi.fn().mockReturnValue({ + dispose: vi.fn(), + }), + visibleTextEditors: [mockTextEditor], + tabGroups: { + all: [mockTabGroup], + close: vi.fn(), + onDidChangeTabs: vi.fn(() => ({ dispose: vi.fn() })), + }, + showErrorMessage: vi.fn(), + }, + workspace: { + workspaceFolders: [ + { + uri: { fsPath: "/mock/workspace/path" }, + name: "mock-workspace", + index: 0, + }, + ], + createFileSystemWatcher: vi.fn(() => ({ + onDidCreate: vi.fn(() => mockDisposable), + onDidDelete: vi.fn(() => mockDisposable), + onDidChange: vi.fn(() => mockDisposable), + dispose: vi.fn(), + })), + fs: { + stat: vi.fn().mockResolvedValue({ type: 1 }), + }, + onDidSaveTextDocument: vi.fn(() => mockDisposable), + getConfiguration: vi.fn(() => ({ get: (key: string, defaultValue: any) => defaultValue })), + }, + env: { + uriScheme: "vscode", + language: "en", + }, + EventEmitter: vi.fn().mockImplementation(function () { + return mockEventEmitter + }), + Disposable: { + from: vi.fn(), + }, + TabInputText: vi.fn(), + RelativePattern: class RelativePattern { + constructor( + public path: string, + public pattern: string, + ) {} + }, + } +}) + +vi.mock("../../mentions", () => ({ + parseMentions: vi.fn().mockImplementation((text) => { + return Promise.resolve({ text: `processed: ${text}`, mode: undefined, contentBlocks: [] }) + }), + openMention: vi.fn(), + getLatestTerminalOutput: vi.fn(), +})) + +vi.mock("../../../integrations/misc/extract-text", () => ({ + extractTextFromFile: vi.fn().mockResolvedValue("Mock file content"), +})) + +vi.mock("../../environment/getEnvironmentDetails", () => ({ + getEnvironmentDetails: vi.fn().mockResolvedValue(""), +})) + +vi.mock("../../ignore/RooIgnoreController") + +vi.mock("../../condense", async (importOriginal) => { + const actual = (await importOriginal()) as any + return { + ...actual, + summarizeConversation: vi.fn().mockResolvedValue({ + messages: [{ role: "user", content: [{ type: "text", text: "continued" }], ts: Date.now() }], + summary: "summary", + cost: 0, + newContextTokens: 1, + }), + } +}) + +vi.mock("../../../utils/storage", () => ({ + getTaskDirectoryPath: vi + .fn() + .mockImplementation((globalStoragePath, taskId) => Promise.resolve(`${globalStoragePath}/tasks/${taskId}`)), + getSettingsDirectoryPath: vi + .fn() + .mockImplementation((globalStoragePath) => Promise.resolve(`${globalStoragePath}/settings`)), +})) + +vi.mock("../../../utils/fs", () => ({ + fileExistsAtPath: vi.fn().mockImplementation((filePath) => { + return filePath.includes("ui_messages.json") || filePath.includes("api_conversation_history.json") + }), +})) + +describe("Task - Streaming Tool Call Handling", () => { + let mockProvider: any + let mockApiConfig: any + let mockOutputChannel: any + let mockExtensionContext: vscode.ExtensionContext + + beforeEach(() => { + NativeToolCallParser.clearAllStreamingToolCalls() + NativeToolCallParser.clearRawChunkState() + + if (!TelemetryService.hasInstance()) { + TelemetryService.createInstance([]) + } + + const storageUri = { + fsPath: path.join(os.tmpdir(), "test-storage"), + } + + mockExtensionContext = { + globalState: { + get: vi.fn().mockImplementation((key: string) => { + if (key === "taskHistory") { + return [] + } + return undefined + }), + update: vi.fn().mockImplementation((_key, _value) => Promise.resolve()), + keys: vi.fn().mockReturnValue([]), + }, + globalStorageUri: storageUri, + workspaceState: { + get: vi.fn().mockImplementation((_key) => undefined), + update: vi.fn().mockImplementation((_key, _value) => Promise.resolve()), + keys: vi.fn().mockReturnValue([]), + }, + secrets: { + get: vi.fn().mockImplementation((_key) => Promise.resolve(undefined)), + store: vi.fn().mockImplementation((_key, _value) => Promise.resolve()), + delete: vi.fn().mockImplementation((_key) => Promise.resolve()), + }, + extensionUri: { + fsPath: "/mock/extension/path", + }, + extension: { + packageJSON: { + version: "1.0.0", + }, + }, + } as unknown as vscode.ExtensionContext + + mockOutputChannel = { + appendLine: vi.fn(), + append: vi.fn(), + clear: vi.fn(), + show: vi.fn(), + hide: vi.fn(), + dispose: vi.fn(), + } + + mockProvider = new ClineProvider( + mockExtensionContext, + mockOutputChannel, + "sidebar", + new ContextProxy(mockExtensionContext), + ) as any + + mockApiConfig = { + apiProvider: "anthropic", + apiModelId: "claude-3-5-sonnet-20241022", + apiKey: "test-api-key", + } + + mockProvider.postMessageToWebview = vi.fn().mockResolvedValue(undefined) + mockProvider.postStateToWebview = vi.fn().mockResolvedValue(undefined) + mockProvider.postStateToWebviewWithoutTaskHistory = vi.fn().mockResolvedValue(undefined) + mockProvider.getTaskWithId = vi.fn().mockResolvedValue(null) + + vi.spyOn(mockProvider, "getState").mockResolvedValue({ + apiConfiguration: mockApiConfig, + autoApprovalEnabled: false, + requestDelaySeconds: 0, + mode: "assistant", + customModes: [], + disabledTools: [], + experiments: {}, + profileThresholds: {}, + } as any) + }) + + afterEach(() => { + NativeToolCallParser.clearAllStreamingToolCalls() + NativeToolCallParser.clearRawChunkState() + }) + + describe("tool_call_partial chunk handling - NativeToolCallParser.processRawChunk", () => { + it("should emit tool_call_start event when processing raw chunk with id and name", () => { + const events = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_123", + name: "read_file", + arguments: '{"path":"a.ts"}', + }) + + expect(events.length).toBeGreaterThan(0) + const startEvent = events.find((e) => e.type === "tool_call_start") + expect(startEvent).toBeDefined() + if (startEvent && startEvent.type === "tool_call_start") { + expect(startEvent.id).toBe("toolu_123") + expect(startEvent.name).toBe("read_file") + } + }) + + it("should emit tool_call_delta events for argument chunks", () => { + const events = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_123", + name: "read_file", + arguments: '{"path":"a.ts"}', + }) + + const deltaEvents = events.filter((e) => e.type === "tool_call_delta") + expect(deltaEvents.length).toBeGreaterThan(0) + if (deltaEvents[0] && deltaEvents[0].type === "tool_call_delta") { + expect(deltaEvents[0].delta).toContain('"path"') + } + }) + + it("should buffer deltas before start event and flush after", () => { + NativeToolCallParser.clearRawChunkState() + + // First chunk with id/name - should emit start + buffered delta + const events1 = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_buffer123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + expect(events1.some((e) => e.type === "tool_call_start")).toBe(true) + expect(events1.some((e) => e.type === "tool_call_delta")).toBe(true) + }) + + it("should handle multiple chunks with same index (stream retry scenario)", () => { + NativeToolCallParser.clearRawChunkState() + + const events1 = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_retry123", + name: "read_file", + arguments: '{"path":"a.ts"}', + }) + + const events2 = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_retry123", + name: "read_file", + arguments: '{"path":"b.ts"}', + }) + + // Both should emit events (the dedup is handled by Task, not NativeToolCallParser) + expect(events1.length).toBeGreaterThan(0) + expect(events2.length).toBeGreaterThan(0) + }) + }) + + describe("NativeToolCallParser streaming state management", () => { + it("should start tracking a streaming tool call and report hasActiveStreamingToolCalls", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + + const id = "toolu_123" + const name = "read_file" + NativeToolCallParser.startStreamingToolCall(id, name) + + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(true) + expect(NativeToolCallParser.getStreamingToolName(NativeToolCallParser.makeStreamingKey(id, name))).toBe( + "read_file", + ) + }) + + it("should accumulate argument deltas via processStreamingChunk", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const id = "toolu_delta_acc" + const name = "execute_command" + NativeToolCallParser.startStreamingToolCall(id, name) + const key = NativeToolCallParser.makeStreamingKey(id, name) + + const chunk1 = NativeToolCallParser.processStreamingChunk(key, '{"command":"echo') + expect(chunk1).toBeDefined() + + const chunk2 = NativeToolCallParser.processStreamingChunk(key, ' "hello"') + expect(chunk2).toBeDefined() + + // Verify accumulated arguments in streaming state + const streamingState = (NativeToolCallParser as any)["streamingToolCalls"].get(key) + expect(streamingState).toBeDefined() + expect(streamingState!.argumentsAccumulator).toContain('"command":"echo') + }) + + it("should finalize tool call and return ToolUse via finalizeStreamingToolCall", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const id = "toolu_final123" + const name = "read_file" + NativeToolCallParser.startStreamingToolCall(id, name) + const key = NativeToolCallParser.makeStreamingKey(id, name) + NativeToolCallParser.processStreamingChunk(key, '{"path":"test.ts"}') + + const result = NativeToolCallParser.finalizeStreamingToolCall(key) + + expect(result).toBeDefined() + expect(result?.type).toBe("tool_use") + expect(result?.name).toBe("read_file") + expect(result?.partial).toBe(false) + // After finalization, should no longer be in streaming state + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + }) + + it("should return null for finalizeStreamingToolCall when arguments are malformed", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const id = "toolu_malformed" + const name = "read_file" + NativeToolCallParser.startStreamingToolCall(id, name) + const key = NativeToolCallParser.makeStreamingKey(id, name) + NativeToolCallParser.processStreamingChunk(key, "{invalid json") + + const result = NativeToolCallParser.finalizeStreamingToolCall(key) + + // finalizeStreamingToolCall uses JSON.parse which will fail on malformed JSON + expect(result).toBeNull() + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + }) + + it("should return null when finalizing unknown tool call id", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const result = NativeToolCallParser.finalizeStreamingToolCall("toolu_unknown::unknown") + expect(result).toBeNull() + }) + + it("should processStreamingChunk return null for unknown tool call id", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const result = NativeToolCallParser.processStreamingChunk("toolu_unknown::unknown", '{"some":"data"}') + expect(result).toBeNull() + }) + + it("should handle multiple sequential streaming tool calls", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + // First tool call + const id1 = "toolu_seq1" + const name1 = "read_file" + NativeToolCallParser.startStreamingToolCall(id1, name1) + const key1 = NativeToolCallParser.makeStreamingKey(id1, name1) + NativeToolCallParser.processStreamingChunk(key1, '{"path":"a.ts"}') + const result1 = NativeToolCallParser.finalizeStreamingToolCall(key1) + expect(result1?.name).toBe("read_file") + + // Second tool call + const id2 = "toolu_seq2" + const name2 = "write_to_file" + NativeToolCallParser.startStreamingToolCall(id2, name2) + const key2 = NativeToolCallParser.makeStreamingKey(id2, name2) + NativeToolCallParser.processStreamingChunk(key2, '{"path":"b.ts","content":"hello"}') + const result2 = NativeToolCallParser.finalizeStreamingToolCall(key2) + expect(result2?.name).toBe("write_to_file") + + // Both should be finalized + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + }) + + it("should handle same toolCallId with different names (MCP tools)", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const id = "toolu_same" + + // First tool with same ID but different name + const name1 = "mcp--server1--read_file" + NativeToolCallParser.startStreamingToolCall(id, name1) + const key1 = NativeToolCallParser.makeStreamingKey(id, name1) + NativeToolCallParser.processStreamingChunk(key1, '{"path":"a.ts"}') + const result1 = NativeToolCallParser.finalizeStreamingToolCall(key1) + expect(result1).toBeDefined() + + // Second tool with same ID but different name + const name2 = "mcp--server2--write_to_file" + NativeToolCallParser.startStreamingToolCall(id, name2) + const key2 = NativeToolCallParser.makeStreamingKey(id, name2) + NativeToolCallParser.processStreamingChunk(key2, '{"path":"b.ts"}') + const result2 = NativeToolCallParser.finalizeStreamingToolCall(key2) + expect(result2).toBeDefined() + + // Both should be finalized (same ID, different names are tracked separately) + expect(NativeToolCallParser.hasActiveStreamingToolCalls()).toBe(false) + }) + }) + + describe("processStreamingChunk partial ToolUse creation", () => { + it("should create partial tool_use with correct structure on start", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const id = "toolu_partial123" + const name = "write_to_file" + NativeToolCallParser.startStreamingToolCall(id, name) + const key = NativeToolCallParser.makeStreamingKey(id, name) + + const partial = NativeToolCallParser.processStreamingChunk(key, '{"path":"output.txt","content":"hello"}') + + expect(partial).toBeDefined() + expect(partial?.type).toBe("tool_use") + expect(partial?.name).toBe("write_to_file") + expect(partial?.partial).toBe(true) + expect(partial?.params).toBeDefined() + }) + + it("should update partial tool_use with accumulated arguments", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const id = "toolu_update123" + const name = "execute_command" + NativeToolCallParser.startStreamingToolCall(id, name) + const key = NativeToolCallParser.makeStreamingKey(id, name) + + const chunk1 = NativeToolCallParser.processStreamingChunk(key, '{"command":"') + expect(chunk1?.params).toBeDefined() + + const chunk2 = NativeToolCallParser.processStreamingChunk(key, 'echo "hello"}') + expect(chunk2?.params).toBeDefined() + // The accumulated arguments should be more complete in chunk2 + if (chunk2?.nativeArgs && typeof chunk2.nativeArgs === "object" && "command" in chunk2.nativeArgs) { + expect((chunk2.nativeArgs as any).command).toContain("echo") + } + }) + + it("should handle severely malformed JSON gracefully", () => { + NativeToolCallParser.clearAllStreamingToolCalls() + + const id = "toolu_fail123" + const name = "read_file" + NativeToolCallParser.startStreamingToolCall(id, name) + const key = NativeToolCallParser.makeStreamingKey(id, name) + + // Partial-json-parser can handle partial JSON like '{"path"' and return a partial result + const partialResult = NativeToolCallParser.processStreamingChunk(key, '{"path"') + + expect(partialResult).toBeDefined() + expect(partialResult?.partial).toBe(true) + // Even severely malformed JSON like '{invalid' gets parsed by partial-json-parser + // It returns an empty object, which is still a valid (though incomplete) result + const veryPartial = NativeToolCallParser.processStreamingChunk("toolu_fail123", "{invalid") + // partial-json-parser handles this gracefully - it may return an empty object or null + // The key point is it doesn't throw an error and the streaming continues + if (veryPartial != null) { + expect(veryPartial.partial).toBe(true) + } else { + // null is also acceptable for severely malformed JSON + expect(veryPartial).toBeNull() + } + }) + }) + + describe("processFinishReason and finalizeRawChunks integration", () => { + it("should emit tool_call_end events when finish_reason is 'tool_calls'", () => { + NativeToolCallParser.clearRawChunkState() + + // First process some raw chunks to populate tracker + NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_finish123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + const events = NativeToolCallParser.processFinishReason("tool_calls") + + expect(events.length).toBeGreaterThan(0) + expect(events[0].type).toBe("tool_call_end") + }) + + it("should finalize remaining raw chunks via finalizeRawChunks", () => { + NativeToolCallParser.clearRawChunkState() + + NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_finalize123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + const events = NativeToolCallParser.finalizeRawChunks() + + expect(events.length).toBeGreaterThan(0) + expect(events[0].type).toBe("tool_call_end") + }) + + it("should clear raw chunk state via clearRawChunkState", () => { + NativeToolCallParser.clearRawChunkState() + + NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_clear123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + NativeToolCallParser.clearRawChunkState() + + const events = NativeToolCallParser.finalizeRawChunks() + expect(events.length).toBe(0) + }) + }) + + describe("tool_call_partial chunk handling - Task integration", () => { + it("should emit tool_call_start event when processing raw chunk with id and name", async () => { + NativeToolCallParser.clearRawChunkState() + + const events = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_123", + name: "read_file", + arguments: '{"path":"a.ts"}', + }) + + // Should emit both start and delta event + expect(events.length).toBeGreaterThan(0) + const startEvent = events.find((e) => e.type === "tool_call_start") + expect(startEvent).toBeDefined() + if (startEvent && startEvent.type === "tool_call_start") { + expect(startEvent.id).toBe("toolu_123") + expect(startEvent.name).toBe("read_file") + } + const deltaEvents = events.filter((e) => e.type === "tool_call_delta") + expect(deltaEvents.length).toBeGreaterThan(0) + }) + + it("should handle duplicate tool_call_partial chunks with same index", async () => { + NativeToolCallParser.clearRawChunkState() + + const events1 = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_dup123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + const events2 = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_dup123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + // Both should emit delta events (dedup is handled by Task, not NativeToolCallParser) + expect(events1.length).toBeGreaterThan(0) + expect(events2.length).toBeGreaterThan(0) + }) + + it("should handle tool_call_delta event without id", async () => { + NativeToolCallParser.clearRawChunkState() + + const events = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_delta123", + name: undefined, + arguments: undefined, + }) + + // Without name, no start event should be emitted + expect(events.length).toBe(0) + }) + + it("should handle tool_call_end via finalizeRawChunks", async () => { + NativeToolCallParser.clearRawChunkState() + + // First process a raw chunk to track the tool call + NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_end123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + // finalizeRawChunks should emit end events for all tracked tools that have started + const events = NativeToolCallParser.finalizeRawChunks() + + expect(events).toHaveLength(1) + expect(events[0]).toEqual({ type: "tool_call_end", id: "toolu_end123", name: "read_file" }) + }) + + it("should handle complete streaming lifecycle: processRawChunk -> finalizeRawChunks", async () => { + NativeToolCallParser.clearRawChunkState() + + // Start + const startEvents = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_lifecycle123", + name: "read_file", + arguments: '{"path":"test.ts"}', + }) + + expect(startEvents.some((e) => e.type === "tool_call_start")).toBe(true) + + // Delta (simulating another chunk with same index) + const deltaEvents = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_lifecycle123", + name: "read_file", + arguments: ',"more":"args"', + }) + + expect(deltaEvents.some((e) => e.type === "tool_call_delta")).toBe(true) + + // End via finalize + const endEvents = NativeToolCallParser.finalizeRawChunks() + + expect(endEvents).toHaveLength(1) + expect(endEvents[0]).toEqual({ type: "tool_call_end", id: "toolu_lifecycle123", name: "read_file" }) + }) + + it("should handle multiple sequential tool calls with different indices", async () => { + NativeToolCallParser.clearRawChunkState() + + // First tool call (index 0) + const events1 = NativeToolCallParser.processRawChunk({ + index: 0, + id: "toolu_multi1", + name: "read_file", + arguments: '{"path":"file1.ts"}', + }) + + expect(events1.some((e) => e.type === "tool_call_start")).toBe(true) + + // Second tool call (index 1) + const events2 = NativeToolCallParser.processRawChunk({ + index: 1, + id: "toolu_multi2", + name: "write_to_file", + arguments: '{"path":"file2.ts","content":"hello"}', + }) + + expect(events2.some((e) => e.type === "tool_call_start")).toBe(true) + + // Finalize both + const endEvents = NativeToolCallParser.finalizeRawChunks() + + expect(endEvents).toHaveLength(2) + const endIds = endEvents.map((e) => e.id) + expect(endIds).toContain("toolu_multi1") + expect(endIds).toContain("toolu_multi2") + }) + }) + + describe("Task.ts compound key dedup and streaming paths", () => { + it("should build and use compound streaming keys via makeStreamingKey", () => { + const key = NativeToolCallParser.makeStreamingKey("toolu_compound123", "read_file") + expect(key).toBe("toolu_compound123::read_file") + + NativeToolCallParser.startStreamingToolCall("toolu_compound123", "read_file") + expect(NativeToolCallParser.getStreamingToolName(key)).toBe("read_file") + }) + + it("should ignore duplicate tool_call_start for same compound key", () => { + const streamingToolCallIndices = new Map() + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}) + const events = [ + { type: "tool_call_start", id: "toolu_dup_compound", name: "read_file" as const }, + { type: "tool_call_start", id: "toolu_dup_compound", name: "read_file" as const }, + ] + const assistantMessageContent: any[] = [] + + for (const event of events) { + const dedupKey = `${event.id}::${event.name}` + if (streamingToolCallIndices.has(dedupKey)) { + console.warn( + `[Task#test] Ignoring duplicate tool_call_start for ID: ${event.id} (tool: ${event.name})`, + ) + continue + } + streamingToolCallIndices.set(dedupKey, assistantMessageContent.length) + assistantMessageContent.push({ type: "tool_use", id: event.id, name: event.name, partial: true }) + } + + expect(assistantMessageContent).toHaveLength(1) + expect(streamingToolCallIndices.has("toolu_dup_compound::read_file")).toBe(true) + expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("[Task#test]")) + warnSpy.mockRestore() + }) + + it("should use compound key for delta events via getStreamingToolCallById", () => { + const streamingToolCallIndices = new Map() + const assistantMessageContent: any[] = [] + + NativeToolCallParser.startStreamingToolCall("toolu_delta_compound", "read_file") + const dedupKey = NativeToolCallParser.makeStreamingKey("toolu_delta_compound", "read_file") + streamingToolCallIndices.set(dedupKey, 0) + assistantMessageContent.push({ + type: "tool_use", + id: "toolu_delta_compound", + name: "read_file", + params: {}, + partial: true, + }) + + const partialToolUse = NativeToolCallParser.processStreamingChunk(dedupKey, '{"path":"') + expect(partialToolUse).not.toBeNull() + + const existingEntry = NativeToolCallParser.getStreamingToolCallById("toolu_delta_compound") + expect(existingEntry).not.toBeNull() + const resolvedKey = existingEntry + ? NativeToolCallParser.makeStreamingKey(existingEntry.id, existingEntry.name) + : undefined + const updatedToolUse = NativeToolCallParser.processStreamingChunk(resolvedKey!, '"test.ts"}') + expect(updatedToolUse).not.toBeNull() + const name = NativeToolCallParser.getStreamingToolName(resolvedKey!) + const toolUseIndex = streamingToolCallIndices.get(`${existingEntry!.id}::${name}`) + expect(toolUseIndex).toBe(0) + assistantMessageContent[toolUseIndex!] = updatedToolUse as any + + expect(assistantMessageContent).toHaveLength(1) + expect(assistantMessageContent[0]?.type).toBe("tool_use") + expect(assistantMessageContent[0]?.name).toBe("read_file") + }) + + it("should finalize and clean up tracking using compound key on tool_call_end", () => { + const streamingToolCallIndices = new Map() + const assistantMessageContent: any[] = [] + const dedupKey = NativeToolCallParser.makeStreamingKey("toolu_cleanup123", "read_file") + + NativeToolCallParser.startStreamingToolCall("toolu_cleanup123", "read_file") + streamingToolCallIndices.set(dedupKey, 0) + assistantMessageContent.push({ type: "tool_use", id: "toolu_cleanup123", name: "read_file", partial: true }) + NativeToolCallParser.processStreamingChunk(dedupKey, '{"path":"test.ts"}') + + const finalToolUse = NativeToolCallParser.finalizeStreamingToolCall(dedupKey) + expect(finalToolUse).not.toBeNull() + assistantMessageContent[0] = finalToolUse as any + streamingToolCallIndices.delete(dedupKey) + + expect(assistantMessageContent).toHaveLength(1) + expect(assistantMessageContent[0]?.type).toBe("tool_use") + expect(assistantMessageContent[0]?.name).toBe("read_file") + expect(assistantMessageContent[0]?.partial).toBe(false) + expect(streamingToolCallIndices.has("toolu_cleanup123::read_file")).toBe(false) + }) + + it("should handle malformed JSON finalization with compound key cleanup", () => { + const streamingToolCallIndices = new Map() + const assistantMessageContent: any[] = [] + const dedupKey = NativeToolCallParser.makeStreamingKey("toolu_malformed123", "read_file") + + NativeToolCallParser.startStreamingToolCall("toolu_malformed123", "read_file") + streamingToolCallIndices.set(dedupKey, 0) + assistantMessageContent.push({ + type: "tool_use", + id: "toolu_malformed123", + name: "read_file", + partial: true, + }) + NativeToolCallParser.processStreamingChunk(dedupKey, "{invalid json") + + const finalToolUse = NativeToolCallParser.finalizeStreamingToolCall(dedupKey) + expect(finalToolUse).toBeNull() + ;(assistantMessageContent[0] as any).partial = false + streamingToolCallIndices.delete(dedupKey) + + expect(assistantMessageContent).toHaveLength(1) + expect(assistantMessageContent[0]?.type).toBe("tool_use") + expect(assistantMessageContent[0]?.name).toBe("read_file") + expect(assistantMessageContent[0]?.partial).toBe(false) + expect(streamingToolCallIndices.has("toolu_malformed123::read_file")).toBe(false) + }) + }) +}) diff --git a/src/core/task/__tests__/duplicate-tool-use-ids.spec.ts b/src/core/task/__tests__/duplicate-tool-use-ids.spec.ts index 33e5afe236..16c401aea5 100644 --- a/src/core/task/__tests__/duplicate-tool-use-ids.spec.ts +++ b/src/core/task/__tests__/duplicate-tool-use-ids.spec.ts @@ -259,4 +259,138 @@ describe("Duplicate tool_use ID Prevention", () => { expect(ids).toEqual(uniqueIds) // All IDs are unique }) }) + + describe("Compound key deduplication (id + name)", () => { + /** + * Tests the compound key fix: when two different tools share the same toolCallId, + * they should both be processed because their (id, name) keys differ. + * + * The streamingToolCallIndices map now uses `${event.id}::${event.name}` as the key + * instead of just `event.id`, allowing different tools with the same call ID to coexist. + */ + + it("should allow two different tools with same toolCallId but different names", () => { + const streamingToolCallIndices = new Map() + const processedEvents: Array<{ id: string; name: string }> = [] + + const processToolCallStart = (id: string, name: string): boolean => { + // Compound key deduplication (new behavior) + const dedupKey = `${id}::${name}` + if (streamingToolCallIndices.has(dedupKey)) { + return false // Skipped as duplicate + } + streamingToolCallIndices.set(dedupKey, processedEvents.length) + processedEvents.push({ id, name }) + return true + } + + // First tool with shared ID + expect(processToolCallStart("toolu_123", "read_file")).toBe(true) + expect(processedEvents).toEqual([{ id: "toolu_123", name: "read_file" }]) + + // Second DIFFERENT tool with same ID should succeed (not dropped) + expect(processToolCallStart("toolu_123", "write_to_file")).toBe(true) + expect(processedEvents).toEqual([ + { id: "toolu_123", name: "read_file" }, + { id: "toolu_123", name: "write_to_file" }, + ]) + + // True duplicate (same ID AND same name) should be rejected + expect(processToolCallStart("toolu_123", "read_file")).toBe(false) + expect(processedEvents).toHaveLength(2) // No change + }) + + it("should reject true duplicates but allow different tools with same callId", () => { + const streamingToolCallIndices = new Map() + const processedEvents: Array<{ id: string; name: string }> = [] + + const processToolCallStart = (id: string, name: string): boolean => { + const dedupKey = `${id}::${name}` + if (streamingToolCallIndices.has(dedupKey)) { + return false + } + streamingToolCallIndices.set(dedupKey, processedEvents.length) + processedEvents.push({ id, name }) + return true + } + + // First call: read_file with toolu_abc + expect(processToolCallStart("toolu_abc", "read_file")).toBe(true) + + // Stream retry of the SAME event should be rejected + expect(processToolCallStart("toolu_abc", "read_file")).toBe(false) + + // Different tool with same callId should succeed + expect(processToolCallStart("toolu_abc", "get_status")).toBe(true) + + // Retry of get_status should also be rejected + expect(processToolCallStart("toolu_abc", "get_status")).toBe(false) + + // Both tools processed + expect(processedEvents).toHaveLength(2) + expect(processedEvents[0].name).toBe("read_file") + expect(processedEvents[1].name).toBe("get_status") + }) + + it("should track correct indices for compound key lookups", () => { + const streamingToolCallIndices = new Map() + let currentIndex = 0 + + const processToolCallStart = (id: string, name: string): number | null => { + const dedupKey = `${id}::${name}` + if (streamingToolCallIndices.has(dedupKey)) { + return null + } + const index = currentIndex + streamingToolCallIndices.set(dedupKey, index) + currentIndex++ + return index + } + + expect(processToolCallStart("toolu_1", "tool_a")).toBe(0) + expect(processToolCallStart("toolu_1", "tool_b")).toBe(1) + expect(processToolCallStart("toolu_2", "tool_a")).toBe(2) + + // Verify compound key lookups return correct indices + expect(streamingToolCallIndices.get("toolu_1::tool_a")).toBe(0) + expect(streamingToolCallIndices.get("toolu_1::tool_b")).toBe(1) + expect(streamingToolCallIndices.get("toolu_2::tool_a")).toBe(2) + + // Old single-key lookup would return wrong index or conflict + // With compound keys, each (id, name) pair has its own entry + expect(streamingToolCallIndices.has("toolu_1")).toBe(false) // No bare key + expect(streamingToolCallIndices.has("toolu_2")).toBe(false) + }) + + it("should clean up after tool_call_end using compound key", () => { + const streamingToolCallIndices = new Map() + + const addTool = (id: string, name: string): void => { + const dedupKey = `${id}::${name}` + streamingToolCallIndices.set(dedupKey, streamingToolCallIndices.size) + } + + const removeTool = (id: string, name: string): boolean => { + const dedupKey = `${id}::${name}` + return streamingToolCallIndices.delete(dedupKey) + } + + addTool("toolu_1", "read_file") + addTool("toolu_1", "write_to_file") + + expect(streamingToolCallIndices.size).toBe(2) + expect(streamingToolCallIndices.has("toolu_1::read_file")).toBe(true) + expect(streamingToolCallIndices.has("toolu_1::write_to_file")).toBe(true) + + // Finalize first tool - should only remove that specific entry + expect(removeTool("toolu_1", "read_file")).toBe(true) + expect(streamingToolCallIndices.size).toBe(1) + expect(streamingToolCallIndices.has("toolu_1::read_file")).toBe(false) + expect(streamingToolCallIndices.has("toolu_1::write_to_file")).toBe(true) + + // Finalize second tool + expect(removeTool("toolu_1", "write_to_file")).toBe(true) + expect(streamingToolCallIndices.size).toBe(0) + }) + }) }) diff --git a/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts b/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts index 4a035447b1..6e9cf50966 100644 --- a/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts +++ b/src/core/tools/__tests__/askFollowupQuestionTool.spec.ts @@ -496,11 +496,14 @@ describe("AskFollowupQuestionTool", () => { it("should build nativeArgs with question and follow_up during streaming", () => { // Start a streaming tool call - NativeToolCallParser.startStreamingToolCall("call_123", "ask_followup_question") + const id1 = "call_123" + const name1 = "ask_followup_question" + NativeToolCallParser.startStreamingToolCall(id1, name1) + const key1 = NativeToolCallParser.makeStreamingKey(id1, name1) // Simulate streaming JSON chunks const chunk1 = '{"question":"What would you like?","follow_up":[{"text":"Option 1","mode":"code"}' - const result1 = NativeToolCallParser.processStreamingChunk("call_123", chunk1) + const result1 = NativeToolCallParser.processStreamingChunk(key1, chunk1) expect(result1).not.toBeNull() expect(result1?.name).toBe("ask_followup_question") @@ -517,14 +520,17 @@ describe("AskFollowupQuestionTool", () => { }) it("should finalize with complete nativeArgs", () => { - NativeToolCallParser.startStreamingToolCall("call_456", "ask_followup_question") + const id2 = "call_456" + const name2 = "ask_followup_question" + NativeToolCallParser.startStreamingToolCall(id2, name2) + const key2 = NativeToolCallParser.makeStreamingKey(id2, name2) // Add complete JSON const completeJson = '{"question":"Choose an option","follow_up":[{"text":"Yes","mode":"code"},{"text":"No","mode":null}]}' - NativeToolCallParser.processStreamingChunk("call_456", completeJson) + NativeToolCallParser.processStreamingChunk(key2, completeJson) - const result = NativeToolCallParser.finalizeStreamingToolCall("call_456") + const result = NativeToolCallParser.finalizeStreamingToolCall(key2) expect(result).not.toBeNull() expect(result?.type).toBe("tool_use") @@ -543,14 +549,17 @@ describe("AskFollowupQuestionTool", () => { }) it("should finalize and forward a non-array follow_up so the tool can report it", () => { - NativeToolCallParser.startStreamingToolCall("call_789", "ask_followup_question") + const id3 = "call_789" + const name3 = "ask_followup_question" + NativeToolCallParser.startStreamingToolCall(id3, name3) + const key3 = NativeToolCallParser.makeStreamingKey(id3, name3) // follow_up arrives as a keyed object instead of an array (the bug repro). const completeJson = '{"question":"How should I proceed?","follow_up":{"0":{"mode":null,"text":"Keep"},"1":{"mode":null,"text":"Remove"}}}' - NativeToolCallParser.processStreamingChunk("call_789", completeJson) + NativeToolCallParser.processStreamingChunk(key3, completeJson) - const result = NativeToolCallParser.finalizeStreamingToolCall("call_789") + const result = NativeToolCallParser.finalizeStreamingToolCall(key3) // The call must NOT be dropped (null) - it should reach the tool with the raw // value so the tool can emit a precise "must be an array" error.