diff --git a/src/core/assistant-message/__tests__/presentAssistantMessage-stale-data.spec.ts b/src/core/assistant-message/__tests__/presentAssistantMessage-stale-data.spec.ts new file mode 100644 index 00000000000..34539d9545e --- /dev/null +++ b/src/core/assistant-message/__tests__/presentAssistantMessage-stale-data.spec.ts @@ -0,0 +1,249 @@ +// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-stale-data.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(), +})) +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + instance: { + captureToolUsage: vi.fn(), + captureConsecutiveMistakeError: vi.fn(), + }, + }, +})) + +/** + * Tests for the fix to ROO-311: write_to_file truncates filenames at special characters + * + * The issue was that during streaming, partial-json parsing could produce truncated + * values (e.g., path "sr" instead of "src/core/prompts/sections/skills.ts"). + * These truncated values would be cloned by presentAssistantMessage before the + * final tool_call_end event updated the array with correct data. + * + * The fix ensures that for non-partial tool_use blocks, we re-read from + * assistantMessageContent to get the final data, not stale partial data. + */ +describe("presentAssistantMessage - Stale Partial Data Fix (ROO-311)", () => { + 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: [], + didCompleteReadingStream: false, + didRejectTool: false, + didAlreadyUseTool: false, + diffEnabled: 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" }), + } + }) + + it("should use fresh data from assistantMessageContent for non-partial tool_use blocks", async () => { + // This test simulates the scenario where: + // 1. A partial tool_use block is added with truncated data + // 2. The block is updated with final data + // 3. presentAssistantMessage should use the final data, not the stale partial data + + const toolCallId = "tool_call_stale_test" + const truncatedPath = "sr" // Truncated by partial-json + const fullPath = "src/core/prompts/sections/skills.ts" // Full path from final JSON + + // Set up the assistantMessageContent with the FINAL data + // (simulating what happens after tool_call_end updates the array) + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: toolCallId, + name: "write_to_file", + params: { + path: fullPath, + content: "file content", + }, + nativeArgs: { + path: fullPath, + content: "file content", + }, + partial: false, // Non-partial = ready for execution + }, + ] + + // Execute presentAssistantMessage + await presentAssistantMessage(mockTask) + + // The block that was processed should have the full path, not truncated + // We can verify this by checking that the tool was called with correct params + // Since we're mocking, we just verify the block in assistantMessageContent has correct data + const processedBlock = mockTask.assistantMessageContent[0] + expect(processedBlock.params.path).toBe(fullPath) + expect(processedBlock.nativeArgs.path).toBe(fullPath) + expect(processedBlock.params.path).not.toBe(truncatedPath) + }) + + it("should re-read block data when processing non-partial tool_use", async () => { + // This test verifies that the fix actually re-reads from assistantMessageContent + // by checking that updates to the array are reflected in the processed block + + const toolCallId = "tool_call_reread_test" + + // Initial setup with partial-like data (but marked as non-partial) + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: toolCallId, + name: "read_file", + params: { + path: "test.ts", + }, + nativeArgs: { + files: [{ path: "test.ts" }], + }, + partial: false, + }, + ] + + // Execute presentAssistantMessage + await presentAssistantMessage(mockTask) + + // Verify the block was processed (not frozen due to stale data) + // The test passes if presentAssistantMessage completes without error + // and the block is still in the expected state + expect(mockTask.assistantMessageContent[0].partial).toBe(false) + }) + + it("should not re-read for partial tool_use blocks (streaming in progress)", async () => { + // Partial blocks should NOT be re-read because they're still being updated + // by the streaming process + + const toolCallId = "tool_call_partial_test" + + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: toolCallId, + name: "write_to_file", + params: { + path: "partial/path", // Partial data + }, + partial: true, // Still streaming + }, + ] + + // Execute presentAssistantMessage + await presentAssistantMessage(mockTask) + + // For partial blocks, the function should return early without processing + // (partial blocks are handled by handlePartial, not execute) + // Verify the block is still partial + expect(mockTask.assistantMessageContent[0].partial).toBe(true) + }) + + it("should handle the case where assistantMessageContent is updated between clone and execution", async () => { + // This test simulates a race condition where: + // 1. Block is cloned with partial data + // 2. Array is updated with final data + // 3. The fix ensures we use the final data + + const toolCallId = "tool_call_race_test" + const partialPath = "src/co" // Truncated + const finalPath = "src/core/prompts/sections/skills.ts" + + // Start with partial-looking data but marked as non-partial + // (simulating the moment right after tool_call_end but before re-read) + mockTask.assistantMessageContent = [ + { + type: "tool_use", + id: toolCallId, + name: "write_to_file", + params: { + path: finalPath, // Final data is already in the array + content: "content", + }, + nativeArgs: { + path: finalPath, + content: "content", + }, + partial: false, + }, + ] + + // Execute presentAssistantMessage + await presentAssistantMessage(mockTask) + + // Verify the final path was used + const block = mockTask.assistantMessageContent[0] + expect(block.params.path).toBe(finalPath) + expect(block.params.path).not.toBe(partialPath) + }) + + it("should handle mcp_tool_use blocks correctly (no re-read needed)", async () => { + // MCP tool use blocks have a different type and should be handled correctly + + const toolCallId = "mcp_tool_call_test" + + // Add getMcpHub mock for MCP tool handling + mockTask.providerRef = { + deref: () => ({ + getState: vi.fn().mockResolvedValue({ + mode: "code", + customModes: [], + }), + getMcpHub: vi.fn().mockReturnValue({ + findServerNameBySanitizedName: vi.fn().mockReturnValue("server"), + }), + }), + } + + mockTask.assistantMessageContent = [ + { + type: "mcp_tool_use", + id: toolCallId, + name: "mcp--server--tool", + serverName: "server", + toolName: "tool", + arguments: { arg: "value" }, + partial: false, + }, + ] + + // Execute presentAssistantMessage + await presentAssistantMessage(mockTask) + + // MCP tool use should be processed without error + // The test passes if no exception is thrown + expect(mockTask.assistantMessageContent[0].type).toBe("mcp_tool_use") + }) +}) diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index 9b1a5fcffb1..2d4435c91b4 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -100,6 +100,13 @@ export async function presentAssistantMessage(cline: Task) { return } + if (block.type === "tool_use" && !block.partial) { + const freshBlock = cline.assistantMessageContent[cline.currentStreamingContentIndex] + if (freshBlock && freshBlock.type === "tool_use" && !freshBlock.partial) { + block = cloneDeep(freshBlock) + } + } + switch (block.type) { case "mcp_tool_use": { // Handle native MCP tool calls (from mcp_serverName_toolName dynamic tools) diff --git a/src/core/tools/WriteToFileTool.ts b/src/core/tools/WriteToFileTool.ts index 7caaeb6d55d..c4e88ee17fb 100644 --- a/src/core/tools/WriteToFileTool.ts +++ b/src/core/tools/WriteToFileTool.ts @@ -146,6 +146,8 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { const partialMessage = JSON.stringify(sharedMessageProps) await task.ask("tool", partialMessage, true).catch(() => {}) await task.diffViewProvider.open(relPath) + } else { + task.diffViewProvider.setRelPath(relPath) } await task.diffViewProvider.update( @@ -187,6 +189,7 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { pushToolResult(message) await task.diffViewProvider.reset() + this.resetPartialState() task.processQueuedMessages() @@ -194,10 +197,14 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { } catch (error) { await handleError("writing file", error as Error) await task.diffViewProvider.reset() + this.resetPartialState() return } } + // Track the last seen path during streaming to detect when the path has stabilized + private lastSeenPartialPath: string | undefined = undefined + override async handlePartial(task: Task, block: ToolUse<"write_to_file">): Promise { const relPath: string | undefined = block.params.path let newContent: string | undefined = block.params.content @@ -217,6 +224,13 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { return } + // During streaming, the partial-json library may return truncated string values + // when chunk boundaries fall mid-value. To avoid creating files at incorrect paths, + // we wait until the path stops changing between consecutive partial blocks before + // creating the file. This ensures we have the complete, final path value. + const pathHasStabilized = this.lastSeenPartialPath !== undefined && this.lastSeenPartialPath === relPath + this.lastSeenPartialPath = relPath + let fileExists: boolean const absolutePath = path.resolve(task.cwd, relPath) @@ -227,12 +241,6 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { task.diffViewProvider.editType = fileExists ? "modify" : "create" } - // Create parent directories early for new files to prevent ENOENT errors - // in subsequent operations (e.g., diffViewProvider.open) - if (!fileExists) { - await createDirectoriesForFile(absolutePath) - } - const isWriteProtected = task.rooProtectedController?.isWriteProtected(relPath) || false const fullPath = absolutePath const isOutsideWorkspace = isPathOutsideWorkspace(fullPath) @@ -248,7 +256,14 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { const partialMessage = JSON.stringify(sharedMessageProps) await task.ask("tool", partialMessage, block.partial).catch(() => {}) - if (newContent) { + // Only create the file and start streaming when the path has stabilized + // (i.e., the same path was seen in consecutive partial blocks) + if (newContent && pathHasStabilized) { + // Create parent directories only when we're sure about the path + if (!fileExists) { + await createDirectoriesForFile(absolutePath) + } + if (!task.diffViewProvider.isEditing) { await task.diffViewProvider.open(relPath) } @@ -259,6 +274,13 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> { ) } } + + /** + * Reset state when the tool finishes (called from execute or on error) + */ + resetPartialState(): void { + this.lastSeenPartialPath = undefined + } } export const writeToFileTool = new WriteToFileTool() diff --git a/src/core/tools/__tests__/writeToFileTool.spec.ts b/src/core/tools/__tests__/writeToFileTool.spec.ts index 3970a99f066..74caa2b8c93 100644 --- a/src/core/tools/__tests__/writeToFileTool.spec.ts +++ b/src/core/tools/__tests__/writeToFileTool.spec.ts @@ -278,10 +278,14 @@ describe("writeToFileTool", () => { ) it.skipIf(process.platform === "win32")( - "creates parent directories early when file does not exist (partial)", + "creates parent directories when path has stabilized (partial)", async () => { + // First call - path not yet stabilized await executeWriteFileTool({}, { fileExists: false, isPartial: true }) + expect(mockedCreateDirectoriesForFile).not.toHaveBeenCalled() + // Second call with same path - path is now stabilized + await executeWriteFileTool({}, { fileExists: false, isPartial: true }) expect(mockedCreateDirectoriesForFile).toHaveBeenCalledWith(absoluteFilePath) }, ) @@ -394,10 +398,14 @@ describe("writeToFileTool", () => { expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() }) - it("streams content updates during partial execution", async () => { + it("streams content updates during partial execution after path stabilizes", async () => { + // First call - path not yet stabilized, no file operations await executeWriteFileTool({}, { isPartial: true }) - expect(mockCline.ask).toHaveBeenCalled() + expect(mockCline.diffViewProvider.open).not.toHaveBeenCalled() + + // Second call with same path - path is now stabilized, file operations proceed + await executeWriteFileTool({}, { isPartial: true }) expect(mockCline.diffViewProvider.open).toHaveBeenCalledWith(testFilePath) expect(mockCline.diffViewProvider.update).toHaveBeenCalledWith(testContent, false) }) @@ -442,11 +450,15 @@ describe("writeToFileTool", () => { expect(mockCline.diffViewProvider.reset).toHaveBeenCalled() }) - it("handles partial streaming errors", async () => { + it("handles partial streaming errors after path stabilizes", async () => { mockCline.diffViewProvider.open.mockRejectedValue(new Error("Open failed")) + // First call - path not yet stabilized, no error yet await executeWriteFileTool({}, { isPartial: true }) + expect(mockHandleError).not.toHaveBeenCalled() + // Second call with same path - path is now stabilized, error occurs + await executeWriteFileTool({}, { isPartial: true }) expect(mockHandleError).toHaveBeenCalledWith("handling partial write_to_file", expect.any(Error)) }) }) diff --git a/src/integrations/editor/DiffViewProvider.ts b/src/integrations/editor/DiffViewProvider.ts index 81e4982124a..ee49fd40ae4 100644 --- a/src/integrations/editor/DiffViewProvider.ts +++ b/src/integrations/editor/DiffViewProvider.ts @@ -46,6 +46,15 @@ export class DiffViewProvider { this.taskRef = new WeakRef(task) } + /** + * Update the relative path without re-opening the diff view. + * This is needed when the diff view was opened during streaming with a partial path + * value, and we now have the complete path to use for the tool result. + */ + setRelPath(relPath: string): void { + this.relPath = relPath + } + async open(relPath: string): Promise { this.relPath = relPath const fileExists = this.editType === "modify"