concurrency-guard

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { TrackedMonitor } from '../types'
5import { REASON_PREFIX, parseTaskEnds, requestedLimits, takeReasonLine, takeReasonPrefix } from './limits'
6import type { Kind } from './limits'
7
8const ALLOW_ONCE = 'Allow once'
9const ALLOW_SESSION = 'No limit this session'
10const DENY = 'Deny'
11// Stored state is JSON, where Infinity would come back as null.
12const UNLIMITED = 1_000_000
13
14const monitors = atom({ plugin: 'concurrency-guard', key: 'monitors' } as const, [])
15const overrides = atom({ plugin: 'concurrency-guard', key: 'overrides' } as const, { subagents: null, monitors: null })
16
17const LABELS: Record<Kind, string> = { subagents: 'subagent', monitors: 'monitor' }
18
19// Several Agent calls in one message are checked at the same moment, before any
20// of them shows up as running. Decisions therefore run one after another, and a
21// call that was let through counts as reserved until it is visibly running.
22const reservations: Record<Kind, Set<string>> = { subagents: new Set(), monitors: new Set() }
23let decisionQueue: Promise<unknown> = Promise.resolve()
24
25function serialized<T>(work: () => Promise<T>): Promise<T> {
26  const run = decisionQueue.then(work, work)
27  decisionQueue = run.catch(() => undefined)
28
29  return run
30}
31
32async function runningMonitors($: EngineInterface): Promise<TrackedMonitor[]> {
33  const now = await $.clock.now()
34  const list = await read($, monitors)
35  const alive = list.filter(i => i.deadline === null || i.deadline > now)
36  if (alive.length !== list.length) {
37    await update($, monitors, () => alive)
38  }
39
40  return alive
41}
42
43async function runningSubagents($: EngineInterface): Promise<number> {
44  const monitorIds = new Set((await read($, monitors)).map(i => i.taskId))
45  const agents = await $.agent.list()
46
47  return agents.filter(i => i.status === 'running' && !monitorIds.has(i.id)).length
48}
49
50async function running($: EngineInterface, kind: Kind): Promise<number> {
51  const visible = kind === 'subagents' ? await runningSubagents($) : (await runningMonitors($)).length
52
53  return visible + reservations[kind].size
54}
55
56// Set from the plugin's options when the module registers.
57let defaults: Record<Kind, number> = { subagents: 4, monitors: 3 }
58
59async function limit($: EngineInterface, kind: Kind): Promise<number> {
60  return (await read($, overrides))[kind] ?? defaults[kind]
61}
62
63async function refreshStatus($: EngineInterface): Promise<void> {
64  const subagents = await running($, 'subagents')
65  const monitorCount = await running($, 'monitors')
66  if (subagents === 0 && monitorCount === 0) {
67    $.ui.status(undefined)
68
69    return
70  }
71  const format = async (count: number, kind: Kind) => {
72    const max = await limit($, kind)
73
74    return `${count}/${max >= UNLIMITED ? '∞' : max}`
75  }
76  $.ui.status(`agents ${await format(subagents, 'subagents')} · monitors ${await format(monitorCount, 'monitors')}`)
77}
78
79// Under the limit the call runs untouched. At the limit it is refused once
80// with instructions; only a retry that states a reason reaches the person.
81async function gate(
82  $: EngineInterface,
83  kind: Kind,
84  task: string,
85  reason: string | null,
86  retryHint: string,
87): Promise<{ deny: string } | null> {
88  const count = await running($, kind)
89  const max = await limit($, kind)
90  if (count < max) {
91    return null
92  }
93  if (reason === null || reason.length === 0) {
94    return {
95      deny: `Concurrency limit reached: ${count} ${kind} are running and the limit is ${max}. Wait for one to finish. Only if another one in parallel is really necessary, ${retryHint} The user will be asked to approve.`,
96    }
97  }
98
99  let answer: string
100  try {
101    answer = await $.ui.ask(
102      `Claude wants to start ${LABELS[kind]} ${count + 1} while the limit is ${max}.\n\nTask: ${task}\nReason: ${reason}\n\nAllow it?`,
103      { options: [ALLOW_ONCE, ALLOW_SESSION, DENY], header: 'Parallel' },
104    )
105  } catch {
106    answer = DENY
107  }
108
109  if (answer === ALLOW_ONCE) {
110    return null
111  }
112  if (answer === ALLOW_SESSION) {
113    await update($, overrides, current => ({ ...current, [kind]: UNLIMITED }))
114    $.ui.toast(`No ${LABELS[kind]} limit for the rest of this session.`)
115
116    return null
117  }
118  const note = answer === DENY ? '' : ` The user answered: "${answer}".`
119
120  return { deny: `The user did not approve another ${LABELS[kind]} beyond the limit of ${max}.${note} Wait for running ones to finish.` }
121}
122
123export const register: Register = (on, options) => {
124  defaults = {
125    subagents: Number(options.maxSubagents ?? 4),
126    monitors: Number(options.maxMonitors ?? 3),
127  }
128
129  on('prompt.submit', async ($, e, next) => {
130    if (e.origin.kind === 'composer') {
131      const requested = requestedLimits(e.text)
132      for (const kind of ['subagents', 'monitors'] as const) {
133        const value = requested[kind]
134        if (value !== undefined && value > defaults[kind]) {
135          await update($, overrides, current => ({ ...current, [kind]: value }))
136          $.ui.toast(`${LABELS[kind]} limit raised to ${value} for this session.`)
137        }
138      }
139    }
140
141    if (e.origin.kind === 'task-notification') {
142      const ended = new Set(parseTaskEnds(e.text).map(i => i.taskId))
143      if (ended.size > 0) {
144        await update($, monitors, list => list.filter(i => !ended.has(i.taskId)))
145      }
146    }
147    await refreshStatus($)
148
149    return next(e)
150  })
151
152  on('tool.call', { tool: 'Agent' }, async ($, e, next) => {
153    // A call missing its parameters is rejected by Claude Code's own validation.
154    if (typeof e.prompt !== 'string' || typeof e.description !== 'string') {
155      return next(e)
156    }
157    const { reason, rest } = takeReasonLine(e.prompt)
158    const refusal = await serialized(async () => {
159      const decision = await gate(
160        $,
161        'subagents',
162        e.description,
163        reason,
164        `retry this exact call with "${REASON_PREFIX} <one sentence why>" as the first line of \`prompt\`.`,
165      )
166      if (decision === null) {
167        reservations.subagents.add(e.tool_use_id)
168      }
169
170      return decision
171    })
172    if (refusal !== null) {
173      return refusal
174    }
175    try {
176      return await next(reason === null ? e : { ...e, prompt: rest })
177    } finally {
178      reservations.subagents.delete(e.tool_use_id)
179      await refreshStatus($)
180    }
181  })
182
183  // A foreground subagent's tool call only returns when it is done; once it has
184  // spawned it is listed as running, so the reservation is no longer needed.
185  on('agent.spawn', async ($, e, next) => {
186    const result = await next(e)
187    reservations.subagents.delete(e.tool_use_id)
188
189    return result
190  })
191
192  on('tool.call', { tool: 'Monitor' }, async ($, e, next) => {
193    if (typeof e.description !== 'string') {
194      return next(e)
195    }
196    const { reason, rest } = takeReasonPrefix(e.description)
197    const refusal = await serialized(async () => {
198      const decision = await gate(
199        $,
200        'monitors',
201        rest,
202        reason,
203        `retry this exact call with \`description\` set to "${REASON_PREFIX} <one sentence why> | <description>".`,
204      )
205      if (decision === null) {
206        reservations.monitors.add(e.tool_use_id)
207      }
208
209      return decision
210    })
211    if (refusal !== null) {
212      return refusal
213    }
214    try {
215      const result = await next(reason === null ? e : { ...e, description: rest })
216      if (result.deny === undefined && result.isError !== true) {
217        const now = await $.clock.now()
218        const deadline = result.result.persistent === true || result.result.timeoutMs === 0 ? null : now + result.result.timeoutMs
219        await update($, monitors, list => [...list, { taskId: result.result.taskId, description: rest, deadline }])
220      }
221
222      return result
223    } finally {
224      reservations.monitors.delete(e.tool_use_id)
225      await refreshStatus($)
226    }
227  })
228
229  on('tool.call', { tool: 'TaskStop' }, async ($, e, next) => {
230    const result = await next(e)
231    const taskId = e.task_id ?? e.shell_id
232    if (taskId !== undefined) {
233      await update($, monitors, list => list.filter(i => i.taskId !== taskId))
234    }
235    await refreshStatus($)
236
237    return result
238  })
239
240  on('turn.complete', async ($, e, next) => {
241    const result = await next(e)
242    await refreshStatus($)
243
244    return result
245  })
246
247  on('session.end', async ($, e, next) => {
248    if (e.reason === 'clear') {
249      await update($, overrides, () => ({ subagents: null, monitors: null }))
250      await update($, monitors, () => [])
251      $.ui.status(undefined)
252    }
253
254    return next(e)
255  })
256}
257

