Merge ace2500418 into cb83656718

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)
+	})
+})

src/core/assistant-message/presentAssistantMessage.ts

@@ -41,6 +41,46 @@ import { codebaseSearchTool } from "../tools/CodebaseSearchTool"
 import { formatResponse } from "../prompts/responses"
 import { sanitizeToolUseId } from "../../utils/tool-id"
 
+/**
+ * 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<string>()
+	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<string>()
+	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
 		}
 
@@ -946,8 +987,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
+			// Last block is complete and it is finished executing.
-			cline.userMessageContentReady = true // Will allow `pWaitFor` to continue.
+			// 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
@@ -964,7 +1011,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
 			}
 		}