diff --git a/cli/src/app.tsx b/cli/src/app.tsx index fc31eef928..4e04caaeba 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 { useByokSelectionStore } from './utils/byok' @@ -167,6 +171,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) @@ -194,6 +205,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 @@ -268,6 +341,12 @@ export const App = ({ onSelectChat={handleResumeChat} onCancelChatHistory={closeChatHistory} onNewChat={handleNewChat} + showUndoHistory={showUndoHistory} + onUndoSelect={handleUndoSelect} + onCancelUndoHistory={handleCancelUndoHistory} + showRedoHistory={showRedoHistory} + onRedoSelect={handleRedoSelect} + onCancelRedoHistory={handleCancelRedoHistory} /> ) } @@ -293,6 +372,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 } /** @@ -345,6 +430,12 @@ const AuthedSurfaceRoutes = ({ onSelectChat, onCancelChatHistory, onNewChat, + showUndoHistory, + onUndoSelect, + onCancelUndoHistory, + showRedoHistory, + onRedoSelect, + onCancelRedoHistory, session, sessionFailure, lastRefund, @@ -418,6 +509,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/__tests__/undo-history-screen.test.tsx b/cli/src/components/__tests__/undo-history-screen.test.tsx new file mode 100644 index 0000000000..94b5e4140b --- /dev/null +++ b/cli/src/components/__tests__/undo-history-screen.test.tsx @@ -0,0 +1,138 @@ +import { afterEach, beforeAll, describe, expect, mock, test } from 'bun:test' +import { createTestRenderer } from '@opentui/core/testing' +import { createRoot, flushSync } from '@opentui/react' +import React from 'react' + +import type { UndoRecord } from '../../state/undo-store' + +const record = (id: string, day: number): UndoRecord => ({ + id, + chatId: 'picker-chat', + hashBefore: `hash-${id}`, + files: [`src/${id}.ts`, 'src/shared.ts'], + message: `prompt ${id}`, + createdAt: new Date(Date.UTC(2026, 0, day, 12)).toISOString(), +}) + +/** Oldest turn first, the way the journal lists them. */ +const THREE = [record('one', 1), record('two', 2), record('three', 3)] + +let undoStack: UndoRecord[] = THREE +let redoStack: UndoRecord[] = [] + +// The journal is not what this file is about: the picker gets a stack with a +// known shape and the assertions are about what it says over it. +mock.module('../../state/undo-store', () => ({ + listUndoEntries: () => undoStack, + listRedoEntries: () => redoStack, +})) + +import { UndoHistoryScreen } from '../undo-history-screen' +import { initializeThemeStore } from '../../hooks/use-theme' + +let cleanupRenderer: (() => void) | undefined + +beforeAll(() => { + initializeThemeStore() +}) + +afterEach(() => { + cleanupRenderer?.() + cleanupRenderer = undefined + undoStack = THREE + redoStack = [] +}) + +const mountPicker = async (mode: 'undo' | 'redo' = 'undo') => { + const setup = await createTestRenderer({ + width: 100, + height: 40, + kittyKeyboard: true, + }) + const root = createRoot(setup.renderer) + cleanupRenderer = () => { + flushSync(() => root.unmount()) + setup.renderer.destroy() + } + + const selections: string[] = [] + flushSync(() => + root.render( + selections.push(recordId)} + onCancel={() => {}} + />, + ), + ) + await setup.renderOnce() + + /** Input is delivered on the render loop and committed by React's + * scheduler, so both have to drain before the next keypress. */ + const settle = async () => { + await setup.renderOnce() + await new Promise((resolve) => setTimeout(resolve, 20)) + await setup.renderOnce() + } + + return Object.assign(setup, { + selections: () => selections, + async press(act: () => void) { + act() + await settle() + }, + }) +} + +describe('UndoHistoryScreen blast radius', () => { + test('lists the files of every turn the undo takes with it', async () => { + const picker = await mountPicker() + + // The newest entry takes nothing with it. + expect(picker.captureCharFrame()).toContain('Files (2)') + + await picker.press(() => picker.mockInput.pressArrow('down')) + await picker.press(() => picker.mockInput.pressArrow('down')) + + // The oldest one reverts all three turns, and shared.ts is listed once. + const frame = picker.captureCharFrame() + expect(frame).toContain('Files (4)') + expect(frame).toContain('src/two.ts') + expect(frame).toContain('src/shared.ts') + }) + + test('stays quiet on the newest turn, which takes nothing with it', async () => { + const picker = await mountPicker() + + const frame = picker.captureCharFrame() + expect(frame).toContain('Select a change to undo') + expect(frame).not.toContain('Also reverts') + }) + + test('says how many newer turns an undo also reverts', async () => { + const picker = await mountPicker() + + await picker.press(() => picker.mockInput.pressArrow('down')) + expect(picker.captureCharFrame()).toContain('Also reverts 1 newer change') + + await picker.press(() => picker.mockInput.pressArrow('down')) + expect(picker.captureCharFrame()).toContain('Also reverts 2 newer changes') + }) + + test('says nothing about changes when the stack has one entry', async () => { + undoStack = [record('only', 1)] + const picker = await mountPicker() + + expect(picker.captureCharFrame()).not.toContain('Also reverts') + }) + + test('warns that a redo discards the newer redo actions', async () => { + redoStack = THREE + const picker = await mountPicker('redo') + + await picker.press(() => picker.mockInput.pressArrow('down')) + await picker.press(() => picker.mockInput.pressArrow('down')) + + expect(picker.captureCharFrame()).toContain('Also discards 2 newer redos') + }) +}) 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' + + // What the action actually touches on disk — not just the selected turn. + // An undo reverts the selected turn and every newer one; a redo brings back + // the turns its undo had reverted. Listing only the focused entry's files + // understated the blast radius the notice below it warns about. + const affectedFiles = useMemo(() => { + if (!focusedRecord) return [] + return mode === 'undo' + ? filesTakenBack(entries, focusedRecord.id) + : filesRestored(focusedRecord) + }, [mode, entries, focusedRecord]) + + const visibleFiles = affectedFiles.slice(0, LAYOUT.MAX_VISIBLE_FILES) + const hiddenFiles = affectedFiles.length - visibleFiles.length + + // The journal takes the selected entry *and everything newer* with it (a + // redo drops the newer redo entries the same way). Say so while the row is + // merely focused, not in the confirmation afterwards: by then the work is + // already gone. + const cascadeCount = focusedRecord + ? newerTurnCount(entries, focusedRecord.id) + : 0 + const cascadeNotice = + cascadeCount === 0 + ? null + : mode === 'undo' + ? `⚠ Also reverts ${cascadeCount} newer ${cascadeCount === 1 ? 'change' : 'changes'}` + : `⚠ Also discards ${cascadeCount} newer ${cascadeCount === 1 ? 'redo' : 'redos'}` + + 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 (${affectedFiles.length})` : 'Files'} + + {visibleFiles.map((file) => ( + + {` • ${file}`} + + ))} + {hiddenFiles > 0 && ( + + {` … and ${hiddenFiles} more`} + + )} + {cascadeNotice && ( + + {` ${cascadeNotice}`} + + )} + + + + {/* 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 2152e52240..5db7b23575 100644 --- a/cli/src/data/slash-commands.ts +++ b/cli/src/data/slash-commands.ts @@ -79,16 +79,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 fe86e341c0..9e0eca111e 100644 --- a/cli/src/hooks/use-send-message.ts +++ b/cli/src/hooks/use-send-message.ts @@ -2,7 +2,13 @@ import { randomUUID } from 'node:crypto' import { useCallback, useEffect, useRef } from 'react' -import { setCurrentChatId } from '../project-files' +import { + getCurrentChatId, + getProjectRoot, + setCurrentChatId, +} from '../project-files' +import { beginUndoTurn, isLatestUndoTurn } from '../state/undo-guards' +import { recordUndoEntry } from '../state/undo-store' import { createStreamController } from './stream-state' import { useChatStore } from '../state/chat-store' import { @@ -18,6 +24,8 @@ import { 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' @@ -337,6 +345,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: { @@ -571,6 +582,13 @@ 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 + // Stamped when the turn starts, presented when it records. The finally + // can run after a newer turn has already started (Esc releases the chain + // lock first), and by then this turn's diff is no longer only its own. + let undoTurnStamp: number | 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 @@ -732,6 +750,23 @@ 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()) + if (undoSnapshotHash) { + undoTurnStamp = beginUndoTurn(runChatId) + } + } + } catch (error) { + logger.debug( + { error }, + '[send-message] Failed to capture undo snapshot', + ) + } + // Open the steering mailbox for this run only once we're committed to // calling run(); the router falls back to the queue before this point. activateSteering(runOwnerId) @@ -847,6 +882,45 @@ export const useSendMessage = ({ }) } } + + // 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 && undoTurnStamp !== null) { + if (!isLatestUndoTurn(runChatId, undoTurnStamp)) { + // A newer turn started before this one finished, so this turn's + // diff would claim files the newer one changed. Drop the turn + // rather than record a change the user never selected. + logger.debug( + {}, + '[send-message] Skipping undo entry: a newer turn started first', + ) + } else { + try { + const undoFiles = await patchSnapshot( + getProjectRoot(), + undoSnapshotHash, + ) + // Checked again after the diff: this await is exactly where a + // newer turn can start editing the worktree. + if ( + undoFiles.length > 0 && + isLatestUndoTurn(runChatId, undoTurnStamp) + ) { + 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-guards.test.ts b/cli/src/state/__tests__/undo-guards.test.ts new file mode 100644 index 0000000000..f2f217414e --- /dev/null +++ b/cli/src/state/__tests__/undo-guards.test.ts @@ -0,0 +1,104 @@ +import { describe, expect, test } from 'bun:test' + +import { + beginUndoTurn, + filesRestored, + filesTakenBack, + isLatestUndoTurn, + newerTurnCount, +} from '../undo-guards' + +import type { UndoRecord } from '../undo-store' + +const record = (id: string): UndoRecord => ({ + id, + chatId: 'undo-guards-chat', + hashBefore: `hash-${id}`, + // Every turn also touches a file another turn touched, so the union is the + // thing under test rather than "all files happen to be distinct". + files: [`${id}.ts`, 'shared.ts'], + message: id, + createdAt: '2026-01-01T00:00:00.000Z', +}) + +// The picker's list: newest first. +const newestFirst = [record('c'), record('b'), record('a')] + +describe('filesTakenBack / filesRestored', () => { + test('an undo touches the selected turn and every newer one', () => { + expect(filesTakenBack(newestFirst, 'c')).toEqual(['c.ts', 'shared.ts']) + expect(filesTakenBack(newestFirst, 'a')).toEqual([ + 'c.ts', + 'shared.ts', + 'b.ts', + 'a.ts', + ]) + }) + + test('lists a file once however many turns touched it', () => { + expect( + filesTakenBack(newestFirst, 'a').filter((file) => file === 'shared.ts'), + ).toHaveLength(1) + }) + + test('is empty for an id the list does not hold', () => { + expect(filesTakenBack(newestFirst, 'gone')).toEqual([]) + }) + + test('a redo brings back the turns its undo had reverted', () => { + const redoable: UndoRecord = { + ...record('x'), + restored: [record('r1'), record('r2')], + } + expect(filesRestored(redoable)).toEqual([ + 'r1.ts', + 'shared.ts', + 'r2.ts', + ]) + }) + + test('a redo without carried turns falls back to its own files', () => { + expect(filesRestored(record('only'))).toEqual(['only.ts', 'shared.ts']) + }) +}) + +describe('newerTurnCount', () => { + test('counts the entries newer than the selected one', () => { + expect(newerTurnCount(newestFirst, 'c')).toBe(0) + expect(newerTurnCount(newestFirst, 'b')).toBe(1) + expect(newerTurnCount(newestFirst, 'a')).toBe(2) + }) + + test('is zero when there is nothing newer to take with it', () => { + expect(newerTurnCount([record('only')], 'only')).toBe(0) + expect(newerTurnCount([], 'a')).toBe(0) + }) + + test('is zero for an id the list does not hold', () => { + expect(newerTurnCount(newestFirst, 'gone')).toBe(0) + }) +}) + +describe('beginUndoTurn / isLatestUndoTurn', () => { + test('a turn that starts and finishes alone may record', () => { + const stamp = beginUndoTurn('guards-alone') + expect(isLatestUndoTurn('guards-alone', stamp)).toBe(true) + }) + + test('a turn whose run outlives the next one may not record', () => { + // Esc releases the chain lock before the interrupted run's `finally`, so + // the next turn starts while the interrupted one is still diffing. + const interrupted = beginUndoTurn('guards-raced') + const next = beginUndoTurn('guards-raced') + expect(isLatestUndoTurn('guards-raced', interrupted)).toBe(false) + expect(isLatestUndoTurn('guards-raced', next)).toBe(true) + }) + + test('stamps are per chat, so another chat cannot invalidate a turn', () => { + const mine = beginUndoTurn('guards-chat-a') + const other = beginUndoTurn('guards-chat-b') + expect(isLatestUndoTurn('guards-chat-a', mine)).toBe(true) + expect(isLatestUndoTurn('guards-chat-b', other)).toBe(true) + expect(isLatestUndoTurn('guards-chat-b', mine)).toBe(false) + }) +}) 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..207ab60098 --- /dev/null +++ b/cli/src/state/__tests__/undo-store.test.ts @@ -0,0 +1,270 @@ +import { afterAll, beforeAll, beforeEach, describe, expect, test } from 'bun:test' + +import { execFileSync } from 'node:child_process' +import { + existsSync, + mkdtempSync, + readdirSync, + readFileSync, + rmSync, + writeFileSync, +} from 'node:fs' +import os from 'node:os' +import path from 'node:path' + +import { getProjectDataDir, setCurrentChatId, setProjectRoot } from '../../project-files' +import { + anchorSnapshot, + listAnchors, + patchSnapshot, + releaseSnapshot, + setSnapshotDirOverrideForTesting, + trackSnapshot, +} from '../../utils/undo-snapshot' +import { + loadUndoState, + peekRedo, + peekUndo, + popRedo, + popUndo, + pushRedo, + pushUndo, + recordUndoEntry, + sweepAnchors, + undoToRecord, +} from '../undo-store' + +const CHAT_ID = 'undo-store-test-chat' + +let projectDir: string +let snapshotRoot: string + +beforeAll(() => { + projectDir = mkdtempSync(path.join(os.tmpdir(), 'undo-store-proj-')) + snapshotRoot = mkdtempSync(path.join(os.tmpdir(), 'undo-store-snapshots-')) + setProjectRoot(projectDir) + setCurrentChatId(CHAT_ID) + // The journal hands snapshots to undo-snapshot, which needs a git project. + setSnapshotDirOverrideForTesting(snapshotRoot) + execFileSync('git', ['init'], { cwd: projectDir, stdio: 'ignore' }) +}) + +// 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 }) + setSnapshotDirOverrideForTesting(undefined) + rmSync(snapshotRoot, { recursive: true, force: true }) +}) + +/** Run the cleanup job's gc now, instead of waiting out its 7-day grace. */ +const pruneSnapshotRepo = (): void => { + const [entry] = readdirSync(snapshotRoot) + if (!entry) return + execFileSync( + 'git', + [ + '--git-dir', + path.join(snapshotRoot, entry), + '--work-tree', + projectDir, + 'gc', + '--quiet', + '--prune=now', + ], + { stdio: ['ignore', 'pipe', 'ignore'] }, + ) +} + +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('a snapshot that is gone', () => { + test('leaves the worktree alone instead of deleting what it cannot restore', async () => { + const kept = path.join(projectDir, 'kept.txt') + writeFileSync(kept, 'the user wrote this\n') + const hash = await trackSnapshot(projectDir) + expect(hash).toBeTruthy() + + // The turn edits the file and the journal records what to revert. + writeFileSync(kept, 'the agent changed it\n') + recordUndoEntry(CHAT_ID, { + hashBefore: hash!, + files: ['kept.txt'], + message: 'a recorded turn', + }) + + // The snapshot is gone — pruned by the cleanup job, or a wiped config dir. + // `checkout` then fails, and "not in the snapshot" must not be read as + // "the turn created this file". + rmSync(snapshotRoot, { recursive: true, force: true }) + + const message = await undoToRecord(CHAT_ID, projectDir, peekUndo(CHAT_ID)!.id) + + expect(message).toBeNull() + expect(existsSync(kept)).toBe(true) + expect(readFileSync(kept, 'utf8')).toBe('the agent changed it\n') + }) +}) + +describe('snapshot anchors', () => { + test('a recorded turn stays restorable after the cleanup job runs', async () => { + const kept = path.join(projectDir, 'anchored.txt') + writeFileSync(kept, 'the user wrote this\n') + const hash = await trackSnapshot(projectDir) + expect(hash).toBeTruthy() + + // The turn edits the file, and the snapshot index moves on with it. + writeFileSync(kept, 'the agent changed it\n') + await patchSnapshot(projectDir, hash!) + recordUndoEntry(CHAT_ID, { + hashBefore: hash!, + files: ['anchored.txt'], + message: 'a recorded turn', + }) + pruneSnapshotRepo() + + const message = await undoToRecord(CHAT_ID, projectDir, peekUndo(CHAT_ID)!.id) + + expect(message).not.toBeNull() + expect(readFileSync(kept, 'utf8')).toBe('the user wrote this\n') + }) + + test('the sweep drops the anchors of a chat whose journal is gone', async () => { + writeFileSync(path.join(projectDir, 'swept.txt'), 'content\n') + const hash = await trackSnapshot(projectDir) + expect(anchorSnapshot(projectDir, 'abandoned-chat', hash!)).toBeTruthy() + expect(listAnchors(projectDir)).toContainEqual({ + key: 'abandoned-chat', + hash: hash!, + }) + + const { released } = sweepAnchors(projectDir) + + expect(released).toBeGreaterThan(0) + expect(listAnchors(projectDir)).not.toContainEqual({ + key: 'abandoned-chat', + hash: hash!, + }) + }) + + test('the sweep re-anchors an entry whose ref went missing', async () => { + writeFileSync(path.join(projectDir, 'reanchor.txt'), 'content\n') + const hash = await trackSnapshot(projectDir) + recordUndoEntry(CHAT_ID, { + hashBefore: hash!, + files: ['reanchor.txt'], + message: 'a recorded turn', + }) + // Pretend the release path ran early, or its write was lost. + expect(releaseSnapshot(projectDir, CHAT_ID, hash!)).toBe(true) + expect(listAnchors(projectDir)).not.toContainEqual({ + key: CHAT_ID, + hash: hash!, + }) + + sweepAnchors(projectDir) + + expect(listAnchors(projectDir)).toContainEqual({ key: CHAT_ID, hash: hash! }) + }) +}) + +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-guards.ts b/cli/src/state/undo-guards.ts new file mode 100644 index 0000000000..4477358a1b --- /dev/null +++ b/cli/src/state/undo-guards.ts @@ -0,0 +1,90 @@ +/** + * Guards shared by the undo picker and the send path. + * + * Both are about a jump-back being honest about what it takes with it: one + * tells the user how far back the action reaches before they commit to it, the + * other keeps a turn's bookkeeping from being written against changes that are + * not its own. + */ + +import type { UndoRecord } from './undo-store' + +/** The files of a set of entries, each one listed once. */ +const unionFiles = (entries: readonly UndoRecord[]): string[] => + Array.from(new Set(entries.flatMap((record) => record.files))) + +/** + * Every file an undo would touch: the selected turn's, plus every newer one's. + * + * The picker used to list only the selected entry's files while the action + * reverts the selected turn *and* everything after it, so the panel understated + * the very blast radius the notice above it was warning about. This is the same + * union `undoToRecord` reverts. + */ +export function filesTakenBack( + newestFirst: readonly UndoRecord[], + recordId: string, +): string[] { + const index = newestFirst.findIndex((record) => record.id === recordId) + if (index === -1) return [] + return unionFiles(newestFirst.slice(0, index + 1)) +} + +/** + * Every file a redo brings back: the turns its undo had reverted, when the + * record carries them. + */ +export function filesRestored(record: UndoRecord): string[] { + if (record.restored && record.restored.length > 0) { + return unionFiles(record.restored) + } + return unionFiles([record]) +} + +/** + * How many entries are newer than `recordId`. + * + * The picker lists entries newest first, so the position of the selected row + * *is* the count: everything above it is newer. The journal reverts the + * selected turn plus every newer one (and an undo of an undo invalidates the + * newer redo entries the same way), which is the part the picker used to leave + * unsaid. + */ +export function newerTurnCount( + newestFirst: readonly UndoRecord[], + recordId: string, +): number { + const index = newestFirst.findIndex((record) => record.id === recordId) + return index > 0 ? index : 0 +} + +/** + * The newest turn that started, per chat. + * + * A turn captures its snapshot before its run and writes its entry after it, + * in the `finally`. Those sit far apart in time and the run can be interrupted + * in between: on Esc the chain lock is released before that `finally` finishes, + * so the next turn can already be editing the worktree when the interrupted + * turn diffs its snapshot. That diff is computed against the worktree as it is + * *then*, so it would report files the newer turn changed, and `/undo` on that + * entry would revert work the user never selected. + * + * So a turn stamps itself when it starts and records only while its stamp is + * still the newest. Past that point its changes can no longer be told apart + * from the next turn's, and the turn is dropped instead of recorded wrong — + * best-effort, like the rest of the feature. + */ +const latestTurnStamp = new Map() +let stampSequence = 0 + +/** Stamp a turn as it starts. Returns what it must present in order to record. */ +export function beginUndoTurn(chatId: string): number { + stampSequence += 1 + latestTurnStamp.set(chatId, stampSequence) + return stampSequence +} + +/** Whether `stamp` is still the newest turn that started for this chat. */ +export function isLatestUndoTurn(chatId: string, stamp: number): boolean { + return latestTurnStamp.get(chatId) === stamp +} 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..a04b78b475 --- /dev/null +++ b/cli/src/state/undo-store.ts @@ -0,0 +1,406 @@ +/** + * 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, + readdirSync, + readFileSync, + writeFileSync, +} from 'node:fs' +import path from 'node:path' + +import { getProjectDataDir, tryGetProjectRoot } from '../project-files' +import { logger } from '../utils/logger' +import { + anchorSnapshot, + diffSnapshot, + isSnapshotAvailable, + listAnchors, + releaseSnapshot, + 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. + } +} + +/** Every hash the journal would need in order to restore from either stack. */ +function referencedHashes(state: UndoState): Set { + const hashes = new Set() + for (const record of [...state.undoStack, ...state.redoStack]) { + if (record.hashBefore) hashes.add(record.hashBefore) + if (record.hashAfter) hashes.add(record.hashAfter) + for (const reverted of record.restored ?? []) { + if (reverted.hashBefore) hashes.add(reverted.hashBefore) + if (reverted.hashAfter) hashes.add(reverted.hashAfter) + } + } + return hashes +} + +/** + * Drop the anchors of the snapshots this chat no longer lists. Snapshots are + * held by refs (see undo-snapshot), so releasing them is what keeps the store + * from growing for the life of the project. + * + * Synchronous, like the rest of the journal's bookkeeping: the file it just + * wrote must not list an entry whose snapshot is still held by a stale ref, + * nor hold a ref no entry lists. + */ +function releaseDroppedAnchors( + chatId: string, + before: Set, + after: Set, +): void { + const projectRoot = tryGetProjectRoot() + if (!projectRoot) return + for (const hash of before) { + if (after.has(hash)) continue + releaseSnapshot(projectRoot, chatId, hash) + } +} + +/** Chat ids that have a data directory on disk, journalled or not. */ +function listJournalChatIds(): string[] { + try { + return readdirSync(path.join(getProjectDataDir(), 'chats'), { + withFileTypes: true, + }) + .filter((entry) => entry.isDirectory()) + .map((entry) => entry.name) + } catch { + // No chats yet, or the directory is unreadable. + return [] + } +} + +/** + * Make the anchors match the journals: hold what a journal lists, drop what + * none lists, and re-anchor an entry whose ref went missing. + * + * The journals are the root set. Releasing on drop covers the entries this + * store hands out, but it runs per chat and cannot know about a chat that was + * deleted: its `undo.json` disappears and the refs its turns created stay + * behind. Sweeping against the journals collects those, and repairs an anchor + * or a release that failed, which is what keeps the store bounded. + */ +export function sweepAnchors(projectRoot: string): { + released: number + anchored: number +} { + const anchors = listAnchors(projectRoot) + const chatIds = new Set(anchors.map((anchor) => anchor.key)) + for (const chatId of listJournalChatIds()) chatIds.add(chatId) + + let released = 0 + let anchored = 0 + for (const chatId of chatIds) { + const needed = referencedHashes(loadUndoState(chatId)) + for (const anchor of anchors) { + if (anchor.key !== chatId || needed.has(anchor.hash)) continue + if (releaseSnapshot(projectRoot, chatId, anchor.hash)) released += 1 + } + for (const hash of needed) { + if (anchors.some((a) => a.key === chatId && a.hash === hash)) continue + if (anchorSnapshot(projectRoot, chatId, hash)) anchored += 1 + } + } + return { released, anchored } +} + +/** The project root the sweep has already run for, so it runs once each. */ +let sweptProjectRoot: string | null = null + +/** + * Sweep once per project root per process, after the current work is done so + * the turn's own path never waits on maintenance. + */ +function sweepAnchorsInBackground(): void { + const projectRoot = tryGetProjectRoot() + if (!projectRoot || sweptProjectRoot === projectRoot) return + sweptProjectRoot = projectRoot + setTimeout(() => { + try { + const { released, anchored } = sweepAnchors(projectRoot) + if (released > 0 || anchored > 0) { + logger.debug( + { released, anchored }, + 'undo-store: swept the snapshot anchors', + ) + } + } catch (error) { + logger.debug({ error }, 'undo-store: anchor sweep failed') + } + }, 0) +} + +/** + * 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) + const referenced = referencedHashes(state) + 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) + releaseDroppedAnchors(chatId, referenced, referencedHashes(state)) + // After the journal lists the turn, so the sweep can never see a snapshot + // no entry is holding yet. + const projectRoot = tryGetProjectRoot() + if (projectRoot) anchorSnapshot(projectRoot, chatId, entry.hashBefore) + sweepAnchorsInBackground() +} + +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) + const referenced = referencedHashes(state) + state.undoStack.push(record) + if (state.undoStack.length > MAX_UNDO_ENTRIES) { + state.undoStack.shift() + } + saveUndoState(chatId, state) + releaseDroppedAnchors(chatId, referenced, referencedHashes(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))) + + // A snapshot can be gone — collected by the cleanup job, or lost with the + // config directory. Abort before touching anything: reverting against a + // snapshot it cannot read is how undoing deletes the files it meant to keep. + if (!(await isSnapshotAvailable(projectRoot, record.hashBefore))) return null + + // 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 + // Held until the redo record that points at it is dropped. + anchorSnapshot(projectRoot, chatId, hashAfter) + + // 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. + const referenced = referencedHashes(state) + 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) + releaseDroppedAnchors(chatId, referenced, referencedHashes(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 a1b82bcd21..49e8f3377b 100644 --- a/cli/src/types/chat.ts +++ b/cli/src/types/chat.ts @@ -230,6 +230,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..03d43710fa --- /dev/null +++ b/cli/src/utils/__tests__/undo-snapshot.test.ts @@ -0,0 +1,257 @@ +import { afterAll, beforeAll, describe, expect, test } from 'bun:test' + +import { execFileSync } from 'node:child_process' +import { + chmodSync, + existsSync, + mkdtempSync, + readdirSync, + rmSync, + writeFileSync, +} from 'node:fs' +import os from 'node:os' +import path from 'node:path' + +import { + anchorSnapshot, + isSnapshotAvailable, + isUndoAvailable, + patchSnapshot, + releaseSnapshot, + restoreSnapshot, + revertFiles, + setSnapshotDirOverrideForTesting, + trackSnapshot, +} from '../undo-snapshot' + +let projectDir: string +let snapshotRoot: string +let plainDir: string + +/** The journal entry the snapshots belong to; anchors are keyed by it. */ +const SNAPSHOT_KEY = 'test-chat' + +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' }) + +/** The snapshot repo the fixture has initialized under the override root. */ +const snapshotRepoDir = (): string => { + const [entry] = readdirSync(snapshotRoot) + if (!entry) throw new Error('no snapshot repo was initialized') + return path.join(snapshotRoot, entry) +} + +/** Run the cleanup job's gc now, instead of waiting out its 7-day grace. */ +const pruneSnapshotRepo = (): void => { + execFileSync( + 'git', + [ + '--git-dir', + snapshotRepoDir(), + '--work-tree', + projectDir, + 'gc', + '--quiet', + '--prune=now', + ], + { stdio: ['ignore', 'pipe', 'ignore'] }, + ) +} + +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('a snapshot whose content is gone', () => { + test('leaves the worktree alone when a blob the tree lists is gone', async () => { + // Make sure the snapshot repo exists before reaching into it. + await trackSnapshot(projectDir) + + // A tree that references a blob which exists nowhere: the tree survives, + // its content does not. This is what borrowing the project's objects + // looks like once the project collects the object it was lending. + const missing = '0123456789abcdef0123456789abcdef01234567' + const tree = execFileSync( + 'git', + [ + '--git-dir', + snapshotRepoDir(), + '--work-tree', + projectDir, + 'mktree', + '--missing', + ], + { + cwd: projectDir, + encoding: 'utf8', + input: `100644 blob ${missing}\torphan.txt\n`, + }, + ).trim() + expect(tree).toBeTruthy() + + writeFileSync(path.join(projectDir, 'orphan.txt'), 'the agent wrote this\n') + + const { restored, deleted } = await revertFiles(projectDir, tree, [ + 'orphan.txt', + ]) + // `git checkout` deletes a worktree file whose blob it cannot read, so + // reverting against this tree would destroy the file instead of restoring + // it. Nothing may be restored, nothing deleted, and the file must survive. + expect(readProjectFile('orphan.txt')).toBe('the agent wrote this\n') + expect(restored).toEqual([]) + expect(deleted).toEqual([]) + expect(existsSync(path.join(projectDir, 'orphan.txt'))).toBe(true) + + // And it must not be advertised as restorable either. + expect(await isSnapshotAvailable(projectDir, tree)).toBe(false) + }) +}) + +describe('snapshot anchoring', () => { + test('git keeps the anchored snapshot and collects the unanchored one', async () => { + writeFileSync(path.join(projectDir, 'lifetime.txt'), 'before the turn\n') + const held = await trackSnapshot(projectDir) + expect(held).toBeTruthy() + expect(anchorSnapshot(projectDir, SNAPSHOT_KEY, held!)).toBeTruthy() + + // Two more turns move the snapshot index on, so nothing in git points at + // either tree any more. + writeFileSync(path.join(projectDir, 'lifetime.txt'), 'after the turn\n') + const unheld = await trackSnapshot(projectDir) + expect(unheld).toBeTruthy() + writeFileSync(path.join(projectDir, 'lifetime.txt'), 'and once more\n') + await trackSnapshot(projectDir) + pruneSnapshotRepo() + + // The anchored one survives; the one no journal holds is collected. + expect(await isSnapshotAvailable(projectDir, held!)).toBe(true) + expect(await isSnapshotAvailable(projectDir, unheld!)).toBe(false) + + // And the anchor is what holds it: drop it and git collects it too. + expect(releaseSnapshot(projectDir, SNAPSHOT_KEY, held!)).toBe(true) + pruneSnapshotRepo() + expect(await isSnapshotAvailable(projectDir, held!)).toBe(false) + expect(await restoreSnapshot(projectDir, held!)).toBe(false) + }) +}) + +describe('failed staging', () => { + test('a failed git add does not hand back a hash', async () => { + // A root user can read anything, so the fixture cannot fail the add there. + if (typeof process.getuid === 'function' && process.getuid() === 0) return + + // git refuses to index an unreadable file, but `write-tree` still returns + // the tree of the stale index: a hash that looks valid and is not. + const blocked = path.join(projectDir, 'blocked.txt') + writeFileSync(blocked, 'unreadable\n') + chmodSync(blocked, 0o000) + try { + expect(await trackSnapshot(projectDir)).toBeNull() + } finally { + chmodSync(blocked, 0o644) + rmSync(blocked, { force: true }) + } + }) +}) + +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 a19e4bc6a6..4be270ac5f 100644 --- a/cli/src/utils/settings.ts +++ b/cli/src/utils/settings.ts @@ -25,6 +25,7 @@ import type { ReasoningEffort } from '@codebuff/common/constants/reasoning-effor 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. @@ -69,6 +70,8 @@ export interface Settings { * moment it renders, not when it is dismissed, so "once" holds however * the launch ends. */ freebucksIntroSeenAt?: string + /** Whether /undo and /redo are enabled. Defaults to true. */ + undo?: boolean } /** @@ -227,6 +230,11 @@ const validateSettings = (parsed: unknown): Settings => { settings.freebucksIntroSeenAt = obj.freebucksIntroSeenAt } + // Validate undo toggle + if (typeof obj.undo === 'boolean') { + settings.undo = obj.undo + } + return settings } @@ -357,3 +365,15 @@ export const hasSeenFreebucksIntro = (): boolean => export const markFreebucksIntroSeen = (): void => { saveSettings({ freebucksIntroSeenAt: new Date().toISOString() }) } + +/** + * 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..6818c0ceec --- /dev/null +++ b/cli/src/utils/undo-snapshot.ts @@ -0,0 +1,751 @@ +/** + * 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 { execFileSync, 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}` + +/** + * Refs that keep a snapshot alive. + * + * `write-tree` returns a bare tree and nothing points at it, so git considers + * it unreachable as soon as the snapshot index moves on with the next turn -- + * and the cleanup job's `gc --prune` collects it while the journal still lists + * the entry. Pointing a ref at it is how git's own docs say to keep an object + * around, and what other snapshot tools do. + * + * The ref is keyed by the journal entry's owner (the chat) so two chats can + * hold the same tree without one of them releasing the other's copy. One ref + * per live entry: `releaseSnapshot` drops it, which is what bounds the store. + */ +const ANCHOR_REF_PREFIX = 'refs/freebuff/undo/' + +const anchorRefFor = (key: string, hash: string): string => + `${ANCHOR_REF_PREFIX}${key}/${hash}` + +/** + * Identity and signing for the anchor commit, set explicitly so the user's own + * global git config (missing `user.email`, `commit.gpgsign` on) cannot fail it. + */ +const ANCHOR_GIT_CONFIG = [ + '-c', + 'user.name=freebuff', + '-c', + 'user.email=freebuff@localhost', + '-c', + 'commit.gpgsign=false', +] + +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. + * @returns false when git refused to stage something, in which case the index + * no longer describes the worktree and must not be turned into a snapshot. + */ +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 true + + 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 true + const pathspecs = `${allowed.map(topLevelLiteral).join('\0')}\0` + const staged = await runGit( + [ + ...base, + 'add', + '--all', + '--sparse', + '--pathspec-from-file=-', + '--pathspec-file-nul', + ], + projectRoot, + { input: pathspecs }, + ) + if (staged.code !== 0) { + // git refused a path (unreadable file, bad name, ...). The index keeps + // the state it had before, and `write-tree` would happily return the + // tree of that stale index: a hash that looks valid and is not. + logger.debug( + { stderr: staged.stderr }, + 'undo-snapshot: git add failed, snapshot would be stale', + ) + return false + } + return true + } catch (error) { + logger.debug({ error }, 'undo-snapshot: failed to stage changes') + return false + } +} + +/** + * Run git and return its trimmed stdout, or null when it failed. + * + * The journal's anchoring is synchronous on purpose. `saveUndoState` already + * writes its file synchronously, and an entry must never be listed before its + * snapshot is anchored -- the sweep reads the anchors and relies on that. + * These calls only touch refs and objects, never the index, so they do not + * contend with the staged operations above, which stay on the async path. + */ +function gitSync(args: string[], projectRoot: string): string | null { + try { + return execFileSync('git', args, { + cwd: projectRoot, + encoding: 'utf8', + stdio: ['ignore', 'pipe', 'ignore'], + }).trim() + } catch (error) { + logger.debug({ error }, 'undo-snapshot: git command failed') + return null + } +} + +/** + * Keep a snapshot out of git's reach by pointing a ref at it, so the cleanup + * job's prune cannot collect a snapshot the journal still lists. + * @returns the ref name, or null when it could not be created. + */ +export function anchorSnapshot( + projectRoot: string, + key: string, + hash: string, +): string | null { + if (!hash) return null + if (!findGitRoot({ cwd: projectRoot })) return null + const snapshotDir = snapshotDirFor(projectRoot) + const base = ['--git-dir', snapshotDir, '--work-tree', projectRoot] + const commit = gitSync( + [ + ...ANCHOR_GIT_CONFIG, + ...base, + 'commit-tree', + hash, + '-m', + 'freebuff undo snapshot', + ], + projectRoot, + ) + if (!commit) return null + const ref = anchorRefFor(key, hash) + // `update-ref` prints nothing on success, so compare against null rather + // than truthiness: an empty string is the happy path here. + return gitSync([...base, 'update-ref', ref, commit], projectRoot) !== null + ? ref + : null +} + +/** + * Drop a snapshot's anchor so the cleanup job can reclaim it. The journal + * calls this when it stops listing an entry (eviction, or the redo stack a new + * turn clears). + */ +export function releaseSnapshot( + projectRoot: string, + key: string, + hash: string, +): boolean { + if (!hash) return false + if (!findGitRoot({ cwd: projectRoot })) return false + const snapshotDir = snapshotDirFor(projectRoot) + return ( + gitSync( + [ + '--git-dir', + snapshotDir, + '--work-tree', + projectRoot, + 'update-ref', + '-d', + anchorRefFor(key, hash), + ], + projectRoot, + ) !== null + ) +} + +/** + * The anchors the snapshot repo is holding, as `{ key, hash }` pairs. The + * journal sweeps its entries against this list. + */ +export function listAnchors( + projectRoot: string, +): { key: string; hash: string }[] { + if (!findGitRoot({ cwd: projectRoot })) return [] + const snapshotDir = snapshotDirFor(projectRoot) + const refs = gitSync( + [ + '--git-dir', + snapshotDir, + '--work-tree', + projectRoot, + 'for-each-ref', + '--format=%(refname)', + ANCHOR_REF_PREFIX, + ], + projectRoot, + ) + if (!refs) return [] + return refs + .split('\n') + .map((line) => line.trim()) + .filter((line) => line.startsWith(ANCHOR_REF_PREFIX)) + .map((ref) => { + const rest = ref.slice(ANCHOR_REF_PREFIX.length) + const separator = rest.indexOf('/') + if (separator <= 0) return null + const key = rest.slice(0, separator) + const hash = rest.slice(separator + 1) + return key && hash && !hash.includes('/') ? { key, hash } : null + }) + .filter((entry): entry is { key: string; hash: string } => entry !== null) +} + +/** + * Capture a snapshot of the project's current state. + * + * The returned tree is not referenced by anything, so it only survives while + * its caller anchors it (`anchorSnapshot`). The cleanup job's prune has a + * grace period of days, which leaves room to anchor after the turn instead of + * on this path. + * + * @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 + // A failed stage leaves the index stale, and `write-tree` would still + // return that stale tree. Refuse rather than hand back a hash that + // would later revert the worktree to the wrong state. + if (!(await stageChanges(snapshotDir, projectRoot, gitRoot))) return null + const result = await runGit( + [...gitArgs(snapshotDir, projectRoot, ['write-tree'])], + projectRoot, + ) + const hash = result.code === 0 ? result.stdout.trim() : '' + if (!hash) return null + void maybePrune(snapshotDir, projectRoot) + return hash + } 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 { + if (!(await stageChanges(snapshotDir, projectRoot, gitRoot))) return [] + 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 + } + }) +} + +/** + * Whether a snapshot can still be restored from. + * + * A stored hash is a tree the journal kept by hand, and it can be gone: the + * cleanup job collected it before this module started anchoring snapshots, or + * the snapshot directory was wiped while `undo.json` stayed behind. Callers + * must ask before reverting. + * + * The tree surviving is not enough on its own. This module borrows the + * project's object database, so the project's own cleanup can collect a blob + * a kept tree still lists. Restoring from such a tree is not merely partial: + * `git checkout` removes the worktree file whose content it cannot read, so a + * snapshot with holes deletes files instead of restoring them. Ask about the + * content too. + * + * Assumes the caller already holds the snapshot repo's lock. + */ +async function snapshotIsComplete( + snapshotDir: string, + projectRoot: string, + hash: string, +): Promise { + const base = gitArgs(snapshotDir, projectRoot, []) + const tree = await runGit([...base, 'cat-file', '-e', hash], projectRoot) + if (tree.code !== 0) return false + // One command reports every object the tree references and cannot find, + // marking each with a leading `?`. + const objects = await runGit( + [...base, 'rev-list', '--objects', '--missing=print', hash], + projectRoot, + ) + if (objects.code !== 0) return false + return !objects.stdout + .split('\n') + .some((line) => line.trimStart().startsWith('?')) +} + +export async function isSnapshotAvailable( + projectRoot: string, + hash: string, +): Promise { + if (!hash) return false + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return false + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + return await snapshotIsComplete(snapshotDir, projectRoot, hash) + } catch (error) { + logger.debug({ error }, 'undo-snapshot: availability check 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 { + // Never start reverting against content that is not all there: the + // checkout below deletes a file whose blob it cannot read, so a snapshot + // with holes would destroy the worktree instead of restoring it. + if (!(await snapshotIsComplete(snapshotDir, projectRoot, hash))) { + logger.debug( + { hash }, + 'undo-snapshot: snapshot incomplete, leaving the worktree alone', + ) + return result + } + 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) { + // The snapshot's tree itself is unreachable, so there is no way to + // tell whether this file was in it. Reading "we could not look" as + // "the turn created it" is how a missing snapshot deletes files the + // user had: leave the worktree alone instead. + logger.debug( + { file, hash }, + 'undo-snapshot: snapshot missing, leaving the file as it is', + ) + continue + } + if (tree.stdout.trim()) { + // It was in the snapshot, but bringing its content back did not work + // and the tree alone cannot say why (a collected object, most + // likely). Leave the file as it is and do not report a restore that + // did not happen: the caller's summary has to stay true. + logger.debug( + { file, hash }, + 'undo-snapshot: content unavailable, leaving the file as it is', + ) + continue + } + // Really was not 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 { + if (!(await stageChanges(snapshotDir, projectRoot, gitRoot))) return '' + const result = await runGit( + [ + ...gitArgs(snapshotDir, projectRoot, [ + 'diff', + '--cached', + '--no-ext-diff', + '--stat', + hash, + '--', + '.', + ]), + ], + projectRoot, + ) + return result.code === 0 ? result.stdout : '' + } catch { + return '' + } + }) +}