turn-budget

1import type { EngineInterface, Register } from 'claude-code'
2
3import {
4  CONTINUE,
5  CONTINUE_QUIETLY,
6  DEFAULTS,
7  STOP,
8  choiceOf,
9  describeLimits,
10  firstMarks,
11  isOver,
12  nextMarks,
13  parse,
14  question,
15  status,
16  stopNote,
17  uncached,
18  HANDOFF,
19} from './budget'
20import type { Limits, Marks, Spend } from './budget'
21import { continueMessage, handoffDoc } from './handoff'
22import type { Todo } from './handoff'
23
24// Before each model request of a turn (subagents' included), checks what the
25// turn has spent; past the limit it asks, before the request is sent, whether
26// to go on. Only between requests, so a running command is never cut.
27
28async function limitsOf($: EngineInterface): Promise<Limits> {
29  return { ...DEFAULTS, ...((await $.store.get('limits')) as Partial<Limits> | undefined) }
30}
31
32type Turn = {
33  id: string
34  start?: number
35  spend: Spend
36  marks: Marks
37  isQuiet: boolean
38  isStopped: boolean
39  /** One question at a time, however many subagents step at once. */
40  asking?: Promise<boolean>
41}
42
43const EDITS = new Set(['Edit', 'MultiEdit', 'Write', 'NotebookEdit'])
44
45// What this session has done, for the handoff: kept for the whole session.
46const tracked = { files: {} as Record<string, number>, todos: [] as Todo[] }
47
48const join = (...parts: string[]) => parts.map((p, i) => (i === 0 ? p.replace(/[\\/]+$/, '') : p)).join('/')
49
50async function gitOut($: EngineInterface, cwd: string, args: string[]) {
51  const ran = await $.process.run(['git', ...args], { cwd, timeoutMs: 10_000 }).catch(() => undefined)
52  return ran && ran.exitCode === 0 ? ran.stdout : undefined
53}
54
55/**
56 * Writes the handoff from exact session data (no model call), stops the turn,
57 * clears the conversation and starts the fresh one with the handoff as its
58 * first message. If the clear is refused, the handoff waits in the prompt box.
59 */
60async function handOff($: EngineInterface, turnId: string | undefined, reason: string) {
61  const cwd = await $.session.cwd()
62  const messages = await $.session.messages().catch(() => [])
63  const prompts = messages.filter(m => m.role === 'user' && m.text.trim() && !m.text.trimStart().startsWith('<')).map(m => m.text)
64  const lastAnswer = [...messages].reverse().find(m => m.role === 'assistant' && m.text.trim())?.text
65  const now = await $.clock.now()
66  const doc = handoffDoc({
67    reason,
68    cwd,
69    prompts,
70    lastAnswer,
71    files: tracked.files,
72    todos: tracked.todos,
73    gitStatus: await gitOut($, cwd, ['status', '--short']),
74    gitDiffStat: await gitOut($, cwd, ['diff', '--stat']),
75    at: new Date(now).toISOString().replace('T', ' ').slice(0, 16),
76  })
77
78  const config = await $.env.get('CLAUDE_CONFIG_DIR')
79  const home = (await $.env.get('USERPROFILE')) ?? (await $.env.get('HOME'))
80  const base = config ?? (home === undefined ? undefined : join(home, '.claude'))
81  const path = base === undefined ? 'not saved' : join(base, 'handoffs', `${new Date(now).toISOString().replace(/[:.]/g, '-')}.md`)
82  if (base !== undefined) await $.fs.write(path, doc).catch(() => undefined)
83
84  if (turnId) await $.turn.abort({ turnId }).catch(() => undefined)
85  $.ui.status(undefined)
86
87  // After this hook returns: a clear and a new prompt cannot start inside it.
88  $.clock.after(300, async () => {
89    const message = continueMessage(doc, path)
90    const cleared = await $.command
91      .run({ command: 'clear' })
92      .then(() => true)
93      .catch(() => false)
94    if (cleared) {
95      tracked.files = {}
96      tracked.todos = []
97      await $.prompt.submit({ text: message })
98      $.ui.toast(`Fresh session started from a handoff. The old conversation is in /resume. Handoff: ${path}`, { timeoutMs: 15_000 })
99    } else {
100      await $.prompt.fill({ text: message, mode: 'replace' }).catch(() => undefined)
101      $.ui.toast(`Handoff written (${path}) and placed in the prompt box. Run /clear, then send it.`, { timeoutMs: 15_000 })
102    }
103  })
104}
105
106// Resolves true to go on, false when the person chose to stop.
107async function check($: EngineInterface, t: Turn, latest: number | undefined): Promise<boolean> {
108  if (t.isStopped) return false
109  if (t.isQuiet || !isOver(t.spend, t.marks)) return true
110  if (t.asking) return t.asking
111
112  if ((await limitsOf($)).onLimit === 'handoff') {
113    t.isStopped = true
114    const used = question(t.spend, t.marks, t.start, latest).replace(/^This turn has used /, '').replace(/ Keep going\?$/, '')
115    await handOff($, t.id, `The previous session was handed off by turn-budget: one turn used ${used}`)
116    return false
117  }
118
119  t.asking = (async () => {
120    let answer: string
121    try {
122      answer = await $.ui.ask(question(t.spend, t.marks, t.start, latest), {
123        options: [HANDOFF, CONTINUE, CONTINUE_QUIETLY, STOP],
124        header: 'Turn budget',
125      })
126    } catch {
127      // Nobody to ask (a -p run, or the dialog was dismissed): never stop a
128      // turn on a question nobody saw; note it and go on.
129      $.ui.status(`${status(t.spend, t.marks)} · over, not asked`)
130      t.isQuiet = true
131      return true
132    }
133    const choice = choiceOf(answer)
134    if (choice === 'handoff') {
135      t.isStopped = true
136      await handOff($, t.id, "The previous session was handed off at the user's request when a turn crossed its budget.")
137      return false
138    }
139    if (choice === 'stop') {
140      t.isStopped = true
141      await $.session
142        .append({ message: { type: 'user', content: [{ type: 'text', text: stopNote(t.spend, t.marks) }] } })
143        .catch(() => undefined)
144      await $.turn.abort({ turnId: t.id }).catch(() => undefined)
145      $.ui.status(undefined)
146      return false
147    }
148    if (choice === 'quiet') t.isQuiet = true
149    else t.marks = nextMarks(t.spend, t.marks, await limitsOf($))
150    return true
151  })()
152
153  try {
154    return await t.asking
155  } finally {
156    t.asking = undefined
157  }
158}
159
160export const register: Register = on => {
161  let latest: number | undefined
162  let turn: Turn | undefined
163
164  const fivePoints = (start: number | undefined) =>
165    start !== undefined && latest !== undefined ? Math.max(0, Math.round((latest - start) * 10) / 10) : undefined
166
167  on('session.start', async ($, e, next) => {
168    await $.command.register({ name: 'turn-budget', description: 'Turn budget: /turn-budget · <points> · tokens <n> · handoff|ask · off|on' })
169    await $.command.register({ name: 'handoff', description: 'Write a handoff from this session (no model call) and continue in a fresh session' })
170    return next(e)
171  })
172
173  // What the handoff needs: the files edited and the latest todo list.
174  on('tool.call', async ($, e, next) => {
175    if (e.agentId === undefined) {
176      const path = (e as { file_path?: unknown }).file_path ?? (e as { notebook_path?: unknown }).notebook_path
177      if (EDITS.has(e.tool) && typeof path === 'string') tracked.files[path] = (tracked.files[path] ?? 0) + 1
178      const todos = (e as { todos?: unknown }).todos
179      if (e.tool === 'TodoWrite' && Array.isArray(todos)) tracked.todos = todos as Todo[]
180    }
181    return next(e)
182  })
183
184  on('command.run', { command: 'handoff' }, async $ => {
185    await handOff($, turn?.id, 'The previous session was handed off with /handoff.')
186    return {}
187  })
188
189  on('session.measure', async ($, e, next) => {
190    const five = e.rateLimits.find(l => l.kind === 'five_hour')
191    if (five) {
192      // The window reset mid-turn: the new reading is the new baseline.
193      if (turn?.start !== undefined && five.percentUsed < turn.start) turn.start = five.percentUsed
194      latest = five.percentUsed
195      if (turn && !turn.isStopped) {
196        turn.spend.points = fivePoints(turn.start)
197        $.ui.status(status(turn.spend, turn.marks))
198      }
199    }
200    return next(e)
201  })
202
203  on('turn.start', async ($, e, next) => {
204    const limits = await limitsOf($)
205    turn = limits.isOff
206      ? undefined
207      : {
208          id: e.turnId,
209          start: latest,
210          spend: { points: latest === undefined ? undefined : 0, tokens: 0, requests: 0 },
211          marks: firstMarks(limits),
212          isQuiet: false,
213          isStopped: false,
214        }
215    $.ui.status(undefined)
216    return next(e)
217  })
218
219  on('turn.complete', async ($, e, next) => {
220    if (e.agentId === undefined) {
221      turn = undefined
222      $.ui.status(undefined)
223    }
224    return next(e)
225  })
226
227  on('turn.step', async function* ($, e, next) {
228    const t = turn
229    if (!t) return yield* next(e)
230
231    if (!(await check($, t, latest))) {
232      // Stopped: answer the step without sending the request.
233      return { turnId: e.turnId, index: e.index, answer: '', toolUses: [], stopReason: null, usage: null }
234    }
235
236    const step = yield* next(e)
237    if (step.usage) {
238      t.spend.tokens += uncached(step.usage)
239      t.spend.requests += 1
240      t.spend.points = fivePoints(t.start)
241      $.ui.status(status(t.spend, t.marks))
242    }
243    return step
244  })
245
246  // Answered with a toast and no text, so the command adds nothing to the conversation.
247  on('command.run', { command: 'turn-budget' }, async ($, e) => {
248    const current = await limitsOf($)
249    const changed = parse(e.args, current)
250    if (changed) await $.store.set('limits', changed)
251    $.ui.toast(describeLimits(changed ?? current), { timeoutMs: 10_000 })
252    return {}
253  })
254}
255

