relay

1import type { AgentInfo, AgentStatus, EngineInterface, Register } from 'claude-code'
2
3const SOFT_LIMIT = 180_000
4const MAX_CYCLES = 10
5
6const stopNote = (why: string) =>
7  `[relay] ${why} ` +
8  'Continue only to the next logical breakpoint, then stop. Do not start new work. ' +
9  'If background agents still run, wait for each to report and record its result first. ' +
10  'Then write a handoff with the writing-handoffs skill. ' +
11  'End your final message with one line: "HANDOFF: <absolute path of the handoff file>". ' +
12  'If the whole task is complete, end with the line "RELAY: DONE" instead.'
13
14const waitNote = (live: AgentInfo[]) =>
15  `[relay] Do not clear yet: ${live.length} background agent(s) still run (${live.map(a => a.description).join('; ')}). ` +
16  '/clear would lose their results. End this turn without the HANDOFF line. ' +
17  'When each agent reports, record its result in the handoff. After the last one reports, end with the HANDOFF line again.'
18
19// Statuses whose result a /clear would lose.
20// ponytail: background Bash tasks are not in $.agent.list(); guard them too if a handoff ever loses one
21const LIVE: AgentStatus[] = ['pending', 'running', 'waiting']
22
23const pastLimit = (tokens: number) =>
24  `The context is at ${Math.round(tokens / 1000)}K tokens, past the ${SOFT_LIMIT / 1000}K soft limit.`
25
26// ponytail: module variables, so a save of this mod or a restart disarms the relay and resets the count
27let isArmed = false
28let cycles = 0
29let isStopAsked = false
30// One cycle the person asked for with /relay now, armed or not.
31let isOnce = false
32// The wait note goes out once per handoff, so a repeated HANDOFF line cannot loop prompts.
33let isWaitAsked = false
34
35function showStatus($: EngineInterface) {
36  const text = isArmed ? `relay ${cycles}/${MAX_CYCLES}` : isOnce ? 'relay once' : undefined
37  $.ui.status(text !== undefined && isStopAsked ? `${text} · stopping` : text)
38}
39
40function disarm($: EngineInterface) {
41  isArmed = false
42  isOnce = false
43  // A stopped relay must not keep denying subagents or holding back the wait note.
44  isStopAsked = false
45  isWaitAsked = false
46  showStatus($)
47}
48
49async function contextTokens($: EngineInterface) {
50  return (await $.session.usage()).context.tokens ?? 0
51}
52
53async function restart($: EngineInterface, path: string) {
54  await $.command.run({ command: 'clear' })
55  await $.prompt.submit({ text: `[relay${isArmed ? ` ${cycles}/${MAX_CYCLES}` : ''}] Continue the task. First read the handoff at ${path}.` })
56  showStatus($)
57}
58
59export const register: Register = on => {
60  on('session.start', async ($, e, next) => {
61    await $.command.register({
62      name: 'relay',
63      description: `Soft stop past ${SOFT_LIMIT / 1000}K context, hand off, /clear, continue (max ${MAX_CYCLES} cycles); now = one cycle at once`,
64      argumentHint: 'on | off | now',
65      immediate: true,
66    })
67
68    return next(e)
69  })
70
71  on('command.run', { command: 'relay' }, async ($, e) => {
72    const arg = e.args.trim()
73    if (arg === 'on') {
74      isArmed = true
75      cycles = 0
76      isStopAsked = false
77    }
78    if (arg === 'off') {
79      disarm($)
80    }
81    if (arg === 'now') {
82      isOnce = true
83      if (!isStopAsked) {
84        isStopAsked = true
85        // Idle, the note starts a turn; mid-turn, the engine queues it into the running turn.
86        $.clock.after(0, () => void $.prompt.submit({ text: stopNote('The owner asked for a handoff now.') }))
87      }
88      showStatus($)
89      return { text: 'Relay: handoff at the next breakpoint, then /clear and continue.' }
90    }
91    showStatus($)
92
93    return { text: isArmed ? `Relay on: cycle ${cycles}/${MAX_CYCLES}, soft limit ${SOFT_LIMIT / 1000}K.` : 'Relay off.' }
94  })
95
96  // Mid-turn: the first tool result past the limit carries the soft stop note to the model.
97  on('tool.call', async ($, e, next) => {
98    // Past the limit, the main loop starts no new subagent: it waits for the live ones, then hands off.
99    if (e.tool === 'Agent' && e.agentId === undefined && (isStopAsked || (isArmed && (await contextTokens($)) >= SOFT_LIMIT))) {
100      return { deny: `[relay] No new subagents past the ${SOFT_LIMIT / 1000}K soft limit. Wait for the running agents, then write the handoff.` }
101    }
102    const ran = await next(e)
103    if (!isArmed || isStopAsked || e.agentId !== undefined || ran.deny !== undefined) {
104      return ran
105    }
106    const tokens = await contextTokens($)
107    if (tokens < SOFT_LIMIT) {
108      return ran
109    }
110    isStopAsked = true
111    showStatus($)
112
113    return { ...ran, context: [...(ran.context ?? []), stopNote(pastLimit(tokens))] }
114  })
115
116  on('turn.complete', async ($, e, next) => {
117    const done = await next(e)
118    if (!(isArmed || isOnce) || e.agentId !== undefined || e.isAborted) {
119      return done
120    }
121
122    if (/^RELAY: DONE\s*$/m.test(e.answer)) {
123      disarm($)
124      $.ui.toast('Relay: task complete.')
125      return done
126    }
127
128    const path = e.answer.match(/^HANDOFF:\s*(\S+)\s*$/m)?.[1]
129    if (path === undefined) {
130      // The turn ended past the limit before any tool call saw it: the turn's end is the breakpoint.
131      const tokens = await contextTokens($)
132      if (!isStopAsked && tokens >= SOFT_LIMIT) {
133        isStopAsked = true
134        showStatus($)
135        $.clock.after(0, () => void $.prompt.submit({ text: stopNote(pastLimit(tokens)) }))
136      }
137      return done
138    }
139
140    const isFile = await $.fs.stat(path).then(
141      stat => stat.kind === 'file',
142      () => false,
143    )
144    if (!isFile) {
145      disarm($)
146      $.ui.toast(`Relay stopped: no handoff file at ${path}`)
147      return done
148    }
149    if (!isOnce && cycles >= MAX_CYCLES) {
150      disarm($)
151      $.ui.toast(`Relay stopped after ${MAX_CYCLES} cycles. Handoff: ${path}`)
152      return done
153    }
154
155    const live = (await $.agent.list()).filter(a => LIVE.includes(a.status))
156    if (live.length > 0) {
157      $.ui.toast(`Relay waits: ${live.length} background agent(s) still run`)
158      if (!isWaitAsked) {
159        isWaitAsked = true
160        $.clock.after(0, () => void $.prompt.submit({ text: waitNote(live) }))
161      }
162      return done
163    }
164
165    if (isArmed) {
166      cycles += 1
167    }
168    isWaitAsked = false
169    isStopAsked = false
170    isOnce = false
171    // After the turn: $.command.run rejects inside a hook the turn waits on.
172    $.clock.after(0, () => void restart($, path))
173
174    return done
175  })
176}
177