ws-mailbox-push-probe

1import type { HookStream, ProcessSpawnChunk, ProcessSpawnResult, Register, ToolCallResult } from 'claude-code'
2
3import { createWatcher, mapWaitExit, parseRecvText, type MailboxEnvelope, type Watcher, type WaitOutcome } from './watcher'
4
5// Probe constants. LAUNCHER is the installed ws plugin's launcher (the
6// `kang-sw-devenv` directory marketplace); pi spawns the same path with
7// `mailbox wait`.
8export const LAUNCHER = '/home/swkang/.claude/plugins/ws-plugin/ws/bin/ws-mcp-launcher.py'
9// Settled design: 0 blocks until mail. Pi uses a finite "10m" so a wedged or
10// orphaned wait self-heals; switch here if the live test shows stale children.
11export const WAIT_TIMEOUT = '0'
12export const RECV_TOOL = 'mcp__plugin_ws_ws__mailbox_recv'
13export const TOOL_NAME = 'mailbox_watch'
14
15export function waitArgv(key: string): string[] {
16  return ['python3', LAUNCHER, 'mailbox', 'wait', '--session-key', key, '--timeout', WAIT_TIMEOUT, '--format', 'json']
17}
18
19/** Consume one spawned `mailbox wait` stream; abort ends it (closing the stream kills the child). */
20async function consumeWait(
21  stream: HookStream<ProcessSpawnChunk, ProcessSpawnResult>,
22  signal: AbortSignal,
23  debug: (text: string) => void,
24): Promise<WaitOutcome> {
25  let aborted = false
26  const onAbort = () => {
27    aborted = true
28    void stream.return(undefined as never).catch(() => {})
29  }
30  signal.addEventListener('abort', onAbort, { once: true })
31  try {
32    const consumed = (async () => {
33      let stderr = ''
34      for (;;) {
35        const step = await stream.next()
36        if (step.done) return { end: step.value, stderr }
37        if (step.value.stream === 'stderr') stderr += step.value.text
38      }
39    })()
40    const stopped = new Promise<'stopped'>((resolve) => {
41      if (signal.aborted) resolve('stopped')
42      signal.addEventListener('abort', () => resolve('stopped'), { once: true })
43    })
44    const raced = await Promise.race([consumed, stopped])
45    if (raced === 'stopped') {
46      consumed.catch(() => {})
47      return 'stopped'
48    }
49    const outcome = mapWaitExit(raced.end?.code ?? null, aborted)
50    if (outcome === 'error') {
51      debug(`ws mailbox: wait exited code=${raced.end?.code} signal=${raced.end?.signal} ${raced.stderr.trim().slice(0, 300)}`)
52    }
53    return outcome
54  } finally {
55    signal.removeEventListener('abort', onAbort)
56  }
57}
58
59/** Interpret a $.tool.call result of the recv tool (pi's createBridgeDrain). */
60export function recvResultToMail(res: ToolCallResult): MailboxEnvelope[] {
61  if (res.deny !== undefined) throw new Error(`recv denied: ${res.deny}`)
62  if (res.isError) throw new Error(`recv error: ${res.text ?? String(res.result)}`)
63  if (typeof res.text === 'string') return parseRecvText(res.text)
64  // Fallback if core gave no `text`: MCP-shaped content blocks in `result`.
65  const content = (res.result as { content?: { type: string; text?: string }[] } | undefined)?.content
66  return parseRecvText(content?.find((b) => b.type === 'text')?.text)
67}
68
69/** Main-loop-only guard and key validation for the bind tool. */
70export function checkWatchCall(agentId: string | undefined, sessionKey: unknown): { key: string } | { deny: string } {
71  if (agentId !== undefined) {
72    return { deny: `${TOOL_NAME} binds the main session only; called from subagent ${agentId}.` }
73  }
74  const key = typeof sessionKey === 'string' ? sessionKey.trim() : ''
75  if (!key) return { deny: `${TOOL_NAME}: session_key (non-empty string) is required.` }
76  return { key }
77}
78
79export const register: Register = (on) => {
80  let watcher: Watcher | undefined
81
82  on('session.start', async ($, e, next) => {
83    watcher?.stop('session.start again')
84    watcher = createWatcher({
85      wait: async (key, signal) => {
86        if (signal.aborted) return 'stopped'
87        const stream = $.process.spawn({ argv: waitArgv(key) })
88        return consumeWait(stream, signal, (text) => $.ui.log(text, { to: 'debug' }))
89      },
90      drain: async (key) => recvResultToMail(await $.tool.call({ tool: RECV_TOOL, session_key: key, format: 'json' })),
91      submit: async (text) => {
92        await $.prompt.submit({ text })
93      },
94      sleep: (ms, signal) =>
95        new Promise<void>((resolve) => {
96          if (signal.aborted) return resolve()
97          const timer = $.clock.after(ms, () => resolve())
98          signal.addEventListener('abort', () => (timer.cancel(), resolve()), { once: true })
99        }),
100      status: (text) => $.ui.status(text),
101      log: (text) => $.ui.log(text),
102    })
103    await $.tool.register({
104      name: TOOL_NAME,
105      description:
106        'Bind ws cross-session mailbox active push to this session. Pass your own ws session_key. ' +
107        'While bound, arriving mail is drained via mailbox.recv and injected as a new turn even when idle. ' +
108        'Same key again is a no-op; a different key rebinds. Main loop only.',
109      inputSchema: {
110        type: 'object',
111        properties: { session_key: { type: 'string', description: 'Your own ws session key.' } },
112        required: ['session_key'],
113      },
114      isDeferred: false,
115    })
116    $.ui.status('ws mailbox: unbound')
117    return next(e)
118  })
119
120  on('tool.call', { tool: 'mcp__ws-mailbox-push-probe__mailbox_watch' }, ($, e) => {
121    const checked = checkWatchCall(e.agentId, e.session_key)
122    if ('deny' in checked) return { deny: checked.deny }
123    if (!watcher) return { deny: `${TOOL_NAME}: session not started yet.` }
124    const outcome = watcher.bind(checked.key)
125    $.ui.log(`ws mailbox: ${outcome} ${checked.key}`)
126    return { result: `mailbox push ${outcome}: session_key=${checked.key}` }
127  }).catch(() => ({ deny: `${TOOL_NAME}: the bind hook failed; see the debug log.` }))
128
129  on('session.end', ($, e, next) => {
130    watcher?.stop(`session.end ${e.reason}`)
131    return next(e)
132  })
133}
134

