replay

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Change } from '../types'
5import { unifiedDiff } from './diff'
6
7const PANE = 'replay'
8const changes = atom({ plugin: 'replay', key: 'changes' } as const, [] as Change[])
9const turn = atom({ plugin: 'replay', key: 'turn' } as const, 0)
10const viewTurn = atom({ plugin: 'replay', key: 'viewTurn' } as const, null)
11
12const FILE_TOOLS = new Set(['Edit', 'Write', 'MultiEdit'])
13const MAX_READ = 4 * 1024 * 1024 - 1
14
15// The file's text; null when it does not exist yet; undefined when it cannot be diffed (too big, unreadable).
16async function snapshot($: EngineInterface, path: string): Promise<string | null | undefined> {
17  const stat = await $.fs.stat(path).catch(() => undefined)
18  if (!stat) return null
19  if (stat.kind !== 'file' || stat.size > MAX_READ) return undefined
20  return $.fs.read(path).then(
21    text => (typeof text === 'string' ? text : undefined),
22    () => undefined,
23  )
24}
25
26async function record(
27  $: EngineInterface,
28  tool: string,
29  file: string,
30  before: string | null | undefined,
31  after: string | null | undefined,
32  limits: { maxChanges: number; maxDiffChars: number },
33) {
34  const at = await read($, turn)
35  const base = { turn: at, tool, file, added: 0, removed: 0, diff: '', omitted: 0 }
36  let entry: Omit<Change, 'id'>
37  if (before === undefined || after === undefined || after === null) {
38    entry = { ...base, note: 'not diffed: the file is over 4 MiB or could not be read' }
39  } else {
40    if (before === after) return
41    const d = unifiedDiff(before ?? '', after, limits.maxDiffChars)
42    entry = { ...base, ...d, ...(before === null ? { note: 'new file' } : {}) }
43  }
44  await update($, changes, list => [...list, { ...entry, id: (list[list.length - 1]?.id ?? 0) + 1 }].slice(-limits.maxChanges))
45}
46
47const turnsOf = (list: readonly Change[]) => [...new Set(list.map(c => c.turn))].sort((a, b) => a - b)
48
49// Moves the view to the previous (-1) or next (+1) turn that has a change.
50async function step($: EngineInterface, by: number) {
51  const turns = turnsOf(await read($, changes))
52  if (turns.length === 0) return
53  const shown = (await read($, viewTurn)) ?? turns[turns.length - 1] ?? 0
54  const at = Math.max(0, turns.indexOf(shown))
55  const next = turns[Math.max(0, Math.min(turns.length - 1, at + by))] ?? shown
56  await update($, viewTurn, () => next)
57}
58
59export const register: Register = (on, options) => {
60  const maxChanges = Math.max(1, typeof options.maxChanges === 'number' ? options.maxChanges : 100)
61  const maxDiffChars = Math.max(500, typeof options.maxDiffChars === 'number' ? options.maxDiffChars : 12000)
62  let cwd = ''
63
64  on('session.start', async ($, e, next) => {
65    cwd = e.cwd
66    await $.command.register({ name: 'replay', description: "Scrub through this session's file changes, turn by turn" })
67    return next(e)
68  })
69
70  on('session.end', async ($, e, next) => {
71    if (e.reason === 'clear') {
72      await update($, changes, () => [])
73      await update($, turn, () => 0)
74      await update($, viewTurn, () => null)
75    }
76    return next(e)
77  })
78
79  on('turn.start', async ($, e, next) => {
80    if (!e.agentId) await update($, turn, n => n + 1)
81    return next(e)
82  })
83
84  // The old file is read before the tool runs and the new one after, so the diff is what is on disk, not what was asked for.
85  on('tool.call', async ($, e, next) => {
86    const tool = String(e.tool)
87    const path = (e as unknown as { file_path?: unknown }).file_path
88    if (!FILE_TOOLS.has(tool) || typeof path !== 'string') return next(e)
89    const before = await snapshot($, path)
90    const ran = await next(e)
91    if (ran.deny !== undefined || ran.isError) return ran
92    const after = await snapshot($, path)
93    const file = cwd && path.startsWith(`${cwd}/`) ? path.slice(cwd.length + 1) : path
94    await record($, tool, file, before, after, { maxChanges, maxDiffChars }).catch(() => undefined)
95    return ran
96  })
97
98  on('command.run', { command: 'replay' }, async $ => {
99    await $.ui.open({ id: PANE, title: 'Replay', focus: true, closeOnEscape: true })
100    return { text: 'Replay pane opened.' }
101  })
102
103  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
104    const { Box, Text, Button, Code } = $.ui.resolve(e)
105    const list = await read($, changes)
106    if (list.length === 0) return <Text dimColor>No file changes yet. Edits and writes show up here as they happen.</Text>
107
108    const turns = turnsOf(list)
109    const picked = await read($, viewTurn)
110    const shown = picked !== null && turns.includes(picked) ? picked : (turns[turns.length - 1] ?? 0)
111    const at = turns.indexOf(shown)
112    const here = list.filter(c => c.turn === shown)
113    const added = here.reduce((sum, c) => sum + c.added, 0)
114    const removed = here.reduce((sum, c) => sum + c.removed, 0)
115    const near = turns.slice(Math.max(0, at - 4), at + 5)
116
117    return (
118      <Box flexDirection="column">
119        <Box>
120          <Button key="prev" label="◀" hotkey="p" onPress={() => step($, -1)} />
121          <Text bold>
122            {' '}
123            turn {shown} ({at + 1} of {turns.length}){' '}
124          </Text>
125          <Button key="next" label="▶" hotkey="n" onPress={() => step($, 1)} />
126          <Text dimColor>
127            {'  '}
128            {here.length} {here.length === 1 ? 'change' : 'changes'} +{added} −{removed}
129          </Text>
130        </Box>
131        <Box>
132          {near.map(t => (
133            <Button key={`turn:${t}`} plain dimColor={t !== shown} onPress={() => update($, viewTurn, () => t)}>
134              {t === shown ? `[${t}]` : String(t)}{' '}
135            </Button>
136          ))}
137        </Box>
138        {here.map(c => (
139          <Box key={`change:${c.id}`} flexDirection="column" marginTop={1}>
140            <Text bold>
141              {c.file} <Text dimColor>{c.tool}{c.note ? `, ${c.note}` : ''} +{c.added} −{c.removed}</Text>
142            </Text>
143            {c.diff !== '' && <Code format="diff" source={c.diff} path={c.file} wrap="truncate-end" />}
144            {c.omitted > 0 && <Text dimColor>… {c.omitted} more {c.omitted === 1 ? 'hunk' : 'hunks'} not kept</Text>}
145          </Box>
146        ))}
147      </Box>
148    )
149  })
150}
151

