mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
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.
This commit is contained in:
Roo Code
parent
67e6a2221b
commit
9789925e2d
2 changed files with 121 additions and 15 deletions
|
|
@ -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,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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", () => {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue