auto-checkpoint

1import type { Register, EngineInterface } from 'claude-code'
2import {
3  absolutePath, buildManifest, checkpointId, filePathsFor, formatListing, parseManifest,
4  pruneList, resolveSelection, riskyBash, sortNewestFirst,
5  type CheckpointRef, type Manifest,
6} from './lib.ts'
7
8type Engine = EngineInterface
9const SNAPSHOT_TOOLS = ['Edit', 'Write', 'MultiEdit', 'NotebookEdit'] as const
10
11let seq = 0
12
13async function baseDir($: Engine): Promise<string> {
14  const home = (await $.env.get('HOME')) ?? ''
15  return `${home}/.claude/checkpoints`
16}
17
18async function sh($: Engine, argv: string[]): Promise<{ exitCode: number; stdout: string }> {
19  const r = await $.process.run(argv, { timeoutMs: 20000 })
20  return { exitCode: r.exitCode, stdout: r.stdout }
21}
22
23async function listRefs($: Engine, base: string): Promise<CheckpointRef[]> {
24  const refs: CheckpointRef[] = []
25  const sessions = await $.fs.list(base).catch(() => [])
26  for (const s of sessions) {
27    if (s.kind !== 'dir') continue
28    const ids = await $.fs.list(`${base}/${s.name}`).catch(() => [])
29    for (const i of ids) if (i.kind === 'dir') refs.push({ session: s.name, id: i.name })
30  }
31  return refs
32}
33
34async function snapshot($: Engine, tool: string, rawPaths: string[]): Promise<void> {
35  const home = (await $.env.get('HOME')) ?? ''
36  const cwd = await $.session.cwd()
37  const base = await baseDir($)
38  const sid = await $.session.id()
39  const now = await $.clock.now()
40  const paths = [...new Set(rawPaths.map(p => absolutePath(p, cwd, home)))]
41  if (paths.length === 0) return
42  const dir = `${base}/${sid}/${checkpointId(now, seq++ % 1000)}`
43  await sh($, ['mkdir', '-p', `${dir}/files`])
44  const items: { path: string; existed: boolean }[] = []
45  for (const p of paths) items.push({ path: p, existed: await $.fs.exists(p).catch(() => false) })
46  const manifest = buildManifest(tool, now, items)
47  for (const entry of manifest.entries) {
48    if (entry.existed && entry.stored !== undefined) {
49      await sh($, ['cp', '-Rp', entry.path, `${dir}/files/${entry.stored}`])
50    }
51  }
52  await $.fs.write(`${dir}/manifest.json`, JSON.stringify(manifest, null, 2))
53  for (const stale of pruneList(await listRefs($, base))) {
54    await sh($, ['rm', '-rf', `${base}/${stale.session}/${stale.id}`])
55  }
56}
57
58async function dirtyTrackedFiles($: Engine): Promise<string[]> {
59  const top = await sh($, ['git', 'rev-parse', '--show-toplevel'])
60  if (top.exitCode !== 0) return []
61  const root = top.stdout.trim()
62  const diff = await sh($, ['git', '-C', root, 'diff', '--name-only', 'HEAD'])
63  if (diff.exitCode !== 0) return []
64  return diff.stdout.split('\n').filter(Boolean).map(f => `${root}/${f}`)
65}
66
67async function loadSession($: Engine): Promise<{ id: string; manifest: Manifest; dir: string }[]> {
68  const base = await baseDir($)
69  const sid = await $.session.id()
70  const ids = (await $.fs.list(`${base}/${sid}`).catch(() => [])).filter(e => e.kind === 'dir')
71  const sorted = sortNewestFirst(ids.map(e => ({ session: sid, id: e.name })))
72  const out: { id: string; manifest: Manifest; dir: string }[] = []
73  for (const ref of sorted) {
74    const dir = `${base}/${sid}/${ref.id}`
75    const text = await $.fs.read(`${dir}/manifest.json`).catch(() => undefined)
76    const manifest = typeof text === 'string' ? parseManifest(text) : undefined
77    if (manifest) out.push({ id: ref.id, manifest, dir })
78  }
79  return out
80}
81
82async function restore($: Engine, dir: string, manifest: Manifest): Promise<string[]> {
83  const lines: string[] = []
84  for (const e of manifest.entries) {
85    if (e.existed && e.stored !== undefined) {
86      const parent = e.path.slice(0, e.path.lastIndexOf('/')) || '/'
87      await sh($, ['mkdir', '-p', parent])
88      await sh($, ['rm', '-rf', e.path])
89      const r = await sh($, ['cp', '-Rp', `${dir}/files/${e.stored}`, e.path])
90      lines.push(r.exitCode === 0 ? `restored ${e.path}` : `FAILED to restore ${e.path}`)
91    } else {
92      await sh($, ['rm', '-rf', e.path])
93      lines.push(`removed ${e.path} (did not exist at checkpoint)`)
94    }
95  }
96  return lines
97}
98
99export const register: Register = on => {
100  for (const tool of SNAPSHOT_TOOLS) {
101    on('tool.call', { tool }, async ($, e, next) => {
102      try {
103        await snapshot($, tool, filePathsFor(tool, e as unknown as Record<string, unknown>))
104      } catch {
105        // checkpointing must never block an edit
106      }
107      return next(e)
108    }).catch(($, e, next) => next(e))
109  }
110
111  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
112    try {
113      const risk = riskyBash(e.command)
114      if (risk) {
115        const paths = [...risk.paths]
116        if (risk.wholeTree) paths.push(...(await dirtyTrackedFiles($)))
117        await snapshot($, 'Bash', paths)
118      }
119    } catch {
120      // checkpointing must never block a command
121    }
122    return next(e)
123  }).catch(($, e, next) => next(e))
124
125  on('session.start', async ($, e, next) => {
126    await $.command.register({
127      name: 'rewind-to',
128      description: 'List file checkpoints, or restore one: /rewind-to <n> | last',
129      argumentHint: '[n|last]',
130    })
131    return next(e)
132  })
133
134  on('command.run', { command: 'rewind-to' }, async ($, e) => {
135    const items = await loadSession($)
136    const sel = resolveSelection(e.args, items.length)
137    if (sel === 'list') return { text: formatListing(items) }
138    if (sel === 'invalid') {
139      return { text: `Nothing to restore for "${e.args}". ${formatListing(items)}` }
140    }
141    const chosen = items[sel]
142    if (!chosen) return { text: 'No such checkpoint.' }
143    const lines = await restore($, chosen.dir, chosen.manifest)
144    return { text: `Rewound to checkpoint ${sel + 1}:\n${lines.join('\n')}` }
145  })
146}
147

