keep-or-undo

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, ToolCallResult } from 'claude-code'
3
4import type { ChangeEntry, ChangeTool, FileEntry, PendingConfirm, Registry, TurnMark } from '../types'
5import { countPatch, diffText, formatPatch, turnRange } from './diff'
6import type { Hunk } from './diff'
7
8type Engine = EngineInterface
9
10const REGISTRY = { plugin: 'keep-or-undo', key: 'registry' } as const
11const TURN = { plugin: 'keep-or-undo', key: 'turn' } as const
12const CONFIRM = { plugin: 'keep-or-undo', key: 'confirm' } as const
13const EXPANDED = { plugin: 'keep-or-undo', key: 'expanded' } as const
14const OWNER = { plugin: 'keep-or-undo', key: 'owner' } as const
15
16const EMPTY: Registry = { files: {}, order: [] }
17const registry = atom(REGISTRY, EMPTY)
18const turnMark = atom(TURN, null)
19const pendingConfirm = atom(CONFIRM, null)
20const expanded = atom(EXPANDED, null)
21const owner = atom(OWNER, null)
22
23const PANE = 'keep-or-undo'
24const PANE_TITLE = "Claude's changes"
25const DAY_MS = 24 * 60 * 60 * 1000
26
27// Every registry operation goes through this queue, so parallel recordings,
28// Keep and Undo never overlap. A module variable: after a reload it starts
29// empty, which is fine.
30let queue: Promise<unknown> = Promise.resolve()
31function serial<T>(work: () => Promise<T>): Promise<T> {
32  const run = queue.then(work, work)
33  queue = run.catch(() => undefined)
34  return run
35}
36
37// Only a cache for `autoOpenPane`: turn.complete has no viewport, the band
38// does.
39let isFullscreen = false
40
41function debug($: Engine, text: string): void {
42  try {
43    $.ui.log(`keep-or-undo: ${text}`, { to: 'debug' })
44  } catch {
45    // logging must never break anything
46  }
47}
48
49function errorText(err: unknown): string {
50  return err instanceof Error ? err.message : String(err)
51}
52
53// ---------------------------------------------------------------------------
54// Paths and files
55
56async function configDir($: Engine): Promise<string> {
57  const config = await $.env.get('CLAUDE_CONFIG_DIR')
58  const home = await $.env.get('HOME')
59  const root = config !== undefined && config !== '' ? config : `${home ?? ''}/.claude`
60  return root.replace(/\/+$/, '')
61}
62
63async function baseDir($: Engine): Promise<string> {
64  return `${await configDir($)}/keep-or-undo`
65}
66
67function safeName(text: string): string {
68  return text.replace(/[^A-Za-z0-9_.-]/g, '_')
69}
70
71async function sessionDir($: Engine, sessionId?: string): Promise<string> {
72  return `${await baseDir($)}/${safeName(sessionId ?? (await $.session.id()))}`
73}
74
75function blobPath(dir: string, ref: string): string {
76  return `${dir}/blobs/${ref}`
77}
78
79/** The file's content: null if it does not exist, undefined if it cannot be read (over 4 MiB). */
80async function readText($: Engine, path: string): Promise<string | null | undefined> {
81  try {
82    if (!(await $.fs.exists(path))) return null
83  } catch {
84    return undefined
85  }
86  try {
87    return await $.fs.read(path)
88  } catch {
89    return undefined
90  }
91}
92
93async function readBlob($: Engine, dir: string, ref: string): Promise<string | undefined> {
94  try {
95    return await $.fs.read(blobPath(dir, ref))
96  } catch {
97    return undefined
98  }
99}
100
101async function realPathOf($: Engine, path: string): Promise<string | undefined> {
102  try {
103    return (await $.fs.stat(path, { resolve: true })).realPath
104  } catch {
105    return undefined
106  }
107}
108
109async function rootOf($: Engine): Promise<string> {
110  const root = await $.session.root()
111  return ((await realPathOf($, root)) ?? root).replace(/\/+$/, '')
112}
113
114function isInside(path: string, dir: string): boolean {
115  return path.startsWith(`${dir}/`) && !path.split('/').includes('..')
116}
117
118function displayPathOf(realPath: string, root: string): string {
119  return isInside(realPath, root) ? realPath.slice(root.length + 1) : realPath
120}
121
122/** Deletes files and folders, only those under `under`. */
123async function removeFiles($: Engine, paths: readonly string[], under: string): Promise<void> {
124  const safe = paths.filter(p => isInside(p, under))
125  if (safe.length === 0) return
126  await $.process.run(['rm', '-rf', '--', ...safe])
127}
128
129/** Deletes a data folder, only if it is under the mod's folder. */
130async function removeDir($: Engine, dir: string): Promise<void> {
131  if (!isInside(dir, await baseDir($))) return
132  await $.process.run(['rm', '-rf', '--', dir])
133}
134
135function refsOf(file: FileEntry): string[] {
136  const refs = new Set<string>()
137  if (file.baseRef !== null) refs.add(file.baseRef)
138  for (const entry of file.entries) {
139    for (const ref of [entry.preRef, entry.postRef, entry.patchRef]) if (ref !== null) refs.add(ref)
140  }
141  return [...refs]
142}
143
144/** The folder of the copy Diff in IDE shows on the left. */
145function ideCopyDir(dir: string, file: FileEntry): string {
146  return `${dir}/ide/${safeName(file.displayPath)}`
147}
148
149async function dropBlobs($: Engine, files: readonly FileEntry[], dir?: string): Promise<void> {
150  if (files.length === 0) return
151  dir ??= await sessionDir($)
152  await removeFiles($, [...files.flatMap(refsOf).map(ref => blobPath(dir, ref)), ...files.map(f => ideCopyDir(dir, f))], dir)
153}
154
155// ---------------------------------------------------------------------------
156// Index: in $.state, copied to $.store and to index.json
157
158type Snapshot = { registry: Registry; updatedAt: number }
159
160function isRegistry(value: unknown): value is Registry {
161  const v = value as Registry | undefined
162  return typeof v === 'object' && v !== null && typeof v.files === 'object' && Array.isArray(v.order)
163}
164
165async function persist($: Engine, reg: Registry): Promise<void> {
166  const sid = await $.session.id()
167  const snapshot: Snapshot = { registry: reg, updatedAt: await $.clock.now() }
168  try {
169    await $.store.set(`sessions/${sid}`, snapshot)
170  } catch (err) {
171    debug($, `store not updated: ${errorText(err)}`)
172  }
173  try {
174    await $.fs.write(`${await sessionDir($, sid)}/index.json`, JSON.stringify(snapshot))
175  } catch (err) {
176    debug($, `index.json not updated: ${errorText(err)}`)
177  }
178}
179
180async function loadSnapshot($: Engine, sid: string): Promise<Registry | undefined> {
181  const stored = (await $.store.get(`sessions/${sid}`)) as Snapshot | undefined
182  if (isRegistry(stored?.registry)) return stored.registry
183  try {
184    const saved = JSON.parse(await $.fs.read(`${await sessionDir($, sid)}/index.json`)) as Snapshot
185    if (isRegistry(saved.registry)) return saved.registry
186  } catch {
187    // no copy on disk
188  }
189  return undefined
190}
191
192/**
193 * The registry in $.state belongs to the current session. After a --resume in
194 * a new process the state is empty; after a /resume or a /clear it belongs to
195 * another session: either way it is reloaded from the store or index.json.
196 * After a mod reload it matches and this does nothing. Call it inside serial.
197 */
198async function ensureOwner($: Engine): Promise<void> {
199  const sid = await $.session.id()
200  const held = await $.state.get(OWNER)
201  if (held.value === sid) return
202  const loaded = await loadSnapshot($, sid)
203  await $.state.set(REGISTRY, loaded ?? EMPTY)
204  await $.state.set(CONFIRM, null)
205  await $.state.set(EXPANDED, null)
206  await $.state.set(OWNER, sid)
207}
208
209/** Does the state belong to the current session? For renders, which never write. */
210async function isOwned($: Engine): Promise<boolean> {
211  return (await read($, owner)) === (await $.session.id())
212}
213
214async function cleanup($: Engine, retentionDays: number): Promise<void> {
215  const cutoff = (await $.clock.now()) - retentionDays * DAY_MS
216  const current = await $.session.id()
217  for (const key of await $.store.keys()) {
218    if (!key.startsWith('sessions/')) continue
219    const sid = key.slice('sessions/'.length)
220    if (sid === current) continue
221    const snapshot = (await $.store.get(key)) as Snapshot | undefined
222    if (typeof snapshot?.updatedAt === 'number' && snapshot.updatedAt >= cutoff) continue
223    await $.store.delete(key)
224    await removeDir($, await sessionDir($, sid))
225  }
226  const base = await baseDir($)
227  if (!(await $.fs.exists(base))) return
228  for (const entry of await $.fs.list(base)) {
229    if (entry.kind !== 'dir' || entry.name === safeName(current)) continue
230    const dir = `${base}/${entry.name}`
231    const index = `${dir}/index.json`
232    const stat = await $.fs.stat((await $.fs.exists(index)) ? index : dir).catch(() => undefined)
233    if (stat !== undefined && stat.mtimeMs < cutoff) await removeDir($, dir)
234  }
235}
236
237function without(reg: Registry, paths: readonly string[]): Registry {
238  const files = { ...reg.files }
239  for (const path of paths) delete files[path]
240  return { files, order: reg.order.filter(p => !paths.includes(p)) }
241}
242
243// ---------------------------------------------------------------------------
244// Recording
245
246type Observed = {
247  tool: ChangeTool
248  path: string
249  toolUseId: string
250  agentId: string | undefined
251}
252
253async function agentLabel($: Engine, agentId: string | undefined): Promise<string | undefined> {
254  if (agentId === undefined) return undefined
255  try {
256    const agent = (await $.agent.list()).find(a => a.id === agentId)
257    if (agent !== undefined) return agent.type || agent.name || agent.description
258  } catch {
259    // fallback below
260  }
261  return agentId.slice(0, 8)
262}
263
264async function record($: Engine, call: Observed, pre: string | null | undefined, r: ToolCallResult): Promise<void> {
265  if (r.deny !== undefined || r.isError === true) return
266  const result = r.result as Record<string, unknown> | undefined
267  if (result === undefined || result === null || result.staged === true) return
268  // NotebookEdit can report a failure in `error` without isError.
269  if (typeof result.error === 'string' && result.error !== '') return
270  await ensureOwner($)
271
272  let before: string | null | undefined
273  let created: boolean
274  let patch: Hunk[] = []
275  if (call.tool === 'NotebookEdit') {
276    // The notebook as it was on disk: original_file is what the tool read.
277    before = pre !== undefined ? pre : typeof result.original_file === 'string' ? result.original_file : undefined
278    created = pre === null
279    const oldSource = typeof result.old_source === 'string' ? result.old_source : ''
280    const newSource = result.edit_mode === 'delete' ? '' : typeof result.new_source === 'string' ? result.new_source : ''
281    patch = diffText(oldSource, newSource)
282  } else {
283    before = typeof result.originalFile === 'string' ? result.originalFile : pre
284    created = call.tool === 'Write' ? result.type === 'create' : pre === null
285    if (Array.isArray(result.structuredPatch)) patch = result.structuredPatch as Hunk[]
286  }
287  if (created) before = null
288
289  // Re-read the content after the change: replaceAll and userModified can alter it.
290  const after = await readText($, call.path)
291  // Nothing changed on disk: nothing to record.
292  if (typeof before === 'string' && before === after) return
293  if (patch.length === 0 && typeof after === 'string' && before !== undefined) {
294    patch = diffText(before ?? '', after)
295  }
296  const counted = countPatch(patch)
297
298  const realPath = (await realPathOf($, call.path)) ?? call.path
299  const root = await rootOf($)
300  const dir = await sessionDir($)
301  const id = safeName(call.toolUseId)
302
303  const preRef = typeof before === 'string' ? `${id}.pre.txt` : null
304  const postRef = typeof after === 'string' ? `${id}.post.txt` : null
305  const patchRef = patch.length > 0 ? `${id}.patch.txt` : null
306  if (preRef !== null) await $.fs.write(blobPath(dir, preRef), before as string)
307  if (postRef !== null) await $.fs.write(blobPath(dir, postRef), after as string)
308  if (patchRef !== null) await $.fs.write(blobPath(dir, patchRef), formatPatch(patch))
309
310  // Requirement 6: if the content before this change is not what the last
311  // recorded change left, someone touched the file in between.
312  const existing = (await read($, registry)).files[realPath]
313  let isExternal = false
314  const last = existing?.entries[existing.entries.length - 1]
315  if (last !== undefined && before !== undefined) {
316    const lastPost = last.postRef === null ? undefined : await readBlob($, dir, last.postRef)
317    if (lastPost !== undefined) isExternal = lastPost !== before
318  }
319
320  const turn = await read($, turnMark)
321  const entry: ChangeEntry = {
322    id: call.toolUseId,
323    turnId: turn?.id ?? null,
324    turnNumber: turn?.number ?? null,
325    tool: call.tool,
326    createdFile: created,
327    preRef,
328    postRef,
329    patchRef,
330    added: counted.added,
331    removed: counted.removed,
332    at: await $.clock.now(),
333  }
334  const agent = await agentLabel($, call.agentId)
335  if (agent !== undefined) entry.agent = agent
336
337  const next = await update($, registry, reg => {
338    const current = reg.files[realPath]
339    if (current === undefined) {
340      const file: FileEntry = {
341        realPath,
342        displayPath: displayPathOf(realPath, root),
343        baseRef: preRef,
344        createdByClaude: created,
345        entries: [entry],
346        absorbedExternal: false,
347        tooLarge: before === undefined || after === undefined,
348      }
349      return { files: { ...reg.files, [realPath]: file }, order: [...reg.order.filter(p => p !== realPath), realPath] }
350    }
351    // baseRef stays: Undo goes back to before the FIRST pending change.
352    const file: FileEntry = {
353      ...current,
354      entries: [...current.entries, entry],
355      absorbedExternal: current.absorbedExternal || isExternal,
356      tooLarge: current.tooLarge || after === undefined,
357    }
358    return { ...reg, files: { ...reg.files, [realPath]: file } }
359  })
360  await update($, turnMark, t => (t === null ? t : { ...t, added: t.added + 1 }))
361  await persist($, next)
362}
363
364/** Observes an editing tool: never blocks it and always returns its result. */
365async function observe($: Engine, call: Observed, run: () => Promise<ToolCallResult>): Promise<ToolCallResult> {
366  let pre: string | null | undefined
367  try {
368    pre = await readText($, call.path)
369  } catch {
370    pre = undefined
371  }
372  const r = await run()
373  try {
374    await serial(() => record($, call, pre, r))
375  } catch (err) {
376    debug($, `recording failed for ${call.path}: ${errorText(err)}`)
377  }
378  return r
379}
380
381// ---------------------------------------------------------------------------
382// Keep, Undo, consistency
383
384async function accept($: Engine, paths: readonly string[]): Promise<void> {
385  await serial(async () => {
386    await ensureOwner($)
387    const reg = await read($, registry)
388    const files = paths.map(p => reg.files[p]).filter((f): f is FileEntry => f !== undefined)
389    const next = await update($, registry, r => without(r, paths))
390    await update($, pendingConfirm, c => (c !== null && c.paths.some(p => paths.includes(p)) ? null : c))
391    await dropBlobs($, files)
392    await persist($, next)
393  })
394}
395
396type Restored = { file: FileEntry; isDeleted: boolean }
397
398function modelMessage(restored: readonly Restored[], isAll: boolean): string {
399  const turnsOf = (file: FileEntry) => {
400    const range = turnRange(file.entries.map(e => e.turnNumber))
401    return range === '' ? '' : range.includes('–') || range.includes(',') ? `in turns ${range}` : `in turn ${range}`
402  }
403  const only = restored[0]
404  if (restored.length === 1 && only !== undefined && !isAll) {
405    const turns = turnsOf(only.file)
406    if (only.isDeleted) {
407      return `The user rejected the creation of \`${only.file.displayPath}\`${turns ? ` (changes ${turns})` : ''}: the file has been deleted. Do not recreate it unless asked.`
408    }
409    return `The user rejected your changes to \`${only.file.displayPath}\`: the file is back to how it was before your changes${turns ? ` ${turns}` : ''}. Re-read it before editing it again.`
410  }
411  const lines = restored.map(({ file, isDeleted }) => {
412    const turns = turnsOf(file)
413    const note = isDeleted ? 'created by you, now deleted' : 'restored'
414    return `- \`${file.displayPath}\` (${note}${turns ? `; changes ${turns}` : ''})`
415  })
416  return [
417    `The user rejected your changes to ${restored.length} files, which are back to how they were before your changes:`,
418    ...lines,
419    'Re-read them before editing them again.',
420  ].join('\n')
421}
422
423type RejectOptions = { isAll: boolean; isConfirmed: boolean }
424
425async function reject($: Engine, paths: readonly string[], options: RejectOptions): Promise<void> {
426  await serial(async () => {
427    await ensureOwner($)
428    const reg = await read($, registry)
429    const dir = await sessionDir($)
430    const root = await rootOf($)
431    const files = paths.map(p => reg.files[p]).filter((f): f is FileEntry => f !== undefined)
432
433    type Plan = { file: FileEntry; current: string | null; base: string | null; isLaterExternal: boolean }
434    const plans: Plan[] = []
435    const settled: FileEntry[] = [] // already at baseRef: just drop them
436    const skipped: string[] = []
437    for (const file of files) {
438      if (file.tooLarge) {
439        skipped.push(`${file.displayPath}: too large, undo not available`)
440        continue
441      }
442      const current = await readText($, file.realPath)
443      const base = file.baseRef === null ? null : await readBlob($, dir, file.baseRef)
444      if (current === undefined || base === undefined) {
445        skipped.push(`${file.displayPath}: content not readable`)
446        continue
447      }
448      if (current === base) {
449        settled.push(file)
450        continue
451      }
452      const last = file.entries[file.entries.length - 1]
453      const lastPost = last?.postRef == null ? undefined : await readBlob($, dir, last.postRef)
454      plans.push({ file, current, base, isLaterExternal: lastPost === undefined || lastPost !== current })
455    }
456
457    if (!options.isConfirmed && plans.some(p => p.isLaterExternal)) {
458      const confirm: PendingConfirm = { paths: [...paths], isAll: options.isAll }
459      await $.state.set(CONFIRM, confirm)
460      return
461    }
462
463    const restored: Restored[] = []
464    const backups: string[] = []
465    const stamp = await $.clock.now()
466    for (const plan of plans) {
467      const { file, current, base } = plan
468      try {
469        if (plan.isLaterExternal && current !== null) {
470          // The index keeps names apart when two paths share a safeName (src/a.ts, src_a.ts).
471          const name = `backup-${stamp}-${backups.length}-${safeName(file.displayPath)}.txt`
472          await $.fs.write(blobPath(dir, name), current)
473          backups.push(blobPath(dir, name))
474        }
475        if (file.createdByClaude) {
476          if (current !== null) {
477            const real = await realPathOf($, file.realPath)
478            if (real !== file.realPath || !isInside(real, root)) {
479              skipped.push(`${file.displayPath}: outside the session folder, not deleting it`)
480              continue
481            }
482            await $.process.run(['rm', '--', real])
483          }
484          restored.push({ file, isDeleted: true })
485        } else {
486          await $.fs.write(file.realPath, base ?? '')
487          restored.push({ file, isDeleted: false })
488        }
489      } catch (err) {
490        skipped.push(`${file.displayPath}: ${errorText(err)}`)
491      }
492    }
493
494    const done = [...settled, ...restored.map(r => r.file)]
495    const donePaths = done.map(f => f.realPath)
496    const next = await update($, registry, r => without(r, donePaths))
497    await $.state.set(CONFIRM, null)
498    await dropBlobs($, done)
499    await persist($, next)
500
501    if (restored.length > 0) {
502      try {
503        await $.session.append({
504          message: { type: 'user', content: [{ type: 'text', text: modelMessage(restored, options.isAll) }] },
505        })
506      } catch (err) {
507        debug($, `could not notify the model: ${errorText(err)}`)
508      }
509    }
510    const notes: string[] = []
511    if (done.length > 0) notes.push(done.length === 1 ? `Changes to ${done[0]!.displayPath} undone` : `Changes to ${done.length} files undone`)
512    if (backups.length > 0) notes.push(`backup of the overwritten content: ${backups.join(', ')}`)
513    if (skipped.length > 0) notes.push(`undo not applied to ${skipped.join('; ')}`)
514    if (notes.length > 0) $.ui.toast(notes.join(' · '), { timeoutMs: backups.length > 0 || skipped.length > 0 ? 12000 : 4000 })
515  })
516}
517
518/** Drops files that are back at baseRef, for example after a native /rewind. */
519async function reconcile($: Engine): Promise<void> {
520  await serial(async () => {
521    await ensureOwner($)
522    const reg = await read($, registry)
523    if (reg.order.length === 0) return
524    const dir = await sessionDir($)
525    const settled: FileEntry[] = []
526    for (const path of reg.order) {
527      const file = reg.files[path]
528      if (file === undefined || file.tooLarge) continue
529      const base = file.baseRef === null ? null : await readBlob($, dir, file.baseRef)
530      if (base === undefined) continue
531      const current = await readText($, file.realPath)
532      if (current === base) settled.push(file)
533    }
534    if (settled.length === 0) return
535    const next = await update($, registry, r => without(r, settled.map(f => f.realPath)))
536    await dropBlobs($, settled)
537    await persist($, next)
538  })
539}
540
541// ---------------------------------------------------------------------------
542// Diff in the IDE. The IDE's own approval tab (openDiff on the `ide` server)
543// is not reachable from a mod: this opens the IDE's diff viewer instead.
544
545/** Is an IDE connected (/ide)? Its MCP server's tools are then listed. */
546async function hasIde($: Engine): Promise<boolean> {
547  try {
548    return (await $.tool.list()).some(t => t.name.startsWith('mcp__ide__'))
549  } catch {
550    return false
551  }
552}
553
554type IdeLock = { pid?: number; ideName?: string; workspaceFolders?: string[] }
555
556/**
557 * The connected IDE's executable, from the lockfiles Claude Code reads in
558 * <config>/ide/: the one of CLAUDE_CODE_SSE_PORT (set in the IDE's terminal),
559 * else the one whose workspace holds the session, the deepest first.
560 */
561async function ideExecutable($: Engine): Promise<string | undefined> {
562  const dir = `${await configDir($)}/ide`
563  const parse = async (name: string) => {
564    try {
565      return JSON.parse(await $.fs.read(`${dir}/${name}`)) as IdeLock
566    } catch {
567      return undefined
568    }
569  }
570  const port = await $.env.get('CLAUDE_CODE_SSE_PORT')
571  let lock = port ? await parse(`${port}.lock`) : undefined
572  if (lock === undefined) {
573    const root = await rootOf($)
574    let depth = -1
575    const names = (await $.fs.exists(dir)) ? (await $.fs.list(dir)).map(e => e.name).filter(n => n.endsWith('.lock')) : []
576    for (const name of names) {
577      const candidate = await parse(name)
578      for (const folder of candidate?.workspaceFolders ?? []) {
579        const f = folder.replace(/\/+$/, '')
580        if ((root === f || isInside(root, f)) && f.length > depth) {
581          depth = f.length
582          lock = candidate
583        }
584      }
585    }
586  }
587  // shortcut: JetBrains only, as its executable forwards `diff a b` to the running instance; VS Code's is Electron's and would need `code --diff`.
588  if (typeof lock?.pid !== 'number' || /code|cursor|windsurf/i.test(lock.ideName ?? '')) return undefined
589  // shortcut: macOS ps prints the full path, Linux only the name; upgrade to /proc/<pid>/exe if Linux is needed.
590  const ps = await $.process.run(['ps', '-o', 'comm=', '-p', String(lock.pid)])
591  const exe = ps.stdout.trim()
592  return ps.exitCode === 0 && exe.startsWith('/') ? exe : undefined
593}
594
595/**
596 * Runs the connected IDE's executable with the arguments `argsOf` gives, or
597 * the reason it cannot. Where it does not run, the row's diff opens in the
598 * pane instead, as a press on the row does without an IDE.
599 */
600async function inIde($: Engine, path: string, what: string, argsOf: (file: FileEntry) => Promise<string[] | string>): Promise<void> {
601  const file = (await read($, registry)).files[path]
602  if (file === undefined) return
603  const fallback = async (reason: string) => {
604    debug($, `${what} in IDE for ${file.displayPath}: ${reason}`)
605    $.ui.toast(`${what} not opened in the IDE (${reason}): diff shown here instead`, { timeoutMs: 6000 })
606    await update($, expanded, () => path)
607  }
608  try {
609    const exe = await ideExecutable($)
610    if (exe === undefined) return await fallback('JetBrains IDEs on macOS only')
611    const args = await argsOf(file)
612    if (typeof args === 'string') return await fallback(args)
613    const run = await $.process.run([exe, ...args])
614    if (run.exitCode !== 0) await fallback(run.stderr.trim() || `exit code ${run.exitCode}`)
615  } catch (err) {
616    await fallback(errorText(err))
617  }
618}
619
620/** Opens base → current in the IDE's diff viewer; the right side is the file itself, so edits there are saved to it. */
621function diffInIde($: Engine, path: string): Promise<void> {
622  return inIde($, path, 'Diff', async file => {
623    const dir = await sessionDir($)
624    const base = file.baseRef === null ? '' : await readBlob($, dir, file.baseRef)
625    if (base === undefined) return 'content before the changes not readable'
626    // A copy under the file's own name: the IDE titles and highlights the left side by it.
627    const left = `${ideCopyDir(dir, file)}/${file.realPath.split('/').pop()}`
628    await $.fs.write(left, base)
629    return ['diff', left, file.realPath]
630  })
631}
632
633/** Opens the file itself in the IDE. */
634function openInIde($: Engine, path: string): Promise<void> {
635  return inIde($, path, 'File', async file => [file.realPath])
636}
637
638async function openPane($: Engine): Promise<void> {
639  try {
640    await reconcile($)
641  } catch (err) {
642    debug($, `consistency check failed: ${errorText(err)}`)
643  }
644  await $.ui.open({ id: PANE, title: PANE_TITLE })
645}
646
647// ---------------------------------------------------------------------------
648// Summary
649
650function totals(reg: Registry): { files: number; added: number; removed: number } {
651  let added = 0
652  let removed = 0
653  for (const path of reg.order) {
654    for (const entry of reg.files[path]?.entries ?? []) {
655      added += entry.added
656      removed += entry.removed
657    }
658  }
659  return { files: reg.order.length, added, removed }
660}
661
662function filesText(count: number): string {
663  return count === 1 ? '1 file' : `${count} files`
664}
665
666function badgesOf(file: FileEntry): string[] {
667  const badges: string[] = []
668  const agents = [...new Set(file.entries.map(e => e.agent).filter((a): a is string => a !== undefined))]
669  for (const agent of agents) badges.push(`subagent: ${agent}`)
670  if (file.createdByClaude) badges.push('new')
671  if (file.absorbedExternal) badges.push('includes external changes')
672  if (file.tooLarge) badges.push('too large, undo not available')
673  return badges
674}
675
676type TreeRow = { kind: 'dir'; key: string; label: string; depth: number } | { kind: 'file'; file: FileEntry; label: string; depth: number }
677
678/** The files as a tree, folders first, alphabetical; a folder holding just one folder shares its row, as IDEs compact them. */
679function treeRows(files: readonly FileEntry[]): TreeRow[] {
680  type Node = { dirs: Map<string, Node>; files: FileEntry[] }
681  const top: Node = { dirs: new Map(), files: [] }
682  for (const file of files) {
683    let at = top
684    for (const part of file.displayPath.split('/').slice(0, -1)) {
685      let next = at.dirs.get(part)
686      if (next === undefined) at.dirs.set(part, (next = { dirs: new Map(), files: [] }))
687      at = next
688    }
689    at.files.push(file)
690  }
691  const rows: TreeRow[] = []
692  const walk = (at: Node, depth: number, prefix: string) => {
693    for (const [name, child] of [...at.dirs].sort(([a], [b]) => a.localeCompare(b))) {
694      let label = name
695      let dir = child
696      for (let only = [...dir.dirs][0]; dir.files.length === 0 && dir.dirs.size === 1 && only !== undefined; only = [...dir.dirs][0]) {
697        label += `/${only[0]}`
698        dir = only[1]
699      }
700      const key = `${prefix}${label}`
701      // An absolute path (a file outside the session folder) starts with an empty name.
702      rows.push({ kind: 'dir', key, label: label === '' ? '/' : label, depth })
703      walk(dir, depth + 1, `${key}/`)
704    }
705    for (const file of [...at.files].sort((a, b) => a.displayPath.localeCompare(b.displayPath))) {
706      rows.push({ kind: 'file', file, label: file.displayPath.split('/').pop() ?? file.displayPath, depth })
707    }
708  }
709  walk(top, 0, '')
710  return rows
711}
712
713type Elements = ReturnType<Engine['ui']['resolve']>
714
715/** A `[ label ]` button in a color of its own: a Button has no color, its Text children do. */
716function actionButton(el: Elements, key: string, label: string, color: string, onPress: () => unknown) {
717  const { Button, Text } = el
718  return (
719    <Button key={key} label={label} plain onPress={() => void onPress()}>
720      <Text color={color} bold>{`[ ${label} ]`}</Text>
721    </Button>
722  )
723}
724
725function confirmText(confirm: PendingConfirm, reg: Registry): string {
726  const only = confirm.paths.length === 1 ? reg.files[confirm.paths[0]!] : undefined
727  return only !== undefined
728    ? `${only.displayPath}: the file has later changes not made by Claude, which will be lost.`
729    : 'Some files have later changes not made by Claude, which will be lost.'
730}
731
732async function confirmReject($: Engine): Promise<void> {
733  const confirm = await read($, pendingConfirm)
734  if (confirm === null) return
735  await reject($, confirm.paths, { isAll: confirm.isAll, isConfirmed: true })
736}
737
738export const register: Register = (on, options) => {
739  const retentionDays = typeof options.retentionDays === 'number' && options.retentionDays > 0 ? options.retentionDays : 14
740  const autoOpenPane = options.autoOpenPane === true
741
742  // --- Recording: never block the tool -------------------------------------
743  on('tool.call', { tool: 'Edit' }, ($, e, next) =>
744    observe($, { tool: 'Edit', path: e.file_path, toolUseId: e.tool_use_id, agentId: e.agentId }, () => next(e)),
745  ).catch(($, e, next) => next(e))
746
747  on('tool.call', { tool: 'Write' }, ($, e, next) =>
748    observe($, { tool: 'Write', path: e.file_path, toolUseId: e.tool_use_id, agentId: e.agentId }, () => next(e)),
749  ).catch(($, e, next) => next(e))
750
751  on('tool.call', { tool: 'NotebookEdit' }, ($, e, next) =>
752    observe($, { tool: 'NotebookEdit', path: e.notebook_path, toolUseId: e.tool_use_id, agentId: e.agentId }, () => next(e)),
753  ).catch(($, e, next) => next(e))
754
755  // --- Lifecycle -------------------------------------------------------------
756  on('session.start', async ($, e, next) => {
757    try {
758      await $.command.register({ name: 'changes', description: "Review Claude's file changes: Keep or Undo" })
759    } catch (err) {
760      debug($, `/changes command not registered: ${errorText(err)}`)
761    }
762    try {
763      await serial(() => ensureOwner($))
764    } catch (err) {
765      debug($, `index not reloaded: ${errorText(err)}`)
766    }
767    // Does not block the first prompt.
768    void cleanup($, retentionDays).catch(err => debug($, `cleanup failed: ${errorText(err)}`))
769    return next(e)
770  })
771
772  // /clear keeps everything: empty registry, copies deleted, files untouched.
773  on('session.end', async ($, e, next) => {
774    if (e.reason === 'clear') {
775      try {
776        await serial(async () => {
777          // Saved copies go; undo backups stay until the periodic cleanup.
778          const dir = await sessionDir($, e.sessionId)
779          const reg = await read($, registry)
780          await dropBlobs($, Object.values(reg.files), dir)
781          await removeFiles($, [`${dir}/index.json`], dir)
782          await $.store.delete(`sessions/${e.sessionId}`)
783          await $.state.set(REGISTRY, EMPTY)
784          await $.state.set(CONFIRM, null)
785          await $.state.set(EXPANDED, null)
786          await $.state.set(TURN, null)
787        })
788      } catch (err) {
789        debug($, `cleanup after /clear failed: ${errorText(err)}`)
790      }
791    }
792    return next(e)
793  })
794
795  on('turn.start', async ($, e, next) => {
796    try {
797      const number = await $.session.turns().catch(() => 0)
798      const mark: TurnMark = { id: e.turnId, number, added: 0 }
799      await $.state.set(TURN, mark)
800      await reconcile($)
801    } catch (err) {
802      debug($, `turn.start: ${errorText(err)}`)
803    }
804    return next(e)
805  })
806
807  on('turn.complete', async ($, e, next) => {
808    if (e.agentId === undefined) {
809      try {
810        const turn = await read($, turnMark)
811        if (turn !== null && turn.id === e.turnId && turn.added > 0) {
812          const sum = totals(await read($, registry))
813          if (sum.files > 0) {
814            $.ui.toast(`${filesText(sum.files)} changed by Claude pending (+${sum.added} −${sum.removed}) · /changes to review`)
815            if (autoOpenPane && isFullscreen) await $.ui.open({ id: PANE, title: PANE_TITLE })
816          }
817        }
818      } catch (err) {
819        debug($, `turn.complete: ${errorText(err)}`)
820      }
821    }
822    return next(e)
823  })
824
825  on('command.run', { command: 'changes' }, async $ => {
826    await openPane($)
827    const sum = totals(await read($, registry))
828    return {
829      text: sum.files === 0 ? 'No pending changes from Claude.' : `${filesText(sum.files)} changed by Claude pending review.`,
830    }
831  })
832
833  // --- Band above the prompt -------------------------------------------------
834  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
835    if (e.viewport?.isFullscreen !== undefined) isFullscreen = e.viewport.isFullscreen
836    if (e.props.hasSurvey || !(await isOwned($))) return next(e)
837    const reg = await read($, registry)
838    const confirm = await read($, pendingConfirm)
839    if (reg.order.length === 0 && confirm === null) return next(e)
840
841    const { Box, Text, Button } = $.ui.resolve(e)
842    const sum = totals(reg)
843    return (
844      <Box key="keep-or-undo" flexDirection="column">
845        {confirm !== null && (
846          <Box key="confirm" flexDirection="row" gap={1}>
847            <Text color="warning" wrap="truncate-end">
848              {confirmText(confirm, reg)}
849            </Text>
850            <Button key="confirm-reject" label="Undo anyway" variant="primary" onPress={() => confirmReject($)} />
851            <Button key="confirm-cancel" label="Cancel" onPress={() => update($, pendingConfirm, () => null)} />
852          </Box>
853        )}
854        {sum.files > 0 && (
855          <Box key="summary" flexDirection="row" gap={1}>
856            <Text>{`${filesText(sum.files)} changed by Claude ·`}</Text>
857            <Text color="diffAdded">+{sum.added}</Text>
858            <Text color="diffRemoved">−{sum.removed}</Text>
859            <Button key="review" label="Review" onPress={() => openPane($)} />
860            <Button key="accept-all" label="Keep all" onPress={async () => accept($, (await read($, registry)).order)} />
861            <Button
862              key="reject-all"
863              label="Undo all"
864              onPress={async () => reject($, (await read($, registry)).order, { isAll: true, isConfirmed: false })}
865            />
866          </Box>
867        )}
868      </Box>
869    )
870  })
871
872  // --- /changes pane ---------------------------------------------------------
873  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
874    const el = $.ui.resolve(e)
875    const { Box, Text, Button, Code } = el
876    const isCurrent = await isOwned($)
877    const reg = isCurrent ? await read($, registry) : EMPTY
878    const confirm = isCurrent ? await read($, pendingConfirm) : null
879    const open = await read($, expanded)
880    const dir = await sessionDir($)
881    const ide = reg.order.length > 0 && (await hasIde($))
882
883    const rows = []
884    for (const row of treeRows(reg.order.map(p => reg.files[p]).filter((f): f is FileEntry => f !== undefined))) {
885      if (row.kind === 'dir') {
886        rows.push(
887          <Box key={`dir:${row.key}`} marginLeft={row.depth * 2}>
888            <Text>
889              <Text dimColor>● </Text>
890              <Text bold>{`${row.label}/`}</Text>
891            </Text>
892          </Box>,
893        )
894        continue
895      }
896      const { file } = row
897      const path = file.realPath
898      const added = file.entries.reduce((n, x) => n + x.added, 0)
899      const removed = file.entries.reduce((n, x) => n + x.removed, 0)
900      const badges = badgesOf(file)
901      const isOpen = open === path
902
903      const sections = []
904      if (isOpen) {
905        for (const entry of file.entries) {
906          const head = [entry.tool, entry.turnNumber !== null ? `turn ${entry.turnNumber}` : undefined, entry.agent ? `subagent ${entry.agent}` : undefined]
907            .filter(Boolean)
908            .join(' · ')
909          const source = entry.patchRef === null ? undefined : await readBlob($, dir, entry.patchRef)
910          sections.push(
911            <Box key={`entry:${entry.id}`} flexDirection="column" marginLeft={2}>
912              <Text dimColor>{`${head} · +${entry.added} −${entry.removed}`}</Text>
913              {source !== undefined && source !== '' ? (
914                <Code source={source} format="diff" path={file.displayPath} />
915              ) : (
916                <Text dimColor>no diff available</Text>
917              )}
918            </Box>,
919          )
920        }
921        if (file.absorbedExternal) {
922          sections.push(
923            <Box key={`external:${path}`} marginLeft={2}>
924              <Text color="warning">This file includes absorbed external changes: Undo reverts those too.</Text>
925            </Box>,
926          )
927        }
928      }
929
930      rows.push(
931        <Box key={`file:${file.displayPath}`} flexDirection="column" marginLeft={row.depth * 2}>
932          <Button
933            key={`open:${file.displayPath}`}
934            plain
935            onPress={() => (ide && !isOpen ? openInIde($, path) : update($, expanded, current => (current === path ? null : path)))}
936          >
937            {/* The name in the link color: a press on it opens the file in the IDE, or the diff here without one (or when that fails). */}
938            <Text color="suggestion">{`${ide && !isOpen ? '→' : isOpen ? '▾' : '▸'} ${row.label}`}</Text> <Text color="diffAdded">+{added}</Text> <Text color="diffRemoved">−{removed}</Text>
939            {badges.length > 0 ? <Text dimColor> [{badges.join('] [')}]</Text> : ''}
940          </Button>
941          {/* Keep and Undo on a row of their own: a long path or many badges must not push them past a narrow pane's edge. */}
942          <Box flexDirection="row" gap={1} marginLeft={2}>
943            {actionButton(el, `accept:${file.displayPath}`, 'Keep', 'success', () => accept($, [path]))}
944            {actionButton(el, `reject:${file.displayPath}`, 'Undo', 'error', () =>
945              reject($, [path], { isAll: false, isConfirmed: false }),
946            )}
947            {ide && actionButton(el, `ide:${file.displayPath}`, 'Diff in IDE', 'suggestion', () => diffInIde($, path))}
948          </Box>
949          {sections}
950        </Box>,
951      )
952    }
953
954    const sum = totals(reg)
955    return (
956      <Box key="keep-or-undo-pane" flexDirection="column" paddingLeft={1}>
957        {confirm !== null && (
958          <Box key="confirm" flexDirection="column">
959            <Text color="warning">{confirmText(confirm, reg)}</Text>
960            <Box flexDirection="row" gap={1}>
961              <Button key="confirm-reject" label="Undo anyway" variant="primary" onPress={() => confirmReject($)} />
962              <Button key="confirm-cancel" label="Cancel" onPress={() => update($, pendingConfirm, () => null)} />
963            </Box>
964          </Box>
965        )}
966        {sum.files === 0 ? (
967          <Text dimColor>No pending changes from Claude.</Text>
968        ) : (
969          <Box key="totals" flexDirection="row" gap={1}>
970            <Text>{`${filesText(sum.files)} ·`}</Text>
971            <Text color="diffAdded">+{sum.added}</Text>
972            <Text color="diffRemoved">−{sum.removed}</Text>
973            {actionButton(el, 'accept-all', 'Keep all', 'success', async () => accept($, (await read($, registry)).order))}
974            {actionButton(el, 'reject-all', 'Undo all', 'error', async () =>
975              reject($, (await read($, registry)).order, { isAll: true, isConfirmed: false }),
976            )}
977          </Box>
978        )}
979        {rows}
980      </Box>
981    )
982  })
983}
984

