mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-06 08:18:39 +00:00
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:
parent
c37aa02b21
commit
93f4fed168
5 changed files with 310 additions and 11 deletions
|
|
@ -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")
|
||||
})
|
||||
})
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue