baton

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register } from 'claude-code'
3
4import type { AgentStats, Bump, Handoff, LoopStats, Window } from '../types'
5
6const main = atom({ plugin: 'baton', key: 'main' } as const, null)
7const agents = atom({ plugin: 'baton', key: 'agents' } as const, {})
8const bumps = atom({ plugin: 'baton', key: 'bumps' } as const, [])
9const phase = atom({ plugin: 'baton', key: 'phase' } as const, 'ok')
10const enginePercent = atom({ plugin: 'baton', key: 'enginePercent' } as const, null)
11const windowState = atom({ plugin: 'baton', key: 'window' } as const, null)
12
13const HANDOFF_TOOL = 'mcp__baton__handoff'
14const STORE_KEY = 'handoff'
15// Wrap up this far into the engine's auto-compaction threshold, so the handoff lands before the summary does.
16const COMPACT_HEADROOM = 0.85
17
18type $ = EngineInterface
19
20// What the next request re-sends (the status line's context_tokens formula).
21const contextOf = (u: ModelUsage) =>
22  u.input_tokens + u.output_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
23
24const k = (n: number) => (Math.abs(n) >= 1000 ? `${(n / 1000).toFixed(Math.abs(n) >= 10_000 ? 0 : 1)}k` : String(n))
25const minutes = (ms: number) => Math.round(ms / 60_000)
26const text = (t: string) => ({ type: 'text' as const, text: t })
27
28const stopText = (reason: string) => `[baton] STOP. ${reason}.
29Do not start new work and do not call more tools unless a single call is needed to finish a sentence of work already in flight.
30Reply now with a HANDOFF for the REMAINING work only, in this shape:
311. Done: one line per finished item, with file paths.
322. Remaining: concrete, ordered items.
333. Next step: the exact first action a fresh agent should take.
344. Gotchas: anything that cost you time.
35Your reply is forwarded to the main agent, which starts a fresh agent on the remaining work.`
36
37const stoppedNote = (a: AgentStats) =>
38  `[baton] Subagent "${a.description}" (${a.id}) was stopped early: ${a.stopReason}. Its result is a HANDOFF of the remaining work, not a finished task. Spawn a fresh subagent whose prompt is that handoff to continue; do not redo what it lists as done.`
39
40const wrapUpNote = (reason: string) => `[baton] ${reason}. Wrap up now:
411. Start no new work. Let running subagents finish; spawn new ones only to continue a stopped one.
422. Update the files your work touched (task plan, docs, memory) so nothing lives only in this conversation.
433. Call the ${HANDOFF_TOOL} tool with a handoff for the REMAINING work only: done, remaining, exact next step, gotchas, files.
444. Then tell the user to start a fresh session in this directory; the handoff is injected there automatically.`
45
46const handoffNote = (h: Handoff) => `[baton] Handoff from your previous session in this directory (written ${new Date(h.writtenAt).toISOString()} at ${h.percent}% context). Continue from it; do not redo what it lists as done. When baton tells you to wrap up, write the next handoff with the ${HANDOFF_TOOL} tool.
47
48${h.text}`
49
50// Filled by register(); top-level so the helpers below can read it.
51const cfg = { mainMax: 70, agentMax: 30, agentMaxMs: 20 * 60_000, maxWorking: 250_000 }
52// ponytail: one window for every loop; subagents on another model read slightly off.
53const win: Window = { model: 200_000, compact: 200_000 }
54
55// The engine's figures: model window, compaction window, auto-compact threshold.
56async function refreshWindow($: $) {
57  try {
58    const u = await $.session.usage({ breakdown: 'summary' })
59    win.model = u.context.window || win.model
60    win.compact = u.context.breakdown?.rawMaxTokens || win.model
61    win.autoCompactAt = u.context.breakdown?.autoCompactThreshold
62  } catch {
63    // not bound yet; keep what we have
64  }
65  await update($, windowState, () => ({ ...win }))
66}
67
68// Where main must wrap up: the tightest of its percent of the compaction window, the headroom
69// before auto-compaction, and the absolute working budget above its baseline. `[limit, why]`.
70function mainLimit(baseline: number): [number, string] {
71  const candidates: [number, string][] = [
72    [Math.round((cfg.mainMax / 100) * win.compact), `${cfg.mainMax}% of the ${k(win.compact)} compaction window`],
73    [baseline + cfg.maxWorking, `${k(cfg.maxWorking)} working tokens above the ${k(baseline)} baseline`],
74  ]
75  if (win.autoCompactAt) {
76    candidates.push([Math.round(COMPACT_HEADROOM * win.autoCompactAt), `${Math.round(COMPACT_HEADROOM * 100)}% of the ${k(win.autoCompactAt)} auto-compaction threshold`])
77  }
78  return candidates.reduce((a, b) => (b[0] < a[0] ? b : a))
79}
80
81// A subagent's working budget: its percent of the compaction window or the absolute cap, whichever is smaller.
82function agentWorkingLimit(): [number, string] {
83  const byPercent = Math.round((cfg.agentMax / 100) * win.compact)
84  return byPercent < cfg.maxWorking ? [byPercent, `${cfg.agentMax}% of the compaction window`] : [cfg.maxWorking, 'the absolute working cap']
85}
86
87async function showStatus($: $) {
88  const m = await read($, main)
89  if (!m) return $.ui.status(undefined)
90  const running = Object.values(await read($, agents)).filter(a => !a.endedAt)
91  const ph = await read($, phase)
92  let line = `baton ctx ${m.percent}% (${k(m.tokens)}/${k(m.window)}) Δ${m.lastDelta >= 0 ? '+' : ''}${k(m.lastDelta)} · work ${k(m.working)}, wrap-up at ${k(m.limit)}`
93  if (running.length) line += ` · agents ${running.length} (max work ${k(Math.max(...running.map(a => a.working)))})`
94  if (ph === 'wrap-up') line += ` · WRAP UP → call ${HANDOFF_TOOL}`
95  if (ph === 'handoff-written') line += ' · handoff saved → start a fresh session'
96  $.ui.status(line)
97}
98
99async function stopAgent($: $, a: AgentStats, reason: string) {
100  const now = await $.clock.now()
101  await update($, agents, all => ({ ...all, [a.id]: { ...all[a.id]!, stoppedAt: now, stopReason: reason } }))
102  try {
103    await $.session.append({ agentId: a.id, message: { type: 'user', content: [text(stopText(reason))] } })
104    $.ui.toast(`baton: stopping subagent "${a.description}": ${reason}`)
105  } catch {
106    // the loop already ended; nothing to stop
107  }
108}
109
110async function checkDurations($: $) {
111  const now = await $.clock.now()
112  for (const a of Object.values(await read($, agents))) {
113    if (a.endedAt || a.stoppedAt || now - a.startedAt < cfg.agentMaxMs) continue
114    await stopAgent($, a, `running ${minutes(now - a.startedAt)} min (limit ${minutes(cfg.agentMaxMs)} min)`)
115  }
116}
117
118export const register: Register = (on, options) => {
119  cfg.mainMax = Number(options.mainWrapUpPercent ?? 70)
120  cfg.agentMax = Number(options.agentMaxPercent ?? 30)
121  cfg.agentMaxMs = Number(options.agentMaxMinutes ?? 20) * 60_000
122  cfg.maxWorking = Number(options.maxWorkingTokens ?? 250_000)
123  const lastTools: Record<string, string[]> = {}
124  let nagAt = 0
125  let handoff: Handoff | null = null
126
127  on('session.start', async ($, e, next) => {
128    await refreshWindow($)
129    await $.tool.register({
130      name: 'handoff',
131      description:
132        'Save a handoff for the next session: what is done, what remains, the exact next step, gotchas and files. Call it when baton tells you to wrap up. The next session started in this directory reads it automatically.',
133      inputSchema: {
134        type: 'object',
135        properties: { text: { type: 'string', description: 'The handoff, markdown, remaining work only.' } },
136        required: ['text'],
137      },
138      isDeferred: false,
139    })
140    await $.command.register({
141      name: 'baton',
142      description: 'Context budget, subagent guard and handoff status',
143      argumentHint: '[clear]',
144    })
145
146    const stored = (await $.store.get(STORE_KEY)) as Handoff | undefined
147    const sessionId = await $.session.id()
148    const isOurs = stored?.cwd === e.cwd && stored.sessionId !== sessionId && (!stored.takenBy || stored.takenBy === sessionId)
149    if (stored && isOurs) {
150      handoff = { ...stored, takenBy: sessionId }
151      if (!stored.takenBy) {
152        await $.store.set(STORE_KEY, handoff)
153        $.ui.toast(`baton: picked up the handoff written ${new Date(stored.writtenAt).toLocaleString()}`)
154      }
155    }
156    $.clock.every(30_000, () => void checkDurations($))
157    await showStatus($)
158    return next(e)
159  })
160
161  on('turn.start', async ($, e, next) => {
162    if (handoff && !handoff.injectedAt) {
163      handoff = { ...handoff, injectedAt: await $.clock.now() }
164      await $.store.set(STORE_KEY, handoff)
165      await $.session.append({ message: { type: 'user', content: [text(handoffNote(handoff))] } })
166    }
167    return next(e)
168  })
169
170  on('agent.spawn', async ($, e, next) => {
171    const r = await next(e)
172    if ('agentId' in r && r.agentId) {
173      const id = r.agentId
174      const now = await $.clock.now()
175      const [limit] = agentWorkingLimit()
176      await update($, agents, all => ({
177        ...all,
178        [id]: {
179          id, description: e.description, type: e.subagentType, startedAt: now,
180          tokens: 0, window: win.model, percent: 0, steps: 0, lastDelta: 0, baseline: 0, working: 0, limit, updatedAt: now,
181        },
182      }))
183    }
184    return r
185  })
186
187  on('turn.step', async function* ($, e, next) {
188    const r = yield* next(e)
189    if (!r.usage) return r
190    const now = await $.clock.now()
191    const key = e.agentId ?? 'main'
192    const tools = lastTools[key] ?? []
193    lastTools[key] = r.toolUses.map(t => t.name)
194    const tokens = contextOf(r.usage)
195    const percent = Math.round((tokens / win.model) * 100)
196    // The first request's size is the loop's baseline (system prompt, tools, memory); work is what grew after it.
197    const baselineOf = (prev: LoopStats | null | undefined) => (prev?.steps ? prev.baseline : tokens)
198    const grow = (prev: LoopStats | null | undefined, limit: number): LoopStats => {
199      const baseline = baselineOf(prev)
200      return { tokens, window: win.model, percent, steps: (prev?.steps ?? 0) + 1, lastDelta: tokens - (prev?.tokens ?? 0), baseline, working: tokens - baseline, limit, updatedAt: now }
201    }
202    let delta = 0
203
204    if (e.agentId) {
205      const id = e.agentId
206      const [workLimit, why] = agentWorkingLimit()
207      const prev = (await read($, agents))[id]
208      const stats: AgentStats = {
209        ...(prev ?? { id, description: '(unlisted)', type: 'unknown', startedAt: now }),
210        ...grow(prev, baselineOf(prev) + workLimit),
211      }
212      delta = stats.lastDelta
213      await update($, agents, cur => ({ ...cur, [id]: stats }))
214      if (!stats.stoppedAt && stats.working >= workLimit) {
215        await stopAgent($, stats, `your working context grew to ${k(stats.working)} tokens (limit ${k(workLimit)}, ${why})`)
216      }
217    } else {
218      const prev = await read($, main)
219      const [limit, why] = mainLimit(baselineOf(prev))
220      const stats = grow(prev, limit)
221      delta = stats.lastDelta
222      await update($, main, () => stats)
223      const ph = await read($, phase)
224      if (ph !== 'handoff-written' && tokens >= limit && percent >= nagAt + 5) {
225        nagAt = percent
226        await update($, phase, () => 'wrap-up')
227        await $.session.append({ message: { type: 'user', content: [text(wrapUpNote(`Your context is at ${k(tokens)} tokens (${percent}%), past the wrap-up limit of ${k(limit)}: ${why}`))] } })
228        $.ui.toast(`baton: context at ${k(tokens)} (${percent}%), wrapping up`)
229      }
230    }
231
232    const bump: Bump = { at: now, agentId: e.agentId, delta, tools: tools.length ? tools : ['(start)'], percentAfter: percent }
233    await update($, bumps, list => [...list, bump].slice(-60))
234    await showStatus($)
235    return r
236  })
237
238  on('turn.complete', async ($, e, next) => {
239    if (e.agentId) {
240      const id = e.agentId
241      const now = await $.clock.now()
242      const a = (await read($, agents))[id]
243      await update($, agents, all => (all[id] ? { ...all, [id]: { ...all[id]!, endedAt: now } } : all))
244      if (a?.stoppedAt) {
245        await $.session.append({ message: { type: 'user', content: [text(stoppedNote(a))] } })
246      }
247      await showStatus($)
248    }
249    return next(e)
250  })
251
252  on('session.measure', async ($, e, next) => {
253    await refreshWindow($)
254    await update($, enginePercent, () => e.context.percent ?? null)
255    return next(e)
256  })
257
258  on('tool.call', { tool: HANDOFF_TOOL }, async ($, e) => {
259    if (e.agentId) return { deny: 'Subagents do not save handoffs; reply with the handoff as your final message instead.' }
260    const body = String(e.text ?? '').trim()
261    if (!body) return { deny: 'A handoff needs text.' }
262    // the tool may run before this step's usage landed in `main`; the engine's figure is current
263    const percent = (await $.session.usage()).context.percent ?? (await read($, main))?.percent ?? 0
264    const h: Handoff = { cwd: await $.session.cwd(), sessionId: await $.session.id(), writtenAt: await $.clock.now(), percent, text: body }
265    await $.store.set(STORE_KEY, h)
266    await update($, phase, () => 'handoff-written')
267    $.ui.toast('baton: handoff saved, start a fresh session to continue')
268    await showStatus($)
269    return { result: 'Handoff saved. Tell the user to start a fresh session in this directory; baton injects the handoff there.' }
270  })
271
272  on('command.run', { command: 'baton' }, async ($, e) => {
273    if (e.args.trim() === 'clear') {
274      await $.store.delete(STORE_KEY)
275      handoff = null
276      return { text: 'baton: stored handoff cleared.' }
277    }
278    const m = await read($, main)
279    const ph = await read($, phase)
280    const eng = await read($, enginePercent)
281    const list = Object.values(await read($, agents))
282    const stored = (await $.store.get(STORE_KEY)) as Handoff | undefined
283    const now = await $.clock.now()
284    const [, mainWhy] = mainLimit(m?.baseline ?? 0)
285    const [agentLimit, agentWhy] = agentWorkingLimit()
286    const lines = [
287      `baton · context budget`,
288      `window: ${k(win.model)} model, ${k(win.compact)} compaction${win.autoCompactAt ? `, auto-compacts at ${k(win.autoCompactAt)}` : ', auto-compaction off'}`,
289      m
290        ? `main: ${m.percent}% (${k(m.tokens)} / ${k(m.window)}), baseline ${k(m.baseline)}, working ${k(m.working)}, ${m.steps} requests${eng === null ? '' : `, engine says ${eng}%`}`
291        : 'main: no model request measured yet',
292      `phase: ${ph}${m ? ` (wrap-up at ${k(m.limit)}: ${mainWhy})` : ''}`,
293      `agents (working limit ${k(agentLimit)}: ${agentWhy}; ${minutes(cfg.agentMaxMs)} min):`,
294      ...(list.length
295        ? list.map(a => {
296            const state = a.endedAt ? 'ended  ' : a.stoppedAt ? 'stopped' : 'running'
297            const reason = a.stopReason ? ` · ${a.stopReason}` : ''
298            return `  ${state} work ${k(a.working).padStart(5)} ${String(minutes((a.endedAt ?? now) - a.startedAt)).padStart(3)} min  "${a.description}"${reason}`
299          })
300        : ['  none']),
301      'biggest context bumps:',
302      ...[...(await read($, bumps))]
303        .sort((x, y) => y.delta - x.delta)
304        .slice(0, 5)
305        .map(b => `  +${k(b.delta).padStart(6)} → ${b.percentAfter}%  ${b.tools.join(', ')}${b.agentId ? '  (agent)' : ''}`),
306      stored
307        ? `handoff: saved ${new Date(stored.writtenAt).toLocaleString()} at ${stored.percent}% · ${stored.takenBy ? `taken by ${stored.takenBy === (await $.session.id()) ? 'this session' : 'an earlier session'}` : 'not yet taken'} · /baton clear removes it`
308        : 'handoff: none stored',
309    ]
310    return { text: lines.join('\n') }
311  })
312}
313

