undo-agent

1// undo-agent: undo the file edits a subagent made.
2//
3//   /undo-agent          list the agents that changed files this session
4//   /undo-agent last     put back the files the most recent agent changed
5//   /undo-agent <id>     the same for one agent by id
6//
7// /rewind restores your own turns' edits, but a subagent's edits aren't
8// captured in checkpoints. This mod snapshots each file before a subagent's
9// Edit, Write, or NotebookEdit, and restores them on request.
10
11const MAX_FILE = 4 * 1024 * 1024 // $.fs.read's limit
12
13type Entry = { file: string; snapshot: string | null; created: boolean }
14type Group = { agentId: string; entries: Entry[]; skipped: string[]; finished: boolean; order: number }
15
16let scratch = ''
17let counter = 0
18const groups = new Map<string, Group>()
19
20function fileOf(e): string | null {
21  const p = e.tool === 'NotebookEdit' ? e.notebook_path : e.file_path
22  return typeof p === 'string' ? p : null
23}
24
25function group(agentId: string): Group {
26  let g = groups.get(agentId)
27  if (!g) {
28    g = { agentId, entries: [], skipped: [], finished: false, order: ++counter }
29    groups.set(agentId, g)
30  }
31  return g
32}
33
34// Save the file's current content before the agent changes it, once per file per agent.
35async function snapshot($, e, next) {
36  const file = fileOf(e)
37  const agentId = typeof e.agentId === 'string' ? e.agentId : null
38  if (!file || !agentId || !scratch) return next(e)
39  const g = group(agentId)
40  if (!g.entries.some((x) => x.file === file) && !g.skipped.includes(file)) {
41    let exists = false
42    try {
43      exists = await $.fs.exists(file)
44    } catch {
45      // Can't tell whether it exists: don't snapshot, and never treat it as created.
46      g.skipped.push(file)
47      return next(e)
48    }
49    if (exists) {
50      let stat
51      try {
52        stat = await $.fs.stat(file, { resolve: false })
53      } catch {
54        g.skipped.push(file)
55        return next(e)
56      }
57      // A symlink or directory is not snapshotted: restoring through a link would write to its target.
58      if (stat.kind !== 'file' || stat.size > MAX_FILE) {
59        g.skipped.push(file)
60        return next(e)
61      }
62      let content: string
63      try {
64        content = await $.fs.read(file)
65      } catch {
66        g.skipped.push(file)
67        return next(e)
68      }
69      const target = scratch + '/' + (++counter) + '.pre'
70      await $.fs.write(target, content)
71      g.entries.push({ file, snapshot: target, created: false })
72    } else {
73      g.entries.push({ file, snapshot: null, created: true })
74    }
75  }
76  return next(e)
77}
78
79async function restore($, g: Group): Promise<string> {
80  const lines: string[] = []
81  for (const entry of g.entries) {
82    try {
83      let kind = 'missing'
84      try {
85        if (await $.fs.exists(entry.file)) kind = (await $.fs.stat(entry.file, { resolve: false })).kind
86      } catch {
87        kind = 'unknown'
88      }
89      if (kind !== 'file' && kind !== 'missing') {
90        lines.push('skipped ' + entry.file + ' (it is now a link, a directory, or unreadable)')
91        continue
92      }
93      if (entry.created) {
94        if (kind === 'file') await $.process.run(['rm', '-f', '--', entry.file])
95        lines.push('removed ' + entry.file)
96      } else if (entry.snapshot) {
97        await $.fs.write(entry.file, await $.fs.read(entry.snapshot))
98        lines.push('restored ' + entry.file)
99      }
100    } catch (error) {
101      lines.push('could not restore ' + entry.file + ': ' + error)
102    }
103  }
104  for (const file of g.skipped) lines.push('skipped ' + file + ' (not snapshotted: a link, a directory, unreadable, or over 4 MiB)')
105  groups.delete(g.agentId)
106  return lines.join('\n')
107}
108
109export function register(on) {
110  on('session.start', async ($, e, next) => {
111    try {
112      const made = await $.process.run(['mktemp', '-d'])
113      scratch = made.stdout.trim()
114    } catch (error) {
115      scratch = ''
116      $.ui.log('undo-agent: could not make a scratch directory, so nothing will be snapshotted: ' + error)
117    }
118    try {
119      await $.command.register({
120        name: 'undo-agent',
121        description: 'Put back the files a subagent changed',
122        argumentHint: '[last | <agent id>]',
123      })
124    } catch (error) {
125      $.ui.log('could not register /undo-agent: ' + error)
126    }
127    return next(e)
128  })
129
130  on('tool.call', { tool: ['Edit', 'Write', 'NotebookEdit'] }, snapshot).catch(async ($, e, next) => {
131    // Snapshotting failed: the edit still runs, and the user is told it can't be undone.
132    if (next.called) return next(e)
133    $.ui.log('undo-agent could not snapshot ' + (fileOf(e) ?? 'a file') + ' (' + next.error.kind + ')')
134    return next(e)
135  })
136
137  on('turn.complete', async ($, e, next) => {
138    if (typeof e.agentId === 'string') {
139      const g = groups.get(e.agentId)
140      if (g && g.entries.length) {
141        g.finished = true
142        $.ui.log('agent ' + e.agentId + ' changed ' + g.entries.length + ' file(s). /undo-agent last puts them back.')
143      }
144    }
145    return next(e)
146  })
147
148  on('command.run', { command: 'undo-agent' }, async ($, e) => {
149    const args = e.args.trim()
150    const all = [...groups.values()].filter((g) => g.entries.length).sort((a, b) => b.order - a.order)
151    if (!args) {
152      if (!all.length) return { text: 'No subagent has changed files this session.' }
153      return {
154        text: all
155          .map((g) => g.agentId + ': ' + g.entries.length + ' file(s)' + (g.finished ? '' : ' (still running)') + '\n  ' + g.entries.map((x) => x.file).join('\n  '))
156          .join('\n'),
157      }
158    }
159    const g = args === 'last' ? all[0] : groups.get(args)
160    if (!g || !g.entries.length) return { text: 'No snapshots for ' + (args === 'last' ? 'any agent' : args) + '. /undo-agent lists them.' }
161    let answer = 'Cancel'
162    try {
163      answer = await $.ui.ask('Put back ' + g.entries.length + ' file(s) changed by agent ' + g.agentId + '? This overwrites their current content.', ['Cancel', 'Restore'])
164    } catch {
165      // Dismissed, or nobody to ask.
166    }
167    if (answer !== 'Restore') return { text: 'Nothing restored.' }
168    return { text: 'Undid agent ' + g.agentId + ':\n' + (await restore($, g)) }
169  }).catch(async () => ({ text: 'undo-agent: the command failed. Nothing was restored.' }))
170}
171