diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index a8b90728b1..7728645db3 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -35,6 +35,7 @@ import { Task } from "../task/Task" import { codebaseSearchTool } from "../tools/codebaseSearchTool" import { experiments, EXPERIMENT_IDS } from "../../shared/experiments" import { applyDiffToolLegacy } from "../tools/applyDiffTool" +import { updateFCOAfterEdit } from "../../services/file-changes/updateAfterEdit" /** * Processes and presents assistant message content to the user interface. @@ -420,6 +421,7 @@ export async function presentAssistantMessage(cline: Task) { case "write_to_file": await checkpointSaveAndMark(cline) await writeToFileTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) + await updateFCOAfterEdit(cline) break case "update_todo_list": await updateTodoListTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) @@ -440,6 +442,7 @@ export async function presentAssistantMessage(cline: Task) { if (isMultiFileApplyDiffEnabled) { await checkpointSaveAndMark(cline) await applyDiffTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) + await updateFCOAfterEdit(cline) } else { await checkpointSaveAndMark(cline) await applyDiffToolLegacy( @@ -450,16 +453,19 @@ export async function presentAssistantMessage(cline: Task) { pushToolResult, removeClosingTag, ) + await updateFCOAfterEdit(cline) } break } case "insert_content": await checkpointSaveAndMark(cline) await insertContentTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) + await updateFCOAfterEdit(cline) break case "search_and_replace": await checkpointSaveAndMark(cline) await searchAndReplaceTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) + await updateFCOAfterEdit(cline) break case "read_file": // Check if this model should use the simplified single-file read tool diff --git a/src/core/checkpoints/__tests__/helpers.ts b/src/core/checkpoints/__tests__/helpers.ts new file mode 100644 index 0000000000..39e9fb0403 --- /dev/null +++ b/src/core/checkpoints/__tests__/helpers.ts @@ -0,0 +1,53 @@ +import { vitest } from "vitest" + +export const createMockTask = (options: { + taskId: string + hasExistingCheckpoints?: boolean + enableCheckpoints?: boolean + provider?: any +}) => { + const mockTask = { + taskId: options.taskId, + instanceId: "test-instance", + rootTask: undefined as any, + parentTask: undefined as any, + taskNumber: 1, + workspacePath: "/mock/workspace", + enableCheckpoints: options.enableCheckpoints ?? true, + checkpointService: null as any, + checkpointServiceInitializing: false, + ongoingCheckpointSaves: new Map(), + clineMessages: options.hasExistingCheckpoints + ? [{ say: "checkpoint_saved", ts: Date.now(), text: "existing-checkpoint-hash" }] + : [], + providerRef: { + deref: () => options.provider || createMockProvider(), + }, + fileContextTracker: {}, + todoList: undefined, + } + + return mockTask +} + +export const createMockProvider = () => ({ + getFileChangeManager: vitest.fn(), + ensureFileChangeManager: vitest.fn(), + log: vitest.fn(), + postMessageToWebview: vitest.fn(), + getGlobalState: vitest.fn(), +}) + +// Mock checkpoint service for testing +export const createMockCheckpointService = () => ({ + saveCheckpoint: vitest.fn().mockResolvedValue({ + commit: "mock-checkpoint-hash", + message: "Mock checkpoint", + }), + restoreCheckpoint: vitest.fn().mockResolvedValue(true), + getDiff: vitest.fn().mockResolvedValue([]), + getCheckpoints: vitest.fn().mockReturnValue([]), + getCurrentCheckpoint: vitest.fn().mockReturnValue("mock-current-checkpoint"), + initShadowGit: vitest.fn().mockResolvedValue(true), + baseHash: "mock-base-hash", +}) diff --git a/src/core/checkpoints/index.ts b/src/core/checkpoints/index.ts index 9f58e9490e..de6c2d9d19 100644 --- a/src/core/checkpoints/index.ts +++ b/src/core/checkpoints/index.ts @@ -166,11 +166,13 @@ async function checkGitInstallation( // Don't throw - allow checkpoint service to continue initializing } + // Note: No initialization checkpoint needed - first checkpoint before file edit serves as baseline if (isCheckpointNeeded) { - log("[Task#getCheckpointService] no checkpoints found, saving initial checkpoint") - checkpointSave(cline, true) + log( + "[Task#getCheckpointService] no checkpoints found, will create baseline checkpoint before first file edit", + ) } else { - log("[Task#getCheckpointService] existing checkpoints found, skipping initial checkpoint") + log("[Task#getCheckpointService] existing checkpoints found, using existing checkpoint as baseline") } } catch (err) { log("[Task#getCheckpointService] caught error in on('initialize'), disabling checkpoints") diff --git a/src/core/tools/attemptCompletionTool.ts b/src/core/tools/attemptCompletionTool.ts index 5074d7f4e8..91af8d862d 100644 --- a/src/core/tools/attemptCompletionTool.ts +++ b/src/core/tools/attemptCompletionTool.ts @@ -89,6 +89,21 @@ export async function attemptCompletionTool( cline.consecutiveMistakeCount = 0 + // Create final checkpoint to capture the last file edit before completion + if (cline.enableCheckpoints) { + try { + await cline.checkpointSave(true) // Force save to capture any final changes + cline.providerRef + .deref() + ?.log("[attemptCompletionTool] Created final checkpoint before task completion") + } catch (error) { + // Non-critical error, don't fail completion + cline.providerRef + .deref() + ?.log(`[attemptCompletionTool] Failed to create final checkpoint: ${error}`) + } + } + // Command execution is permanently disabled in attempt_completion // Users must use execute_command tool separately before attempt_completion await cline.say("completion_result", result, undefined, false) diff --git a/src/core/tools/writeToFileTool.ts b/src/core/tools/writeToFileTool.ts index 187236ed73..e82eab92bc 100644 --- a/src/core/tools/writeToFileTool.ts +++ b/src/core/tools/writeToFileTool.ts @@ -304,13 +304,6 @@ export async function writeToFileTool( // Get the formatted response message const message = await cline.diffViewProvider.pushToolWriteResult(cline, cline.cwd, !fileExists) - // Track file as edited by LLM for FCO - try { - await cline.fileContextTracker.trackFileContext(relPath.toString(), "roo_edited") - } catch (error) { - console.error("Failed to track file edit in context:", error) - } - pushToolResult(message) await cline.diffViewProvider.reset() diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index ba5941df7b..429cd23478 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -813,7 +813,9 @@ export class ClineProvider return task } - public async createTaskWithHistoryItem(historyItem: HistoryItem & { rootTask?: Task; parentTask?: Task }) { + public async createTaskWithHistoryItem( + historyItem: HistoryItem & { rootTask?: Task; parentTask?: Task; preservedFCOState?: any }, + ) { await this.removeClineFromStack() // If the history item has a saved mode, restore it and its associated API configuration @@ -894,6 +896,37 @@ export class ClineProvider `[subtasks] ${task.parentTask ? "child" : "parent"} task ${task.taskId}.${task.instanceId} instantiated`, ) + // Restore preserved FCO state if provided (from task abort/cancel) + if (historyItem.preservedFCOState) { + try { + const fileChangeManager = await this.ensureFileChangeManager() + if (fileChangeManager && historyItem.preservedFCOState.files) { + // Restore the file changes in FileChangeManager + fileChangeManager.setFiles(historyItem.preservedFCOState.files) + + // Send restored FCO state to webview + const filteredChangeset = await fileChangeManager.getLLMOnlyChanges( + task.taskId, + task.fileContextTracker, + ) + + if (filteredChangeset.files.length > 0) { + this.postMessageToWebview({ + type: "filesChanged", + filesChanged: filteredChangeset, + }) + + this.log( + `[createTaskWithHistoryItem] Restored FCO state with ${filteredChangeset.files.length} LLM-only file changes`, + ) + } + } + } catch (error) { + this.log(`[createTaskWithHistoryItem] Failed to restore FCO state: ${error}`) + // Non-critical error, don't fail task creation + } + } + return task } @@ -1297,6 +1330,18 @@ export class ClineProvider const rootTask = cline.rootTask const parentTask = cline.parentTask + // Preserve FCO state before aborting task to prevent FCO from disappearing + let preservedFCOState: any = undefined + try { + const fileChangeManager = this.getFileChangeManager() + if (fileChangeManager) { + preservedFCOState = fileChangeManager.getChanges() + this.log(`[cancelTask] Preserved FCO state with ${preservedFCOState.files.length} files`) + } + } catch (error) { + this.log(`[cancelTask] Failed to preserve FCO state: ${error}`) + } + cline.abortTask() await pWaitFor( @@ -1323,7 +1368,7 @@ export class ClineProvider } // Clears task again, so we need to abortTask manually above. - await this.createTaskWithHistoryItem({ ...historyItem, rootTask, parentTask }) + await this.createTaskWithHistoryItem({ ...historyItem, rootTask, parentTask, preservedFCOState }) } async updateCustomInstructions(instructions?: string) { diff --git a/src/services/file-changes/updateAfterEdit.ts b/src/services/file-changes/updateAfterEdit.ts new file mode 100644 index 0000000000..168c49afd8 --- /dev/null +++ b/src/services/file-changes/updateAfterEdit.ts @@ -0,0 +1,100 @@ +import { Task } from "../../core/task/Task" +import { getCheckpointService } from "../../core/checkpoints" +import { FileChangeType } from "@roo-code/types" +import { FileChangeManager } from "./FileChangeManager" + +/** + * Updates FCO immediately after a file edit without changing checkpoint timing. + * This provides immediate visibility of changes while preserving rollback safety. + */ +export async function updateFCOAfterEdit(task: Task): Promise { + const provider = task.providerRef.deref() + if (!provider) { + return + } + + try { + const fileChangeManager = provider.getFileChangeManager() + const checkpointService = await getCheckpointService(task) + + if (!fileChangeManager || !checkpointService || !task.taskId || !task.fileContextTracker) { + return + } + + // Get current baseline for FCO + const baseline = fileChangeManager.getChanges().baseCheckpoint + + // Calculate diff from baseline to current working directory state + // We use the checkpointService to get a diff from baseline to HEAD (current state) + try { + const changes = await checkpointService.getDiff({ + from: baseline, + to: "HEAD", // Current working directory state + }) + + if (!changes || changes.length === 0) { + // No changes detected, keep current FCO state + return + } + + // Convert checkpoint service changes to FileChange format + const fileChanges = changes.map((change: any) => { + const type = ( + change.paths.newFile ? "create" : change.paths.deletedFile ? "delete" : "edit" + ) as FileChangeType + + // Calculate line differences + let linesAdded = 0 + let linesRemoved = 0 + + if (type === "create") { + linesAdded = change.content.after ? change.content.after.split("\n").length : 0 + linesRemoved = 0 + } else if (type === "delete") { + linesAdded = 0 + linesRemoved = change.content.before ? change.content.before.split("\n").length : 0 + } else { + const lineDifferences = FileChangeManager.calculateLineDifferences( + change.content.before || "", + change.content.after || "", + ) + linesAdded = lineDifferences.linesAdded + linesRemoved = lineDifferences.linesRemoved + } + + return { + uri: change.paths.relative, + type, + fromCheckpoint: baseline, + toCheckpoint: "HEAD", // This represents current state, not an actual checkpoint + linesAdded, + linesRemoved, + } + }) + + // Update FileChangeManager with the new files + fileChangeManager.setFiles(fileChanges) + + // Get LLM-only changes for the webview (filters out accepted/rejected files) + const filteredChangeset = await fileChangeManager.getLLMOnlyChanges(task.taskId, task.fileContextTracker) + + // Send updated changes to webview only if there are changes to show + if (filteredChangeset.files.length > 0) { + provider.postMessageToWebview({ + type: "filesChanged", + filesChanged: filteredChangeset, + }) + + provider.log( + `[updateFCOAfterEdit] Updated FCO with ${filteredChangeset.files.length} LLM-only file changes`, + ) + } + } catch (diffError) { + // If we can't calculate diff (e.g., baseline is invalid), don't update FCO + provider.log(`[updateFCOAfterEdit] Failed to calculate diff from ${baseline} to HEAD: ${diffError}`) + } + } catch (error) { + // Non-critical error, don't throw - just log and continue + provider?.log(`[updateFCOAfterEdit] Error updating FCO after edit: ${error}`) + } +}