edit-undo

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { SnapshotMeta } from '../types'
5import {
6  MAX_FILE_BYTES,
7  addSkipped,
8  displayPath,
9  done,
10  evictions,
11  fileRows,
12  formatBytes,
13  isLossy,
14  listText,
15  parseArgs,
16  plan,
17  plural,
18  preview,
19  resolvePath,
20  rowLabel,
21  stackOf,
22  utf8Bytes,
23} from './snapshots'
24import type { Plan } from './snapshots'
25
26const PANE = 'edit-undo'
27const TITLE = 'Edit undo'
28const RM_MS = 10_000
29const QUEUE_MS = 2_000
30
31const index = atom({ plugin: 'edit-undo', key: 'index' } as const, [])
32const skipped = atom({ plugin: 'edit-undo', key: 'skipped' } as const, [])
33const confirm = atom({ plugin: 'edit-undo', key: 'confirm' } as const, null)
34const forceable = atom({ plugin: 'edit-undo', key: 'forceable' } as const, null)
35const armed = atom({ plugin: 'edit-undo', key: 'armed' } as const, null)
36const notice = atom({ plugin: 'edit-undo', key: 'notice' } as const, null)
37const before = { plugin: 'edit-undo', key: 'before' } as const
38
39type Captured = { kind: 'missing' } | { kind: 'text'; text: string } | { kind: 'skip'; reason: string }
40
41async function sha256(text: string): Promise<string> {
42  const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(text))
43
44  return [...new Uint8Array(digest)].map(b => b.toString(16).padStart(2, '0')).join('')
45}
46
47// The file as it stands: absent, its text, or a reason it is not snapshotted.
48async function capture($: EngineInterface, path: string): Promise<Captured> {
49  const stat = await $.fs.stat(path).catch(() => null)
50  if (stat === null) return (await $.fs.exists(path).catch(() => true)) ? { kind: 'skip', reason: 'could not be read' } : { kind: 'missing' }
51  if (stat.kind !== 'file') return { kind: 'skip', reason: 'not a regular file' }
52  if (stat.size > MAX_FILE_BYTES) return { kind: 'skip', reason: `over 1 MB (${formatBytes(stat.size)}), not snapshotted` }
53  const text = await $.fs.read(path).catch(() => null)
54  if (typeof text !== 'string') return { kind: 'skip', reason: 'could not be read' }
55  if (isLossy(text)) return { kind: 'skip', reason: 'not UTF-8 text, not snapshotted' }
56
57  return { kind: 'text', text }
58}
59
60// The hash of the file now, null when absent, undefined when it cannot be read.
61async function currentHash($: EngineInterface, path: string): Promise<string | null | undefined> {
62  if (!(await $.fs.exists(path).catch(() => true))) return null
63  const text = await $.fs.read(path).catch(() => null)
64
65  return typeof text === 'string' ? sha256(text) : undefined
66}
67
68async function skip($: EngineInterface, path: string, reason: string) {
69  const time = await $.clock.now()
70  await update($, skipped, list => addSkipped(list, { path, reason, time }))
71}
72
73// Runs once the tool's result went back: records the content the tool left and pushes the snapshot.
74async function record($: EngineInterface, path: string, tool: string, pre: { kind: 'missing' } | { kind: 'text'; text: string }) {
75  const after = await currentHash($, path)
76  if (after === undefined) return
77  const time = await $.clock.now()
78  const id = `${time.toString(36)}-${Math.random().toString(36).slice(2, 10)}`
79  const meta: SnapshotMeta = {
80    id,
81    path,
82    time,
83    tool,
84    existed: pre.kind === 'text',
85    bytes: pre.kind === 'text' ? utf8Bytes(pre.text) : 0,
86    afterHash: after,
87  }
88  try {
89    await $.state.set({ ...before, id }, pre.kind === 'text' ? pre.text : '')
90  } catch {
91    await skip($, path, 'too large to keep in session state, not snapshotted')
92    return
93  }
94  let dropped: string[] = []
95  await update($, index, list => {
96    const next = [...list, meta]
97    dropped = evictions(next)
98    return next.filter(one => !dropped.includes(one.id))
99  })
100  for (const gone of dropped) await $.state.set({ ...before, id: gone }, '').catch(() => undefined)
101}
102
103// Each record waits for the previous one, and a new tool call waits for both: the hash read after a call is that call's.
104let queue: Promise<void> = Promise.resolve()
105let busy = 0
106
107async function snapshotAround<T extends { deny?: unknown; isError?: boolean }>($: EngineInterface, tool: string, given: string, agentId: string | undefined, run: () => Promise<T>): Promise<T> {
108  if (agentId !== undefined || given === '') return run()
109  if (busy > 0) await Promise.race([queue, $.clock.sleep(QUEUE_MS).catch(() => undefined)])
110  const path = resolvePath(given, await $.session.cwd())
111  const pre = await capture($, path)
112  const ran = await run()
113  if (ran.deny !== undefined || ran.isError === true) return ran
114  if (pre.kind === 'skip') {
115    $.clock.after(0, () => void skip($, path, pre.reason).catch(() => undefined))
116    return ran
117  }
118  busy += 1
119  const previous = queue
120  queue = new Promise<void>(resolve => {
121    $.clock.after(0, () => {
122      void previous
123        .then(() => record($, path, tool, pre))
124        .catch(() => undefined)
125        .finally(() => {
126          busy -= 1
127          resolve()
128        })
129    })
130  })
131
132  return ran
133}
134
135async function planFor($: EngineInterface, path: string, force: boolean): Promise<{ p: Plan; left: number }> {
136  const stack = stackOf(await read($, index), path)
137  const hash = await currentHash($, path)
138  if (hash === undefined) return { p: { refuse: 'the file cannot be read now', isForceable: false }, left: 0 }
139
140  return { p: plan(stack[stack.length - 1], hash, force), left: Math.max(0, stack.length - 1) }
141}
142
143// Restores (or deletes) the file from its latest snapshot, then pops it. Only ever called from a command or a button.
144// `expected` is the snapshot the preview described: a newer edit since then is not undone on that confirm.
145async function undo($: EngineInterface, path: string, force: boolean, expected: string): Promise<{ text: string; isDone: boolean; isStale?: boolean }> {
146  const cwd = await $.session.cwd()
147  const { p } = await planFor($, path, force)
148  if ('refuse' in p) return { text: done(p, path, cwd), isDone: false }
149  if (p.snapshot.id !== expected) return { text: `Not undone: ${displayPath(path, cwd)} was edited again since the preview.`, isDone: false, isStale: true }
150  try {
151    if (p.action === 'restore') {
152      const { value } = await $.state.get({ ...before, id: p.snapshot.id })
153      if (typeof value !== 'string') return { text: `Cannot undo ${displayPath(path, cwd)}: the snapshot is gone.`, isDone: false }
154      await $.fs.write(path, value)
155    } else {
156      const stat = await $.fs.stat(path)
157      if (stat.kind !== 'file') return { text: `Cannot undo ${displayPath(path, cwd)}: not a regular file.`, isDone: false }
158      await $.process.run(['rm', '--', path], { timeoutMs: RM_MS })
159      if (await $.fs.exists(path).catch(() => true)) return { text: `Cannot undo ${displayPath(path, cwd)}: the file could not be deleted.`, isDone: false }
160    }
161  } catch {
162    return { text: `Cannot undo ${displayPath(path, cwd)}: the file could not be written.`, isDone: false }
163  }
164  await update($, index, list => list.filter(one => one.id !== p.snapshot.id))
165  await $.state.set({ ...before, id: p.snapshot.id }, '').catch(() => undefined)
166  const text = done(p, path, cwd)
167  $.ui.toast(text)
168
169  return { text, isDone: true }
170}
171
172async function resetConfirm($: EngineInterface, message: string | null) {
173  await update($, confirm, () => null)
174  await update($, forceable, () => null)
175  await update($, notice, () => message)
176}
177
178// Two-step: the first press shows what will happen and arms the button, the second does it.
179async function press($: EngineInterface, path: string, force: boolean) {
180  const key = `${force ? 'force' : 'undo'}:${path}`
181  const pending = await read($, confirm)
182  let lead = ''
183  if (pending !== null && pending.key === key) {
184    const out = await undo($, path, force, pending.snapshot)
185    if (out.isStale !== true) {
186      await resetConfirm($, out.text)
187      return
188    }
189    lead = `${out.text} `
190  }
191  const cwd = await $.session.cwd()
192  const { p, left } = await planFor($, path, force)
193  const text = lead + preview(p, path, cwd, await $.clock.now(), left)
194  if ('refuse' in p) {
195    await update($, confirm, () => null)
196    await update($, forceable, () => (p.isForceable ? path : null))
197  } else {
198    await update($, confirm, () => ({ key, snapshot: p.snapshot.id }))
199    if (!force) await update($, forceable, () => null)
200  }
201  await update($, notice, () => ('refuse' in p && p.isForceable ? `${text} "Force undo" restores it anyway.` : text))
202}
203
204async function fileCommand($: EngineInterface, given: string, isConfirm: boolean, force: boolean): Promise<string> {
205  const cwd = await $.session.cwd()
206  const path = resolvePath(given, cwd)
207  const key = `${force ? 'force' : 'undo'}:${path}`
208  let lead = ''
209  if (isConfirm) {
210    const pending = await read($, armed)
211    if (pending === null || pending.key !== key) {
212      return `Run /undo ${given}${force ? ' force' : ''} first to see what it will do, then add "confirm".`
213    }
214    await update($, armed, () => null)
215    const out = await undo($, path, force, pending.snapshot)
216    if (out.isStale !== true) return out.text
217    lead = `${out.text} `
218  }
219  const { p, left } = await planFor($, path, force)
220  const text = lead + preview(p, path, cwd, await $.clock.now(), left)
221  if ('refuse' in p) {
222    await update($, armed, () => null)
223    return p.isForceable ? `${text} Run /undo ${given} force to restore it anyway.` : text
224  }
225  await update($, armed, () => ({ key, snapshot: p.snapshot.id }))
226
227  return `${text} Run /undo ${given}${force ? ' force' : ''} confirm to do it.`
228}
229
230export const register: Register = on => {
231  on('session.start', async ($, e, next) => {
232    await $.command.register({
233      name: 'undo',
234      description: 'Undo Claude’s file edits of this session: /undo (pane) | /undo <path> [force] [confirm] | /undo list | /undo clear',
235      argumentHint: '[<path> [force] [confirm] | list | clear | close]',
236    })
237
238    return next(e)
239  })
240
241  on('tool.call', { tool: 'Edit' }, ($, e, next) => snapshotAround($, 'Edit', e.file_path, e.agentId, () => next(e)))
242  on('tool.call', { tool: 'Write' }, ($, e, next) => snapshotAround($, 'Write', e.file_path, e.agentId, () => next(e)))
243  on('tool.call', { tool: 'NotebookEdit' }, ($, e, next) => snapshotAround($, 'NotebookEdit', e.notebook_path, e.agentId, () => next(e)))
244
245  on('command.run', { command: 'undo' }, async ($, e) => {
246    const args = parseArgs(e.args)
247    const cwd = await $.session.cwd()
248    if (args.kind === 'close') {
249      await $.ui.close({ id: PANE })
250      return { text: 'Undo pane closed.' }
251    }
252    if (args.kind === 'list') return { text: listText(await read($, index), await read($, skipped), cwd, await $.clock.now()) }
253    if (args.kind === 'clear') {
254      const list = await read($, index)
255      await update($, index, () => [])
256      await update($, skipped, () => [])
257      await update($, armed, () => null)
258      await resetConfirm($, null)
259      for (const one of list) await $.state.set({ ...before, id: one.id }, '').catch(() => undefined)
260      return { text: `Cleared ${plural(list.length, 'snapshot')}.` }
261    }
262    if (args.kind === 'file') return { text: await fileCommand($, args.path, args.confirm, args.force) }
263    await resetConfirm($, null)
264    await $.ui.open({ id: PANE, title: TITLE })
265    const count = fileRows(await read($, index)).length
266
267    return { text: `Undo pane opened (${plural(count, 'file')}).` }
268  })
269
270  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
271    const { Box, Button, Text } = $.ui.resolve(e)
272    const rows = fileRows(await read($, index))
273    const skips = await read($, skipped)
274    const pending = await read($, confirm)
275    const canForce = await read($, forceable)
276    const message = await read($, notice)
277    const cwd = await $.session.cwd()
278    const now = await $.clock.now()
279
280    return (
281      <Box flexDirection="column">
282        <Text bold>{rows.length === 0 ? 'No edits to undo yet.' : plural(rows.length, 'file')}</Text>
283        {message !== null && <Text>{message}</Text>}
284        {rows.map(row => (
285          <Box flexDirection="column">
286            <Text>{rowLabel(row, cwd, now)}</Text>
287            <Button key={`undo:${row.path}`} label={pending?.key === `undo:${row.path}` ? 'Confirm undo?' : 'Undo'} onPress={() => press($, row.path, false)} />
288            {canForce === row.path && (
289              <Button key={`force:${row.path}`} label={pending?.key === `force:${row.path}` ? 'Confirm force undo?' : 'Force undo'} onPress={() => press($, row.path, true)} />
290            )}
291          </Box>
292        ))}
293        {skips.map(one => (
294          <Text dimColor>{`${displayPath(one.path, cwd)}: ${one.reason}`}</Text>
295        ))}
296      </Box>
297    )
298  })
299}
300

