context-restore

1import type { EngineInterface, Register } from 'claude-code'
2import { baseName, bodyOfFile, changedLog, changedNote, currentLog, currentText, rereadLog, rereadNote, rulePathsOf, sectionKey, skillDirOf, skillFileOf, skillFilesIn, statusText, type Line } from './restore.ts'
3
4const ENABLED_KEY = 'enabled'
5const CONSUMER = 'context-restore'
6const USAGE = 'expects nothing (the status), on or off'
7
8/**
9 * The on/off setting, the directory the session started in, the host's config directory, when the session
10 * started, every rules file it read with the time it was last written and its text, every file the main
11 * loop's Read tool read with the time it was last written then, and the last thing the mod did.
12 */
13type State = { enabled: boolean; root: string; config: string; startedAt: number; rules: Map<string, { at: number; text: string }>; reads: Map<string, number>; last?: string }
14
15/**
16 * Reads the on/off setting from the store, which every window shares, so a change made in another
17 * window applies here at the next hook that acts on it.
18 */
19async function readSettings($: EngineInterface, state: State): Promise<void> {
20  state.enabled = (await $.store.get(ENABLED_KEY)) !== false
21}
22
23function errorText(err: unknown): string {
24  return err instanceof Error ? err.message : String(err)
25}
26
27/** What the mod did, told to the person: the sidebar while it is open, else one transcript line. */
28async function toPerson($: EngineInterface, state: State, line: Line): Promise<void> {
29  state.last = line.text
30  try {
31    if (await $.sidebar.set({ consumer: CONSUMER, key: sectionKey(line.text), title: 'context restored', lines: [line], until: 'stream' })) return
32  } catch {
33    // The sidebar mod is not installed.
34  }
35  $.ui.log(line.text)
36}
37
38/** The last write of a file, or undefined when it is not there. */
39async function mtimeOf($: EngineInterface, path: string): Promise<number | undefined> {
40  if (!(await $.fs.exists(path))) return undefined
41  return (await $.fs.stat(path)).mtimeMs
42}
43
44/** The directories a plugin is installed in, from the host's install record. */
45async function pluginDirs($: EngineInterface, state: State, plugin: string): Promise<string[]> {
46  const record = `${state.config}/plugins/installed_plugins.json`
47  if (!(await $.fs.exists(record))) return []
48  const all = (JSON.parse(String(await $.fs.read(record))) as { plugins?: Record<string, { installPath?: unknown }[]> }).plugins ?? {}
49  return Object.entries(all)
50    .filter(([id]) => id.startsWith(`${plugin}@`))
51    .flatMap(([, rows]) => rows.map(r => r.installPath).filter((p): p is string => typeof p === 'string'))
52}
53
54/**
55 * The file a command came from: a plugin's `commands/<name>.md` for `<plugin>:<name>`, else the project's
56 * or the person's `commands/<name>.md`. A built-in command has none.
57 */
58async function commandFileOf($: EngineInterface, state: State, name: string): Promise<string | undefined> {
59  const colon = name.indexOf(':')
60  const candidates = colon > 0
61    ? (await pluginDirs($, state, name.slice(0, colon))).map(dir => `${dir}/commands/${name.slice(colon + 1)}.md`)
62    : [`${state.root}/.claude/commands/${name}.md`, `${state.config}/commands/${name}.md`]
63  for (const path of candidates) if (await $.fs.exists(path)) return path
64  return undefined
65}
66
67/** A file's text as the engine hands it to the model. */
68async function readBody($: EngineInterface, path: string): Promise<string> {
69  const text = String(await $.fs.read(path))
70  return path.endsWith('/SKILL.md') ? bodyOfFile(text, path.slice(0, -'/SKILL.md'.length)) : bodyOfFile(text)
71}
72
73/**
74 * The text one call of a skill or command should carry: the file's current text when the engine handed an
75 * older copy it loaded at the start, else undefined. A skill names its directory, a command is looked up.
76 */
77async function fresherText($: EngineInterface, state: State, name: string, text: string): Promise<string | undefined> {
78  const file = skillFileOf(text) ?? (await commandFileOf($, state, name))
79  if (file === undefined) return undefined
80  const at = await mtimeOf($, file)
81  if (at === undefined) return undefined
82  return currentText(text, await readBody($, file), file, at > state.startedAt)
83}
84
85/**
86 * The files of a skill's directory the main loop read and that were written since, relative to the directory.
87 * A file read again takes a new time, so it is named until the model reads it again.
88 */
89async function changedSinceRead($: EngineInterface, state: State, dir: string): Promise<string[]> {
90  const out: string[] = []
91  for (const file of skillFilesIn(state.reads.keys(), dir)) {
92    const at = await mtimeOf($, `${dir}/${file}`)
93    if (at !== undefined && at !== state.reads.get(`${dir}/${file}`)) out.push(file)
94  }
95  return out
96}
97
98/** A skill's text with the note to read its changed files again, or the text unchanged when none changed. */
99async function withRereadNote($: EngineInterface, state: State, skill: string, text: string): Promise<string> {
100  const dir = skillDirOf(text)
101  const files = dir === undefined ? [] : await changedSinceRead($, state, dir)
102  if (files.length === 0) return text
103  await toPerson($, state, rereadLog(skill, files))
104  return `${text.trimEnd()}\n\n${rereadNote(files)}`
105}
106
107/** The text one skill or command call should carry: the file's current text, then the note on changed files the model read. */
108async function callText($: EngineInterface, state: State, skill: string, text: string): Promise<string> {
109  const fresher = await fresherText($, state, skill, text)
110  if (fresher !== undefined) await toPerson($, state, currentLog(skill))
111  return withRereadNote($, state, skill, fresher ?? text)
112}
113
114/** Records a file the main loop's Read tool read, with the time it was last written. */
115async function recordRead($: EngineInterface, state: State, path: string): Promise<void> {
116  const at = await mtimeOf($, path)
117  if (at === undefined) state.reads.delete(path)
118  else state.reads.set(path, at)
119}
120
121/** Records the rules files and the global CLAUDE.md the session read, with the time each was last written and its text. */
122async function recordRules($: EngineInterface, state: State, text: string): Promise<void> {
123  for (const path of rulePathsOf(text, `${state.config}/CLAUDE.md`)) {
124    const at = await mtimeOf($, path)
125    if (at !== undefined) state.rules.set(path, { at, text: String(await $.fs.read(path)) })
126  }
127}
128
129/** One rules file that changed on disk: its label, its path and its new text. */
130type Change = { label: string; path: string; text: string }
131
132/**
133 * The rules files whose text changed on disk since the session read them; each record takes the new time
134 * and text. A file written again with the same text (an editor's save, a checkout) sends nothing, because
135 * its whole text would reach the model for no change (measured: a touched CLAUDE.md sent 21 KB).
136 */
137async function changedRules($: EngineInterface, state: State): Promise<Change[]> {
138  const out: Change[] = []
139  for (const [path, seen] of state.rules) {
140    const at = await mtimeOf($, path)
141    if (at === undefined || at === seen.at) continue
142    const text = String(await $.fs.read(path))
143    state.rules.set(path, { at, text })
144    if (text !== seen.text) out.push({ label: baseName(path), path, text })
145  }
146  return out
147}
148
149/** The notes for every rules file that changed on disk, and one line to the person naming them. */
150async function changeNotes($: EngineInterface, state: State): Promise<string[]> {
151  let changes: Change[]
152  try {
153    changes = await changedRules($, state)
154  } catch (err) {
155    $.ui.log(`a changed rules file was not read, so the model keeps its earlier text: ${errorText(err)}`)
156    return []
157  }
158  if (changes.length === 0) return []
159  await toPerson($, state, changedLog(changes.map(c => c.label)))
160  return changes.map(c => changedNote(c.label, c.path, c.text))
161}
162
163async function setEnabled($: EngineInterface, state: State, on: boolean): Promise<string> {
164  state.enabled = on
165  await $.store.set(ENABLED_KEY, on)
166  return on ? 'on: a call of a changed skill or command gets its current text, and a changed rules file reaches the model' : 'off: the engine\'s text stays as it is'
167}
168
169async function runCommand($: EngineInterface, state: State, args: string): Promise<string> {
170  const word = args.trim()
171  if (word === 'on' || word === 'off') return setEnabled($, state, word === 'on')
172  if (word !== '') return USAGE
173  await readSettings($, state)
174  return statusText(state.enabled, state.rules.size, state.last)
175}
176
177export const register: Register = on => {
178  const state: State = { enabled: true, root: '', config: '', startedAt: 0, rules: new Map(), reads: new Map() }
179
180  on('session.start', async ($, e, next) => {
181    const r = await next(e)
182    state.startedAt = await $.clock.now()
183    await readSettings($, state)
184    state.root = e.cwd
185    state.config = (await $.env.get('CLAUDE_CONFIG_DIR')) || `${(await $.env.get('HOME')) ?? ''}/.claude`
186    await $.command.register({ name: 'context-restore', description: 'The current text of a changed skill, command or rules file: status, on, off (context-restore)', argumentHint: '[on | off]', immediate: true })
187    return r
188  })
189
190  // The engine prints the plugin name in front of command text and log lines, so the texts do not repeat it.
191  on('command.run', { command: 'context-restore' }, async ($, e) => ({ text: await runCommand($, state, String(e.args ?? '')) }))
192
193  // The engine loads each skill and command once and hands that copy at every call, also after its file changed.
194  on('skill.prompt', async ($, e, next) => {
195    const r = await next(e)
196    await readSettings($, state)
197    if (!state.enabled) return r
198    try {
199      const text = await callText($, state, e.skill, r.text)
200      return text === r.text ? r : { ...r, text }
201    } catch (err) {
202      $.ui.log(`the file of ${e.skill} was not read, so the call keeps the engine's text: ${errorText(err)}`)
203      return r
204    }
205  })
206
207  // A subagent's reads live in its own context, so only the main loop's count.
208  on('tool.call', { tool: 'Read' }, async ($, e, next) => {
209    const r = await next(e)
210    if (e.agentId !== undefined || r.deny !== undefined || r.isError === true) return r
211    await readSettings($, state)
212    if (!state.enabled) return r
213    try {
214      await recordRead($, state, e.file_path)
215    } catch (err) {
216      $.ui.log(`the write time of ${e.file_path} was not read, so a change to it is not noticed: ${errorText(err)}`)
217    }
218    return r
219  })
220
221  on('prompt.attachment', { type: 'instructions' }, async ($, e, next) => {
222    await readSettings($, state)
223    if (state.enabled) await recordRules($, state, e.text)
224    return next(e)
225  })
226
227  // The notes go to the model alone; the person reads one line naming the files.
228  on('prompt.submit', async ($, e, next) => {
229    // No rules file was recorded, so there is nothing to compare.
230    if (state.rules.size === 0) return next(e)
231    await readSettings($, state)
232    if (!state.enabled) return next(e)
233    const notes = await changeNotes($, state)
234    return notes.length === 0 ? next(e) : next({ ...e, context: [...(e.context ?? []), ...notes] })
235  })
236}
237

