mm

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, HookFailure, PromptSubmitResult, Register, Timer } from 'claude-code'
3
4import type { MmWatch } from '../types'
5import { doneMessage, isFlagForm, splitArgs } from './args'
6import type { ModelResult } from './args'
7
8const POLL_MS = 5000
9const RUN_LINE = /^RUN (\S+)/m
10const FOLLOW_UP = /^(NEXT|THEN) /
11const MMRUN_START = /\bmmrun\s+(review|run)\b/
12const CONTEXT_MAX = 32_000
13const CUT_NOTE = `\n\n[mm plugin: cut to ${CONTEXT_MAX} characters; \`mmrun result <model> <rid>\` prints the whole result]`
14const RESULT_TRIES = 3
15
16type Waiting = MmWatch & { message: string }
17
18const watches = atom({ plugin: 'mm', key: 'watches' } as const, [] as MmWatch[])
19
20// Module timers die with the module on reload; session.start fires again then and restarts it.
21let poller: Timer | null = null
22let isPolling = false
23// From a prompt entering until the main thread stops with no background task left to wake it.
24let isBusy = false
25// Runs started whose watch could not be written yet; the next poll writes them.
26const unwatched: Pick<MmWatch, 'rid' | 'kind'>[] = []
27// Ticks in a row `mmrun result` failed, by run.
28const resultFailures = new Map<string, number>()
29// Bumped by each /clear: work that started before one must not watch or deliver its runs.
30let clears = 0
31
32function logFailure($: EngineInterface, event: string, error: HookFailure): void {
33  $.ui.log(`mm: ${event} hook failed (${error.kind}): ${error.message ?? 'no message'}`, { to: 'debug' })
34}
35
36/** Runs work no hook awaits (a timer tick, a submit in flight), logging its failure. */
37function detach($: EngineInterface, what: string, work: () => Promise<unknown>): void {
38  void work().catch((error: unknown) => $.ui.log(`mm: ${what} failed: ${String(error)}`, { to: 'debug' }))
39}
40
41/** Claims the messages of ended runs: removed from `watches` in one write, so each is handed to one taker. */
42async function takeMessages($: EngineInterface): Promise<Waiting[]> {
43  if (!(await read($, watches)).some(one => one.message !== undefined)) {
44    return []
45  }
46
47  let taken: Waiting[] = []
48
49  await update($, watches, list => {
50    taken = list.filter((one): one is Waiting => one.message !== undefined)
51
52    return list.filter(one => one.message === undefined)
53  })
54
55  return taken
56}
57
58/** Puts claimed messages back for the next poll or prompt, replacing a watch of the same run; not across a /clear since `at`. */
59async function putBack($: EngineInterface, taken: Waiting[], at: number): Promise<void> {
60  if (taken.length > 0 && at === clears) {
61    await update($, watches, list => [...list.filter(one => !taken.some(back => back.rid === one.rid)), ...taken])
62  }
63}
64
65/** Submits each waiting message as a prompt of its own, without waiting for its turn; one not accepted goes back. */
66async function flush($: EngineInterface): Promise<void> {
67  for (const watch of await takeMessages($)) {
68    detach($, 'submitting a finished run', async () => {
69      const at = clears
70      let r: PromptSubmitResult
71
72      try {
73        r = await $.prompt.submit({ text: watch.message })
74      } catch (error) {
75        await putBack($, [watch], at)
76        throw error
77      }
78
79      if (r.drop !== undefined) {
80        await putBack($, [watch], at)
81      }
82    })
83  }
84}
85
86async function home($: EngineInterface): Promise<string> {
87  return (await $.env.get('HOME')) ?? ''
88}
89
90async function readText($: EngineInterface, path: string): Promise<string | null> {
91  try {
92    return String(await $.fs.read(path))
93  } catch {
94    return null
95  }
96}
97
98/** Each model's status file in the run directory; 'gone' when the directory is, null when a listing or read failed. */
99async function statusesOf($: EngineInterface, dir: string): Promise<{ model: string; status: string }[] | 'gone' | null> {
100  let models: string[]
101
102  try {
103    models = (await $.fs.list(dir)).filter(entry => entry.name.endsWith('.status')).map(entry => entry.name.slice(0, -'.status'.length))
104  } catch {
105    return (await $.fs.exists(dir).catch(() => true)) ? null : 'gone'
106  }
107
108  const found: { model: string; status: string }[] = []
109
110  for (const model of models.sort()) {
111    const status = await readText($, `${dir}/${model}.status`)
112
113    if (status === null) {
114      return null
115    }
116
117    found.push({ model, status: status.trim() })
118  }
119
120  return found
121}
122
123/** Each model's result; null while a failed `mmrun result` is still to be tried again, a failure note after RESULT_TRIES ticks. */
124async function resultsOf($: EngineInterface, mmrun: string, watch: MmWatch, models: { model: string; status: string }[]): Promise<ModelResult[] | null> {
125  const results: ModelResult[] = []
126  let isFailed = false
127
128  for (const { model, status } of models) {
129    let output = ''
130
131    if (status === 'DONE') {
132      const r = await $.process
133        .run([mmrun, 'result', model, watch.rid, ...(watch.kind === 'review' ? ['--top'] : [])])
134        .catch((error: unknown) => ({ exitCode: -1, stdout: '', stderr: String(error) }))
135
136      isFailed ||= r.exitCode !== 0
137      output =
138        r.exitCode === 0
139          ? r.stdout.trim()
140          : `mm plugin: could not read the result of ${model} in ${watch.rid} after ${RESULT_TRIES} tries (${r.stderr.split('\n')[0] ?? ''}); \`mmrun result ${model} ${watch.rid}\` prints it.`
141    }
142
143    results.push({ model, status, output })
144  }
145
146  const failures = isFailed ? (resultFailures.get(watch.rid) ?? 0) + 1 : 0
147
148  if (failures > 0 && failures < RESULT_TRIES) {
149    resultFailures.set(watch.rid, failures)
150
151    return null
152  }
153
154  resultFailures.delete(watch.rid)
155
156  return results
157}
158
159/** One pass over the watches: queues the message of each run whose models have all ended, then forgets it. */
160async function pollWatches($: EngineInterface): Promise<void> {
161  if (isPolling) {
162    return
163  }
164
165  isPolling = true
166
167  try {
168    const mmrun = `${$.plugin.root}/bin/mmrun`
169
170    if (unwatched.length > 0) {
171      const at = clears
172      const adding = unwatched.splice(0)
173      const startedAt = await $.clock.now()
174
175      try {
176        await update($, watches, list =>
177          at === clears ? [...list, ...adding.filter(one => !list.some(other => other.rid === one.rid)).map(one => ({ ...one, startedAt }))] : list,
178        )
179      } catch (error) {
180        unwatched.push(...adding)
181        throw error
182      }
183    }
184
185    for (const watch of (await read($, watches)).filter(one => one.message === undefined)) {
186      // A dead worker leaves RUNNING behind until `mmrun status` rewrites it to STALE.
187      await $.process.run([mmrun, 'status', watch.rid])
188
189      const models = await statusesOf($, `${await home($)}/.claude/mmruns/${watch.rid}`)
190
191      if (models === null || (models !== 'gone' && (models.length === 0 || models.some(one => one.status === 'RUNNING')))) {
192        continue
193      }
194
195      const results = models === 'gone' ? null : await resultsOf($, mmrun, watch, models)
196
197      if (models !== 'gone' && results === null) {
198        continue
199      }
200
201      const text = results === null ? null : doneMessage(watch.kind, watch.rid, results)
202
203      await update($, watches, list =>
204        text === null ? list.filter(one => one.rid !== watch.rid) : list.map(one => (one.rid === watch.rid ? { ...one, message: text } : one)),
205      )
206    }
207
208    if (!isBusy) {
209      await flush($)
210    }
211  } finally {
212    isPolling = false
213  }
214}
215
216function ensurePoller($: EngineInterface): void {
217  if (poller === null) {
218    poller = $.clock.every(POLL_MS, () => detach($, 'poll', () => pollWatches($)))
219  }
220}
221
222/** Starts `mmrun <kind> …` in the session root and watches the run it prints; review without --models runs `reviewModels`. */
223async function start($: EngineInterface, kind: MmWatch['kind'], words: string[], reviewModels: string): Promise<{ text: string; context?: string[] }> {
224  const models = kind === 'review' && reviewModels !== '' && !words.includes('--models') ? ['--models', reviewModels] : []
225  const r = await $.process.run([`${$.plugin.root}/bin/mmrun`, kind, ...words, ...models, '--dir', await $.session.root()])
226
227  if (r.exitCode !== 0) {
228    return { text: `mmrun ${kind} failed: ${r.stderr.split('\n')[0] ?? ''}` }
229  }
230
231  const text = r.stdout
232    .split('\n')
233    .filter(line => !FOLLOW_UP.test(line))
234    .join('\n')
235    .trim()
236  const rid = RUN_LINE.exec(r.stdout)?.[1]
237
238  if (rid === undefined) {
239    return { text }
240  }
241
242  const note =
243    kind === 'review'
244      ? `Review ${rid} was started by /mm:review. When it finishes, the mm plugin sends each model's critical/major findings; do not run \`mmrun wait\` yourself.`
245      : `Run ${rid} was started by /mm:run. When it finishes, the mm plugin sends the results; do not run \`mmrun wait\` yourself.`
246
247  // The run is started: a failure from here on must not reach the hook's catch, whose markdown fallback starts another.
248  try {
249    const watch: MmWatch = { rid, kind, startedAt: await $.clock.now() }
250
251    await update($, watches, list => [...list, watch])
252  } catch (error) {
253    $.ui.log(`mm: watching ${rid} failed, the next poll tries again: ${String(error)}`, { to: 'debug' })
254    unwatched.push({ rid, kind })
255  }
256
257  ensurePoller($)
258
259  return { text, context: [note] }
260}
261
262/** Whether the words start mmrun from the mod: options only, and for run both --model and --task. */
263function isStartable(kind: MmWatch['kind'], words: string[]): boolean {
264  return isFlagForm(words) && (kind === 'review' || (words.includes('--model') && words.includes('--task')))
265}
266
267export const register: Register = (on, options) => {
268  const reviewModels = String(options.reviewModels)
269
270  on('session.start', async ($, e, next) => {
271    poller?.cancel()
272    poller = null
273    ensurePoller($)
274
275    return next(e)
276  }).catch(($, e, next) => {
277    logFailure($, 'session.start', next.error)
278
279    return next(e)
280  })
281
282  // /clear ends the conversation the runs would report to: forget them (they keep running, as Claude Code's background tasks do).
283  on('session.end', async ($, e, next) => {
284    if (e.reason === 'clear') {
285      clears++
286      unwatched.splice(0)
287      resultFailures.clear()
288      await update($, watches, () => [])
289    }
290
291    return next(e)
292  }).catch(($, e, next) => {
293    logFailure($, 'session.end', next.error)
294
295    return next(e)
296  })
297
298  // Options only start mmrun here; anything with natural language goes to commands/review.md.
299  on('command.run', { command: 'mm:review' }, async ($, e, next) => {
300    const words = splitArgs(e.args)
301
302    return isFlagForm(words) ? start($, 'review', words, reviewModels) : next(e)
303  }).catch(($, e, next) => {
304    logFailure($, 'command.run', next.error)
305
306    return next(e)
307  })
308
309  on('command.run', { command: 'mm:run' }, async ($, e, next) => {
310    const words = splitArgs(e.args)
311
312    return isStartable('run', words) ? start($, 'run', words, reviewModels) : next(e)
313  }).catch(($, e, next) => {
314    logFailure($, 'command.run', next.error)
315
316    return next(e)
317  })
318
319  // The model's own Skill call to the same commands. Never hook tool.call on Bash (claude-code#92533).
320  on('tool.call', { tool: 'Skill' }, async ($, e, next) => {
321    if (e.tool !== 'Skill' || e.agentId !== undefined || (e.skill !== 'mm:review' && e.skill !== 'mm:run')) {
322      return next(e)
323    }
324
325    const kind = e.skill === 'mm:review' ? 'review' : 'run'
326    const words = splitArgs(e.args ?? '')
327
328    if (words.length === 0 || !isStartable(kind, words)) {
329      return next(e)
330    }
331
332    const r = await start($, kind, words, reviewModels)
333    const result = { success: true, commandName: e.skill }
334
335    return r.context === undefined
336      ? { result, context: [`mmrun failed to start: ${r.text}; start it with Bash as /${e.skill} describes instead`] }
337      : { result, context: [...r.context, `mmrun output:\n${r.text}`] }
338  }).catch(($, e, next) => {
339    logFailure($, 'tool.call', next.error)
340
341    return next(e)
342  })
343
344  // Runs the model starts through Bash, so their end is handed over as /mm:review's are.
345  on('classic.PostToolUse', async ($, e, next) => {
346    const input = typeof e.tool_input === 'object' && e.tool_input !== null ? (e.tool_input as Record<string, unknown>) : {}
347    const output = typeof e.tool_response === 'object' && e.tool_response !== null ? (e.tool_response as Record<string, unknown>) : {}
348    const kind = e.agent_id === undefined && e.tool_name === 'Bash' && typeof input.command === 'string' ? MMRUN_START.exec(input.command)?.[1] : undefined
349    const rid = typeof output.stdout === 'string' ? RUN_LINE.exec(output.stdout)?.[1] : undefined
350
351    if ((kind === 'review' || kind === 'run') && rid !== undefined) {
352      const watch: MmWatch = { rid, kind, startedAt: await $.clock.now() }
353
354      await update($, watches, list => (list.some(one => one.rid === rid) ? list : [...list, watch]))
355      ensurePoller($)
356    }
357
358    return next(e)
359  }).catch(($, e, next) => {
360    logFailure($, 'classic.PostToolUse', next.error)
361
362    return next(e)
363  })
364
365  // Waiting messages ride the prompt that enters first, so none is also submitted. Not prompt.context: that spends the cache.
366  // A prompt that does not enter (dropped, or a hook beneath failing) puts its messages back.
367  on('prompt.submit', async ($, e, next) => {
368    const wasBusy = isBusy
369
370    isBusy = true
371
372    const at = clears
373    const taken = await takeMessages($)
374    const text = taken.map(one => one.message).join('\n\n')
375    let r: PromptSubmitResult
376
377    try {
378      r = await next(
379        taken.length === 0 ? e : { ...e, context: [...(e.context ?? []), text.length > CONTEXT_MAX ? text.slice(0, CONTEXT_MAX - CUT_NOTE.length) + CUT_NOTE : text] },
380      )
381    } catch (error) {
382      await putBack($, taken, at)
383      throw error
384    }
385
386    if (r.drop !== undefined) {
387      isBusy = wasBusy
388      await putBack($, taken, at)
389    }
390
391    return r
392  }).catch(($, e, next) => {
393    logFailure($, 'prompt.submit', next.error)
394
395    return next(e)
396  })
397
398  // An interrupted or failed turn may end with no Stop; it leaves nothing running that would hand the messages over.
399  on('turn.complete', async ($, e, next) => {
400    if (e.agentId === undefined && e.reason !== 'answer') {
401      isBusy = false
402    }
403
404    return next(e)
405  }).catch(($, e, next) => {
406    logFailure($, 'turn.complete', next.error)
407
408    return next(e)
409  })
410
411  // Background work still in flight wakes the session again, so the messages wait for that turn. The next poll
412  // submits them: the host refuses a submit from a Stop hook.
413  on('classic.Stop', async ($, e, next) => {
414    if (e.agent_id === undefined) {
415      isBusy = (e.background_tasks ?? []).length > 0
416    }
417
418    return next(e)
419  }).catch(($, e, next) => {
420    logFailure($, 'classic.Stop', next.error)
421
422    return next(e)
423  })
424}
425