1// Diffs and counts, with no dependencies: the module has no Node.
2
3export type Hunk = {
4  oldStart: number
5  oldLines: number
6  newStart: number
7  newLines: number
8  lines: readonly string[]
9}
10
11export type Counted = { added: number; removed: number }
12
13/** Sums the '+' and '-' lines of the engine's patches. */
14export function countPatch(patch: readonly Hunk[]): Counted {
15  let added = 0
16  let removed = 0
17  for (const hunk of patch) {
18    for (const line of hunk.lines) {
19      if (line.startsWith('+')) added += 1
20      else if (line.startsWith('-')) removed += 1
21    }
22  }
23  return { added, removed }
24}
25
26/** The patches as unified text, which `Code format="diff"` reads. */
27export function formatPatch(patch: readonly Hunk[]): string {
28  return patch
29    .map(h => [`@@ -${h.oldStart},${h.oldLines} +${h.newStart},${h.newLines} @@`, ...h.lines].join('\n'))
30    .join('\n')
31}
32
33function splitLines(text: string): string[] {
34  if (text === '') return []
35  const lines = text.split('\n')
36  if (lines[lines.length - 1] === '') lines.pop()
37  return lines
38}
39
40// Above this threshold (lines × lines of the changed span) no LCS:
41// the whole span counts as removed and then added.
42const LCS_LIMIT = 4_000_000
43
44/**
45 * A line diff between two texts, as a single hunk with 3 lines of context.
46 * Used where the engine gives no patch: NotebookEdit, a Write that creates.
47 */
48export function diffText(before: string, after: string): Hunk[] {
49  const a = splitLines(before)
50  const b = splitLines(after)
51  let head = 0
52  while (head < a.length && head < b.length && a[head] === b[head]) head += 1
53  let tail = 0
54  while (tail < a.length - head && tail < b.length - head && a[a.length - 1 - tail] === b[b.length - 1 - tail]) tail += 1
55  const midA = a.slice(head, a.length - tail)
56  const midB = b.slice(head, b.length - tail)
57  if (midA.length === 0 && midB.length === 0) return []
58
59  const ops: string[] = []
60  if (midA.length * midB.length <= LCS_LIMIT) {
61    const n = midA.length
62    const m = midB.length
63    // One flat buffer, row i at i * width: no object per row.
64    const width = m + 1
65    const table = new Uint32Array((n + 1) * width)
66    for (let i = n - 1; i >= 0; i -= 1) {
67      const row = i * width
68      const below = row + width
69      for (let j = m - 1; j >= 0; j -= 1) {
70        table[row + j] = midA[i] === midB[j] ? table[below + j + 1]! + 1 : Math.max(table[below + j]!, table[row + j + 1]!)
71      }
72    }
73    let i = 0
74    let j = 0
75    while (i < n && j < m) {
76      if (midA[i] === midB[j]) {
77        ops.push(` ${midA[i]}`)
78        i += 1
79        j += 1
80      } else if (table[(i + 1) * width + j]! >= table[i * width + j + 1]!) {
81        ops.push(`-${midA[i]}`)
82        i += 1
83      } else {
84        ops.push(`+${midB[j]}`)
85        j += 1
86      }
87    }
88    while (i < n) ops.push(`-${midA[i++]}`)
89    while (j < m) ops.push(`+${midB[j++]}`)
90  } else {
91    for (const line of midA) ops.push(`-${line}`)
92    for (const line of midB) ops.push(`+${line}`)
93  }
94
95  const before3 = a.slice(Math.max(0, head - 3), head).map(l => ` ${l}`)
96  const after3 = a.slice(a.length - tail, Math.min(a.length, a.length - tail + 3)).map(l => ` ${l}`)
97  const lines = [...before3, ...ops, ...after3]
98  const start = head - before3.length
99  const oldLines = lines.filter(l => !l.startsWith('+')).length
100  const newLines = lines.filter(l => !l.startsWith('-')).length
101  return [
102    {
103      oldStart: oldLines === 0 ? start : start + 1,
104      oldLines,
105      newStart: newLines === 0 ? start : start + 1,
106      newLines,
107      lines,
108    },
109  ]
110}
111
112/** "3", "3–5", "2, 4–6": the turn numbers for the message to the model. */
113export function turnRange(numbers: readonly (number | null)[]): string {
114  const sorted = [...new Set(numbers.filter((n): n is number => n !== null))].sort((x, y) => x - y)
115  const parts: string[] = []
116  let i = 0
117  while (i < sorted.length) {
118    let j = i
119    while (j + 1 < sorted.length && sorted[j + 1] === sorted[j]! + 1) j += 1
120    parts.push(i === j ? `${sorted[i]}` : `${sorted[i]}–${sorted[j]}`)
121    i = j + 1
122  }
123  return parts.join(', ')
124}
125

