fix: prevent write_to_file from creating files at truncated paths during streaming

During streaming, the partial-json library may return truncated string values when
chunk boundaries fall mid-value (e.g., 'sr' instead of 'src/core/prompts/sections/skills.ts').

This fix:
1. Waits for the path to stabilize (same value in consecutive partial blocks) before creating files
2. Updates diffViewProvider.relPath in execute() when isEditing is already true
3. Re-reads fresh block data for non-partial tool_use blocks

Fixes ROO-311
This commit is contained in:
daniel-lxs 2025-12-30 23:24:13 -05:00
parent c37aa02b21
commit 93f4fed168
No known key found for this signature in database
GPG key ID: 21C74479048B3AA6
5 changed files with 310 additions and 11 deletions

View file

@ -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")
})
})

View file

@ -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)

View file

@ -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<void> {
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()

View file

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

View file

@ -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<void> {
this.relPath = relPath
const fileExists = this.editType === "modify"