diff --git a/cli/src/app.tsx b/cli/src/app.tsx index f86b12977c..51a607293c 100644 --- a/cli/src/app.tsx +++ b/cli/src/app.tsx @@ -4,6 +4,7 @@ import { useShallow } from 'zustand/react/shallow' import { Chat } from './chat' import { ChatHistoryScreen } from './components/chat-history-screen' +import { UndoHistoryScreen } from './components/undo-history-screen' import { ChatRuntimeProvider } from './contexts/chat-runtime-context' import { FreebuffSupersededScreen } from './components/freebuff-superseded-screen' import { LoginModal } from './components/login-modal' @@ -13,10 +14,13 @@ import { useAuthQuery } from './hooks/use-auth-query' import { useAuthState } from './hooks/use-auth-state' import { useFreebuffSession } from './hooks/use-freebuff-session' import { useTerminalFocus } from './hooks/use-terminal-focus' -import { getProjectRoot, startNewChat } from './project-files' +import { getCurrentChatId, getProjectRoot, startNewChat } from './project-files' import { useChatHistoryStore } from './state/chat-history-store' +import { useUndoHistoryStore } from './state/undo-history-store' +import { redoToRecord, undoToRecord } from './state/undo-store' import { stopActiveRun } from './utils/active-run' import { useChatStore } from './state/chat-store' +import { getSystemMessage } from './utils/message-history' import type { TopBannerType } from './types/store' import { IS_FREEBUFF } from './utils/constants' import { findGitRoot } from './utils/git' @@ -160,6 +164,13 @@ export const App = ({ // Chat history state from store const { showChatHistory, closeChatHistory } = useChatHistoryStore() + const { + showUndoHistory, + closeUndoHistory, + showRedoHistory, + closeRedoHistory, + } = useUndoHistoryStore() + // State to track which chat to resume (set when user selects from history) const [resumeChatId, setResumeChatId] = useState(null) @@ -187,6 +198,68 @@ export const App = ({ setResumeChatId(null) }, [closeChatHistory, resetChatStore]) + // Undo/redo pickers: revert (or restore) the selected record against the + // project, then surface the result as a system message in the current chat. + const handleUndoSelect = useCallback( + async (recordId: string) => { + stopActiveRun('undo-history') + closeUndoHistory() + try { + const message = + (await undoToRecord(getCurrentChatId(), projectRoot, recordId)) ?? + 'Could not undo — the snapshot store is unavailable.' + useChatStore.getState().setMessages((prev) => [ + ...prev, + getSystemMessage(message, { commandResult: 'undo' }), + ]) + } catch (error) { + useChatStore.getState().setMessages((prev) => [ + ...prev, + getSystemMessage('Could not undo the selected change.', { + commandResult: 'undo', + }), + ]) + } + setInputFocused(true) + }, + [closeUndoHistory, projectRoot, setInputFocused], + ) + + const handleRedoSelect = useCallback( + async (recordId: string) => { + stopActiveRun('redo-history') + closeRedoHistory() + try { + const message = + (await redoToRecord(getCurrentChatId(), projectRoot, recordId)) ?? + 'Could not redo — the snapshot store is unavailable.' + useChatStore.getState().setMessages((prev) => [ + ...prev, + getSystemMessage(message, { commandResult: 'redo' }), + ]) + } catch (error) { + useChatStore.getState().setMessages((prev) => [ + ...prev, + getSystemMessage('Could not redo the selected change.', { + commandResult: 'redo', + }), + ]) + } + setInputFocused(true) + }, + [closeRedoHistory, projectRoot, setInputFocused], + ) + + const handleCancelUndoHistory = useCallback(() => { + closeUndoHistory() + setInputFocused(true) + }, [closeUndoHistory, setInputFocused]) + + const handleCancelRedoHistory = useCallback(() => { + closeRedoHistory() + setInputFocused(true) + }, [closeRedoHistory, setInputFocused]) + // Determine effective continueChat values const effectiveContinueChat = continueChat || resumeChatId !== null const effectiveContinueChatId = resumeChatId ?? continueChatId @@ -260,6 +333,12 @@ export const App = ({ onSelectChat={handleResumeChat} onCancelChatHistory={closeChatHistory} onNewChat={handleNewChat} + showUndoHistory={showUndoHistory} + onUndoSelect={handleUndoSelect} + onCancelUndoHistory={handleCancelUndoHistory} + showRedoHistory={showRedoHistory} + onRedoSelect={handleRedoSelect} + onCancelRedoHistory={handleCancelRedoHistory} /> ) } @@ -285,6 +364,12 @@ interface AuthedSurfaceProps { onSelectChat: (chatId: string) => void onCancelChatHistory: () => void onNewChat: () => void + showUndoHistory: boolean + onUndoSelect: (recordId: string) => void + onCancelUndoHistory: () => void + showRedoHistory: boolean + onRedoSelect: (recordId: string) => void + onCancelRedoHistory: () => void } /** @@ -327,6 +412,12 @@ const AuthedSurfaceRoutes = ({ onSelectChat, onCancelChatHistory, onNewChat, + showUndoHistory, + onUndoSelect, + onCancelUndoHistory, + showRedoHistory, + onRedoSelect, + onCancelRedoHistory, session, sessionError, }: AuthedSurfaceProps & { @@ -381,6 +472,26 @@ const AuthedSurfaceRoutes = ({ ) } + if (showUndoHistory) { + return ( + + ) + } + + if (showRedoHistory) { + return ( + + ) + } + return ( { + const { streamingAgents, isChainInProgress } = useChatStore.getState() + if (streamingAgents.size > 0 || isChainInProgress) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage( + '⏳ Please wait for the current task to finish before undoing.', + ), + ]) + return + } + if (!isUndoEnabled()) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage('Undo is disabled in settings.'), + ]) + return + } + const projectRoot = tryGetProjectRoot() + if (!projectRoot) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage('No project is open to undo.'), + ]) + return + } + if (!isUndoAvailable(projectRoot)) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage( + 'Undo requires the project to be a git repository.', + ), + ]) + return + } + if (listUndoEntries(getCurrentChatId()).length === 0) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage('Nothing to undo.'), + ]) + return + } + + // Always open the picker: the user chooses which recorded turn to + // revert (Enter reverts that turn and everything newer). + params.saveToHistory(params.inputValue.trim()) + clearInput(params) + return { openUndoHistory: true } + }, + }), + defineCommand({ + name: 'redo', + aliases: ['r'], + handler: (params) => { + const { streamingAgents, isChainInProgress } = useChatStore.getState() + if (streamingAgents.size > 0 || isChainInProgress) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage( + '⏳ Please wait for the current task to finish before redoing.', + ), + ]) + return + } + if (!isUndoEnabled()) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage('Redo is disabled in settings.'), + ]) + return + } + const projectRoot = tryGetProjectRoot() + if (!projectRoot) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage('No project is open to redo.'), + ]) + return + } + if (listRedoEntries(getCurrentChatId()).length === 0) { + params.setMessages((prev) => [ + ...prev, + getSystemMessage('Nothing to redo.'), + ]) + return + } + + // Always open the picker: the user chooses which undone change to + // restore. + params.saveToHistory(params.inputValue.trim()) + clearInput(params) + return { openRedoHistory: true } + }, + }), defineCommandWithArgs({ name: 'interview', handler: (params, args) => { diff --git a/cli/src/components/blocks/command-result-block.tsx b/cli/src/components/blocks/command-result-block.tsx new file mode 100644 index 0000000000..b533f42680 --- /dev/null +++ b/cli/src/components/blocks/command-result-block.tsx @@ -0,0 +1,159 @@ +import { TextAttributes } from '@opentui/core' +import { memo } from 'react' + +import { useTheme } from '../../hooks/use-theme' + +import type { ReactNode } from 'react' + +interface CommandResultBlockProps { + content: string + commandResult: 'undo' | 'redo' +} + +/** + * Renders /undo and /redo confirmation messages with per-part colors so they + * stand out from plain agent output without relying on non-serializable + * render blocks: + * + * - heading: bold, in the accent color (amber for undo, green for redo) + * - icons (↺ / 🗑): accent color; filenames: default foreground; " (deleted)": muted + * - diff stat bar (e.g. "app.js | 10 ++++++---"): " | " in the accent color, + * and the + / - runs colored git-style (green additions, red deletions) + * - diff summary: "(+)" green, "(-)" red + * + * The content format is produced by `undoToRecord` / `redoToRecord` in + * `cli/src/state/undo-store.ts`. + */ +export const CommandResultBlock = memo( + ({ content, commandResult }: CommandResultBlockProps) => { + const theme = useTheme() + const accent = commandResult === 'undo' ? theme.warning : theme.success + + const renderBar = (bar: string, keyPrefix: string): ReactNode[] => { + // Split the stat bar into contiguous + and - runs so each gets its own + // color (git convention: additions green, deletions red). + const nodes: ReactNode[] = [] + let run = '' + for (let i = 0; i <= bar.length; i++) { + const next = i < bar.length ? bar[i] : '' + // End the run before adding a different char so runs stay exact + // (e.g. "+++++++++----" splits into "+++++++++" and "----"). + if (run && (i === bar.length || next !== run[0])) { + nodes.push( + + {run} + , + ) + run = '' + } + if (i < bar.length) run += bar[i] + } + return nodes + } + + const renderLine = (rawLine: string, idx: number): ReactNode => { + const line = rawLine + if (!line.trim()) return null + + // Heading: **Undid the last change:** / **Redid the last change (2 files):** + const headingMatch = /^\*\*(.+)\*\*$/u.exec(line) + if (headingMatch) { + return ( + + + {headingMatch[1]} + + + ) + } + + // File lines: " ↺ app.js" / " 🗑 notas.md (deleted)" + // The u flag matters: 🗑 is an astral code point (surrogate pair) and + // without it a character class treats it as two lone surrogates. + const fileMatch = /^ {2}([↺🗑]) (.*)$/u.exec(line) + if (fileMatch) { + const rest = fileMatch[2]! + const deletedMatch = /^(.*) \(deleted\)$/u.exec(rest) + return ( + + {fileMatch[1]} + + {` ${deletedMatch ? deletedMatch[1] : rest}`} + + {deletedMatch && (deleted)} + + ) + } + + // Diff stat line: "app.js | 10 ++++++---" + const statMatch = /^\s*(.+?)\s+\|\s+(\d+)\s+([+-]+)\s*$/u.exec(line) + if (statMatch) { + return ( + + {statMatch[1]} + {' | '} + {` ${statMatch[2]} `} + {renderBar(statMatch[3]!, `bar-${idx}`)} + + ) + } + + // Diff summary: "1 file changed, 2 insertions(+), 5 deletions(-)" + if (/^\s*\d+ files? changed/u.test(line)) { + const parts = line.trim().split(/(\([+-]\))/) + return ( + + {parts.map((part, pIdx) => { + if (part === '(+)') { + return ( + + {part} + + ) + } + if (part === '(-)') { + return ( + + {part} + + ) + } + return ( + + {part} + + ) + })} + + ) + } + + // Binary stat lines: " Bin 0 -> 123 bytes" + if (/^\s*Bin\b/u.test(line)) { + return ( + + {line.trim()} + + ) + } + + // Fallback: plain text. + return ( + + {line} + + ) + } + + return ( + + {content.split('\n').map((line, idx) => renderLine(line, idx))} + + ) + }, +) + +CommandResultBlock.displayName = 'CommandResultBlock' diff --git a/cli/src/components/message-block.tsx b/cli/src/components/message-block.tsx index adbd6fd488..623d6f544c 100644 --- a/cli/src/components/message-block.tsx +++ b/cli/src/components/message-block.tsx @@ -2,6 +2,7 @@ import { TextAttributes } from '@opentui/core' import { memo, useState } from 'react' import { BlocksRenderer } from './blocks/blocks-renderer' +import { CommandResultBlock } from './blocks/command-result-block' import { UserContentWithCopyButton } from './blocks/user-content-copy' import { Button } from './button' import { FileAttachmentCard } from './file-attachment-card' @@ -273,7 +274,12 @@ export const MessageBlock = memo(({ )} - {blocks ? ( + {metadata?.commandResult ? ( + + ) : blocks ? ( void + onCancel: () => void +} + +export const UndoHistoryScreen: React.FC = ({ + mode, + onSelect, + onCancel, +}) => { + const theme = useTheme() + const { terminalWidth, terminalHeight } = useTerminalLayout() + const contentWidth = terminalWidth - LAYOUT.CONTENT_PADDING + + // Load the stack once at mount, newest first for display. + const entries = useMemo(() => { + const chatId = getCurrentChatId() + const stack = + mode === 'undo' ? listUndoEntries(chatId) : listRedoEntries(chatId) + return [...stack].reverse() + }, [mode]) + + const isCompactMode = terminalHeight < LAYOUT.COMPACT_MODE_THRESHOLD + const isNarrowWidth = terminalWidth < LAYOUT.NARROW_WIDTH_THRESHOLD + const [cancelHovered, setCancelHovered] = useState(false) + + // Format: "[time] [n files] [prompt title]" + // reservedWidth accounts for: time col, files col, 2 gaps, list border (2), + // scrollbar (1), and button padding (2). + const reservedWidth = + LAYOUT.TIME_COL_WIDTH + + LAYOUT.FILES_COL_WIDTH + + LAYOUT.GAP_WIDTH * 2 + + 5 + const maxPromptWidth = Math.max(20, contentWidth - reservedWidth) + + const truncateText = (text: string, maxLen: number): string => { + const singleLine = text.replace(/\n/g, ' ').trim() + if (singleLine.length <= maxLen) return singleLine + return singleLine.slice(0, maxLen - 1) + '…' + } + + const padRight = (text: string, width: number): string => { + // Count code points so emoji/wide chars don't break padding + const len = Array.from(text).length + if (len >= width) return text + return text + ' '.repeat(width - len) + } + + const items: SelectableListItem[] = useMemo( + () => + entries.map((record) => { + const time = padRight( + formatRelativeTime(new Date(record.createdAt)), + LAYOUT.TIME_COL_WIDTH, + ) + const fileCount = padRight( + `${record.files.length} file${record.files.length === 1 ? '' : 's'}`, + LAYOUT.FILES_COL_WIDTH, + ) + const title = padRight( + truncateText(record.message, maxPromptWidth), + maxPromptWidth, + ) + return { + id: record.id, + label: `${time}${' '.repeat(LAYOUT.GAP_WIDTH)}${fileCount}${' '.repeat(LAYOUT.GAP_WIDTH)}${title}`, + // Keep the original prompt + files for search filtering. + secondary: `${record.message} ${record.files.join(' ')}`, + hideSecondary: true, + } + }), + [entries, maxPromptWidth], + ) + + const filterByPromptAndFiles = useCallback( + (item: SelectableListItem, query: string) => + (item.secondary ?? '').toLowerCase().includes(query.toLowerCase()), + [], + ) + + const { + searchQuery, + setSearchQuery, + focusedIndex, + setFocusedIndex, + filteredItems, + handleFocusChange, + } = useSearchableList({ + items, + filterFn: filterByPromptAndFiles, + }) + + // The record behind the currently focused row, for the files panel. + const focusedRecord: UndoRecord | undefined = useMemo(() => { + const focused = filteredItems[focusedIndex] + if (!focused) return undefined + return entries.find((record) => record.id === focused.id) + }, [filteredItems, focusedIndex, entries]) + + const handleKeyIntercept = useCallback( + (key: { + name?: string + sequence?: string + shift?: boolean + ctrl?: boolean + meta?: boolean + option?: boolean + }) => { + if (key.name === 'escape') { + if (searchQuery.length > 0) { + setSearchQuery('') + } else { + onCancel() + } + return true + } + if (key.name === 'up') { + setFocusedIndex((prev) => Math.max(0, prev - 1)) + return true + } + if (key.name === 'down') { + const maxIndex = Math.max(0, filteredItems.length - 1) + setFocusedIndex((prev) => Math.min(maxIndex, prev + 1)) + return true + } + if (isPlainEnterKey(key)) { + const focused = filteredItems[focusedIndex] + if (focused) { + onSelect(focused.id) + } + return true + } + if (key.name === 'c' && key.ctrl) { + onCancel() + return true + } + return false + }, + [ + searchQuery, + setSearchQuery, + setFocusedIndex, + filteredItems, + focusedIndex, + onSelect, + onCancel, + ], + ) + + const actionVerb = mode === 'undo' ? 'undo' : 'redo' + const emptyMessage = + entries.length === 0 + ? mode === 'undo' + ? 'Nothing to undo yet' + : 'Nothing to redo yet' + : searchQuery + ? 'No matching changes' + : 'No changes found' + + const visibleFiles = + focusedRecord?.files.slice(0, LAYOUT.MAX_VISIBLE_FILES) ?? [] + const hiddenFiles = + (focusedRecord?.files.length ?? 0) - visibleFiles.length + + return ( + + + {/* Title */} + {!isCompactMode && ( + + + {mode === 'undo' + ? 'Select a change to undo' + : 'Select a change to redo'} + + + )} + + {/* Search input */} + + setSearchQuery(text)} + onSubmit={() => {}} + onPaste={() => {}} + onKeyIntercept={handleKeyIntercept} + placeholder="Search changes..." + focused={true} + maxHeight={1} + minHeight={1} + cursorPosition={searchQuery.length} + /> + + + {/* Change list - grows to fill remaining space */} + + onSelect(item.id)} + onFocusChange={handleFocusChange} + emptyMessage={emptyMessage} + /> + + + {/* Files preview for the focused entry */} + + + {focusedRecord + ? `Files (${focusedRecord.files.length})` + : 'Files'} + + {visibleFiles.map((file) => ( + + {` • ${file}`} + + ))} + {hiddenFiles > 0 && ( + + {` … and ${hiddenFiles} more`} + + )} + + + + {/* Bottom bar */} + + + + + {`↑↓ navigate · Enter ${actionVerb} · Click to ${actionVerb} · Esc cancel`} + + + + {!isNarrowWidth && ( + + + + )} + + + + ) +} diff --git a/cli/src/data/slash-commands.ts b/cli/src/data/slash-commands.ts index ff6ad09870..56968a73cd 100644 --- a/cli/src/data/slash-commands.ts +++ b/cli/src/data/slash-commands.ts @@ -76,16 +76,16 @@ const ALL_SLASH_COMMANDS: SlashCommand[] = [ description: 'Create a starter knowledge.md file', implicitCommand: true, }, - // { - // id: 'undo', - // label: 'undo', - // description: 'Undo the last change made by the assistant', - // }, - // { - // id: 'redo', - // label: 'redo', - // description: 'Redo the most recent undone change', - // }, + { + id: 'undo', + label: 'undo', + description: 'Undo the last change made by the assistant', + }, + { + id: 'redo', + label: 'redo', + description: 'Redo the most recent undone change', + }, { id: 'usage', label: 'usage', diff --git a/cli/src/hooks/use-send-message.ts b/cli/src/hooks/use-send-message.ts index 39065317ee..eea6b0b68c 100644 --- a/cli/src/hooks/use-send-message.ts +++ b/cli/src/hooks/use-send-message.ts @@ -2,7 +2,12 @@ import { randomUUID } from 'node:crypto' import { useCallback, useEffect, useRef } from 'react' -import { setCurrentChatId } from '../project-files' +import { + getCurrentChatId, + getProjectRoot, + setCurrentChatId, +} from '../project-files' +import { recordUndoEntry } from '../state/undo-store' import { createStreamController } from './stream-state' import { useChatStore } from '../state/chat-store' import { @@ -13,6 +18,8 @@ import { getCodebuffClient } from '../utils/codebuff-client' import { AGENT_MODE_TO_COST_MODE, IS_FREEBUFF } from '../utils/constants' import { createEventHandlerState } from '../utils/create-event-handler-state' import { createRunConfig } from '../utils/create-run-config' +import { isUndoEnabled } from '../utils/settings' +import { patchSnapshot, trackSnapshot } from '../utils/undo-snapshot' import { getAgentIdForMode } from '../utils/freebuff-agent-selection' import { loadAgentDefinitions } from '../utils/local-agent-registry' import { logger } from '../utils/logger' @@ -296,6 +303,9 @@ export const useSendMessage = ({ const abortController = new AbortController() const runChatDir = resolveCurrentChatDir() const runChatIsCurrent = () => resolveCurrentChatDir() === runChatDir + // The chat that started this run — undo entries must follow it even if + // the user switches chats (/new, /history) while the run is in flight. + const runChatId = getCurrentChatId() let latestRunStateSnapshot: RunState = previousRunStateRef.current ?? { traceSessionId: randomUUID(), output: { @@ -526,6 +536,9 @@ export const useSendMessage = ({ // called at the start of sendMessage to ensure they happen synchronously // before any async work, so the router can correctly detect busy state. let actualCredits: number | undefined + // Snapshot hash captured before the run; consumed in the finally block + // to record an undo entry for this turn. + let undoSnapshotHash: string | null = null // Checkpoint the turn to disk immediately so that killing the process // (closed terminal, crash) can't lose the user's prompt, then keep the @@ -633,6 +646,20 @@ export const useSendMessage = ({ }, '[send-message] Sending message with sdk run config', ) + // Undo support: snapshot the project before the turn so /undo can + // restore the pre-turn state. Best-effort — a failure only disables + // undo for this turn. + try { + if (isUndoEnabled()) { + undoSnapshotHash = await trackSnapshot(getProjectRoot()) + } + } catch (error) { + logger.debug( + { error }, + '[send-message] Failed to capture undo snapshot', + ) + } + const runState = await client.run(runConfig) // Only adopt and persist the result while this run's chat is still @@ -708,6 +735,29 @@ export const useSendMessage = ({ logger.debug({ error }, '[send-message] Ignoring error after abort') } } finally { + // Undo: record the turn's file changes (success, error, or abort) so + // the user can revert them with /undo. No-op when nothing changed. + if (undoSnapshotHash) { + try { + const undoFiles = await patchSnapshot( + getProjectRoot(), + undoSnapshotHash, + ) + if (undoFiles.length > 0) { + recordUndoEntry(runChatId, { + hashBefore: undoSnapshotHash, + files: undoFiles, + message: content, + }) + } + } catch (error) { + logger.debug( + { error }, + '[send-message] Failed to record undo entry', + ) + } + } + // Stop exit-flushing this run's checkpoint; the final state (or last // checkpoint, on error) has been saved above. Owner-guarded so an // aborted run resolving late can't clear a newer run's provider. diff --git a/cli/src/state/__tests__/undo-store.test.ts b/cli/src/state/__tests__/undo-store.test.ts new file mode 100644 index 0000000000..d177c8052c --- /dev/null +++ b/cli/src/state/__tests__/undo-store.test.ts @@ -0,0 +1,136 @@ +import { afterAll, beforeAll, beforeEach, describe, expect, test } from 'bun:test' + +import { existsSync, mkdtempSync, rmSync } from 'node:fs' +import os from 'node:os' +import path from 'node:path' + +import { getProjectDataDir, setCurrentChatId, setProjectRoot } from '../../project-files' +import { + loadUndoState, + peekRedo, + peekUndo, + popRedo, + popUndo, + pushRedo, + pushUndo, + recordUndoEntry, +} from '../undo-store' + +const CHAT_ID = 'undo-store-test-chat' + +let projectDir: string + +beforeAll(() => { + projectDir = mkdtempSync(path.join(os.tmpdir(), 'undo-store-proj-')) + setProjectRoot(projectDir) + setCurrentChatId(CHAT_ID) +}) + +// Each test starts from a clean journal for this chat. +beforeEach(() => { + const chatDir = path.join(getProjectDataDir(), 'chats', CHAT_ID) + rmSync(chatDir, { recursive: true, force: true }) + setCurrentChatId(CHAT_ID) +}) + +afterAll(() => { + const chatDir = path.join(getProjectDataDir(), 'chats', CHAT_ID) + rmSync(chatDir, { recursive: true, force: true }) + rmSync(projectDir, { recursive: true, force: true }) +}) + +describe('recordUndoEntry', () => { + test('records an entry and ignores empty ones', () => { + recordUndoEntry(CHAT_ID, { + hashBefore: 'abc123', + files: ['a.txt', 'b.txt'], + message: 'fix the bug', + }) + const state = loadUndoState(CHAT_ID) + expect(state.undoStack).toHaveLength(1) + expect(state.undoStack[0]).toMatchObject({ + chatId: CHAT_ID, + hashBefore: 'abc123', + files: ['a.txt', 'b.txt'], + message: 'fix the bug', + }) + expect(peekUndo(CHAT_ID)?.hashBefore).toBe('abc123') + + // No hash or no files → nothing recorded. + recordUndoEntry(CHAT_ID, { hashBefore: '', files: [], message: 'x' }) + recordUndoEntry(CHAT_ID, { hashBefore: 'def', files: [], message: 'x' }) + expect(loadUndoState(CHAT_ID).undoStack).toHaveLength(1) + }) + + test('clears the redo stack when a new entry arrives', () => { + pushRedo(CHAT_ID, { + id: 'r1', + chatId: CHAT_ID, + hashBefore: 'old', + hashAfter: 'new', + files: ['a.txt'], + message: 'redo me', + createdAt: new Date().toISOString(), + }) + expect(peekRedo(CHAT_ID)).not.toBeNull() + + recordUndoEntry(CHAT_ID, { + hashBefore: 'xyz', + files: ['c.txt'], + message: 'new turn', + }) + expect(peekRedo(CHAT_ID)).toBeNull() + }) +}) + +describe('undo/redo stack operations', () => { + test('popUndo returns the most recent record and persists', () => { + recordUndoEntry(CHAT_ID, { + hashBefore: 'first', + files: ['one.txt'], + message: 'first turn', + }) + recordUndoEntry(CHAT_ID, { + hashBefore: 'second', + files: ['two.txt'], + message: 'second turn', + }) + + const record = popUndo(CHAT_ID) + expect(record?.hashBefore).toBe('second') + expect(peekUndo(CHAT_ID)?.hashBefore).toBe('first') + // Reload from disk to confirm the pop persisted. + expect(loadUndoState(CHAT_ID).undoStack).toHaveLength(1) + }) + + test('pushUndo restores a record and popRedo cycles', () => { + recordUndoEntry(CHAT_ID, { + hashBefore: 'cycle-hash', + files: ['x.txt'], + message: 'cycle', + }) + const record = popUndo(CHAT_ID)! + pushRedo(CHAT_ID, { ...record, hashAfter: 'after-state' }) + + const redoRecord = popRedo(CHAT_ID) + expect(redoRecord?.hashBefore).toBe('cycle-hash') + expect(redoRecord?.hashAfter).toBe('after-state') + + pushUndo(CHAT_ID, { ...redoRecord!, hashAfter: undefined }) + expect(peekUndo(CHAT_ID)?.hashBefore).toBe('cycle-hash') + }) + + test('returns null from empty stacks', () => { + setCurrentChatId('empty-chat') + expect(popUndo('empty-chat')).toBeNull() + expect(popRedo('empty-chat')).toBeNull() + setCurrentChatId(CHAT_ID) + }) +}) + +describe('corrupt file handling', () => { + test('loads an empty state for a nonexistent chat', () => { + expect(loadUndoState('never-existed').undoStack).toEqual([]) + expect(loadUndoState('never-existed').redoStack).toEqual([]) + }) +}) diff --git a/cli/src/state/undo-history-store.ts b/cli/src/state/undo-history-store.ts new file mode 100644 index 0000000000..d46f518030 --- /dev/null +++ b/cli/src/state/undo-history-store.ts @@ -0,0 +1,32 @@ +import { create } from 'zustand' + +interface UndoHistoryStoreState { + showUndoHistory: boolean + showRedoHistory: boolean +} + +interface UndoHistoryStoreActions { + openUndoHistory: () => void + closeUndoHistory: () => void + openRedoHistory: () => void + closeRedoHistory: () => void + reset: () => void +} + +type UndoHistoryStore = UndoHistoryStoreState & UndoHistoryStoreActions + +const initialState: UndoHistoryStoreState = { + showUndoHistory: false, + showRedoHistory: false, +} + +export const useUndoHistoryStore = create()((set) => ({ + ...initialState, + + openUndoHistory: () => set({ showUndoHistory: true }), + closeUndoHistory: () => set({ showUndoHistory: false }), + openRedoHistory: () => set({ showRedoHistory: true }), + closeRedoHistory: () => set({ showRedoHistory: false }), + + reset: () => set(initialState), +})) diff --git a/cli/src/state/undo-store.ts b/cli/src/state/undo-store.ts new file mode 100644 index 0000000000..327935657a --- /dev/null +++ b/cli/src/state/undo-store.ts @@ -0,0 +1,267 @@ +/** + * Per-chat undo/redo journal. + * + * Each assistant turn that changed files records an entry here: the snapshot + * hash captured before the turn plus the files it changed. `/undo` pops the + * most recent entry and reverts those files via the snapshot repo; `/redo` + * restores the state captured at undo time. Persisted as `undo.json` inside + * the chat's data directory so it survives restarts and follows the chat + * across `/history` resumes. + */ + +import { randomUUID } from 'node:crypto' +import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs' +import path from 'node:path' + +import { getProjectDataDir } from '../project-files' +import { + diffSnapshot, + restoreSnapshot, + revertFiles, + trackSnapshot, +} from '../utils/undo-snapshot' + +export type UndoRecord = { + id: string + chatId: string + /** Snapshot (git tree hash) captured before the assistant turn. */ + hashBefore: string + /** Files the turn changed, relative to the project root. */ + files: string[] + /** The user message that started the turn. */ + message: string + createdAt: string + /** Snapshot captured at /undo time; set only while the record is redoable. */ + hashAfter?: string + /** + * Turns reverted by an undo action that jumped back past them. Stored on + * the redo record so /redo can restore both the files and the stacks. + */ + restored?: UndoRecord[] +} + +export type UndoState = { + undoStack: UndoRecord[] + redoStack: UndoRecord[] +} + +const MAX_UNDO_ENTRIES = 20 +const MAX_MESSAGE_CHARS = 120 + +/** Keep the journal small: first line of the prompt, truncated. */ +function truncateMessage(message: string): string { + const firstLine = (message.split('\n')[0] ?? '').trim() + if (firstLine.length <= MAX_MESSAGE_CHARS) return firstLine + return `${firstLine.slice(0, MAX_MESSAGE_CHARS - 1)}…` +} + +function chatDirFor(chatId: string): string { + return path.join(getProjectDataDir(), 'chats', chatId) +} + +function undoFilePath(chatId: string): string { + return path.join(chatDirFor(chatId), 'undo.json') +} + +export function loadUndoState(chatId: string): UndoState { + try { + const file = undoFilePath(chatId) + if (existsSync(file)) { + const parsed = JSON.parse(readFileSync(file, 'utf8')) as Partial + return { + undoStack: Array.isArray(parsed.undoStack) ? parsed.undoStack : [], + redoStack: Array.isArray(parsed.redoStack) ? parsed.redoStack : [], + } + } + } catch { + // Corrupt or unreadable — treat as empty rather than breaking commands. + } + return { undoStack: [], redoStack: [] } +} + +export function saveUndoState(chatId: string, state: UndoState): void { + try { + const dir = chatDirFor(chatId) + mkdirSync(dir, { recursive: true }) + writeFileSync(undoFilePath(chatId), JSON.stringify(state, null, 2)) + } catch { + // Best-effort; undo is a convenience feature. + } +} + +/** + * Record a completed turn. Clears the redo stack: new edits invalidate redo. + * No-ops when the snapshot is missing or nothing changed. + */ +export function recordUndoEntry( + chatId: string, + entry: { hashBefore: string; files: string[]; message: string }, +): void { + if (!entry.hashBefore || entry.files.length === 0) return + const state = loadUndoState(chatId) + state.undoStack.push({ + id: randomUUID(), + chatId, + hashBefore: entry.hashBefore, + files: entry.files, + message: truncateMessage(entry.message), + createdAt: new Date().toISOString(), + }) + if (state.undoStack.length > MAX_UNDO_ENTRIES) { + state.undoStack.shift() + } + state.redoStack = [] + saveUndoState(chatId, state) +} + +export function peekUndo(chatId: string): UndoRecord | null { + const stack = loadUndoState(chatId).undoStack + return stack[stack.length - 1] ?? null +} + +export function popUndo(chatId: string): UndoRecord | null { + const state = loadUndoState(chatId) + const record = state.undoStack.pop() ?? null + if (record) saveUndoState(chatId, state) + return record +} + +/** Push a record back onto the undo stack (used by /redo). Does not touch redo. */ +export function pushUndo(chatId: string, record: UndoRecord): void { + const state = loadUndoState(chatId) + state.undoStack.push(record) + if (state.undoStack.length > MAX_UNDO_ENTRIES) { + state.undoStack.shift() + } + saveUndoState(chatId, state) +} + +export function peekRedo(chatId: string): UndoRecord | null { + const stack = loadUndoState(chatId).redoStack + return stack[stack.length - 1] ?? null +} + +export function popRedo(chatId: string): UndoRecord | null { + const state = loadUndoState(chatId) + const record = state.redoStack.pop() ?? null + if (record) saveUndoState(chatId, state) + return record +} + +export function pushRedo(chatId: string, record: UndoRecord): void { + const state = loadUndoState(chatId) + state.redoStack.push(record) + saveUndoState(chatId, state) +} + +/** The undo stack for a chat, oldest turn first. */ +export function listUndoEntries(chatId: string): UndoRecord[] { + return loadUndoState(chatId).undoStack +} + +/** The redo stack for a chat, oldest action first. */ +export function listRedoEntries(chatId: string): UndoRecord[] { + return loadUndoState(chatId).redoStack +} + +/** + * Undo back to a specific recorded turn (the OpenCode model): reverts that + * turn's files AND everything the agent changed in newer turns, restoring the + * project to the snapshot captured before the selected turn. The reverted + * turns move onto the redo stack so /redo can restore the exact state that + * was left behind. Returns the confirmation message, or null when the + * snapshot store is unavailable (in which case nothing is changed). + */ +export async function undoToRecord( + chatId: string, + projectRoot: string, + recordId: string, +): Promise { + const state = loadUndoState(chatId) + const index = state.undoStack.findIndex((record) => record.id === recordId) + if (index === -1) return null + const record = state.undoStack[index]! + // The selected turn plus every newer one — all of it is reverted. + const affected = state.undoStack.slice(index) + const files = Array.from(new Set(affected.flatMap((r) => r.files))) + + // Capture the current state first so /redo can restore exactly what was + // undone. If the snapshot store is unavailable, abort without mutating the + // stacks (mirrors the old inline handler's behavior). + const hashAfter = await trackSnapshot(projectRoot) + if (!hashAfter) return null + + // Diff of what the affected turns changed (computed before reverting). + const diffStat = await diffSnapshot(projectRoot, record.hashBefore) + const { restored, deleted } = await revertFiles( + projectRoot, + record.hashBefore, + files, + ) + + state.undoStack = state.undoStack.slice(0, index) + state.redoStack.push({ ...record, hashAfter, restored: affected }) + saveUndoState(chatId, state) + + const undone = [ + ...restored.map((file) => ` ↺ ${file}`), + ...deleted.map((file) => ` 🗑 ${file} (deleted)`), + ].join('\n') + if (!undone) return 'Could not undo the selected change.' + const turns = affected.length + const heading = + turns === 1 + ? '**Undid the last change:**' + : `**Undid ${turns} change(s) back to: ${truncateMessage(record.message)}**` + return `${heading}\n${undone}${diffStat ? `\n\n${diffStat}` : ''}` +} + +/** + * Redo a specific undo action: restores the project to the state captured + * when that undo ran and moves the reverted turns back onto the undo stack. + * Redo actions newer than the selected one are invalidated by the jump. + * Returns the confirmation message, or null when the restore fails. + */ +export async function redoToRecord( + chatId: string, + projectRoot: string, + recordId: string, +): Promise { + const state = loadUndoState(chatId) + const index = state.redoStack.findIndex((record) => record.id === recordId) + if (index === -1) return null + const record = state.redoStack[index]! + if (!record.hashAfter) return null + + // Diff of what this redo restores (computed before the tree changes). + const diffStat = await diffSnapshot(projectRoot, record.hashAfter) + const ok = await restoreSnapshot(projectRoot, record.hashAfter) + if (!ok) return null + + // Drop newer redo actions — jumping back invalidates them. + state.redoStack = state.redoStack.slice(0, index) + if (record.restored && record.restored.length > 0) { + state.undoStack.push(...record.restored) + } else { + state.undoStack.push({ ...record, hashAfter: undefined, restored: undefined }) + } + saveUndoState(chatId, state) + + const turns = record.restored?.length ?? 1 + const files = record.restored + ? Array.from(new Set(record.restored.flatMap((r) => r.files))) + : record.files + // After the restore, files present in the tree came back; files missing + // were removed by the restored state (the agent had deleted them). + const fileLines = files.map((file) => + existsSync(path.join(projectRoot, file)) + ? ` ↺ ${file}` + : ` 🗑 ${file} (deleted)`, + ) + const fileWord = files.length === 1 ? 'file' : 'files' + const heading = + turns === 1 + ? `**Redid the last change (${files.length} ${fileWord}):**` + : `**Redid ${turns} change(s) (${files.length} ${fileWord}):**` + return `${heading}\n${fileLines.join('\n')}${diffStat ? `\n\n${diffStat}` : ''}` +} diff --git a/cli/src/types/chat.ts b/cli/src/types/chat.ts index 77e0862931..61ed246052 100644 --- a/cli/src/types/chat.ts +++ b/cli/src/types/chat.ts @@ -161,6 +161,9 @@ export type AgentMessage = { export type ChatMessageMetadata = { /** Working directory where a bash command was executed */ bashCwd?: string + /** UI-only marker for /undo or /redo confirmation messages so they render + * with a distinct color instead of looking like plain agent output. */ + commandResult?: 'undo' | 'redo' /** UI-only marker for a response created in this process. Restored messages * strip it so ads are never fetched retroactively into settled history. */ allowInlineAds?: boolean diff --git a/cli/src/utils/__tests__/undo-snapshot.test.ts b/cli/src/utils/__tests__/undo-snapshot.test.ts new file mode 100644 index 0000000000..55a98f8e52 --- /dev/null +++ b/cli/src/utils/__tests__/undo-snapshot.test.ts @@ -0,0 +1,128 @@ +import { afterAll, beforeAll, describe, expect, test } from 'bun:test' + +import { execFileSync } from 'node:child_process' +import { existsSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs' +import os from 'node:os' +import path from 'node:path' + +import { + isUndoAvailable, + patchSnapshot, + restoreSnapshot, + revertFiles, + setSnapshotDirOverrideForTesting, + trackSnapshot, +} from '../undo-snapshot' + +let projectDir: string +let snapshotRoot: string +let plainDir: string + +const gitInProject = (args: string[]): string => + execFileSync('git', args, { + cwd: projectDir, + encoding: 'utf8', + // Keep expected stderr noise (e.g. the "ambiguous HEAD" probe) out of the + // test output. + stdio: ['ignore', 'pipe', 'ignore'], + }) + +const readProjectFile = (file: string): string => + execFileSync('cat', [path.join(projectDir, file)], { encoding: 'utf8' }) + +beforeAll(() => { + projectDir = mkdtempSync(path.join(os.tmpdir(), 'undo-snapshot-proj-')) + snapshotRoot = mkdtempSync(path.join(os.tmpdir(), 'undo-snapshot-store-')) + plainDir = mkdtempSync(path.join(os.tmpdir(), 'undo-snapshot-plain-')) + setSnapshotDirOverrideForTesting(snapshotRoot) + // A real (but commit-less) git repository. write-tree does not need commits + // or user identity, so this is all the fixture requires. + execFileSync('git', ['init'], { cwd: projectDir }) +}) + +afterAll(() => { + setSnapshotDirOverrideForTesting(undefined) + rmSync(projectDir, { recursive: true, force: true }) + rmSync(snapshotRoot, { recursive: true, force: true }) + rmSync(plainDir, { recursive: true, force: true }) +}) + +describe('isUndoAvailable', () => { + test('is true for a git repository', () => { + expect(isUndoAvailable(projectDir)).toBe(true) + }) + + test('is false for a directory without git', () => { + expect(isUndoAvailable(plainDir)).toBe(false) + }) +}) + +describe('trackSnapshot', () => { + test('returns null for a non-git directory', async () => { + expect(await trackSnapshot(plainDir)).toBeNull() + }) + + test('returns a hash and detects modified, added, and deleted files', async () => { + writeFileSync(path.join(projectDir, 'a.txt'), 'hello\n') + writeFileSync(path.join(projectDir, 'c.txt'), 'keep me\n') + const hash = await trackSnapshot(projectDir) + expect(hash).toBeTruthy() + + // Modify a tracked-in-snapshot file, add a new file, delete another. + writeFileSync(path.join(projectDir, 'a.txt'), 'hello world\n') + writeFileSync(path.join(projectDir, 'b.txt'), 'new file\n') + rmSync(path.join(projectDir, 'c.txt')) + + const changed = await patchSnapshot(projectDir, hash!) + expect(changed.sort()).toEqual(['a.txt', 'b.txt', 'c.txt']) + }) +}) + +describe('restoreSnapshot', () => { + test('overwrites the worktree with the tracked state', async () => { + writeFileSync(path.join(projectDir, 'restore.txt'), 'original\n') + const hash = await trackSnapshot(projectDir) + + writeFileSync(path.join(projectDir, 'restore.txt'), 'changed by agent\n') + + const ok = await restoreSnapshot(projectDir, hash!) + expect(ok).toBe(true) + expect(readProjectFile('restore.txt')).toBe('original\n') + }) +}) + +describe('revertFiles', () => { + test('restores a modified file and deletes a file created after the snapshot', async () => { + writeFileSync(path.join(projectDir, 'revert.txt'), 'v1\n') + const hash = await trackSnapshot(projectDir) + + writeFileSync(path.join(projectDir, 'revert.txt'), 'v2 by agent\n') + writeFileSync(path.join(projectDir, 'agent-created.txt'), 'agent made this\n') + + const { restored, deleted } = await revertFiles(projectDir, hash!, [ + 'revert.txt', + 'agent-created.txt', + ]) + // revert.txt existed in the snapshot and is restored; agent-created.txt did + // not exist in the snapshot (the agent created it) and is deleted. + expect(restored).toEqual(['revert.txt']) + expect(deleted).toEqual(['agent-created.txt']) + expect(readProjectFile('revert.txt')).toBe('v1\n') + expect(existsSync(path.join(projectDir, 'agent-created.txt'))).toBe(false) + }) +}) + +describe('isolation from the real repository', () => { + test('never stages or commits anything in the project git repo', async () => { + // Fresh fixture area. + writeFileSync(path.join(projectDir, 'isolated.txt'), 'content\n') + await trackSnapshot(projectDir) + writeFileSync(path.join(projectDir, 'isolated.txt'), 'content v2\n') + await trackSnapshot(projectDir) + + // The real repo must have no staged changes and no commits of ours. + const staged = gitInProject(['diff', '--cached', '--name-only']).trim() + expect(staged).toBe('') + expect(() => gitInProject(['rev-parse', 'HEAD'])).toThrow() + }) +}) diff --git a/cli/src/utils/active-run.ts b/cli/src/utils/active-run.ts index 9ef9ff396b..345802f90f 100644 --- a/cli/src/utils/active-run.ts +++ b/cli/src/utils/active-run.ts @@ -4,6 +4,8 @@ export type ActiveRunStopReason = | 'logout' | 'new-chat' | 'history-resume' + | 'undo-history' + | 'redo-history' | 'session-transition' | 'process-exit' @@ -24,6 +26,8 @@ export const ACTIVE_RUN_QUEUE_POLICIES = { logout: 'clear-and-block', 'new-chat': 'clear-and-block', 'history-resume': 'clear-and-block', + 'undo-history': 'clear-and-block', + 'redo-history': 'clear-and-block', 'session-transition': 'clear-and-block', 'process-exit': 'preserve-and-block', } satisfies Record diff --git a/cli/src/utils/message-history.ts b/cli/src/utils/message-history.ts index 11c3497bf5..bbd5fe501d 100644 --- a/cli/src/utils/message-history.ts +++ b/cli/src/utils/message-history.ts @@ -5,7 +5,14 @@ import { getConfigDir } from './auth' import { formatTimestamp } from './helpers' import { logger } from './logger' -import type { ChatMessage, ContentBlock, FileAttachment, ImageAttachment, TextAttachment } from '../types/chat' +import type { + ChatMessage, + ChatMessageMetadata, + ContentBlock, + FileAttachment, + ImageAttachment, + TextAttachment, +} from '../types/chat' const MAX_HISTORY_SIZE = 1000 @@ -35,6 +42,7 @@ export function getUserMessage( export function getSystemMessage( content: string | ContentBlock[], + metadata?: ChatMessageMetadata, ): ChatMessage { return { id: `sys-${Date.now()}`, @@ -48,6 +56,7 @@ export function getSystemMessage( blocks: content, }), timestamp: formatTimestamp(), + ...(metadata ? { metadata } : {}), } } diff --git a/cli/src/utils/settings.ts b/cli/src/utils/settings.ts index ee99dfb00a..9a1bbc1cad 100644 --- a/cli/src/utils/settings.ts +++ b/cli/src/utils/settings.ts @@ -16,6 +16,7 @@ import type { AgentMode } from './constants' const DEFAULT_SETTINGS: Settings = { mode: 'DEFAULT' as const, adsEnabled: true, + undo: true, } // Note: The old FREE mode has been renamed back to LITE; migrate on load. @@ -38,6 +39,8 @@ export interface Settings { * first-time onboarding suggested prompts so they only show to brand-new * users and quietly retire afterwards. */ hasSubmittedFirstPrompt?: boolean + /** Whether /undo and /redo are enabled. Defaults to true. */ + undo?: boolean } /** @@ -146,6 +149,11 @@ const validateSettings = (parsed: unknown): Settings => { settings.hasSubmittedFirstPrompt = obj.hasSubmittedFirstPrompt } + // Validate undo toggle + if (typeof obj.undo === 'boolean') { + settings.undo = obj.undo + } + return settings } @@ -224,3 +232,15 @@ export const markFirstPromptSubmitted = (): void => { if (loadSettings().hasSubmittedFirstPrompt === true) return saveSettings({ hasSubmittedFirstPrompt: true }) } + +/** + * Whether /undo and /redo are enabled. Defaults to true. + */ +export const isUndoEnabled = (): boolean => loadSettings().undo !== false + +/** + * Enable or disable /undo and /redo. + */ +export const setUndoEnabled = (enabled: boolean): void => { + saveSettings({ undo: enabled }) +} diff --git a/cli/src/utils/undo-snapshot.ts b/cli/src/utils/undo-snapshot.ts new file mode 100644 index 0000000000..d154dec6aa --- /dev/null +++ b/cli/src/utils/undo-snapshot.ts @@ -0,0 +1,494 @@ +/** + * Undo snapshot store. + * + * Tracks the filesystem state before and after each assistant turn using a + * dedicated, hidden git repository per project — the same approach as + * OpenCode's snapshot service (`packages/opencode/src/snapshot/index.ts`). + * + * The snapshot repository never touches the project's real `.git`: every git + * invocation runs with `--git-dir --work-tree `. + * A snapshot is the tree hash produced by `git write-tree` after staging the + * project's changes, so restoring one is a matter of `git read-tree` + + * `git checkout-index`. The source repository's object database is reused via + * `objects/info/alternates` (with the source index copied) so the first + * snapshot stays fast even on huge repositories. + * + * All operations are best-effort: a failure disables undo for that turn + * rather than breaking the chat. + */ + +import { spawn } from 'node:child_process' +import { createHash } from 'node:crypto' +import { + copyFileSync, + existsSync, + mkdirSync, + readFileSync, + rmSync, + statSync, + writeFileSync, +} from 'node:fs' +import path from 'node:path' + +import { getConfigDir } from './config-dir' +import { findGitRoot } from './git' +import { logger } from './logger' + +/** Files larger than this are excluded from snapshots (same limit as OpenCode). */ +const MAX_SNAPSHOT_FILE_BYTES = 2 * 1024 * 1024 + +/** Absolute-path pathspec magic used with `--pathspec-from-file`. */ +const topLevelLiteral = (file: string): string => `:(top,literal)${file}` + +type GitResult = { + code: number + stdout: string + stderr: string +} + +async function runGit( + args: string[], + cwd: string, + options?: { input?: string; env?: Record }, +): Promise { + return new Promise((resolve) => { + const child = spawn('git', args, { + cwd, + env: options?.env ? { ...process.env, ...options.env } : process.env, + stdio: ['pipe', 'pipe', 'pipe'], + }) + let stdout = '' + let stderr = '' + child.stdout.setEncoding('utf8') + child.stderr.setEncoding('utf8') + child.stdout.on('data', (chunk: string) => { + stdout += chunk + }) + child.stderr.on('data', (chunk: string) => { + stderr += chunk + }) + child.on('error', (error) => { + resolve({ + code: 1, + stdout, + stderr: error instanceof Error ? error.message : String(error), + }) + }) + child.on('close', (code) => { + resolve({ code: code ?? 1, stdout, stderr }) + }) + if (options?.input) { + child.stdin.write(options.input) + } + child.stdin.end() + }) +} + +/** Serialize git operations per snapshot repo so concurrent tracks can't corrupt the index. */ +const locks = new Map>() + +/** Hourly GC keeps snapshot repos from growing without bound (prune 7 days). */ +const GC_INTERVAL_MS = 60 * 60 * 1000 +const GC_PRUNE = '7.days' +const lastGcByDir = new Map() + +function withLock(key: string, fn: () => Promise): Promise { + const previous = locks.get(key) ?? Promise.resolve() + const next = previous.then(fn, fn) + locks.set( + key, + next.catch(() => { + // Swallow so a failed op never wedges the chain for later ones. + }), + ) + return next +} + +/** Where snapshot repos live. Test-overridable via setSnapshotDirOverrideForTesting. */ +let snapshotDirOverride: string | undefined + +export function setSnapshotDirOverrideForTesting(dir: string | undefined): void { + snapshotDirOverride = dir +} + +function getSnapshotBaseDir(): string { + return snapshotDirOverride ?? path.join(getConfigDir(), 'undo-snapshots') +} + +function snapshotDirFor(projectRoot: string): string { + const key = createHash('sha1').update(projectRoot).digest('hex').slice(0, 12) + return path.join(getSnapshotBaseDir(), `${path.basename(projectRoot)}-${key}`) +} + +/** git args that point every command at the snapshot repo, not the real one. */ +const gitArgs = ( + snapshotDir: string, + projectRoot: string, + command: string[], +): string[] => ['--git-dir', snapshotDir, '--work-tree', projectRoot, ...command] + +/** Whether undo can work for this project at all (it must be a git repo). */ +export function isUndoAvailable(projectRoot: string): boolean { + return findGitRoot({ cwd: projectRoot }) !== null +} + +async function ensureInitialized( + snapshotDir: string, + projectRoot: string, +): Promise { + if (existsSync(snapshotDir)) return true + try { + mkdirSync(snapshotDir, { recursive: true }) + await runGit(['init'], projectRoot, { + env: { GIT_DIR: snapshotDir, GIT_WORK_TREE: projectRoot }, + }) + const configArgs = ['--git-dir', snapshotDir, 'config'] + await runGit([...configArgs, 'core.autocrlf', 'false'], projectRoot) + await runGit([...configArgs, 'core.longpaths', 'true'], projectRoot) + await runGit([...configArgs, 'feature.manyFiles', 'true'], projectRoot) + await runGit([...configArgs, 'index.version', '4'], projectRoot) + await runGit([...configArgs, 'core.untrackedCache', 'true'], projectRoot) + await seedFromSourceRepo(snapshotDir, projectRoot) + return true + } catch (error) { + logger.debug({ error }, 'undo-snapshot: failed to initialize snapshot repo') + return false + } +} + +/** + * Reuse the source repo's object database (and index) so already-hashed file + * content does not need re-hashing on the first snapshot. Best-effort. + */ +async function seedFromSourceRepo( + snapshotDir: string, + projectRoot: string, +): Promise { + try { + const common = await runGit( + ['rev-parse', '--path-format=absolute', '--git-common-dir'], + projectRoot, + ) + if (common.code !== 0) return + const source = common.stdout.trim() + if (!source || !existsSync(source)) return + + const sourceObjects = path.join(source, 'objects') + if (!existsSync(sourceObjects)) return + + const alternatesDir = path.join(snapshotDir, 'objects', 'info') + mkdirSync(alternatesDir, { recursive: true }) + writeFileSync(path.join(alternatesDir, 'alternates'), `${sourceObjects}\n`) + + const sourceIndex = path.join(source, 'index') + if (existsSync(sourceIndex)) { + copyFileSync(sourceIndex, path.join(snapshotDir, 'index')) + } + } catch (error) { + logger.debug({ error }, 'undo-snapshot: failed to seed from source repo') + } +} + +/** Mirror the source repo's info/exclude plus blocked (oversized) files into the snapshot repo. */ +function syncExcludes( + snapshotDir: string, + projectRoot: string, + gitRoot: string | null, + blocked: string[], +): void { + try { + const lines: string[] = [] + const sourceExclude = gitRoot + ? path.join(gitRoot, '.git', 'info', 'exclude') + : null + if (sourceExclude && existsSync(sourceExclude)) { + lines.push(...readFileSync(sourceExclude, 'utf8').split('\n')) + } + for (const file of blocked) { + lines.push(`/${file.replaceAll('\\', '/')}`) + } + mkdirSync(path.join(snapshotDir, 'info'), { recursive: true }) + writeFileSync( + path.join(snapshotDir, 'info', 'exclude'), + `${lines.filter((line) => line.trim()).join('\n')}\n`, + ) + } catch { + // Best-effort. + } +} + +/** + * Stage every changed, added, and deleted file (respecting the project's own + * ignore rules and skipping files over the size limit) into the snapshot + * index, so a following `write-tree` reflects the current state. + */ +async function stageChanges( + snapshotDir: string, + projectRoot: string, + gitRoot: string | null, +): Promise { + try { + const base = gitArgs(snapshotDir, projectRoot, []) + const [diff, others] = await Promise.all([ + runGit( + [...base, 'diff-files', '--name-only', '-z', '--', '.'], + projectRoot, + ), + runGit( + [ + ...base, + 'ls-files', + '--full-name', + '--others', + '--exclude-standard', + '-z', + '--', + '.', + ], + projectRoot, + ), + ]) + const changed = diff.stdout.split('\0').filter(Boolean) + const untracked = others.stdout.split('\0').filter(Boolean) + const all = Array.from(new Set([...changed, ...untracked])) + if (all.length === 0) return + + const allowed: string[] = [] + const blocked: string[] = [] + for (const file of all) { + try { + const stat = statSync(path.join(projectRoot, file)) + if (stat.isFile() && stat.size > MAX_SNAPSHOT_FILE_BYTES) { + blocked.push(file) + } else { + allowed.push(file) + } + } catch { + // Deleted or unreadable — still stage the removal. + allowed.push(file) + } + } + + syncExcludes(snapshotDir, projectRoot, gitRoot, blocked) + + if (allowed.length === 0) return + const pathspecs = `${allowed.map(topLevelLiteral).join('\0')}\0` + await runGit( + [ + ...base, + 'add', + '--all', + '--sparse', + '--pathspec-from-file=-', + '--pathspec-file-nul', + ], + projectRoot, + { input: pathspecs }, + ) + } catch (error) { + logger.debug({ error }, 'undo-snapshot: failed to stage changes') + } +} + +/** + * Capture a snapshot of the project's current state. + * @returns the snapshot hash, or null when unavailable/failed. + */ +export async function trackSnapshot(projectRoot: string): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return null + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + if (!(await ensureInitialized(snapshotDir, projectRoot))) return null + await stageChanges(snapshotDir, projectRoot, gitRoot) + const result = await runGit( + [...gitArgs(snapshotDir, projectRoot, ['write-tree'])], + projectRoot, + ) + const hash = result.code === 0 ? result.stdout.trim() : '' + if (hash) { + void maybePrune(snapshotDir, projectRoot) + } + return hash || null + } catch (error) { + logger.debug({ error }, 'undo-snapshot: track failed') + return null + } + }) +} + +/** + * Run `git gc --prune=7.days` at most once per hour per snapshot repo. + * Fire-and-forget; never throws. + */ +async function maybePrune(snapshotDir: string, projectRoot: string): Promise { + const last = lastGcByDir.get(snapshotDir) ?? 0 + if (Date.now() - last < GC_INTERVAL_MS) return + lastGcByDir.set(snapshotDir, Date.now()) + const result = await runGit( + ['--git-dir', snapshotDir, 'gc', `--prune=${GC_PRUNE}`], + projectRoot, + ) + if (result.code !== 0) { + logger.debug({ stderr: result.stderr }, 'undo-snapshot: gc failed') + } +} + +/** + * List the files that changed since the given snapshot. + * @returns project-relative file paths (empty when nothing changed). + */ +export async function patchSnapshot( + projectRoot: string, + hash: string, +): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return [] + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + await stageChanges(snapshotDir, projectRoot, gitRoot) + const result = await runGit( + [ + ...gitArgs(snapshotDir, projectRoot, [ + 'diff', + '--cached', + '--no-ext-diff', + '--name-only', + hash, + '--', + '.', + ]), + ], + projectRoot, + ) + if (result.code !== 0) return [] + return result.stdout + .split('\n') + .map((line) => line.trim()) + .filter(Boolean) + } catch (error) { + logger.debug({ error }, 'undo-snapshot: patch failed') + return [] + } + }) +} + +/** Restore the entire worktree to a snapshot. Returns false on failure. */ +export async function restoreSnapshot( + projectRoot: string, + hash: string, +): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return false + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + const base = gitArgs(snapshotDir, projectRoot, []) + const read = await runGit([...base, 'read-tree', hash], projectRoot) + if (read.code !== 0) return false + const checkout = await runGit( + [...base, 'checkout-index', '-a', '-f'], + projectRoot, + ) + return checkout.code === 0 + } catch (error) { + logger.debug({ error }, 'undo-snapshot: restore failed') + return false + } + }) +} + +/** + * Restore a specific set of files to a snapshot. Files that did not exist in + * the snapshot (created after it) are deleted. + * @returns the files that were restored and the files that were deleted. + */ +export async function revertFiles( + projectRoot: string, + hash: string, + files: string[], +): Promise<{ restored: string[]; deleted: string[] }> { + const empty = { restored: [] as string[], deleted: [] as string[] } + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot || files.length === 0) return empty + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + const result = { restored: [] as string[], deleted: [] as string[] } + try { + const base = gitArgs(snapshotDir, projectRoot, []) + const root = path.resolve(projectRoot) + for (const file of files) { + // Defense in depth: never touch a path that escapes the project. + const resolved = path.resolve(root, file) + if (resolved !== root && !resolved.startsWith(`${root}${path.sep}`)) { + continue + } + const checkout = await runGit( + [...base, 'checkout', hash, '--', file], + projectRoot, + ) + if (checkout.code === 0) { + result.restored.push(file) + continue + } + const tree = await runGit( + [...base, 'ls-tree', hash, '--', file], + projectRoot, + ) + if (tree.code === 0 && tree.stdout.trim()) { + // Existed in the snapshot but the checkout failed; leave it be. + result.restored.push(file) + } else { + // Did not exist in the snapshot — the turn created it. + try { + rmSync(resolved, { force: true }) + result.deleted.push(file) + } catch { + // Keep going with the remaining files. + } + } + } + return result + } catch (error) { + logger.debug({ error }, 'undo-snapshot: revert failed') + return result + } + }) +} + +/** + * Compact diff stat of everything that changed since a snapshot (used in the + * /undo result message). Empty string when there is nothing or on failure. + */ +export async function diffSnapshot( + projectRoot: string, + hash: string, +): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return '' + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + await stageChanges(snapshotDir, projectRoot, gitRoot) + const result = await runGit( + [ + ...gitArgs(snapshotDir, projectRoot, [ + 'diff', + '--cached', + '--no-ext-diff', + '--stat', + hash, + '--', + '.', + ]), + ], + projectRoot, + ) + return result.code === 0 ? result.stdout : '' + } catch { + return '' + } + }) +}