turn-checkpoint

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Turn } from '../types'
5
6const turns = atom({ plugin: 'turn-checkpoint', key: 'turns' } as const, [])
7
8const WRITERS = new Set(['Edit', 'Write', 'MultiEdit', 'NotebookEdit'])
9
10const size = (list: Turn[]) =>
11  list.reduce((sum, t) => sum + t.snaps.reduce((s, one) => s + (one.original?.length ?? 0), 0), 0)
12
13export const register: Register = (on, options) => {
14  const maxTurns = Math.max(1, Number(options.maxTurns ?? 20))
15  const maxBytes = Math.max(0.1, Number(options.maxMegabytes ?? 2)) * 1_000_000
16  let counter = 0
17
18  on('session.start', async ($, e, next) => {
19    await $.command.register({ name: 'undo-turn', description: 'Restore the files the last turn changed' })
20    await $.command.register({ name: 'checkpoints', description: 'List the turns that changed files' })
21
22    return next(e)
23  })
24
25  on('prompt.submit', async ($, e, next) => {
26    if (e.command === undefined) {
27      counter += 1
28      const no = counter
29      const at = await $.clock.now()
30      await update($, turns, list => [...list, { no, at, prompt: e.text.slice(0, 60), snaps: [] }].slice(-maxTurns))
31    }
32
33    return next(e)
34  })
35
36  on('tool.call', async ($, e, next) => {
37    if (WRITERS.has(e.tool)) {
38      const input = e as unknown as Record<string, unknown>
39      const path = String(input.file_path ?? input.notebook_path ?? '')
40      const list = await read($, turns)
41      const current = list[list.length - 1]
42      if (path !== '' && current !== undefined && !current.snaps.some(one => one.path === path)) {
43        let original: string | null = null
44        try {
45          original = await $.fs.read(path)
46        } catch {
47          // The file does not exist yet: the turn creates it.
48        }
49        await update($, turns, all =>
50          all.map(t => (t.no === current.no ? { ...t, snaps: [...t.snaps, { path, original }] } : t)),
51        )
52        await update($, turns, all => {
53          let kept = all
54          while (kept.length > 1 && size(kept) > maxBytes) kept = kept.slice(1)
55          return kept
56        })
57      }
58    }
59
60    return next(e)
61  })
62
63  on('command.run', { command: 'checkpoints' }, async $ => {
64    const list = (await read($, turns)).filter(t => t.snaps.length > 0)
65    if (list.length === 0) {
66      return { text: 'No turn has changed a file yet.' }
67    }
68
69    return {
70      text: list
71        .map(t => `#${t.no} "${t.prompt}": ${t.snaps.map(s => s.path.split('/').pop()).join(', ')}`)
72        .join('\n'),
73    }
74  })
75
76  on('command.run', { command: 'undo-turn' }, async $ => {
77    const list = await read($, turns)
78    const target = [...list].reverse().find(t => t.snaps.length > 0)
79    if (target === undefined) {
80      return { text: 'Nothing to undo: no turn has changed a file.' }
81    }
82    const restored: string[] = []
83    const created: string[] = []
84    for (const snap of target.snaps) {
85      if (snap.original === null) {
86        created.push(snap.path)
87      } else {
88        await $.fs.write(snap.path, snap.original)
89        restored.push(snap.path)
90      }
91    }
92    await update($, turns, all => all.filter(t => t.no !== target.no))
93
94    return {
95      text: [
96        `Undid turn #${target.no} ("${target.prompt}").`,
97        `Restored ${restored.length}: ${restored.join(', ') || 'none'}`,
98        created.length > 0 ? `Created by that turn and left in place (delete them yourself): ${created.join(', ')}` : '',
99      ]
100        .filter(Boolean)
101        .join('\n'),
102    }
103  })
104}
105

1export type Snap = { path: string; original: string | null }
2export type Turn = { no: number; at: number; prompt: string; snaps: Snap[] }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'turn-checkpoint': { turns: Turn[] }
7  }
8}
9