1import type { SkippedFile, SnapshotMeta } from '../types'
2
3export const MAX_PER_FILE = 20
4export const MAX_FILE_BYTES = 1024 * 1024
5export const MAX_TOTAL_BYTES = 50 * 1024 * 1024
6export const MAX_SNAPSHOTS = 2000
7export const MAX_SKIPPED = 50
8
9export type Limits = { perFile: number; totalBytes: number; count: number }
10export const LIMITS: Limits = { perFile: MAX_PER_FILE, totalBytes: MAX_TOTAL_BYTES, count: MAX_SNAPSHOTS }
11
12const encoder = new TextEncoder()
13
14export const utf8Bytes = (text: string) => encoder.encode(text).length
15
16// A text read of a file that is not UTF-8 replaces bytes with U+FFFD: writing it back would damage the file.
17export const isLossy = (text: string) => text.includes('�')
18
19// An absolute path with `.`, `..` and repeated slashes resolved; relative paths are taken under `cwd`.
20export function resolvePath(path: string, cwd: string): string {
21  const joined = path.startsWith('/') ? path : `${cwd.replace(/\/+$/, '')}/${path}`
22  const parts: string[] = []
23  for (const part of joined.split('/')) {
24    if (part === '' || part === '.') continue
25    if (part === '..') parts.pop()
26    else parts.push(part)
27  }
28
29  return `/${parts.join('/')}`
30}
31
32// The path relative to `cwd` when it is inside it, otherwise as is.
33export function displayPath(path: string, cwd: string): string {
34  const root = cwd.replace(/\/+$/, '')
35
36  return root !== '' && path.startsWith(`${root}/`) ? path.slice(root.length + 1) : path
37}
38
39// Which snapshots to drop so the index fits: per file, then in count and size, oldest first.
40export function evictions(index: readonly SnapshotMeta[], limits: Limits = LIMITS): string[] {
41  const byAge = [...index].sort((a, b) => a.time - b.time)
42  const drop = new Set<string>()
43  const perPath = new Map<string, SnapshotMeta[]>()
44  for (const one of byAge) perPath.set(one.path, [...(perPath.get(one.path) ?? []), one])
45  for (const list of perPath.values()) for (const one of list.slice(0, Math.max(0, list.length - limits.perFile))) drop.add(one.id)
46  let count = index.length - drop.size
47  let bytes = index.reduce((sum, one) => sum + (drop.has(one.id) ? 0 : one.bytes), 0)
48  for (const one of byAge) {
49    if (count <= limits.count && bytes <= limits.totalBytes) break
50    if (drop.has(one.id)) continue
51    drop.add(one.id)
52    count -= 1
53    bytes -= one.bytes
54  }
55
56  return [...drop]
57}
58
59// The snapshots of one path, oldest first; the last is the one an undo restores.
60export const stackOf = (index: readonly SnapshotMeta[], path: string) =>
61  index.filter(one => one.path === path).sort((a, b) => a.time - b.time)
62
63export type FileRow = { path: string; count: number; last: SnapshotMeta }
64
65// One row per file, most recently changed first.
66export function fileRows(index: readonly SnapshotMeta[]): FileRow[] {
67  const rows = new Map<string, FileRow>()
68  for (const one of index) {
69    const row = rows.get(one.path)
70    if (row === undefined) rows.set(one.path, { path: one.path, count: 1, last: one })
71    else rows.set(one.path, { path: one.path, count: row.count + 1, last: one.time >= row.last.time ? one : row.last })
72  }
73
74  return [...rows.values()].sort((a, b) => b.last.time - a.last.time)
75}
76
77export function addSkipped(list: readonly SkippedFile[], entry: SkippedFile): SkippedFile[] {
78  return [...list.filter(one => one.path !== entry.path), entry].slice(-MAX_SKIPPED)
79}
80
81export type Plan = { action: 'restore' | 'delete'; snapshot: SnapshotMeta } | { refuse: string; isForceable: boolean }
82
83// What undoing `top` does given the file's current hash (null: the file is absent now).
84export function plan(top: SnapshotMeta | undefined, currentHash: string | null, force: boolean): Plan {
85  if (top === undefined) return { refuse: 'no snapshot to undo', isForceable: false }
86  const isChanged = currentHash !== top.afterHash
87  if (!top.existed) {
88    return isChanged
89      ? { refuse: 'the file did not exist before and has changed since Claude wrote it: it is not deleted', isForceable: false }
90      : { action: 'delete', snapshot: top }
91  }
92  if (isChanged && !force) return { refuse: 'the file has changed since Claude last edited it', isForceable: true }
93
94  return { action: 'restore', snapshot: top }
95}
96
97export type UndoArgs =
98  | { kind: 'pane' }
99  | { kind: 'close' }
100  | { kind: 'list' }
101  | { kind: 'clear' }
102  | { kind: 'file'; path: string; confirm: boolean; force: boolean }
103
104// `/undo`, `/undo list|clear|close`, `/undo <path> [confirm] [force]` (the flags read off the end, in any order).
105export function parseArgs(raw: string): UndoArgs {
106  const text = raw.trim()
107  if (text === '') return { kind: 'pane' }
108  if (text === 'list' || text === 'clear' || text === 'close') return { kind: text }
109  let rest = text
110  let confirm = false
111  let force = false
112  for (;;) {
113    const flag = /\s+(confirm|force)$/.exec(rest)
114    if (flag === null || (flag[1] === 'confirm' ? confirm : force)) break
115    if (flag[1] === 'confirm') confirm = true
116    else force = true
117    rest = rest.slice(0, flag.index)
118  }
119
120  return { kind: 'file', path: unquote(rest), confirm, force }
121}
122
123const unquote = (text: string) => (/^(['"]).*\1$/.test(text) ? text.slice(1, -1) : text)
124
125export function age(ms: number): string {
126  const s = Math.max(0, Math.round(ms / 1000))
127  if (s < 60) return `${s}s ago`
128  if (s < 3600) return `${Math.floor(s / 60)}m ago`
129  if (s < 86_400) return `${Math.floor(s / 3600)}h ago`
130
131  return `${Math.floor(s / 86_400)}d ago`
132}
133
134export function formatBytes(bytes: number): string {
135  if (bytes < 1024) return `${bytes} B`
136  if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`
137
138  return `${(bytes / (1024 * 1024)).toFixed(1)} MB`
139}
140
141export const plural = (n: number, word: string) => `${n} ${word}${n === 1 ? '' : 's'}`
142
143export function rowLabel(row: FileRow, cwd: string, now: number): string {
144  return `${displayPath(row.path, cwd)} · ${plural(row.count, 'snapshot')} · last ${row.last.tool} ${age(now - row.last.time)}`
145}
146
147// What an undo will do, said before it is confirmed.
148export function preview(p: Plan, path: string, cwd: string, now: number, left: number): string {
149  const shown = displayPath(path, cwd)
150  if ('refuse' in p) return `Cannot undo ${shown}: ${p.refuse}.`
151  const when = `${p.snapshot.tool} ${age(now - p.snapshot.time)}`
152
153  return p.action === 'delete'
154    ? `Undo will delete ${shown}: it did not exist before ${when}.`
155    : `Undo will restore ${shown} to its content before ${when} (${plural(left, 'snapshot')} left after).`
156}
157
158export function done(p: Plan, path: string, cwd: string): string {
159  const shown = displayPath(path, cwd)
160  if ('refuse' in p) return `Cannot undo ${shown}: ${p.refuse}.`
161
162  return p.action === 'delete' ? `Undo: deleted ${shown}` : `Undo: restored ${shown} (before ${p.snapshot.tool})`
163}
164
165export function listText(index: readonly SnapshotMeta[], skipped: readonly SkippedFile[], cwd: string, now: number): string {
166  const rows = fileRows(index)
167  const lines = rows.length === 0 ? ['No snapshots in this session.'] : [`${plural(rows.length, 'file')}, ${plural(index.length, 'snapshot')}, ${formatBytes(index.reduce((s, one) => s + one.bytes, 0))}:`]
168  for (const row of rows) lines.push(`  ${rowLabel(row, cwd, now)}`)
169  if (skipped.length > 0) {
170    lines.push('Not snapshotted:')
171    for (const one of skipped) lines.push(`  ${displayPath(one.path, cwd)}: ${one.reason}`)
172  }
173
174  return lines.join('\n')
175}
176

1// One snapshot of a file taken before a main-agent Edit/Write/NotebookEdit; its content is kept apart, in `before`.
2export type SnapshotMeta = {
3  id: string
4  path: string
5  time: number
6  tool: string
7  existed: boolean // false: the file did not exist before the tool call
8  bytes: number // UTF-8 size of the content before
9  afterHash: string | null // SHA-256 of the content the tool left; null when the file was absent after
10}
11
12export type SkippedFile = { path: string; reason: string; time: number }
13
14// An armed two-step undo: which action on which file, and the snapshot its preview described.
15export type Armed = { key: string; snapshot: string }
16
17declare module 'claude-code' {
18  interface PluginState {
19    'edit-undo': {
20      index: SnapshotMeta[]
21      before: StateFamily<string>
22      skipped: SkippedFile[]
23      confirm: Armed | null
24      forceable: string | null
25      armed: Armed | null
26      notice: string | null
27    }
28  }
29}
30