1/** Context fill of one model loop, measured from the API's own usage figures. */
2export type LoopStats = {
3  /** Tokens the next request re-sends: input + cache read + cache write + output of the last response. */
4  tokens: number
5  /** The model's context window, in tokens. */
6  window: number
7  /** tokens over window, whole percent. */
8  percent: number
9  /** Model requests counted so far. */
10  steps: number
11  /** Growth since the previous request. */
12  lastDelta: number
13  /** The loop's first request: system prompt, tools, memory. Fixed overhead, not work. */
14  baseline: number
15  /** tokens minus baseline: what the work added. */
16  working: number
17  /** The token count at which baton acts on this loop (wrap-up for main, stop for a subagent). */
18  limit: number
19  updatedAt: number
20}
21
22/** The engine's context figures baton measures against. */
23export type Window = {
24  /** The model's context window. */
25  model: number
26  /** The auto-compaction window; smaller than `model` when one is configured. */
27  compact: number
28  /** Where the engine auto-compacts; absent when auto-compaction is off. */
29  autoCompactAt?: number
30}
31
32export type AgentStats = LoopStats & {
33  id: string
34  description: string
35  type: string
36  startedAt: number
37  /** Set once baton told the agent to stop and hand off. */
38  stoppedAt?: number
39  stopReason?: string
40  endedAt?: number
41}
42
43/** One context growth step, attributed to the tool calls that preceded it. */
44export type Bump = {
45  at: number
46  agentId?: string
47  delta: number
48  tools: string[]
49  percentAfter: number
50}
51
52export type Phase = 'ok' | 'wrap-up' | 'handoff-written'
53
54/** What a session leaves for the next one, kept in $.store under "handoff". */
55export type Handoff = {
56  cwd: string
57  sessionId: string
58  writtenAt: number
59  percent: number
60  text: string
61  /** The session that picked it up; a handoff is injected into exactly one session. */
62  takenBy?: string
63  /** When that session's first turn injected it; a reload never injects it again. */
64  injectedAt?: number
65}
66
67declare module 'claude-code' {
68  interface PluginState {
69    baton: {
70      main: LoopStats | null
71      agents: Record<string, AgentStats>
72      bumps: Bump[]
73      phase: Phase
74      /** The engine's own status-line percent, from session.measure, for cross-checking. */
75      enginePercent: number | null
76      window: Window | null
77    }
78  }
79}
80