diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-parallel-tools.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-parallel-tools.spec.ts new file mode 100644 index 00000000000..0ec2f4447ac --- /dev/null +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-parallel-tools.spec.ts @@ -0,0 +1,307 @@ +// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-parallel-tools.spec.ts + +import { describe, it, expect, beforeEach, vi } from "vitest" +import { presentAssistantMessage } from "../presentAssistantMessage" + +// Mock dependencies +vi.mock("../../task/Task") +vi.mock("../../tools/validateToolUse", () => ({ + validateToolUse: vi.fn(), + isValidToolName: vi.fn(() => true), +})) +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureToolUsage: vi.fn(), + captureConsecutiveMistakeError: vi.fn(), + }, + }, +})) + +// Mock the tool handlers to avoid complex setup +vi.mock("../../tools/ListFilesTool", () => ({ + listFilesTool: { + handle: vi.fn().mockImplementation(async (cline, block, callbacks) => { + // Simulate async tool execution - tool result is pushed asynchronously + await Promise.resolve() + callbacks.pushToolResult("list_files result") + }), + }, +})) + +vi.mock("../../tools/ReadFileTool", () => ({ + readFileTool: { + handle: vi.fn().mockImplementation(async (cline, block, callbacks) => { + // Simulate async tool execution - tool result is pushed asynchronously + await Promise.resolve() + callbacks.pushToolResult("read_file result") + }), + getReadFileToolDescription: vi.fn(() => "[read_file]"), + }, +})) + +describe("presentAssistantMessage - Parallel Tool Execution Timing", () => { + let mockTask: any + + beforeEach(() => { + // Create a mock Task with minimal properties needed for testing + mockTask = { + taskId: "test-task-id", + instanceId: "test-instance", + abort: false, + presentAssistantMessageLocked: false, + presentAssistantMessageHasPendingUpdates: false, + currentStreamingContentIndex: 0, + assistantMessageContent: [], + userMessageContent: [], + userMessageContentReady: false, + didCompleteReadingStream: false, + didRejectTool: false, + didAlreadyUseTool: false, + consecutiveMistakeCount: 0, + clineMessages: [], + api: { + getModel: () => ({ id: "test-model", info: {} }), + }, + browserSession: { + closeBrowser: vi.fn().mockResolvedValue(undefined), + }, + recordToolUsage: vi.fn(), + recordToolError: vi.fn(), + toolRepetitionDetector: { + check: vi.fn().mockReturnValue({ allowExecution: true }), + }, + providerRef: { + deref: () => ({ + getState: vi.fn().mockResolvedValue({ + mode: "code", + customModes: [], + }), + }), + }, + say: vi.fn().mockResolvedValue(undefined), + ask: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }), + } + + // Add pushToolResultToUserContent method + mockTask.pushToolResultToUserContent = vi.fn().mockImplementation((toolResult: any) => { + const existingResult = mockTask.userMessageContent.find( + (block: any) => block.type === "tool_result" && block.tool_use_id === toolResult.tool_use_id, + ) + if (existingResult) { + return false + } + mockTask.userMessageContent.push(toolResult) + return true + }) + }) + + it("should NOT set userMessageContentReady until all tool_results are collected for parallel tools", async () => { + // Set up multiple tool_use blocks (parallel tool calls) + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: "tool_call_1", + name: "list_files", + params: { path: "/test" }, + partial: false, + }, + { + type: "tool_use", + id: "tool_call_2", + name: "read_file", + params: { path: "/test/file.txt" }, + partial: false, + }, + ] + + mockTask.didCompleteReadingStream = true + + // Process first tool + await presentAssistantMessage(mockTask) + + // After processing first tool, userMessageContentReady should NOT be true + // because tool_call_2 doesn't have a tool_result yet + // Note: Due to how the mock is set up, the first tool should push its result + // but the second tool hasn't been processed yet + expect(mockTask.userMessageContent.length).toBeGreaterThanOrEqual(0) + + // If only one tool_result exists for two tool_use blocks, userMessageContentReady should be false + if (mockTask.userMessageContent.length === 1) { + expect(mockTask.userMessageContentReady).toBe(false) + } + }) + + it("should set userMessageContentReady when all tool_results are collected", async () => { + // Set up a single tool_use block + const toolCallId = "tool_call_single" + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: toolCallId, + name: "list_files", + params: { path: "/test" }, + partial: false, + }, + ] + + mockTask.didCompleteReadingStream = true + + await presentAssistantMessage(mockTask) + + // After the tool executes and pushes its result, userMessageContentReady should be true + // because there's 1 tool_use and 1 tool_result + const toolResultCount = mockTask.userMessageContent.filter((b: any) => b.type === "tool_result").length + + if (toolResultCount === 1) { + expect(mockTask.userMessageContentReady).toBe(true) + } + }) + + it("should handle text-only content without waiting for tool_results", async () => { + // Set up a text-only content block (no tools) + mockTask.assistantMessageContent = [ + { + type: "text", + content: "Hello, this is a text response", + partial: false, + }, + ] + + mockTask.didCompleteReadingStream = true + + await presentAssistantMessage(mockTask) + + // With no tool_use blocks, userMessageContentReady should be true after processing text + expect(mockTask.userMessageContentReady).toBe(true) + }) + + it("should wait for tool_results even when didRejectTool is true", async () => { + // Set up multiple tool_use blocks + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: "tool_call_rejected_1", + name: "list_files", + params: { path: "/test" }, + partial: false, + }, + { + type: "tool_use", + id: "tool_call_rejected_2", + name: "read_file", + params: { path: "/test/file.txt" }, + partial: false, + }, + ] + + mockTask.didRejectTool = true + mockTask.didCompleteReadingStream = true + + await presentAssistantMessage(mockTask) + + // When didRejectTool is true, error tool_results should be pushed for each tool + // Both should have tool_results (skipped messages) + const toolResults = mockTask.userMessageContent.filter((b: any) => b.type === "tool_result") + + // The function should have pushed error tool_results for rejected tools + expect(toolResults.length).toBeGreaterThan(0) + + // If all tool_results are collected, userMessageContentReady should be true + const toolUseCount = mockTask.assistantMessageContent.filter( + (b: any) => b.type === "tool_use" || b.type === "mcp_tool_use", + ).length + + if (toolResults.length >= toolUseCount) { + expect(mockTask.userMessageContentReady).toBe(true) + } + }) + + it("should not set userMessageContentReady if stream is not complete", async () => { + // Set up a tool_use block + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: "tool_call_stream", + name: "list_files", + params: { path: "/test" }, + partial: false, + }, + ] + + // Stream is NOT complete + mockTask.didCompleteReadingStream = false + + await presentAssistantMessage(mockTask) + + // Even if the tool executed, userMessageContentReady should NOT be true + // because the stream hasn't completed yet (more content may arrive) + // Note: The fix specifically checks both conditions + expect(mockTask.userMessageContentReady).toBe(false) + }) + + it("should handle mcp_tool_use blocks the same as tool_use blocks", async () => { + // Set up an mcp_tool_use block (MCP tool) + mockTask.assistantMessageContent = [ + { + type: "mcp_tool_use", + id: "mcp_tool_call_1", + name: "mcp_server_tool", + serverName: "test_server", + toolName: "test_tool", + arguments: {}, + partial: false, + }, + ] + + mockTask.didRejectTool = true // Use rejection to get a simple tool_result + mockTask.didCompleteReadingStream = true + + await presentAssistantMessage(mockTask) + + // The mcp_tool_use should be treated similarly - needs tool_result before ready + const toolResults = mockTask.userMessageContent.filter((b: any) => b.type === "tool_result") + const toolUseCount = mockTask.assistantMessageContent.filter( + (b: any) => b.type === "tool_use" || b.type === "mcp_tool_use", + ).length + + // If all tool_results are collected, userMessageContentReady should be true + if (toolResults.length >= toolUseCount) { + expect(mockTask.userMessageContentReady).toBe(true) + } + }) + + it("should correctly count mixed tool_use and mcp_tool_use blocks", async () => { + // Set up mixed tool blocks + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: "regular_tool_1", + name: "list_files", + params: { path: "/test" }, + partial: false, + }, + { + type: "mcp_tool_use", + id: "mcp_tool_1", + name: "mcp_server_tool", + serverName: "test_server", + toolName: "test_tool", + arguments: {}, + partial: false, + }, + ] + + mockTask.didRejectTool = true // Simplify by using rejection + mockTask.didCompleteReadingStream = true + + await presentAssistantMessage(mockTask) + + // Both tool types should require tool_results + const toolUseCount = mockTask.assistantMessageContent.filter( + (b: any) => b.type === "tool_use" || b.type === "mcp_tool_use", + ).length + + expect(toolUseCount).toBe(2) + }) +}) diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index 5b9e1c4840e..03fe0039e43 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -41,6 +41,46 @@ import { codebaseSearchTool } from "../tools/CodebaseSearchTool" import { formatResponse } from "../prompts/responses" +/** + * Checks if all tool_use blocks in the assistant message have corresponding tool_result + * blocks in the user message content. This is critical for parallel tool execution - + * we must wait for all tool results before signaling that the message is ready. + * + * @param cline - The Task instance + * @returns true if all tool_use blocks have matching tool_result blocks, false otherwise + */ +function areAllToolResultsCollected(cline: Task): boolean { + // Count tool_use and mcp_tool_use blocks in assistant message + const toolUseIds = new Set() + for (const block of cline.assistantMessageContent) { + if ((block.type === "tool_use" || block.type === "mcp_tool_use") && (block as any).id) { + toolUseIds.add((block as any).id) + } + } + + // If no tool_use blocks, we're ready + if (toolUseIds.size === 0) { + return true + } + + // Count tool_result blocks in user message content + const toolResultIds = new Set() + for (const block of cline.userMessageContent) { + if (block.type === "tool_result" && (block as any).tool_use_id) { + toolResultIds.add((block as any).tool_use_id) + } + } + + // Check if every tool_use has a corresponding tool_result + for (const toolUseId of toolUseIds) { + if (!toolResultIds.has(toolUseId)) { + return false + } + } + + return true +} + /** * Processes and presents assistant message content to the user interface. * @@ -76,7 +116,8 @@ export async function presentAssistantMessage(cline: Task) { // streaming could finish. If streaming is finished, and we're out of // bounds then this means we already presented/executed the last // content block and are ready to continue to next request. - if (cline.didCompleteReadingStream) { + // CRITICAL: Also verify all tool results are collected for parallel tool execution + if (cline.didCompleteReadingStream && areAllToolResultsCollected(cline)) { cline.userMessageContentReady = true } @@ -984,8 +1025,14 @@ export async function presentAssistantMessage(cline: Task) { // streaming is finished then we set `userMessageContentReady` to // true when out of bounds. This gracefully allows the stream to // continue on and all potential content blocks be presented. - // Last block is complete and it is finished executing - cline.userMessageContentReady = true // Will allow `pWaitFor` to continue. + // Last block is complete and it is finished executing. + // CRITICAL: For parallel tool execution, we must verify all tool results + // are collected before signaling ready. Without this check, the message + // queue could proceed before all tool_result blocks are in userMessageContent. + // Also verify the stream is complete - more content may still arrive. + if (cline.didCompleteReadingStream && areAllToolResultsCollected(cline)) { + cline.userMessageContentReady = true // Will allow `pWaitFor` to continue. + } } // Call next block if it exists (if not then read stream will call it @@ -1002,7 +1049,8 @@ export async function presentAssistantMessage(cline: Task) { } else { // CRITICAL FIX: If we're out of bounds and the stream is complete, set userMessageContentReady // This handles the case where assistantMessageContent is empty or becomes empty after processing - if (cline.didCompleteReadingStream) { + // Also verify all tool results are collected for parallel tool execution + if (cline.didCompleteReadingStream && areAllToolResultsCollected(cline)) { cline.userMessageContentReady = true } }