1// The budget arithmetic: pure, so the tests exercise it directly.
2
3export type Limits = {
4  /** Session points (5-hour window) one turn may use before it asks. */
5  points: number
6  /** Uncached input tokens (uncached + cache-written) one turn may send before it asks. */
7  tokens: number
8  isOff: boolean
9  /** At the limit: hand off to a fresh session by itself (default), or ask first. */
10  onLimit: 'handoff' | 'ask'
11}
12
13export const DEFAULTS: Limits = { points: 5, tokens: 500_000, isOff: false, onLimit: 'handoff' }
14
15export type Spend = {
16  /** Session points this turn has used, from Claude Code's readings; undefined off a subscription. */
17  points?: number
18  tokens: number
19  requests: number
20}
21
22/** The next marks to ask at; each Continue raises the crossed one by its own step. */
23export type Marks = { points: number; tokens: number }
24
25export const firstMarks = (limits: Limits): Marks => ({ points: limits.points, tokens: limits.tokens })
26
27export const isOver = (spend: Spend, marks: Marks): boolean =>
28  (spend.points !== undefined && spend.points >= marks.points) || spend.tokens >= marks.tokens
29
30/** Which limit the turn crossed: the one to name when asking and in the note. Points first when both are. */
31export type Crossed = 'points' | 'tokens'
32
33export const crossed = (spend: Spend, marks: Marks): Crossed =>
34  spend.points !== undefined && spend.points >= marks.points ? 'points' : 'tokens'
35
36export const nextMarks = (spend: Spend, marks: Marks, limits: Limits): Marks => {
37  let { points, tokens } = marks
38  while (spend.points !== undefined && spend.points >= points) points += limits.points
39  while (spend.tokens >= tokens) tokens += limits.tokens
40  return { points, tokens }
41}
42
43export type ModelUsageLike = {
44  input_tokens: number
45  cache_creation_input_tokens?: number
46}
47
48/** Cache reads are left out: they cost a fraction of new input. */
49export const uncached = (usage: ModelUsageLike) => usage.input_tokens + (usage.cache_creation_input_tokens ?? 0)
50
51const n = (value: number) => value.toLocaleString('en-US')
52
53export const question = (spend: Spend, marks: Marks, start: number | undefined, now: number | undefined) => {
54  const used =
55    crossed(spend, marks) === 'points' && start !== undefined && now !== undefined
56      ? `${spend.points} ${spend.points === 1 ? 'point' : 'points'} of your session (${start}% → ${now}%, limit +${marks.points})`
57      : `${n(spend.tokens)} uncached input tokens (limit ${n(marks.tokens)})`
58  return `This turn has used ${used} over ${spend.requests} ${spend.requests === 1 ? 'request' : 'requests'}. Keep going?`
59}
60
61export const HANDOFF = 'Hand off to a fresh session'
62export const CONTINUE = 'Continue'
63export const CONTINUE_QUIETLY = "Don't ask again"
64export const STOP = 'Stop here'
65
66export type Choice = 'handoff' | 'continue' | 'quiet' | 'stop'
67
68/** Anything but a Continue label stops: an answer that is not a clear yes protects the session. */
69export const choiceOf = (answer: string): Choice =>
70  answer === HANDOFF ? 'handoff' : answer === CONTINUE ? 'continue' : answer === CONTINUE_QUIETLY ? 'quiet' : 'stop'
71
72export const stopNote = (spend: Spend, marks: Marks) =>
73  `[turn-budget] The user stopped the previous turn at ${crossed(spend, marks) === 'points' ? `+${spend.points} session points` : `${n(spend.tokens)} uncached input tokens`}. Ask before continuing that work.`
74
75export const status = (spend: Spend, marks: Marks) =>
76  spend.points !== undefined ? `turn budget ${spend.points}/${marks.points} pts` : `turn budget ${n(spend.tokens)}/${n(marks.tokens)} tok`
77
78/** `/turn-budget 8`, `/turn-budget tokens 1000000`, `/turn-budget off|on`; undefined for anything else. */
79export const parse = (args: string, limits: Limits): Limits | undefined => {
80  const [a, b] = args.trim().toLowerCase().split(/\s+/)
81  if (a === 'off' || a === 'on') return { ...limits, isOff: a === 'off' }
82  if (a === 'handoff' || a === 'ask') return { ...limits, onLimit: a }
83  if (a === 'tokens' && b && /^\d+$/.test(b) && Number(b) > 0) return { ...limits, tokens: Number(b) }
84  if (a && /^\d+(\.\d+)?$/.test(a) && Number(a) > 0) return { ...limits, points: Number(a) }
85  return undefined
86}
87
88export const describeLimits = (limits: Limits) =>
89  limits.isOff
90    ? 'Turn budget is off. /turn-budget on to turn it back on.'
91    : `Turn budget: at +${limits.points} session points or ${n(limits.tokens)} uncached input tokens in one turn, it ${limits.onLimit === 'handoff' ? 'hands off to a fresh session by itself' : 'asks first'}. /turn-budget <points> · tokens <n> · handoff|ask · off`
92

