fork-probe

1import type { Register } from 'claude-code'
2
3// What the probe saw, newest last; module state, so a reload starts it over
4type Entry = { at: number; what: string; data?: unknown }
5
6const RUN = 'mcp__fork-probe__run'
7const REPORT = 'mcp__fork-probe__report'
8
9export const register: Register = on => {
10  const log: Entry[] = []
11  const t0 = Date.now()
12  const note = (what: string, data?: unknown) => log.push({ at: Date.now() - t0, what, data })
13  const forks = new Set<string>()
14  let isDropping = false
15
16  on('session.start', async ($, e, next) => {
17    const r = await next(e)
18    await $.tool.register({
19      name: 'run',
20      description:
21        'fork-probe: spawns a background fork of this conversation with $.agent.spawn({ subagentType: "fork" }) and returns at once. Pass a nonce you just wrote; the fork is asked to repeat it.',
22      inputSchema: {
23        type: 'object',
24        properties: {
25          nonce: { type: 'string' },
26          dropNotice: {
27            type: 'boolean',
28            description: "drop the fork's completion notice before it reaches the main loop",
29          },
30          subagentType: { type: 'string', description: 'default "fork"' },
31        },
32        required: ['nonce'],
33      },
34      isDeferred: false,
35    })
36    await $.tool.register({
37      name: 'report',
38      description: 'fork-probe: returns everything the probe recorded so far, as JSON.',
39      isDeferred: false,
40    })
41    return r
42  })
43
44  on('tool.call', { tool: RUN }, async ($, e) => {
45    const input = e as unknown as { nonce: string; dropNotice?: boolean; subagentType?: string }
46    isDropping = input.dropNotice === true
47    const subagentType = input.subagentType ?? 'fork'
48    const prompt = `Reply with only the nonce argument of the most recent ${RUN} tool call in this conversation, then the word DONE. If you cannot see one, reply NO-CONTEXT.`
49    note('spawn.request', { subagentType, dropNotice: isDropping })
50    try {
51      const r = await $.agent.spawn({ subagentType, prompt, description: 'fork probe' })
52      note('spawn.result', r)
53      if (r.agentId !== undefined) forks.add(r.agentId)
54      return { result: JSON.stringify({ spawned: r }) }
55    } catch (err) {
56      note('spawn.rejected', String(err))
57      return { result: JSON.stringify({ rejected: String(err) }) }
58    }
59  })
60
61  on('tool.call', { tool: REPORT }, async () => ({ result: JSON.stringify(log, null, 1) }))
62
63  // What the engine tells the hooks about the spawn: is it a fork, in the background
64  on('agent.spawn', async ($, e, next) => {
65    if (e.description === 'fork probe')
66      note('agent.spawn', { subagentType: e.subagentType, fork: e.fork, background: e.background, model: e.model })
67    return next(e)
68  })
69
70  // How many messages the fork's first request carries: a fork that inherited sends the whole transcript
71  on('turn.step', async function* ($, e, next) {
72    if (e.agentId !== undefined && forks.has(e.agentId))
73      note('fork.step', { index: e.index, messageCount: e.messageCount, model: e.model })
74    return yield* next(e)
75  })
76
77  on('turn.complete', async ($, e, next) => {
78    if (e.agentId === undefined || !forks.has(e.agentId)) return next(e)
79    const r = await next(e)
80    const msgs = await $.session.messages({ agentId: e.agentId }).catch(err => String(err))
81    note('fork.complete', {
82      reason: e.reason,
83      text: r.text.slice(0, 300),
84      usage: e.usage,
85      transcript: Array.isArray(msgs)
86        ? { count: msgs.length, last: msgs.slice(-2).map(m => ({ role: m.role, text: m.text.slice(0, 120) })) }
87        : msgs,
88    })
89    // The hand-back may carry the agent's final text: when dropping, hand back a stub instead
90    return isDropping ? { ...r, text: 'fork-probe: reply saved to its branch' } : r
91  })
92
93  // Rows the main conversation keeps: does the hand-back pass here, and through which door
94  on('session.append', async ($, e, next) => {
95    const raw = JSON.stringify(e)
96    if (e.agentId !== undefined || ![...forks].some(id => raw.includes(id))) return next(e)
97    const isDelivery = e.door === 'delivery'
98    note('append', { door: e.door, rewritten: isDelivery && isDropping, head: raw.slice(0, 300) })
99    if (!isDelivery || !isDropping) return next(e)
100    // The hand-back is the agent's own final text; this row is the one place it can be cut
101    const text = 'fork-probe: a background reply finished and was kept off this conversation'
102    return next({ ...e, message: { ...e.message, content: [{ type: 'text', text }] } } as typeof e)
103  })
104
105  // A subagent's hand-back may travel as a message: session.send from its loop, session.receive into main
106  on('session.send', async ($, e, next) => {
107    const from = (e as unknown as { agentId?: string }).agentId
108    if (from !== undefined && forks.has(from))
109      note('send', { from, text: String((e as unknown as { text?: string }).text ?? '').slice(0, 200) })
110    return next(e)
111  })
112
113  on('session.receive', async ($, e, next) => {
114    const isOurs =
115      [...forks].some(id => e.text.includes(id)) || e.text.includes('NO-CONTEXT') || e.text.includes('DONE')
116    note('receive', {
117      origin: e.origin,
118      agentId: e.agentId ?? null,
119      isOurs,
120      dropped: isOurs && isDropping,
121      text: e.text.slice(0, 300),
122    })
123    if (isOurs && isDropping) return { consumed: 'fork-probe consumed the background agent’s hand-back' }
124    return next(e)
125  })
126
127  // A background task's notice reaches the main loop as a prompt of origin task-notification
128  on('prompt.submit', async ($, e, next) => {
129    const kind = (e.origin as { kind?: string } | undefined)?.kind
130    if (kind !== 'task-notification') return next(e)
131    const isOurs = [...forks].some(id => e.text.includes(id)) || e.text.includes('fork probe')
132    note('notice', { isOurs, intoTurn: e.turnId ?? null, dropped: isOurs && isDropping, text: e.text.slice(0, 400) })
133    if (isOurs && isDropping) return { drop: 'fork-probe dropped the fork’s completion notice' }
134    return next(e)
135  })
136}
137