1/** How the engine expands a skill or command file for the model, and the texts this mod writes. */
2
3const BASE_DIR = /^Base directory for this skill: ([^\n]+)/
4
5/** The directory a skill's text names in its first line, or undefined for a text that names none. */
6export function skillDirOf(text: string): string | undefined {
7  return BASE_DIR.exec(text)?.[1]
8}
9const FRONTMATTER = /^---\n[\s\S]*?\n---\n+/
10const RULE_HEADER = /^Contents of (.+?) \(/gm
11/** What the engine fills in when it expands a file: its arguments, a `${...}` variable, or a shell command's output. */
12const PLACEHOLDER = /\$ARGUMENTS|\$\d|\$\{|!`/
13/** The placeholders a file cannot be read back from: a positional argument, a variable, a shell command's output. */
14const OPAQUE_PLACEHOLDER = /\$\d|\$\{|!`/
15const ARGUMENTS = '$ARGUMENTS'
16/** What the engine adds after a file with no placeholder when the call carries arguments (measured on 2.1.280). */
17const ARGUMENTS_TAIL = '\n\nARGUMENTS: '
18
19/** The SKILL.md a skill's text names in its first line, or undefined for a text that names none. */
20export function skillFileOf(text: string): string | undefined {
21  const dir = skillDirOf(text)
22  return dir === undefined ? undefined : `${dir}/SKILL.md`
23}
24
25/** A skill or command file as the engine hands it to the model: no frontmatter, a skill with its directory first. */
26export function bodyOfFile(fileText: string, dir?: string): string {
27  const body = fileText.replace(FRONTMATTER, '')
28  return dir === undefined ? body : `Base directory for this skill: ${dir}\n\n${body}`
29}
30
31/** The note that follows the engine's text when the file changed and holds text the engine fills in. */
32export function appendedNote(path: string, body: string): string {
33  return `context-restore: ${path} changed on disk after this session loaded it. Its current text follows, with its placeholders not filled in. It replaces the instructions above; the arguments above still apply.\n\n${body}`
34}
35
36/** The engine's text of a file with no placeholder, split before the `ARGUMENTS:` part it adds for a call with arguments. */
37function splitArguments(engineText: string): { head: string; tail: string } {
38  const at = engineText.lastIndexOf(ARGUMENTS_TAIL)
39  return at < 0 ? { head: engineText, tail: '' } : { head: engineText.slice(0, at), tail: engineText.slice(at) }
40}
41
42/**
43 * Whether the engine's text is the file with each `$ARGUMENTS` filled in: every other part word for word,
44 * each `$ARGUMENTS` any text.
45 */
46export function fitsArguments(engineText: string, fileBody: string): boolean {
47  const parts = fileBody.trimEnd().split(ARGUMENTS).map(p => p.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'))
48  return new RegExp(`^${parts.join('[\\s\\S]*')}$`).test(engineText.trimEnd())
49}
50
51/** The engine's filled text with the file's current text after it. */
52function withNote(engineText: string, fileBody: string, path: string): string {
53  return `${engineText.trimEnd()}\n\n${appendedNote(path, fileBody)}`
54}
55
56/**
57 * The text the model reads for one call of a skill or command whose file holds placeholders. A file whose only
58 * placeholder is `$ARGUMENTS` is compared as a template, so the note follows exactly when the engine's copy
59 * differs. Any other placeholder cannot be read back, so the note follows once the file was written after
60 * the session started.
61 */
62function placeholderText(engineText: string, fileBody: string, path: string, writtenSinceStart: boolean): string | undefined {
63  if (!OPAQUE_PLACEHOLDER.test(fileBody)) return fitsArguments(engineText, fileBody) ? undefined : withNote(engineText, fileBody, path)
64  return writtenSinceStart ? withNote(engineText, fileBody, path) : undefined
65}
66
67/**
68 * The text the model reads for one call of a skill or command, when the engine's copy is older than the file.
69 * A file with no placeholder is compared as it is, and its text takes the place of a different engine text,
70 * with the call's arguments kept. A file with placeholders keeps the engine's filled text, with the new file
71 * text after it (`placeholderText`). Undefined: the engine's text is current.
72 */
73export function currentText(engineText: string, fileBody: string, path: string, writtenSinceStart: boolean): string | undefined {
74  if (PLACEHOLDER.test(fileBody)) return placeholderText(engineText, fileBody, path, writtenSinceStart)
75  const { head, tail } = splitArguments(engineText)
76  if (fileBody.trimEnd() === head.trimEnd()) return undefined
77  return tail === '' ? fileBody : `${fileBody.trimEnd()}\n${tail}`
78}
79
80/**
81 * The watched files an `instructions` attachment carries, by the `Contents of <path> (` line each starts
82 * with: every rules file, and the user's global CLAUDE.md. A project's CLAUDE.md is not watched.
83 */
84export function rulePathsOf(text: string, globalFile: string): string[] {
85  return [...text.matchAll(RULE_HEADER)].map(m => m[1] ?? '').filter(p => p.includes('/rules/') || p === globalFile)
86}
87
88/** The files a skill's directory holds besides its SKILL.md, out of `paths`, relative to the directory. */
89export function skillFilesIn(paths: Iterable<string>, dir: string): string[] {
90  return [...paths].filter(p => p.startsWith(`${dir}/`) && p !== `${dir}/SKILL.md`).map(p => p.slice(dir.length + 1))
91}
92
93/** The line a skill call ends with when files of its directory the model read changed on disk since. */
94export function rereadNote(files: readonly string[]): string {
95  return `context-restore: these files of this skill's base directory changed on disk after this session read them, so the copies read earlier are out of date; read them again before you use them: ${files.join(', ')}`
96}
97
98/** How the sidebar colours a line or a part of one. */
99type Tone = 'ok' | 'warn' | 'error' | 'dim'
100export type Part = { text: string; kind?: Tone }
101/** A sidebar line, as the sidebar mod's contract names it; `text` holds the whole line for a sidebar that draws no parts. */
102export type Line = { text: string; kind?: Tone; parts?: Part[] }
103
104const part = (text: string, kind: Tone | undefined): Part => (kind === undefined ? { text } : { text, kind })
105
106/** A line made of parts, its `text` their texts joined. */
107const partsLine = (parts: Part[]): Line => ({ text: parts.map(p => p.text).join(''), parts })
108
109/**
110 * A line the person reads: what the mod did about files that changed on disk, faint, and which files, in
111 * `tone`. Its `text` is the transcript line.
112 */
113function onDisk(done: string, names: readonly string[], tone?: Tone, tail: Part[] = []): Line {
114  return partsLine([part(`changed on disk${done}: `, 'dim'), part(names.join(', '), tone), ...tail])
115}
116
117/** The line the person reads when a skill call asked the model to read changed files again; the files are yellow. */
118export function rereadLog(skill: string, files: readonly string[]): Line {
119  return onDisk(' since the model read it, the call asks to read again', files, 'warn', [part(` (${skill})`, 'dim')])
120}
121
122/** The last part of a path, as the texts name a file. */
123export function baseName(path: string): string {
124  return path.slice(path.lastIndexOf('/') + 1)
125}
126
127/** The line the person reads when a call got the file's current text in place of the engine's older copy. */
128export function currentLog(name: string): Line {
129  return onDisk(', the call got the current text', [name])
130}
131
132/** The line the person reads when rules files changed on disk: what the model was handed again. */
133export function changedLog(labels: readonly string[]): Line {
134  return onDisk(', the new text went to the model', labels)
135}
136
137/** The note the model reads for one rules file that changed on disk after the session read it. */
138export function changedNote(label: string, path: string, text: string): string {
139  return `context-restore: ${label} (${path}) changed on disk after this session read it. Its current text follows and replaces the earlier one; follow it from now on.\n\n${text}`
140}
141
142/** A sidebar section key: the subject cut to what the sidebar takes. */
143export function sectionKey(text: string): string {
144  return text.replace(/[^A-Za-z0-9._:-]+/g, '-').slice(0, 64) || 'restore'
145}
146
147/** The `/context-restore` answer: the setting, what is watched, and the last thing done. */
148export function statusText(enabled: boolean, rules: number, last: string | undefined): string {
149  const done = last === undefined ? '' : ` · last: ${last}`
150  return `${enabled ? 'on' : 'off'} · every skill and command call is checked against its file, ${rules} rules file(s) watched${done}`
151}
152