1// The handoff document: built from exact session data, with no model call, so a
2// fresh session can pick the work up without re-sending the old context.
3
4export type Todo = { content: string; status: string }
5
6export type HandoffInput = {
7  /** Why the old session stopped, in one line. */
8  reason: string
9  cwd: string
10  /** The person's prompts, oldest first. */
11  prompts: string[]
12  /** The last thing Claude said before the stop. */
13  lastAnswer?: string
14  /** Files Claude edited in the old session, with how many edits each. */
15  files: Record<string, number>
16  todos: Todo[]
17  gitStatus?: string
18  gitDiffStat?: string
19  at: string
20}
21
22const PROMPTS = 5
23const PROMPT_CHARS = 800
24const ANSWER_CHARS = 1_200
25const GIT_LINES = 40
26const MAX_CHARS = 12_000
27
28const clip = (text: string, n: number) => {
29  const t = text.trim()
30  return t.length > n ? `${t.slice(0, n - 1)}…` : t
31}
32
33const lines = (text: string | undefined, n: number) => {
34  // Leading spaces matter: in `git status --short`, " M" is not "M ".
35  const all = (text ?? '').split('\n').filter(l => l.trim() !== '')
36  return all.length > n ? [...all.slice(0, n), `… ${all.length - n} more`] : all
37}
38
39const MARK: Record<string, string> = { completed: '[x]', in_progress: '[~]', pending: '[ ]' }
40
41export const handoffDoc = (h: HandoffInput) => {
42  const out: string[] = [
43    `# Handoff from a previous session`,
44    '',
45    `${h.reason} Written ${h.at} in \`${h.cwd}\`, without a model call: everything below is taken from the session as it was.`,
46    '',
47    '## What the user asked (latest last, in full; earlier ones shortened)',
48    // The latest prompt is what to continue: never cut, its newlines kept.
49    ...h.prompts
50      .slice(-PROMPTS)
51      .map((p, i, all) => (i === all.length - 1 ? `${i + 1} (latest). ${p.trim()}` : `${i + 1}. ${clip(p, PROMPT_CHARS).replace(/\n+/g, ' ')}`)),
52  ]
53  if (h.todos.length) out.push('', '## Todo list when it stopped', ...h.todos.map(t => `- ${MARK[t.status] ?? '[ ]'} ${t.content}`))
54  const files = Object.entries(h.files).sort((a, b) => b[1] - a[1])
55  if (files.length) out.push('', '## Files edited', ...files.map(([f, n]) => `- \`${f}\` (${n} ${n === 1 ? 'edit' : 'edits'})`))
56  const status = lines(h.gitStatus, GIT_LINES)
57  if (status.length) out.push('', '## git status --short', '```', ...status, '```')
58  const diff = lines(h.gitDiffStat, GIT_LINES)
59  if (diff.length) out.push('', '## git diff --stat', '```', ...diff, '```')
60  if (h.lastAnswer?.trim()) out.push('', '## Last thing Claude said', ...clip(h.lastAnswer, ANSWER_CHARS).split('\n').map(l => `> ${l}`))
61  out.push(
62    '',
63    '## Next',
64    'Continue the latest request. Read the files above as you need them rather than all at once, check the todo list, and ask the user if anything here is unclear.',
65  )
66  const doc = out.join('\n')
67  // The cap leaves the latest prompt out of the count: it comes before the
68  // rest, so the cut never reaches it.
69  const max = MAX_CHARS + (h.prompts.at(-1)?.trim().length ?? 0)
70  return doc.length > max ? `${doc.slice(0, max - 40)}\n\n… (handoff cut at ${max.toLocaleString('en-US')} characters)` : doc
71}
72
73/** The one message that starts the fresh session: the handoff itself, so nothing needs reading first. */
74export const continueMessage = (doc: string, path: string) =>
75  `Continue the work from my previous session. The handoff is below (also saved at ${path}).\n\n${doc}`
76