1// Pure helpers for auto-checkpoint: target extraction, manifest, pruning, listing.
2
3export const MAX_CHECKPOINTS = 200
4
5export type ManifestEntry = {
6  path: string
7  existed: boolean
8  /** name of the copy inside the checkpoint folder (files/<stored>), when existed */
9  stored?: string
10}
11
12export type Manifest = {
13  version: 1
14  createdAt: number
15  tool: string
16  entries: ManifestEntry[]
17}
18
19export type CheckpointRef = { session: string; id: string }
20
21const FILE_TOOLS: Record<string, string> = {
22  Edit: 'file_path',
23  Write: 'file_path',
24  MultiEdit: 'file_path',
25  NotebookEdit: 'notebook_path',
26}
27
28/** Paths a file-editing tool call is about to change. */
29export function filePathsFor(tool: string, input: Record<string, unknown>): string[] {
30  const key = FILE_TOOLS[tool]
31  if (key === undefined) return []
32  const v = input[key]
33  return typeof v === 'string' && v !== '' ? [v] : []
34}
35
36export type RiskyBash = { paths: string[]; wholeTree: boolean }
37
38function splitWords(s: string): string[] {
39  const out: string[] = []
40  const re = /"([^"]*)"|'([^']*)'|(\S+)/g
41  let m: RegExpExecArray | null
42  while ((m = re.exec(s)) !== null) out.push(m[1] ?? m[2] ?? m[3] ?? '')
43  return out
44}
45
46/**
47 * Says which paths a shell command may destroy. `wholeTree` means the tracked
48 * changes of the repo are at risk (git reset --hard, git checkout -- .).
49 * Returns undefined when the command is not risky.
50 */
51export function riskyBash(command: string): RiskyBash | undefined {
52  const paths: string[] = []
53  let wholeTree = false
54  let risky = false
55  for (const segment of command.split(/&&|\|\||;|\|/)) {
56    const words = splitWords(segment.trim())
57    while (words[0] === 'sudo' || words[0] === 'command' || /^\w+=/.test(words[0] ?? '')) words.shift()
58    const cmd = words[0]
59    const args = words.slice(1)
60    if (cmd === 'rm' || cmd === 'mv') {
61      risky = true
62      paths.push(...args.filter(a => !a.startsWith('-') && !/[*?]/.test(a)))
63    } else if (cmd === 'git' && args[0] === 'reset' && args.includes('--hard')) {
64      risky = true
65      wholeTree = true
66    } else if (cmd === 'git' && args[0] === 'checkout' && args.includes('--')) {
67      risky = true
68      const rest = args.slice(args.indexOf('--') + 1)
69      if (rest.length === 0 || rest.includes('.')) wholeTree = true
70      else paths.push(...rest)
71    } else if (cmd === 'git' && (args[0] === 'clean' || args[0] === 'restore')) {
72      risky = true
73      if (args[0] === 'restore') wholeTree = true
74    }
75  }
76  return risky ? { paths: [...new Set(paths)], wholeTree } : undefined
77}
78
79/** Resolves a path against cwd without touching the file system. */
80export function absolutePath(p: string, cwd: string, home: string): string {
81  let full = p
82  if (full === '~') full = home
83  else if (full.startsWith('~/')) full = `${home}/${full.slice(2)}`
84  else if (!full.startsWith('/')) full = `${cwd.replace(/\/$/, '')}/${full}`
85  const parts: string[] = []
86  for (const seg of full.split('/')) {
87    if (seg === '' || seg === '.') continue
88    if (seg === '..') parts.pop()
89    else parts.push(seg)
90  }
91  return `/${parts.join('/')}`
92}
93
94/** Sortable checkpoint id: zero-padded epoch milliseconds plus a sequence tail. */
95export function checkpointId(nowMs: number, seq = 0): string {
96  return `${String(nowMs).padStart(15, '0')}-${String(seq).padStart(3, '0')}`
97}
98
99export function buildManifest(
100  tool: string,
101  nowMs: number,
102  items: { path: string; existed: boolean }[],
103): Manifest {
104  return {
105    version: 1,
106    createdAt: nowMs,
107    tool,
108    entries: items.map((it, i) => ({
109      path: it.path,
110      existed: it.existed,
111      ...(it.existed ? { stored: String(i) } : {}),
112    })),
113  }
114}
115
116/** Newest first. */
117export function sortNewestFirst(refs: CheckpointRef[]): CheckpointRef[] {
118  return [...refs].sort((a, b) =>
119    a.id !== b.id ? (a.id < b.id ? 1 : -1) : a.session < b.session ? 1 : a.session > b.session ? -1 : 0,
120  )
121}
122
123/** Checkpoints to delete so only the newest `keep` remain, across sessions. */
124export function pruneList(refs: CheckpointRef[], keep = MAX_CHECKPOINTS): CheckpointRef[] {
125  return sortNewestFirst(refs).slice(keep)
126}
127
128/** Resolves `/rewind-to` arguments against a newest-first list of `count`. */
129export function resolveSelection(arg: string, count: number): number | 'list' | 'invalid' {
130  const a = arg.trim().toLowerCase()
131  if (a === '') return 'list'
132  if (a === 'last') return count > 0 ? 0 : 'invalid'
133  if (!/^\d+$/.test(a)) return 'invalid'
134  const n = Number(a)
135  return n >= 1 && n <= count ? n - 1 : 'invalid'
136}
137
138export function formatListing(items: { id: string; manifest: Manifest }[]): string {
139  if (items.length === 0) return 'No checkpoints in this session yet.'
140  const lines = items.map((it, i) => {
141    const when = new Date(it.manifest.createdAt).toISOString().replace('T', ' ').slice(0, 19)
142    const names = it.manifest.entries.map(e => e.path.split('/').pop() ?? e.path)
143    const shown = names.slice(0, 3).join(', ') + (names.length > 3 ? `, +${names.length - 3} more` : '')
144    return `${String(i + 1).padStart(3)}. ${when}Z  ${it.manifest.tool}  ${shown}`
145  })
146  return `Checkpoints (newest first). Restore with /rewind-to <n> or /rewind-to last.\n${lines.join('\n')}`
147}
148
149export function parseManifest(text: string): Manifest | undefined {
150  try {
151    const m = JSON.parse(text) as Manifest
152    return m && m.version === 1 && Array.isArray(m.entries) ? m : undefined
153  } catch {
154    return undefined
155  }
156}
157