1export type ChangeTool = 'Edit' | 'Write' | 'NotebookEdit'
2
3/** A recorded change: one successful call to Edit, Write or NotebookEdit. */
4export type ChangeEntry = {
5  /** The call's tool_use_id: also the prefix of its blobs. */
6  id: string
7  /** The main turn it happened in (`turn.start`), if known. */
8  turnId: string | null
9  /** The turn number (`$.session.turns()`), for messages to the model. */
10  turnNumber: number | null
11  tool: ChangeTool
12  /** The subagent that made the change, if any. */
13  agent?: string
14  createdFile: boolean
15  /** Blob with the content BEFORE the change; null if the file did not exist or cannot be read. */
16  preRef: string | null
17  /** Blob with the content AFTER the change; null if the file cannot be read (over 4 MiB). */
18  postRef: string | null
19  /** Blob with the unified diff to show; null if there is none. Kept out of the index to stay under the store's 4 MiB. */
20  patchRef: string | null
21  added: number
22  removed: number
23  at: number
24}
25
26export type FileEntry = {
27  realPath: string
28  /** Relative to session.root when under it, otherwise the real path. */
29  displayPath: string
30  /** Blob with the content before the FIRST pending change: where Undo goes back to. Null if the file was created. */
31  baseRef: string | null
32  /** The file did not exist before the first pending change. */
33  createdByClaude: boolean
34  entries: ChangeEntry[]
35  /** The file contains absorbed external changes (requirement 6). */
36  absorbedExternal: boolean
37  /** Some content could not be read (over 4 MiB): undo is not available. */
38  tooLarge: boolean
39}
40
41export type Registry = {
42  files: Record<string, FileEntry>
43  /** The realPaths in order of first change. */
44  order: string[]
45}
46
47export type TurnMark = {
48  id: string
49  number: number
50  /** Changes recorded in this turn. */
51  added: number
52}
53
54/** An undo request waiting for confirmation because of later external changes. */
55export type PendingConfirm = {
56  paths: string[]
57  /** Pressed from "Undo all": one message to the model. */
58  isAll: boolean
59}
60
61declare module 'claude-code' {
62  interface PluginState {
63    'keep-or-undo': {
64      registry: Registry
65      turn: TurnMark | null
66      confirm: PendingConfirm | null
67      expanded: string | null
68      /** The session the registry belongs to: reloaded after a session change. */
69      owner: string | null
70    }
71  }
72}
73