1// Options of `mmrun review` and `mmrun run` that take the next word as their value; any other `--x` is a switch.
2const VALUED = new Set([
3  '--base',
4  '--commit',
5  '--paths',
6  '--focus',
7  '--notes-file',
8  '--max',
9  '--models',
10  '--dir',
11  '--grok-model',
12  '--codex-model',
13  '--agy-model',
14  '--patch',
15  '--model',
16  '--task',
17  '--model-name',
18])
19
20/** One model's end of a run: its status file, and `mmrun result` output when it is DONE. */
21export type ModelResult = { model: string; status: string; output: string }
22
23/** Splits like a shell: blanks separate words, '…' is literal, "…" and a bare \ escape; no expansion. */
24export function splitArgs(args: string): string[] {
25  const words: string[] = []
26  let word = ''
27  let isWord = false
28  let quote: '' | "'" | '"' = ''
29
30  for (let i = 0; i < args.length; i++) {
31    const ch = args[i] ?? ''
32
33    if (quote === "'") {
34      if (ch === "'") {
35        quote = ''
36      } else {
37        word += ch
38      }
39    } else if (quote === '"') {
40      if (ch === '"') {
41        quote = ''
42      } else if (ch === '\\' && '"\\$`'.includes(args[i + 1] ?? '')) {
43        word += args[++i]
44      } else {
45        word += ch
46      }
47    } else if (/\s/.test(ch)) {
48      if (isWord) {
49        words.push(word)
50        word = ''
51        isWord = false
52      }
53    } else {
54      isWord = true
55
56      if (ch === "'" || ch === '"') {
57        quote = ch
58      } else if (ch === '\\' && i + 1 < args.length) {
59        word += args[++i]
60      } else {
61        word += ch
62      }
63    }
64  }
65
66  if (isWord) {
67    words.push(word)
68  }
69
70  return words
71}
72
73/** Options only: the first word is an option, and every other word is the value of the option before it. */
74export function isFlagForm(words: string[]): boolean {
75  if (words.length === 0) {
76    return true
77  }
78
79  return words.every((word, i) => (i === 0 ? word.startsWith('--') : word.startsWith('--') || VALUED.has(words[i - 1] ?? '')))
80}
81
82const JUDGE = [
83  'Judge rules (as in step 4 of /mm:review):',
84  '- Settle each finding by reading the code it points at; drop those with confidence below 80, and mention the dropped ones only in one line at the end: `Dropped N low-confidence findings`.',
85  '- A finding both models report is high confidence; one that only a single model reports, settle by reading the code yourself: no voting, and do not drop it because the other model did not report it.',
86  '- Act only on critical/major.',
87  '- Coming back empty-handed is fine: with no finding at 80 or above, say plainly there are no issues; do not pad.',
88  '- When the findings show `Not expanded` greater than 0, tell the user and let them decide whether to upgrade to --exhaustive.',
89]
90
91/** The message handed to the model when every model of a watched run has ended. */
92export function doneMessage(kind: 'review' | 'run', rid: string, results: ModelResult[]): string {
93  // The engine refuses a submitted text that begins with `/`: it would run as a command.
94  const head = kind === 'review' ? `mm plugin: /mm:review review ${rid} has finished` : `mm plugin: /mm:run run ${rid} has finished`
95  const sections = results.map(one =>
96    one.status === 'DONE'
97      ? `## ${one.model}(${one.status})\n${one.output}`
98      : `## ${one.model}(${one.status})\nNo result. To troubleshoot, read the last lines: \`tail -20 ~/.claude/mmruns/${rid}/${one.model}.raw\`; do not read the whole raw.`,
99  )
100  const tail =
101    kind === 'review'
102      ? JUDGE
103      : [
104          'Next (as in step 4 of /mm:run):',
105          ...results.map(one => `- Read the patch first: \`~/.claude/mmruns/${rid}/${one.model}.patch\`, against the task's acceptance items.`),
106          `- For a cross review, \`mmrun review --patch ${rid}\`.`,
107          `- If satisfied, \`mmrun apply ${rid}\`; if not, \`mmrun discard ${rid}\`.`,
108          '- On `status: blocked`, take the questions to the user; do not guess the answers and delegate again.',
109        ]
110
111  return [head, ...sections, tail.join('\n')].join('\n\n')
112}
113

1/** An mmrun run started by /mm:review or /mm:run, watched until every model has ended. */
2export type MmWatch = {
3  rid: string
4  kind: 'review' | 'run'
5  /** Epoch ms, from `$.clock.now()` when the run was started. */
6  startedAt: number
7  /** The run's conclusions once every model ended, waiting to be handed to the model; absent while it runs. */
8  message?: string
9}
10
11declare module 'claude-code' {
12  interface PluginState {
13    mm: {
14      watches: MmWatch[]
15    }
16  }
17}
18