1export type Op = { t: ' ' | '-' | '+'; s: string }
2
3export function splitLines(text: string): string[] {
4  if (text === '') return []
5  const lines = text.split('\n')
6  if (lines[lines.length - 1] === '') lines.pop()
7  return lines
8}
9
10// Line diff: common head and tail are trimmed, the middle is an LCS table; a middle too big for one is a plain replacement.
11export function diffOps(a: string[], b: string[]): Op[] {
12  let head = 0
13  while (head < a.length && head < b.length && a[head] === b[head]) head += 1
14  let tail = 0
15  while (tail < a.length - head && tail < b.length - head && a[a.length - 1 - tail] === b[b.length - 1 - tail]) tail += 1
16  const x = a.slice(head, a.length - tail)
17  const y = b.slice(head, b.length - tail)
18  const ops: Op[] = a.slice(0, head).map(s => ({ t: ' ', s }))
19
20  if (x.length * y.length > 2_000_000) {
21    for (const s of x) ops.push({ t: '-', s })
22    for (const s of y) ops.push({ t: '+', s })
23  } else {
24    const w = y.length + 1
25    const lcs = new Uint32Array((x.length + 1) * w)
26    for (let i = x.length - 1; i >= 0; i -= 1)
27      for (let j = y.length - 1; j >= 0; j -= 1)
28        lcs[i * w + j] = x[i] === y[j] ? (lcs[(i + 1) * w + j + 1] ?? 0) + 1 : Math.max(lcs[(i + 1) * w + j] ?? 0, lcs[i * w + j + 1] ?? 0)
29    let i = 0
30    let j = 0
31    while (i < x.length || j < y.length) {
32      if (i < x.length && j < y.length && x[i] === y[j]) {
33        ops.push({ t: ' ', s: x[i] ?? '' })
34        i += 1
35        j += 1
36      } else if (i < x.length && (j === y.length || (lcs[(i + 1) * w + j] ?? 0) >= (lcs[i * w + j + 1] ?? 0))) {
37        ops.push({ t: '-', s: x[i] ?? '' })
38        i += 1
39      } else {
40        ops.push({ t: '+', s: y[j] ?? '' })
41        j += 1
42      }
43    }
44  }
45  for (const s of a.slice(a.length - tail)) ops.push({ t: ' ', s })
46  return ops
47}
48
49export type Hunk = { text: string; added: number; removed: number }
50
51// Groups changes that lie within 2 * context lines of each other, as `diff -u` does.
52export function toHunks(ops: Op[], context = 3): Hunk[] {
53  const changed: number[] = []
54  ops.forEach((op, i) => op.t !== ' ' && changed.push(i))
55  const groups: [number, number][] = []
56  for (const i of changed) {
57    const last = groups[groups.length - 1]
58    if (last && i - last[1] <= 2 * context) last[1] = i
59    else groups.push([i, i])
60  }
61  // Line numbers of every op's first old and new line.
62  const oldAt: number[] = []
63  const newAt: number[] = []
64  let o = 1
65  let n = 1
66  for (const op of ops) {
67    oldAt.push(o)
68    newAt.push(n)
69    if (op.t !== '+') o += 1
70    if (op.t !== '-') n += 1
71  }
72  return groups.map(([first, last]) => {
73    const from = Math.max(0, first - context)
74    const to = Math.min(ops.length - 1, last + context)
75    const slice = ops.slice(from, to + 1)
76    const oldCount = slice.filter(op => op.t !== '+').length
77    const newCount = slice.filter(op => op.t !== '-').length
78    const oldStart = oldCount === 0 ? (oldAt[from] ?? 1) - 1 : (oldAt[from] ?? 1)
79    const newStart = newCount === 0 ? (newAt[from] ?? 1) - 1 : (newAt[from] ?? 1)
80    return {
81      text: [`@@ -${oldStart},${oldCount} +${newStart},${newCount} @@`, ...slice.map(op => op.t + op.s)].join('\n'),
82      added: slice.filter(op => op.t === '+').length,
83      removed: slice.filter(op => op.t === '-').length,
84    }
85  })
86}
87
88// One hunk cut to its first lines, its header counts redone so the Code element still parses it.
89export function clipHunk(text: string, maxChars: number): string {
90  const [header = '', ...lines] = text.split('\n')
91  const m = /^@@ -(\d+),\d+ \+(\d+),\d+ @@$/.exec(header)
92  if (!m) return text
93  const kept: string[] = []
94  let size = 0
95  for (const line of lines) {
96    if (size + line.length > maxChars && kept.length > 0) break
97    kept.push(line)
98    size += line.length + 1
99  }
100  const oldCount = kept.filter(l => l[0] !== '+').length
101  const newCount = kept.filter(l => l[0] !== '-').length
102  return [`@@ -${Number(m[1])},${oldCount} +${Number(m[2])},${newCount} @@`, ...kept].join('\n')
103}
104
105export type Diffed = { diff: string; added: number; removed: number; omitted: number }
106
107// The whole change as hunks, cut at a hunk boundary once `maxChars` is passed (the first hunk always stays).
108export function unifiedDiff(before: string, after: string, maxChars: number): Diffed {
109  const hunks = toHunks(diffOps(splitLines(before), splitLines(after)))
110  const kept: string[] = []
111  let size = 0
112  for (const h of hunks) {
113    if (kept.length > 0 && size + h.text.length > maxChars) break
114    kept.push(h.text)
115    size += h.text.length + 1
116  }
117  const first = kept[0]
118  const isClipped = first !== undefined && first.length > maxChars
119  if (isClipped) kept[0] = clipHunk(first, maxChars)
120  return {
121    diff: kept.join('\n'),
122    added: hunks.reduce((sum, h) => sum + h.added, 0),
123    removed: hunks.reduce((sum, h) => sum + h.removed, 0),
124    omitted: hunks.length - kept.length + (isClipped ? 1 : 0),
125  }
126}
127

1export type Change = {
2  id: number
3  // The prompt-turn the edit happened in, counted from 1 since the session started.
4  turn: number
5  tool: string
6  file: string
7  added: number
8  removed: number
9  // Unified-diff hunks (no ---/+++ header), as the Code element draws them; empty when `note` says why.
10  diff: string
11  // Hunks left out to keep the change small.
12  omitted: number
13  note?: string
14}
15
16declare module 'claude-code' {
17  interface PluginState {
18    replay: { changes: Change[]; turn: number; viewTurn: number | null }
19  }
20}
21