exit-memory-update

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, Timer } from 'claude-code'
3
4import type { Countdown } from '../types'
5
6const PLUGIN = 'claude-md-management'
7const SKILL = 'revise-claude-md'
8
9const YES = 'Yes, update and exit'
10const NO = 'No, just exit'
11const CANCEL = 'Cancel'
12
13// How long the "exiting" band waits after the update, so its summary can be read.
14const COUNTDOWN_SECONDS = 5
15
16const FALLBACK_PROMPT = `The session is about to end. Before it does, update CLAUDE.md with the learnings from this session.
17
181. Review this conversation for durable, non-obvious knowledge a future session would need: commands that work (build, test, lint, run), project conventions, architecture facts, gotchas and pitfalls hit, preferences the user stated.
192. Read the existing CLAUDE.md in the project root (create it if it does not exist). Skip anything it already says, anything derivable from the code itself, and anything that only mattered to this conversation.
203. Make concise, targeted edits under the fitting sections; don't rewrite unrelated content.
214. Finish with a short summary of what you added or changed, or say that there was nothing worth recording.`
22
23const countdown = atom({ plugin: 'exit-memory-update', key: 'countdown' } as const, null)
24
25const isClaudeMd = (path: string) => /(^|[\\/])CLAUDE\.md$/.test(path)
26
27const describeEdits = (edits: number) => `CLAUDE.md updated (${edits} ${edits === 1 ? 'edit' : 'edits'}).`
28
29const UPDATING = 'Updating CLAUDE.md, then exiting · Esc to stay'
30
31// The countdown's ticking timer, while the band counts down to exit.
32let tick: Timer | undefined
33
34async function stopCountdown($: EngineInterface) {
35  tick?.cancel()
36  tick = undefined
37  await update($, countdown, () => null)
38}
39
40async function exitNow($: EngineInterface) {
41  await stopCountdown($)
42  $.command.run({ command: 'exit' }).catch(() => {})
43}
44
45async function stay($: EngineInterface) {
46  await stopCountdown($)
47  $.ui.toast("Staying in the session. Run /exit when you're ready.")
48}
49
50export const register: Register = on => {
51  // Set while the update turn runs; when it finishes we run /exit for real.
52  let isExitPending = false
53  // Successful Edit/Write calls on a CLAUDE.md: this session's, and the update turn's.
54  let sessionEdits = 0
55  let updateEdits = 0
56
57  on('session.start', async ($, e, next) => {
58    isExitPending = false
59    sessionEdits = 0
60    updateEdits = 0
61    await stopCountdown($)
62    return next(e)
63  })
64
65  on('tool.call', async ($, e, next) => {
66    const result = await next(e)
67    const isEdit = (e.tool === 'Edit' || e.tool === 'Write') && isClaudeMd(e.file_path)
68    if (isEdit && !result.deny && !result.isError) {
69      sessionEdits++
70      if (isExitPending) updateEdits++
71    }
72    return result
73  })
74
75  on('command.run', { command: 'exit' }, async ($, e, next) => {
76    // Our own follow-up /exit, or a headless/SDK run: exit as asked.
77    if (e.origin?.kind === 'plugin' && e.origin.name === $.plugin.name) return next(e)
78    if (e.origin?.kind === 'sdk') return next(e)
79    // The band is up (counting down or waiting): /exit again just skips the wait.
80    if (tick || (await read($, countdown))) {
81      await stopCountdown($)
82      return next(e)
83    }
84    if ((await $.session.turns()) === 0) return next(e)
85
86    const commands = await $.command.list()
87    const skill = commands.find(
88      c => c.name === `${PLUGIN}:${SKILL}` || (c.plugin === PLUGIN && c.name.endsWith(SKILL)),
89    )
90
91    const question = sessionEdits > 0
92      ? 'CLAUDE.md was already updated this session. Update it again with the latest learnings before exiting?'
93      : "Update CLAUDE.md with this session's learnings before exiting?"
94    const how = skill ? `Runs /${skill.name}.` : `Claude edits it directly (${PLUGIN} isn't installed).`
95
96    let answer: string
97    try {
98      answer = await $.ui.ask(`${question} ${how} To steer it, pick Type something and say what to focus on.`, {
99        header: 'CLAUDE.md',
100        options: [YES, NO, CANCEL],
101      })
102    } catch {
103      return { text: 'Exit cancelled.' }
104    }
105
106    if (answer === NO) return next(e)
107    // Anything typed under "Type something" is guidance for the update, not a cancel.
108    const guidance = answer === YES || answer === CANCEL ? '' : answer.trim()
109    if (answer !== YES && !guidance) return { text: 'Exit cancelled.' }
110
111    isExitPending = true
112    updateEdits = 0
113    $.ui.status(UPDATING)
114    // A command.run hook can't start a command or prompt itself; do it from a timer once /exit returns.
115    $.clock.after(0, () => {
116      const started = skill
117        ? $.command.run({ command: skill.name, args: guidance })
118        : $.prompt.submit({
119            text: guidance ? `${FALLBACK_PROMPT}\n\nThe user asked to focus on: ${guidance}` : FALLBACK_PROMPT,
120          })
121      started.catch(err => {
122        isExitPending = false
123        $.ui.status(undefined)
124        $.ui.toast(`Couldn't start the CLAUDE.md update: ${String(err)}`)
125      })
126    })
127
128    const focus = guidance ? ` (focus: ${guidance})` : ''
129    return {
130      text: skill
131        ? `Updating CLAUDE.md with /${skill.name}${focus}; the session exits when it's done.`
132        : `${PLUGIN} isn't installed; updating CLAUDE.md directly${focus}. The session exits when it's done.`,
133    }
134  })
135
136  on('turn.complete', async ($, e, next) => {
137    const result = await next(e)
138    if (!isExitPending) return result
139    isExitPending = false
140    $.ui.status(undefined)
141
142    if (e.reason !== 'answer') {
143      $.ui.toast('CLAUDE.md update did not finish, so the session stays open. Run /exit again when ready.')
144      return result
145    }
146
147    const edits = updateEdits
148    // Nothing applied yet: the update most likely ended by proposing edits and asking for approval.
149    // Wait for the reply instead of exiting under it; the band still offers Exit now and Stay.
150    if (edits === 0) {
151      await update($, countdown, () => ({ edits }))
152      return result
153    }
154
155    // Leave the summary on screen for a few seconds; the band offers Exit now and Stay.
156    await update($, countdown, () => ({ seconds: COUNTDOWN_SECONDS, edits }))
157    tick?.cancel()
158    tick = $.clock.every(1000, async () => {
159      const left = ((await read($, countdown))?.seconds ?? 1) - 1
160      if (left <= 0) return exitNow($)
161      await update($, countdown, c => (c ? { ...c, seconds: left } : c))
162    })
163    return result
164  })
165
166  on('prompt.submit', async ($, e, next) => {
167    const state = await read($, countdown)
168    if (state && state.seconds === undefined) {
169      // A reply to the update's proposed edits: carry on, and exit once that turn is done.
170      await stopCountdown($)
171      isExitPending = true
172      $.ui.status(UPDATING)
173    } else if (tick) {
174      // Typing a new prompt during the countdown means the person wants to stay.
175      await stopCountdown($)
176    }
177    return next(e)
178  })
179
180  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
181    const state: Countdown | null = await read($, countdown)
182    if (!state || e.props.hasSurvey) return next(e)
183
184    const { Box, Button, Text } = $.ui.resolve(e)
185    return (
186      <Box flexDirection="row" gap={2}>
187        {state.seconds === undefined ? (
188          <Text>
189            No changes to CLAUDE.md yet. <Text dimColor>Reply to Claude to apply them; the session exits after.</Text>
190          </Text>
191        ) : (
192          <Text>
193            {describeEdits(state.edits)} <Text dimColor>Exiting in {state.seconds}s</Text>
194          </Text>
195        )}
196        <Button key="exit" label="Exit now" hotkey="1" plain onPress={() => exitNow($)} />
197        <Button key="stay" label="Stay" hotkey="2" plain onPress={() => stay($)} />
198      </Box>
199    )
200  })
201}
202

1// `seconds` counts down to exit; without it, the band waits for the person's reply to the update.
2export type Countdown = { seconds?: number; edits: number }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'exit-memory-update': { countdown: Countdown | null }
7  }
8}
9