From 9789925e2d76ea75a37b62c1bb6709d2558fb2f2 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Thu, 7 Aug 2025 09:37:23 +0000 Subject: [PATCH] feat: add checkpoints before file edits for increased checkpoint frequency - Save checkpoints BEFORE file editing operations (write_to_file, apply_diff, insert_content, search_and_replace) - Maintain existing checkpoints AFTER file edits - Update checkpointSaveAndMark function to handle before/after timing - Add comprehensive tests for new checkpoint behavior This addresses user feedback requesting checkpoints before edits to allow reverting to the state just before making changes. --- .../presentAssistantMessage.ts | 34 ++++-- .../__tests__/checkpoint-timing.spec.ts | 102 ++++++++++++++++-- 2 files changed, 121 insertions(+), 15 deletions(-) diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index 9e3e34ec60..1e8d1a6b22 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -410,9 +410,11 @@ export async function presentAssistantMessage(cline: Task) { switch (block.name) { case "write_to_file": + // Save checkpoint BEFORE file edit + await checkpointSaveAndMark(cline, "before") await writeToFileTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) // Save checkpoint AFTER file edit - await checkpointSaveAndMark(cline) + await checkpointSaveAndMark(cline, "after") break case "update_todo_list": await updateTodoListTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) @@ -430,6 +432,9 @@ export async function presentAssistantMessage(cline: Task) { ) } + // Save checkpoint BEFORE file edit + await checkpointSaveAndMark(cline, "before") + if (isMultiFileApplyDiffEnabled) { await applyDiffTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) } else { @@ -443,18 +448,22 @@ export async function presentAssistantMessage(cline: Task) { ) } // Save checkpoint AFTER file edit - await checkpointSaveAndMark(cline) + await checkpointSaveAndMark(cline, "after") break } case "insert_content": + // Save checkpoint BEFORE file edit + await checkpointSaveAndMark(cline, "before") await insertContentTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) // Save checkpoint AFTER file edit - await checkpointSaveAndMark(cline) + await checkpointSaveAndMark(cline, "after") break case "search_and_replace": + // Save checkpoint BEFORE file edit + await checkpointSaveAndMark(cline, "before") await searchAndReplaceTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) // Save checkpoint AFTER file edit - await checkpointSaveAndMark(cline) + await checkpointSaveAndMark(cline, "after") break case "read_file": await readFileTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) @@ -586,16 +595,25 @@ export async function presentAssistantMessage(cline: Task) { /** * save checkpoint and mark done in the current streaming task. * @param task The Task instance to checkpoint save and mark. + * @param timing Whether this is a "before" or "after" checkpoint for file edits * @returns */ -async function checkpointSaveAndMark(task: Task) { - if (task.currentStreamingDidCheckpoint) { +async function checkpointSaveAndMark(task: Task, timing?: "before" | "after") { + // For "before" checkpoints, always save regardless of currentStreamingDidCheckpoint + // For "after" checkpoints or no timing specified, use the original logic + if (timing !== "before" && task.currentStreamingDidCheckpoint) { return } try { await task.checkpointSave(true) - task.currentStreamingDidCheckpoint = true + // Only mark as done for "after" checkpoints or when no timing is specified + if (timing !== "before") { + task.currentStreamingDidCheckpoint = true + } } catch (error) { - console.error(`[Task#presentAssistantMessage] Error saving checkpoint: ${error.message}`, error) + console.error( + `[Task#presentAssistantMessage] Error saving checkpoint (${timing || "default"}): ${error.message}`, + error, + ) } } diff --git a/src/core/checkpoints/__tests__/checkpoint-timing.spec.ts b/src/core/checkpoints/__tests__/checkpoint-timing.spec.ts index 3198dbfac6..48dc39588d 100644 --- a/src/core/checkpoints/__tests__/checkpoint-timing.spec.ts +++ b/src/core/checkpoints/__tests__/checkpoint-timing.spec.ts @@ -216,8 +216,8 @@ describe("Checkpoint Timing", () => { vi.clearAllMocks() }) - describe("Checkpoint after file edits", () => { - it("should save checkpoint AFTER write_to_file tool execution", async () => { + describe("Checkpoint before and after file edits", () => { + it("should save checkpoint BEFORE and AFTER write_to_file tool execution", async () => { // Setup assistant message content with write_to_file tool mockTask.assistantMessageContent = [ { @@ -241,8 +241,10 @@ describe("Checkpoint Timing", () => { // Execute presentAssistantMessage await presentAssistantMessage(mockTask) - // Verify checkpoint was saved after the tool execution + // Verify checkpoint was saved twice (before and after the tool execution) + expect(mockTask.checkpointSave).toHaveBeenCalledTimes(2) expect(mockTask.checkpointSave).toHaveBeenCalledWith(true) + // Note: currentStreamingDidCheckpoint is only set to true after the "after" checkpoint expect(mockTask.currentStreamingDidCheckpoint).toBe(true) }) @@ -251,7 +253,7 @@ describe("Checkpoint Timing", () => { // through the other file editing tools (write_to_file, insert_content, search_and_replace) // which all follow the same pattern of saving checkpoints after file edits. - it("should save checkpoint AFTER insert_content tool execution", async () => { + it("should save checkpoint BEFORE and AFTER insert_content tool execution", async () => { // Setup assistant message content with insert_content tool mockTask.assistantMessageContent = [ { @@ -276,12 +278,14 @@ describe("Checkpoint Timing", () => { // Execute presentAssistantMessage await presentAssistantMessage(mockTask) - // Verify checkpoint was saved after the tool execution + // Verify checkpoint was saved twice (before and after the tool execution) + expect(mockTask.checkpointSave).toHaveBeenCalledTimes(2) expect(mockTask.checkpointSave).toHaveBeenCalledWith(true) + // Note: currentStreamingDidCheckpoint is only set to true after the "after" checkpoint expect(mockTask.currentStreamingDidCheckpoint).toBe(true) }) - it("should save checkpoint AFTER search_and_replace tool execution", async () => { + it("should save checkpoint BEFORE and AFTER search_and_replace tool execution", async () => { // Setup assistant message content with search_and_replace tool mockTask.assistantMessageContent = [ { @@ -306,10 +310,94 @@ describe("Checkpoint Timing", () => { // Execute presentAssistantMessage await presentAssistantMessage(mockTask) - // Verify checkpoint was saved after the tool execution + // Verify checkpoint was saved twice (before and after the tool execution) + expect(mockTask.checkpointSave).toHaveBeenCalledTimes(2) expect(mockTask.checkpointSave).toHaveBeenCalledWith(true) + // Note: currentStreamingDidCheckpoint is only set to true after the "after" checkpoint expect(mockTask.currentStreamingDidCheckpoint).toBe(true) }) + + it("should handle checkpoint errors gracefully for file edit tools", async () => { + // Setup assistant message content with write_to_file tool + mockTask.assistantMessageContent = [ + { + type: "tool_use", + name: "write_to_file", + params: { + path: "test.txt", + content: "test content", + }, + partial: false, + }, + ] + + // Mock the write_to_file tool execution + const writeToFileModule = await import("../../tools/writeToFileTool") + vi.spyOn(writeToFileModule, "writeToFileTool").mockImplementation(async () => { + // Simulate tool execution + return undefined + }) + + // Mock checkpointSave to fail on first call (before) and succeed on second (after) + mockTask.checkpointSave + .mockRejectedValueOnce(new Error("Checkpoint before failed")) + .mockResolvedValueOnce(undefined) + + // Mock console.error to verify error logging + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + + // Execute presentAssistantMessage + await presentAssistantMessage(mockTask) + + // Verify checkpoint was attempted twice + expect(mockTask.checkpointSave).toHaveBeenCalledTimes(2) + + // Verify error was logged for the "before" checkpoint failure + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining("Error saving checkpoint (before)"), + expect.any(Error), + ) + + // Clean up + consoleErrorSpy.mockRestore() + }) + + it("should not set currentStreamingDidCheckpoint for 'before' checkpoints", async () => { + // Setup assistant message content with write_to_file tool + mockTask.assistantMessageContent = [ + { + type: "tool_use", + name: "write_to_file", + params: { + path: "test.txt", + content: "test content", + }, + partial: false, + }, + ] + + // Mock the write_to_file tool execution to track when it's called + const writeToFileModule = await import("../../tools/writeToFileTool") + let toolExecuted = false + vi.spyOn(writeToFileModule, "writeToFileTool").mockImplementation(async () => { + // At this point, the "before" checkpoint should have been saved + // but currentStreamingDidCheckpoint should still be false + expect(mockTask.checkpointSave).toHaveBeenCalledTimes(1) + expect(mockTask.currentStreamingDidCheckpoint).toBe(false) + toolExecuted = true + return undefined + }) + + // Execute presentAssistantMessage + await presentAssistantMessage(mockTask) + + // Verify tool was executed + expect(toolExecuted).toBe(true) + + // After execution, currentStreamingDidCheckpoint should be true (from "after" checkpoint) + expect(mockTask.currentStreamingDidCheckpoint).toBe(true) + expect(mockTask.checkpointSave).toHaveBeenCalledTimes(2) + }) }) describe("Checkpoint before new prompts", () => {