1export const REASON_PREFIX = 'Over-limit reason:'
2
3export type Kind = 'subagents' | 'monitors'
4
5const NUMBER_WORDS: Record<string, number> = {
6  zwei: 2, drei: 3, vier: 4, fünf: 5, sechs: 6, sieben: 7, acht: 8, neun: 9, zehn: 10, elf: 11, zwölf: 12,
7  two: 2, three: 3, four: 4, five: 5, six: 6, seven: 7, eight: 8, nine: 9, ten: 10, eleven: 11, twelve: 12,
8}
9
10const COUNT = `(\\d+|${Object.keys(NUMBER_WORDS).join('|')})`
11
12// "Nutze 8 Subagents", "starte fünf Subagenten", "use up to six agents" — a
13// count directly before the noun, optionally with "parallel" in between.
14const REQUEST_PATTERNS: Record<Kind, RegExp> = {
15  subagents: new RegExp(`(?<![\\p{L}\\d])${COUNT}\\s+(?:parallele?n?\\s+|parallel\\s+)?(?:sub-?agent(?:s|en)?|agent(?:s|en)?)(?![\\p{L}])`, 'iu'),
16  monitors: new RegExp(`(?<![\\p{L}\\d])${COUNT}\\s+(?:parallele?n?\\s+|parallel\\s+)?monitor(?:s|e|en)?(?![\\p{L}])`, 'iu'),
17}
18
19export function requestedLimits(text: string): Partial<Record<Kind, number>> {
20  const limits: Partial<Record<Kind, number>> = {}
21  for (const kind of ['subagents', 'monitors'] as const) {
22    const match = REQUEST_PATTERNS[kind].exec(text)
23    const token = match?.[1]?.toLowerCase()
24    const value = token === undefined ? Number.NaN : (NUMBER_WORDS[token] ?? Number(token))
25    if (Number.isInteger(value) && value > 0) {
26      limits[kind] = value
27    }
28  }
29
30  return limits
31}
32
33export type Reasoned = { reason: string | null; rest: string }
34
35// Subagents carry the reason as the first line of their prompt.
36export function takeReasonLine(prompt: string): Reasoned {
37  const [first = '', ...others] = prompt.split('\n')
38  if (!first.trim().startsWith(REASON_PREFIX)) {
39    return { reason: null, rest: prompt }
40  }
41
42  return { reason: first.trim().slice(REASON_PREFIX.length).trim(), rest: others.join('\n').replace(/^\n+/, '') }
43}
44
45// Monitors carry it in their one-line description: "Over-limit reason: <why> | <description>".
46export function takeReasonPrefix(description: string): Reasoned {
47  const trimmed = description.trim()
48  if (!trimmed.startsWith(REASON_PREFIX)) {
49    return { reason: null, rest: description }
50  }
51  const body = trimmed.slice(REASON_PREFIX.length)
52  const separator = body.lastIndexOf('|')
53  if (separator === -1) {
54    return { reason: body.trim(), rest: body.trim() }
55  }
56
57  return { reason: body.slice(0, separator).trim(), rest: body.slice(separator + 1).trim() }
58}
59
60export type TaskEnd = { taskId: string; status: string }
61
62// Background tasks report their end as a <task-notification> delivered like a prompt.
63export function parseTaskEnds(text: string): TaskEnd[] {
64  const ends: TaskEnd[] = []
65  const pattern = /<task-notification>([\s\S]*?)<\/task-notification>/g
66  for (const match of text.matchAll(pattern)) {
67    const body = match[1] ?? ''
68    const taskId = /<task-id>\s*([^<\s]+)\s*<\/task-id>/.exec(body)?.[1]
69    const status = /<status>\s*([^<\s]+)\s*<\/status>/.exec(body)?.[1]
70    if (taskId !== undefined && status !== undefined && ['completed', 'failed', 'killed', 'stopped', 'timeout'].includes(status)) {
71      ends.push({ taskId, status })
72    }
73  }
74
75  return ends
76}
77

1export type TrackedMonitor = { taskId: string; description: string; deadline: number | null }
2
3export type LimitOverrides = { subagents: number | null; monitors: number | null }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'concurrency-guard': {
8      monitors: TrackedMonitor[]
9      overrides: LimitOverrides
10    }
11  }
12}
13