1/**
2 * Pure mailbox-push controller, mirroring agents-plugin-pi/src/mailbox-waiter.ts.
3 *
4 * `mailbox wait` only PEEKS (exit 0 = mail present, 3 = timed out, 130 =
5 * interrupted, other = error) and is used purely as the block-until-mail
6 * signal. The mailbox recv tool is the authoritative DRAIN: on a wake the
7 * controller drains through recv and submits exactly what recv removed as ONE
8 * prompt, so a re-arm never re-delivers the same mail (no hot loop, no double
9 * delivery). An empty drain after a mail wake backs off instead of spinning.
10 *
11 * Every side effect is an injected edge so the loop is testable without a
12 * live mailbox or engine.
13 */
14
15export type WaitOutcome = 'mail' | 'timeout' | 'stopped' | 'error'
16
17export interface MailboxEnvelope {
18  from?: string
19  reply_to?: string
20  content: string
21  sent_at?: string
22}
23
24export interface WatcherEdges {
25  /** Run one `mailbox wait` for `key`; must resolve `stopped` soon after `signal` aborts. */
26  wait: (key: string, signal: AbortSignal) => Promise<WaitOutcome>
27  /** Drain via mailbox recv; throws on a tool-level error. */
28  drain: (key: string) => Promise<MailboxEnvelope[]>
29  /** Submit one prompt that wakes the idle session. */
30  submit: (text: string) => Promise<void>
31  /** Resolve after `ms`, or early (resolve, not reject) when `signal` aborts. */
32  sleep: (ms: number, signal: AbortSignal) => Promise<void>
33  status: (text: string | undefined) => void
34  log: (text: string) => void
35}
36
37export type BindResult = 'bound' | 'unchanged' | 'rebound'
38
39export const BACKOFF_MIN_MS = 1_000
40export const BACKOFF_MAX_MS = 60_000
41export const EMPTY_DRAIN_BACKOFF_MS = 2_000
42
43export function nextBackoff(prev: number): number {
44  return prev <= 0 ? BACKOFF_MIN_MS : Math.min(prev * 2, BACKOFF_MAX_MS)
45}
46
47/** Map a finished `mailbox wait` child to an outcome (pi's mapMailboxWaitExit). */
48export function mapWaitExit(code: number | null, aborted: boolean): WaitOutcome {
49  if (aborted) return 'stopped'
50  if (code === 0) return 'mail'
51  if (code === 3) return 'timeout'
52  // 130 (external SIGINT/SIGTERM) is deliberately an error, not stopped: a
53  // stop we did not cause must not silently end push for the session.
54  return 'error'
55}
56
57/** Parse recv's `format: "json"` text (a bare array of envelopes); malformed -> []. */
58export function parseRecvText(text: string | undefined): MailboxEnvelope[] {
59  if (!text) return []
60  let parsed: unknown
61  try {
62    parsed = JSON.parse(text)
63  } catch {
64    return []
65  }
66  if (!Array.isArray(parsed)) return []
67  return parsed.filter(
68    (c): c is MailboxEnvelope =>
69      !!c && typeof c === 'object' && typeof (c as { content?: unknown }).content === 'string',
70  )
71}
72
73/** One prompt for the whole drained batch, never one per envelope. */
74export function formatBatch(key: string, mail: readonly MailboxEnvelope[]): string {
75  const lines = [
76    `[ws mailbox push] ${mail.length} message(s) arrived for session_key ${key}.`,
77    'They were already drained by mailbox.recv (do not call recv for these). Reply with mailbox.send to the handle shown.',
78  ]
79  mail.forEach((m, i) => {
80    const handle = m.from ?? m.reply_to ?? '(unknown)'
81    lines.push('', `--- ${i + 1}/${mail.length} from ${handle}${m.sent_at ? ` at ${m.sent_at}` : ''} ---`, m.content)
82  })
83  return lines.join('\n')
84}
85
86export interface Watcher {
87  bind: (key: string) => BindResult
88  stop: (reason: string) => void
89  current: () => string | undefined
90}
91
92export function createWatcher(edges: WatcherEdges): Watcher {
93  let key: string | undefined
94  let ctrl: AbortController | undefined
95
96  async function loop(myKey: string, signal: AbortSignal): Promise<void> {
97    let backoff = 0
98    let emptyStreak = 0
99    while (!signal.aborted) {
100      edges.status(`ws mailbox: armed ${myKey}`)
101      let outcome: WaitOutcome
102      try {
103        outcome = await edges.wait(myKey, signal)
104      } catch (err) {
105        edges.log(`ws mailbox: wait failed to start: ${String(err)}`)
106        outcome = 'error'
107      }
108      if (signal.aborted || outcome === 'stopped') return
109      if (outcome === 'timeout') {
110        backoff = 0
111        continue
112      }
113      if (outcome === 'error') {
114        backoff = nextBackoff(backoff)
115        edges.status(`ws mailbox: wait error, retry in ${backoff / 1000}s`)
116        edges.log(`ws mailbox: wait error for ${myKey}; backing off ${backoff}ms`)
117        await edges.sleep(backoff, signal)
118        continue
119      }
120      let mail: MailboxEnvelope[]
121      try {
122        mail = await edges.drain(myKey)
123      } catch (err) {
124        backoff = nextBackoff(backoff)
125        edges.status(`ws mailbox: drain error, retry in ${backoff / 1000}s`)
126        edges.log(`ws mailbox: recv failed for ${myKey}: ${String(err)}; backing off ${backoff}ms`)
127        await edges.sleep(backoff, signal)
128        continue
129      }
130      backoff = 0
131      if (mail.length === 0) {
132        // wait saw mail but recv drained nothing (raced, or recv cannot reach
133        // a queue wait can see): never re-arm instantly, or the peek spins.
134        emptyStreak += 1
135        const delay = Math.min(EMPTY_DRAIN_BACKOFF_MS * emptyStreak, BACKOFF_MAX_MS)
136        edges.log(`ws mailbox: wake with empty drain for ${myKey} (x${emptyStreak}); sleeping ${delay}ms`)
137        await edges.sleep(delay, signal)
138        continue
139      }
140      emptyStreak = 0
141      // Recv already removed this mail: deliver it even if a rebind raced in.
142      edges.log(`ws mailbox: drained ${mail.length} message(s) for ${myKey}; submitting one prompt`)
143      try {
144        await edges.submit(formatBatch(myKey, mail))
145      } catch (err) {
146        edges.log(`ws mailbox: prompt.submit failed (mail already drained): ${String(err)}`)
147      }
148    }
149  }
150
151  return {
152    bind(next) {
153      if (next === key && ctrl && !ctrl.signal.aborted) return 'unchanged'
154      const had = key !== undefined
155      ctrl?.abort()
156      key = next
157      const mine = new AbortController()
158      ctrl = mine
159      void loop(next, mine.signal).catch((err) => edges.log(`ws mailbox: loop crashed: ${String(err)}`))
160      return had ? 'rebound' : 'bound'
161    },
162    stop(reason) {
163      if (!ctrl) return
164      ctrl.abort()
165      ctrl = undefined
166      key = undefined
167      edges.status(undefined)
168      edges.log(`ws mailbox: stopped (${reason})`)
169    },
170    current: () => key,
171  }
172}
173