From fd8176c4d787ae780aca7d1c587745de25410b45 Mon Sep 17 00:00:00 2001 From: Dhruv Sharma <237243628+Dhruv-Sharma29@users.noreply.github.com> Date: Fri, 6 Nov 2026 14:08:44 +0530 Subject: [PATCH] feat(work): add undo for accepted AI Work results Work Mode let users accept AI results but gave no way to back them out from the Work panel. Project Work was worse: it rewrites several tabs through React state, so editor Undo could not revert it at all. - Active document: an "Undo AI Changes" button restores the pre-Work text as one undoable editor edit. It refuses if the document was edited since, pointing the user to editor Undo instead. - Open project: accepting a proposal snapshots every affected tab. "Undo accepted changes" restores them and closes tabs the result opened. Tabs edited or saved after accepting are left alone and listed. Closes #462 --- frontend/src/app/page.tsx | 4 +- frontend/src/components/AIPanel.tsx | 13 ++++ frontend/src/components/ProjectWorkReview.tsx | 1 + frontend/src/hooks/useAIActions.ts | 25 ++++++- frontend/src/hooks/useEditor.ts | 15 ++++ frontend/src/hooks/useProjectWork.ts | 18 ++++- frontend/src/lib/project-work-undo.mjs | 75 +++++++++++++++++++ frontend/tests/project-work-undo.test.mjs | 44 +++++++++++ 8 files changed, 190 insertions(+), 5 deletions(-) create mode 100644 frontend/src/lib/project-work-undo.mjs create mode 100644 frontend/tests/project-work-undo.test.mjs diff --git a/frontend/src/app/page.tsx b/frontend/src/app/page.tsx index ed846b0e..12abcefc 100644 --- a/frontend/src/app/page.tsx +++ b/frontend/src/app/page.tsx @@ -47,7 +47,7 @@ export default function HomePage() { const { isDesktopRuntime, backendStatus, backendMessage, showPackagedBackendStatus } = useTauriBackend(editor); const projectRoot = project.projectRoot; const fileIO = useFileIO({ editor, isDesktopRuntime, backendStatus, monacoRef }); - const { selectedText, syncSelectedText, applyAction: handleAIApplyAction, executePrompt: executeAIPrompt, executeWork, isWorking } = useAIActions({ + const { selectedText, syncSelectedText, applyAction: handleAIApplyAction, executePrompt: executeAIPrompt, executeWork, isWorking, undoWork, canUndoWork } = useAIActions({ documentId: editor.activeTab.id, documentText: editor.activeTab.content, editorRef: monacoRef, onDocumentChange: editor.handleContentChange, }); const selectPanel = useCallback((panel: NonNullable) => { @@ -191,6 +191,8 @@ export default function HomePage() { initialTab={aiPanelInitialTab} executeWork={executeWork} isWorking={isWorking} + undoWork={undoWork} + canUndoWork={canUndoWork} /> ), [PANELS.CITATION]: ( diff --git a/frontend/src/components/AIPanel.tsx b/frontend/src/components/AIPanel.tsx index 8134c904..2b24c434 100644 --- a/frontend/src/components/AIPanel.tsx +++ b/frontend/src/components/AIPanel.tsx @@ -22,6 +22,9 @@ type Props = { initialTab?: Tab; executeWork?: (instruction: string) => Promise; isWorking?: boolean; + /** Revert the last AI Work result applied to the active document. */ + undoWork?: () => void; + canUndoWork?: boolean; }; const LANGUAGES = [ @@ -172,6 +175,8 @@ export default function AIPanel({ initialTab, executeWork, isWorking, + undoWork, + canUndoWork, }: Props) { const currentDocumentRef = useRef({ documentId, documentText }); currentDocumentRef.current = { documentId, documentText }; @@ -929,6 +934,14 @@ export default function AIPanel({ > {isWorking || projectWork.busy ? "Working..." : workScope === "project" ? "Plan Project Changes" : "Apply Changes"} + {workScope === "document" && canUndoWork && undoWork && ( + + )} ) : tab === "translate" ? (
diff --git a/frontend/src/components/ProjectWorkReview.tsx b/frontend/src/components/ProjectWorkReview.tsx index df964ecd..ac16f4db 100644 --- a/frontend/src/components/ProjectWorkReview.tsx +++ b/frontend/src/components/ProjectWorkReview.tsx @@ -10,6 +10,7 @@ export default function ProjectWorkReview({ projectRoot, work }: Props) {

Project: {projectRoot}

The AI discovers paths first and reads only files it requests. It cannot write to disk.

{work.message &&

{work.message}

} + {work.canUndo && } {work.task &&

{work.task.summary}

{work.task.file_count} files · {work.task.steps} steps · {work.task.inspected.length} inspected

diff --git a/frontend/src/hooks/useAIActions.ts b/frontend/src/hooks/useAIActions.ts index 82567331..5080f17e 100644 --- a/frontend/src/hooks/useAIActions.ts +++ b/frontend/src/hooks/useAIActions.ts @@ -3,7 +3,7 @@ import { useCallback, useLayoutEffect, useRef, useState, type RefObject } from "react"; import type { editor as MonacoEditor } from "monaco-editor"; import { AI } from "@/lib/api"; -import { applyWorkResultIfFresh } from "@/lib/work-result-guard.mjs"; +import { applyWorkResultIfFresh, isWorkResultFresh } from "@/lib/work-result-guard.mjs"; import { applyWorkEdit } from "@/lib/apply-work-edit.mjs"; type UseAIActionsOptions = { @@ -21,6 +21,8 @@ export function useAIActions({ }: UseAIActionsOptions) { const [selectedText, setSelectedText] = useState(""); const [isWorking, setIsWorking] = useState(false); + // The last AI Work result applied here, so the Work panel can offer to undo it. + const [lastWork, setLastWork] = useState<{ documentId: string; before: string; after: string } | null>(null); const currentDocumentRef = useRef({ documentId, content: documentText }); useLayoutEffect(() => { @@ -121,7 +123,9 @@ export function useAIActions({ // preserves the old behaviour when no editor/model is available. (content) => applyWorkEdit(editorRef.current, onDocumentChange, content) ); - if (!applied) { + if (applied) { + setLastWork({ documentId, before: requestSnapshot.content, after: data.modified_content }); + } else { alert( "The document changed while AI Work was running. No AI changes were applied. Run the task again on the current document." ); @@ -134,6 +138,21 @@ export function useAIActions({ } }, [documentId, documentText, editorRef, onDocumentChange]); + const undoWork = useCallback(() => { + if (!lastWork) return; + const stillAIResult = isWorkResultFresh( + { documentId: lastWork.documentId, content: lastWork.after }, + currentDocumentRef.current + ); + if (!stillAIResult) { + alert("The document was edited after AI Work was applied. Use Undo (Cmd/Ctrl+Z) in the editor to step back through those edits."); + return; + } + // Restored as its own undoable edit, so editor Undo (Cmd/Ctrl+Z) can bring the AI result back. + applyWorkEdit(editorRef.current, onDocumentChange, lastWork.before); + setLastWork(null); + }, [editorRef, lastWork, onDocumentChange]); + return { selectedText, syncSelectedText, @@ -141,5 +160,7 @@ export function useAIActions({ executePrompt, executeWork, isWorking, + undoWork, + canUndoWork: !isWorking && lastWork?.documentId === documentId, }; } diff --git a/frontend/src/hooks/useEditor.ts b/frontend/src/hooks/useEditor.ts index 796fb798..7bc5a6d6 100644 --- a/frontend/src/hooks/useEditor.ts +++ b/frontend/src/hooks/useEditor.ts @@ -15,6 +15,7 @@ import { resolveTabClose } from "@/lib/tabLifecycle.mjs"; import { settleSavedTab } from "@/lib/save-settle.mjs"; import { parentDirOf, resolvePreviewBaseDir } from "@/lib/preview-base-dir.mjs"; import { createLatestRequestGuard } from "@/lib/latest-request.mjs"; +import { revertProjectWork, type ProjectWorkUndoEntry } from "@/lib/project-work-undo.mjs"; export type Tab = { id: string; @@ -391,6 +392,19 @@ export function useEditor() { if (activeChange) refreshPreview(activeChange.content, activeChange.file_path); }, [activeTab.filePath, refreshPreview]); + /** Revert an accepted Project Work result; returns the paths left untouched. */ + const revertProjectChanges = useCallback((record: ProjectWorkUndoEntry[]) => { + const { tabs: reverted, skipped } = revertProjectWork(tabs, record); + const remaining = reverted.length ? reverted : [makeTab(nextTabId())]; + const active = remaining.find((tab) => tab.id === activeTabId) ?? remaining[remaining.length - 1]; + setTabs(remaining); + setActiveTabId(active.id); + if (active.id !== activeTabId || active.content !== activeTab.content) { + refreshPreview(active.content, active.previewBaseDir ?? active.filePath ?? undefined); + } + return skipped; + }, [tabs, activeTabId, activeTab.content, refreshPreview]); + // ── recent files ─────────────────────────────────────────────────────────── const loadRecentFiles = useCallback(async () => { const { entries } = await Files.getRecent(); @@ -435,6 +449,7 @@ export function useEditor() { closeTab, closeAllTabs, applyProjectChanges, + revertProjectChanges, loadRecentFiles, refreshPreview, exportAs, diff --git a/frontend/src/hooks/useProjectWork.ts b/frontend/src/hooks/useProjectWork.ts index d8e1b669..2a4ad88d 100644 --- a/frontend/src/hooks/useProjectWork.ts +++ b/frontend/src/hooks/useProjectWork.ts @@ -8,6 +8,7 @@ import { projectTabsAreFresh, taskMatchesProject, } from "@/lib/project-work-guard.mjs"; +import { captureProjectWorkUndo, type ProjectWorkUndoEntry } from "@/lib/project-work-undo.mjs"; type Editor = ReturnType; @@ -20,6 +21,7 @@ export function useProjectWork(editor: Editor, projectRoot: string | null) { const requestSnapshot = useRef(editor.tabs); const projectRootRef = useRef(projectRoot); const runGuard = useRef(createProjectWorkRunGuard()); + const [undoRecord, setUndoRecord] = useState(null); useLayoutEffect(() => { currentTabs.current = editor.tabs; }, [editor.tabs]); @@ -35,6 +37,7 @@ export function useProjectWork(editor: Editor, projectRoot: string | null) { if (projectRootRef.current === projectRoot) return; projectRootRef.current = projectRoot; discard(); + setUndoRecord(null); }, [discard, projectRoot]); const run = useCallback(async (instruction: string) => { @@ -84,12 +87,14 @@ export function useProjectWork(editor: Editor, projectRoot: string | null) { if (!projectTabsAreFresh(requestSnapshot.current, currentTabs.current)) { throw new Error("Open documents changed while Project Work was running. Run the task again."); } + const record = captureProjectWorkUndo(currentTabs.current, task.changes); editor.applyProjectChanges(task.changes); + setUndoRecord(record); await AI.discardProjectWork(task.id); if (!isCurrentRun()) return; taskId.current = null; setTask(null); - setMessage("Changes were staged in editor tabs. Review and save each file."); + setMessage("Changes were staged in editor tabs. Review and save each file, or undo them."); } catch (error) { if (isCurrentRun()) setMessage(error instanceof Error ? error.message : String(error)); } finally { @@ -97,5 +102,14 @@ export function useProjectWork(editor: Editor, projectRoot: string | null) { } }, [busy, editor, task]); - return { task, busy, message, run, apply, discard }; + const undo = useCallback(() => { + if (!undoRecord || busy) return; + const skipped = editor.revertProjectChanges(undoRecord); + setUndoRecord(null); + setMessage(skipped.length + ? `Undid Project Work. These files were edited or saved afterwards and were left as they are:\n${skipped.join("\n")}` + : "Undid Project Work. The affected tabs are back to their previous content."); + }, [busy, editor, undoRecord]); + + return { task, busy, message, run, apply, discard, undo, canUndo: !!undoRecord && !busy }; } diff --git a/frontend/src/lib/project-work-undo.mjs b/frontend/src/lib/project-work-undo.mjs new file mode 100644 index 00000000..9cecb515 --- /dev/null +++ b/frontend/src/lib/project-work-undo.mjs @@ -0,0 +1,75 @@ +/** + * Undo support for accepted Project Work results. + * + * Accepting Project Work rewrites the content of several editor tabs (and may + * open new ones). Those writes go through React state, not one Monaco model, so + * the editor's own Undo cannot step back over them. Instead we snapshot every + * tab the result is about to touch, and revert from that snapshot on request. + * + * @typedef {{ id: string, filePath: string | null, content: string, dirty: boolean }} TabLike + * @typedef {{ file_path: string, content: string }} ChangeLike + * @typedef {{ + * filePath: string, + * applied: string, + * before: { content: string, dirty: boolean } | null, + * }} ProjectWorkUndoEntry + */ + +/** + * Record the state of every tab a Project Work result is about to change. + * `before` is null for files the result opens as new tabs. + * + * @param {TabLike[]} tabs + * @param {ChangeLike[]} changes + * @returns {ProjectWorkUndoEntry[]} + */ +export function captureProjectWorkUndo(tabs, changes) { + return changes.map((change) => { + const tab = tabs.find((candidate) => candidate.filePath === change.file_path); + return { + filePath: change.file_path, + applied: change.content, + before: tab ? { content: tab.content, dirty: tab.dirty } : null, + }; + }); +} + +/** + * Restore the tabs recorded by `captureProjectWorkUndo`. + * + * A tab is only reverted while it still holds exactly the AI content, so undo + * never throws away edits the user made afterwards. Such tabs, and new tabs + * that were already saved to disk, are left alone and reported in `skipped`. + * A reverted tab whose AI content was saved stays dirty, because the file on + * disk no longer matches the restored text. + * + * @template {TabLike} T + * @param {T[]} tabs + * @param {ProjectWorkUndoEntry[]} record + * @returns {{ tabs: T[], reverted: string[], skipped: string[] }} + */ +export function revertProjectWork(tabs, record) { + let next = tabs; + const reverted = []; + const skipped = []; + + for (const entry of record) { + const tab = next.find((candidate) => candidate.filePath === entry.filePath); + if (!tab) continue; // Closed since apply: nothing left to revert. + if (tab.content !== entry.applied || (!entry.before && !tab.dirty)) { + skipped.push(entry.filePath); + continue; + } + const before = entry.before; + next = before + ? next.map((candidate) => + candidate === tab + ? { ...candidate, content: before.content, dirty: tab.dirty ? before.dirty : true } + : candidate + ) + : next.filter((candidate) => candidate !== tab); + reverted.push(entry.filePath); + } + + return { tabs: next, reverted, skipped }; +} diff --git a/frontend/tests/project-work-undo.test.mjs b/frontend/tests/project-work-undo.test.mjs new file mode 100644 index 00000000..2ebebce9 --- /dev/null +++ b/frontend/tests/project-work-undo.test.mjs @@ -0,0 +1,44 @@ +import { test } from "node:test"; +import assert from "node:assert/strict"; +import { captureProjectWorkUndo, revertProjectWork } from "../src/lib/project-work-undo.mjs"; + +const tab = (id, filePath, content, dirty = false) => ({ id, filePath, content, dirty }); + +test("undo restores edited tabs and closes tabs that Project Work opened", () => { + const before = [tab("a", "/p/a.md", "A")]; + const changes = [ + { file_path: "/p/a.md", content: "A by AI" }, + { file_path: "/p/new.md", content: "New by AI" }, + ]; + const record = captureProjectWorkUndo(before, changes); + const afterApply = [tab("a", "/p/a.md", "A by AI", true), tab("n", "/p/new.md", "New by AI", true)]; + + const { tabs, reverted, skipped } = revertProjectWork(afterApply, record); + + assert.deepEqual(tabs, [tab("a", "/p/a.md", "A", false)]); + assert.deepEqual(reverted, ["/p/a.md", "/p/new.md"]); + assert.deepEqual(skipped, []); +}); + +test("undo never discards edits the user made after accepting", () => { + const record = captureProjectWorkUndo([tab("a", "/p/a.md", "A")], [{ file_path: "/p/a.md", content: "AI" }]); + const edited = [tab("a", "/p/a.md", "AI plus user typing", true)]; + + const { tabs, skipped } = revertProjectWork(edited, record); + + assert.equal(tabs, edited); + assert.deepEqual(skipped, ["/p/a.md"]); +}); + +test("a saved AI result stays dirty after undo, and a saved new file is kept", () => { + const record = captureProjectWorkUndo( + [tab("a", "/p/a.md", "A")], + [{ file_path: "/p/a.md", content: "AI" }, { file_path: "/p/new.md", content: "New" }] + ); + const saved = [tab("a", "/p/a.md", "AI", false), tab("n", "/p/new.md", "New", false)]; + + const { tabs, skipped } = revertProjectWork(saved, record); + + assert.deepEqual(tabs, [tab("a", "/p/a.md", "A", true), tab("n", "/p/new.md", "New", false)]); + assert.deepEqual(skipped, ["/p/new.md"]); +});