Skip to content

Commit 16c4d48

Browse files
easonliang28easonliang28
committed
fix(write-to-file): guard diff reset cleanup
1 parent 0966556 commit 16c4d48

2 files changed

Lines changed: 105 additions & 77 deletions

File tree

‎src/core/tools/WriteToFileTool.ts‎

Lines changed: 89 additions & 73 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,12 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> {
5757
this.partialStreamFailuresByTaskId.delete(key)
5858
}
5959

60+
private async resetDiffViewAfterWrite(task: Task): Promise<void> {
61+
await task.diffViewProvider.reset().catch((resetError) => {
62+
console.error("Error resetting write_to_file diff view:", resetError)
63+
})
64+
}
65+
6066
override resetPartialState(): void {
6167
super.resetPartialState()
6268
this.partialStreamFailuresByTaskId.clear()
@@ -147,92 +153,104 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> {
147153
EXPERIMENT_IDS.PREVENT_FOCUS_DISRUPTION,
148154
)
149155

150-
if (isPreventFocusDisruptionEnabled) {
151-
task.diffViewProvider.editType = fileExists ? "modify" : "create"
152-
if (fileExists) {
153-
const absolutePath = path.resolve(task.cwd, relPath)
154-
task.diffViewProvider.originalContent = await fs.readFile(absolutePath, "utf-8")
156+
try {
157+
if (isPreventFocusDisruptionEnabled) {
158+
task.diffViewProvider.editType = fileExists ? "modify" : "create"
159+
if (fileExists) {
160+
const absolutePath = path.resolve(task.cwd, relPath)
161+
task.diffViewProvider.originalContent = await fs.readFile(absolutePath, "utf-8")
162+
} else {
163+
task.diffViewProvider.originalContent = ""
164+
}
165+
166+
let unified = fileExists
167+
? formatResponse.createPrettyPatch(relPath, task.diffViewProvider.originalContent, newContent)
168+
: convertNewFileToUnifiedDiff(newContent, relPath)
169+
unified = sanitizeUnifiedDiff(unified)
170+
const completeMessage = JSON.stringify({
171+
...sharedMessageProps,
172+
content: unified,
173+
diffStats: computeDiffStats(unified) || undefined,
174+
} satisfies ClineSayTool)
175+
176+
const didApprove = await askApproval("tool", completeMessage, undefined, isWriteProtected)
177+
178+
if (!didApprove) {
179+
return
180+
}
181+
182+
await task.diffViewProvider.saveDirectly(
183+
relPath,
184+
newContent,
185+
false,
186+
diagnosticsEnabled,
187+
writeDelayMs,
188+
)
155189
} else {
156-
task.diffViewProvider.originalContent = ""
190+
if (!task.diffViewProvider.isEditing) {
191+
const partialMessage = JSON.stringify(sharedMessageProps)
192+
await task.ask("tool", partialMessage, true).catch(() => {})
193+
await task.diffViewProvider.open(relPath)
194+
}
195+
196+
await task.diffViewProvider.update(
197+
everyLineHasLineNumbers(newContent) ? stripLineNumbers(newContent) : newContent,
198+
true,
199+
)
200+
201+
await delay(300)
202+
task.diffViewProvider.scrollToFirstDiff()
203+
204+
let unified = fileExists
205+
? formatResponse.createPrettyPatch(relPath, task.diffViewProvider.originalContent, newContent)
206+
: convertNewFileToUnifiedDiff(newContent, relPath)
207+
unified = sanitizeUnifiedDiff(unified)
208+
const completeMessage = JSON.stringify({
209+
...sharedMessageProps,
210+
content: unified,
211+
diffStats: computeDiffStats(unified) || undefined,
212+
} satisfies ClineSayTool)
213+
214+
const didApprove = await askApproval("tool", completeMessage, undefined, isWriteProtected)
215+
216+
if (!didApprove) {
217+
await task.diffViewProvider.revertChanges()
218+
return
219+
}
220+
221+
await task.diffViewProvider.saveChanges(diagnosticsEnabled, writeDelayMs)
157222
}
158223

159-
let unified = fileExists
160-
? formatResponse.createPrettyPatch(relPath, task.diffViewProvider.originalContent, newContent)
161-
: convertNewFileToUnifiedDiff(newContent, relPath)
162-
unified = sanitizeUnifiedDiff(unified)
163-
const completeMessage = JSON.stringify({
164-
...sharedMessageProps,
165-
content: unified,
166-
diffStats: computeDiffStats(unified) || undefined,
167-
} satisfies ClineSayTool)
168-
169-
const didApprove = await askApproval("tool", completeMessage, undefined, isWriteProtected)
170-
171-
if (!didApprove) {
172-
return
224+
if (relPath) {
225+
await task.fileContextTracker.trackFileContext(relPath, "roo_edited" as RecordSource)
173226
}
174227

175-
await task.diffViewProvider.saveDirectly(relPath, newContent, false, diagnosticsEnabled, writeDelayMs)
176-
} else {
177-
if (!task.diffViewProvider.isEditing) {
178-
const partialMessage = JSON.stringify(sharedMessageProps)
179-
await task.ask("tool", partialMessage, true).catch(() => {})
180-
await task.diffViewProvider.open(relPath)
181-
}
228+
task.didEditFile = true
182229

183-
await task.diffViewProvider.update(
184-
everyLineHasLineNumbers(newContent) ? stripLineNumbers(newContent) : newContent,
185-
true,
186-
)
230+
const message = await task.diffViewProvider.pushToolWriteResult(task, task.cwd, !fileExists)
187231

188-
await delay(300)
189-
task.diffViewProvider.scrollToFirstDiff()
232+
pushToolResult(message)
190233

191-
let unified = fileExists
192-
? formatResponse.createPrettyPatch(relPath, task.diffViewProvider.originalContent, newContent)
193-
: convertNewFileToUnifiedDiff(newContent, relPath)
194-
unified = sanitizeUnifiedDiff(unified)
195-
const completeMessage = JSON.stringify({
196-
...sharedMessageProps,
197-
content: unified,
198-
diffStats: computeDiffStats(unified) || undefined,
199-
} satisfies ClineSayTool)
234+
await this.resetDiffViewAfterWrite(task)
200235

201-
const didApprove = await askApproval("tool", completeMessage, undefined, isWriteProtected)
236+
task.processQueuedMessages()
202237

203-
if (!didApprove) {
204-
await task.diffViewProvider.revertChanges()
205-
return
206-
}
207-
208-
await task.diffViewProvider.saveChanges(diagnosticsEnabled, writeDelayMs)
238+
return
239+
} finally {
240+
this.resetTaskPartialState(task)
209241
}
210-
211-
if (relPath) {
212-
await task.fileContextTracker.trackFileContext(relPath, "roo_edited" as RecordSource)
213-
}
214-
215-
task.didEditFile = true
216-
217-
const message = await task.diffViewProvider.pushToolWriteResult(task, task.cwd, !fileExists)
218-
219-
pushToolResult(message)
220-
221-
await task.diffViewProvider.reset()
222-
this.resetTaskPartialState(task)
223-
224-
task.processQueuedMessages()
225-
226-
return
227242
} catch (error) {
228243
// Finalize any open partial tool message so the UI spinner doesn't get stuck.
229244
// The partial ask fired during streaming (handlePartial) or early in execute sets
230245
// partial: true on the webview message; without this, the spinner persists even
231246
// after the error bubble appears.
232-
await task.finalizePartialToolAsk()
233-
await handleError("writing file", error as Error)
234-
await task.diffViewProvider.reset()
235-
this.resetTaskPartialState(task)
247+
try {
248+
await task.finalizePartialToolAsk()
249+
await handleError("writing file", error as Error)
250+
await this.resetDiffViewAfterWrite(task)
251+
} finally {
252+
this.resetTaskPartialState(task)
253+
}
236254
return
237255
}
238256
}
@@ -315,9 +333,7 @@ export class WriteToFileTool extends BaseTool<"write_to_file"> {
315333
// partial tool message each time.
316334
this.partialStreamFailuresByTaskId.add(partialStreamFailureKey)
317335
await task.finalizePartialToolAsk(partialMessage)
318-
await task.diffViewProvider.reset().catch((resetError) => {
319-
console.error("Error resetting write_to_file diff view after partial failure:", resetError)
320-
})
336+
await this.resetDiffViewAfterWrite(task)
321337
}
322338
}
323339
}

‎src/core/tools/__tests__/writeToFileTool.spec.ts‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,7 @@ describe("writeToFileTool", () => {
191191
mockCline.finalizePartialToolAsk = vi.fn().mockResolvedValue(undefined)
192192
mockCline.recordToolError = vi.fn()
193193
mockCline.sayAndCreateMissingParamError = vi.fn().mockResolvedValue("Missing param error")
194+
mockCline.processQueuedMessages = vi.fn()
194195

195196
mockAskApproval = vi.fn().mockResolvedValue(true)
196197
mockHandleError = vi.fn().mockResolvedValue(undefined)
@@ -396,6 +397,20 @@ describe("writeToFileTool", () => {
396397
// Should process normally without issues
397398
expect(mockCline.consecutiveMistakeCount).toBe(0)
398399
})
400+
401+
it("does not report a successful write as failed when final diff reset rejects", async () => {
402+
const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {})
403+
mockCline.diffViewProvider.reset.mockRejectedValue(new Error("reset failed"))
404+
405+
await executeWriteFileTool({}, { fileExists: false })
406+
407+
expect(mockHandleError).not.toHaveBeenCalled()
408+
expect(mockPushToolResult).toHaveBeenCalledWith("Tool result message")
409+
expect(mockCline.didEditFile).toBe(true)
410+
expect(consoleErrorSpy).toHaveBeenCalledWith("Error resetting write_to_file diff view:", expect.any(Error))
411+
412+
consoleErrorSpy.mockRestore()
413+
})
399414
})
400415

401416
describe("partial block handling", () => {
@@ -632,10 +647,7 @@ describe("writeToFileTool", () => {
632647
expect(mockCline.finalizePartialToolAsk).toHaveBeenCalled()
633648
expect(mockCline.diffViewProvider.reset).toHaveBeenCalled()
634649
expect(mockHandleError).not.toHaveBeenCalled()
635-
expect(consoleErrorSpy).toHaveBeenCalledWith(
636-
"Error resetting write_to_file diff view after partial failure:",
637-
expect.any(Error),
638-
)
650+
expect(consoleErrorSpy).toHaveBeenCalledWith("Error resetting write_to_file diff view:", expect.any(Error))
639651

640652
consoleErrorSpy.mockRestore()
641653
})

0 commit comments

Comments
 (0)