time-budget

1import type { EngineInterface, Register, Timer, TurnStepChunk, TurnStepResult } from 'claude-code'
2import type { Report, Unit } from '../types'
3import {
4  CHECKPOINTS, GRACE_MS, LAST, MAX_RESUMES, TOOL, budgetFor, budgetLine, calibrationFactor, crossed, dueAt, gate, haltReason,
5  isReport, mins, parsePairs, projectedMin, reportText, statusText, type Pair,
6} from './budget'
7import { AWAKE_DEF, register as registerAwake } from './awake'
8
9// time-budget: the user's rule, enforced rather than requested --
10// "작업을 시작하기 전에 시간을 추정하고, 보고하라. 추정 시간 기준 1/3, 1/2, 2/3,
11// 5/6 지점에 중간 보고하라. 추정 시간을 넘을 것 같다면 즉시 작업을 중단하고 보고하라."
12//
13// A unit is one clm tracker task on the main thread, or one subagent run. A
14// main thread with no doing task is ungated; once a task has a budget, each
15// checkpoint marks a report due and every other tool waits on it. A report
16// that projects past the budget, or the budget running out, denies every tool
17// but the report, and GRACE_MS after 6/6 a main turn still running is aborted.
18// The calibration history is time-budget.py's file, read and written in the
19// same shape, one pair per unit.
20
21const UNITS = { plugin: 'time-budget', key: 'units' } as const
22const PARKED = { plugin: 'time-budget', key: 'parked' } as const
23const CODEX = 'codex-subagent:'
24const MAIN_ABORT = 'main-abort'
25const STATUS_REFRESH_MS = 60_000
26
27// Unlimited mode is the user's unattended run ("맥북 꺼두고 자동사냥하고 싶음"):
28// no budget stops the work, and a turn that ends with clm tasks still open is
29// resumed by a plugin prompt, up to MAX_RESUMES in a row without the user.
30const UNLIMITED = { plugin: 'time-budget', key: 'unlimited' } as const
31const RESUMES = { plugin: 'time-budget', key: 'resumes' } as const
32const UNLIMITED_ENV = 'CLAUDE_TIME_BUDGET_UNLIMITED'
33const UNLIMITED_COMMAND = 'budget-unlimited'
34const OPEN_STATUSES = new Set(['todo', 'doing'])
35
36const minutes = { type: 'number', exclusiveMinimum: 0 }
37const TOOLS = [
38  {
39    name: 'estimate',
40    description: 'Open a new clm tracker task and declare how long it will take, or pass task with an existing issue id to resume or re-estimate that task. The harness scales the estimate by how far past estimates ran, shows the budget to the user, and asks for a report at 1/3, 1/2, 2/3 and 5/6 of it.',
41    inputSchema: {
42      type: 'object',
43      properties: {
44        minutes: { ...minutes, description: 'Your estimate for the whole task, in minutes' },
45        scope: { type: 'string', description: 'What the task covers, in one line' },
46        steps: { type: 'array', items: { type: 'string' }, description: 'The planned steps, in order' },
47        task: { type: 'string', description: 'An existing clm issue id to resume' },
48      },
49      required: ['minutes', 'scope', 'steps'],
50    },
51  },
52  {
53    name: 'report',
54    description: 'File a progress report against the time budget: what is verified done, what is open, and the extra minutes the open items need. The harness projects the finish from it.',
55    inputSchema: {
56      type: 'object',
57      properties: {
58        done: { type: 'array', items: { type: 'object', properties: { item: { type: 'string' }, check: { type: 'string', description: 'The check that proved it done' } }, required: ['item', 'check'] } },
59        open: { type: 'array', items: { type: 'object', properties: { item: { type: 'string' }, next: { type: 'string', description: 'The next action on it' } }, required: ['item', 'next'] } },
60        more_min: { type: 'number', minimum: 0, description: 'Extra minutes the open items need beyond the budget' },
61      },
62      required: ['done', 'open'],
63    },
64  },
65  {
66    name: 'grant',
67    description: 'Give a running subagent more minutes from now on its time budget, after reading its report.',
68    inputSchema: {
69      type: 'object',
70      properties: { agent_id: { type: 'string', description: 'The subagent id its report names' }, minutes },
71      required: ['agent_id', 'minutes'],
72    },
73  },
74] as const
75
76type Ctx = EngineInterface
77type S = {
78  units: Record<string, Unit>
79  parked: Record<string, Unit>
80  timers: Map<string, Timer>
81  handedBack: Map<string, boolean>
82  extending: Set<string>
83  runningTurn?: string
84  pairs?: Pair[]
85  calibrationWrite?: Promise<void>
86  refresh?: Timer
87  unlimited: boolean
88  envUnlimited: boolean
89  resumes: number
90}
91
92const isUnlimited = (s: S) => s.unlimited || s.envUnlimited
93
94const errText = (err: unknown) => (err instanceof Error ? err.message : String(err))
95
96function log($: Ctx, text: string) {
97  $.ui.log(`[time-budget] ${text}`, { to: 'debug' })
98}
99
100async function save($: Ctx, s: S) {
101  try {
102    await $.state.set(UNITS, structuredClone(s.units))
103    await $.state.set(PARKED, structuredClone(s.parked))
104  } catch (err) {
105    log($, `state.set failed: ${errText(err)}`)
106  }
107}
108
109async function showStatus($: Ctx, s: S) {
110  const u = mainEntry(s)?.[1]
111  try {
112    $.ui.status(u ? statusText(u, await $.clock.now()) : undefined)
113  } catch (err) {
114    log($, `ui.status failed: ${errText(err)}`)
115  }
116}
117
118function mainEntry(s: S): [string, Unit] | undefined {
119  return Object.entries(s.units).find(([, unit]) => unit.kind === 'main')
120}
121
122async function calibrationFile($: Ctx) {
123  const dir = (await $.env.get('TIME_BUDGET_DIR')) ?? `${(await $.env.get('HOME')) ?? '.'}/.local/share/claude-time-budget`
124  return `${dir}/calibration.jsonl`
125}
126
127async function loadPairs($: Ctx, s: S): Promise<Pair[]> {
128  if (s.pairs) return s.pairs
129  const path = await calibrationFile($)
130  try {
131    s.pairs = (await $.fs.exists(path)) ? parsePairs((await $.fs.read(path)) as string) : []
132  } catch (err) {
133    log($, `cannot read ${path}: ${errText(err)}; calibration falls back to 1.0`)
134    s.pairs = []
135  }
136  return s.pairs
137}
138
139// One pair per unit, logged when the unit closes; `end` is when its last turn ended.
140async function logPair($: Ctx, s: S, u: Unit, end: number) {
141  if (u.estimateMin === undefined || u.budgetMin === undefined) return
142  const pair: Pair = {
143    kind: u.kind, agent_type: u.agentType ?? null, estimate_min: u.estimateMin, budget_min: u.budgetMin,
144    actual_min: Math.round(((end - u.start - (u.pausedMs ?? 0)) / 60_000) * 100) / 100, at: Math.floor(end / 1000),
145  }
146  const write = async () => {
147    const path = await calibrationFile($)
148    try {
149      const prev = (await $.fs.exists(path)) ? ((await $.fs.read(path)) as string) : ''
150      const existing = parsePairs(prev)
151      const next = [...existing, pair]
152      const text = `${next.map(p => JSON.stringify(p)).join('\n')}\n`
153      // A random temp avoids same-second collisions; separate sessions can still lose an update between read and rename.
154      const temp = `${path}.tmp-${pair.at}-${Math.random().toString(36).slice(2)}`
155      await $.fs.write(temp, text)
156      const moved = await $.process.run(['mv', temp, path], { timeoutMs: 5_000 })
157      if (moved.exitCode !== 0) throw new Error(`mv ${temp} ${path} failed (${moved.exitCode}): ${moved.stderr}`)
158      s.pairs = next
159    } catch (err) {
160      log($, `cannot append ${JSON.stringify(pair)} to ${path}: ${errText(err)}`)
161    }
162  }
163  const pending = s.calibrationWrite ?? Promise.resolve()
164  s.calibrationWrite = pending.catch(() => undefined).then(write)
165  await s.calibrationWrite
166}
167
168function disarm(s: S, key: string) {
169  s.timers.get(key)?.cancel()
170  s.timers.delete(key)
171}
172
173function resume(s: S, key: string, u: Unit) {
174  disarm(s, key)
175  if (u.kind === 'main') disarm(s, MAIN_ABORT)
176  s.handedBack.delete(key)
177  u.halted = undefined
178  u.haltedAt = undefined
179  u.reportDue = undefined
180  u.report = undefined
181  u.reportAt = undefined
182  u.unread = undefined
183}
184
185// A main turn still running GRACE_MS after the stop has had its chance to report.
186function armAbort($: Ctx, s: S, key: string) {
187  disarm(s, MAIN_ABORT)
188  if (isUnlimited(s)) return
189  s.timers.set(MAIN_ABORT, $.clock.after(GRACE_MS, () => {
190    const turnId = s.runningTurn
191    if (isUnlimited(s) || !s.units[key]?.halted || !turnId) return
192    $.turn.abort({ turnId }).catch(err => log($, `turn.abort(${turnId}) failed: ${errText(err)}`))
193  }))
194}
195
196// Moves the unit to the checkpoint the clock has reached; the timer and every tool call run it.
197async function advance($: Ctx, s: S, key: string, now: number) {
198  const u = s.units[key]
199  if (!u || u.budgetMin === undefined || u.halted) return
200  const c = crossed(now - u.start, u.budgetMin)
201  let changed = false
202  if (c > u.fired) {
203    u.fired = c
204    if (c < LAST) u.reportDue = CHECKPOINTS[c - 1]!.label
205    changed = true
206  }
207  const reason = haltReason(u, now, isUnlimited(s))
208  if (reason) {
209    u.halted = reason
210    u.haltedAt = now
211    u.reportDue = undefined
212    if (u.kind === 'main') armAbort($, s, key)
213    changed = true
214  }
215  if (!changed) return
216  await save($, s)
217  if (u.kind === 'main') await showStatus($, s)
218}
219
220async function arm($: Ctx, s: S, key: string) {
221  disarm(s, key)
222  const u = s.units[key]
223  if (!u || u.budgetMin === undefined || u.halted || u.fired >= LAST) return
224  const delay = Math.max(0, u.start + dueAt(u.fired, u.budgetMin) - (await $.clock.now()))
225  s.timers.set(key, $.clock.after(delay, () => {
226    void (async () => {
227      await advance($, s, key, await $.clock.now())
228      await arm($, s, key)
229    })().catch(err => log($, `checkpoint for ${key} failed: ${errText(err)}`))
230  }))
231}
232
233async function closeMain($: Ctx, s: S, now: number, key = mainEntry(s)?.[0]) {
234  if (!key) return
235  const u = s.units[key]
236  if (!u) return
237  disarm(s, key)
238  if (u.kind === 'main') disarm(s, MAIN_ABORT)
239  delete s.units[key]
240  // Its clm issue can be resumed by id, so the start and budget it ran under are kept for that.
241  if (u.kind === 'main') s.parked[key] = { ...u, calibrated: true }
242  if (!u.calibrated) await logPair($, s, u, u.lastEnd ?? now)
243}
244
245async function syncMainTask($: Ctx, s: S, now: number) {
246  const entry = mainEntry(s)
247  if (!entry) return
248  const [key] = entry
249  try {
250    const issue = (await $.clm.issues({ issue: key }))[0]
251    if (issue?.status === 'done' || issue?.status === 'dropped') {
252      await closeMain($, s, now, key)
253      await save($, s)
254    }
255  } catch (err) {
256    log($, `clm task status failed for ${key}: ${errText(err)}`)
257  }
258}
259
260function extend(s: S, key: string, u: Unit, extra: number, now: number): boolean {
261  if (u.budgetMin === undefined) return false
262  u.pausedMs = (u.pausedMs ?? 0) + Math.max(0, now - (u.lastEnd ?? now))
263  const fromNow = (now - u.start) / 60_000 + extra
264  // A halted unit restarts from now; a running one is never shortened by a grant.
265  u.budgetMin = u.halted ? fromNow : Math.max(u.budgetMin, fromNow)
266  resume(s, key, u)
267  u.fired = crossed(now - u.start, u.budgetMin)
268  return true
269}
270
271async function extendMain($: Ctx, s: S, extra: number): Promise<boolean> {
272  const entry = mainEntry(s)
273  if (!entry) return false
274  const [key, u] = entry
275  if (!u.halted || u.budgetMin === undefined || s.extending.has(key)) return false
276  s.extending.add(key)
277  const halted = u.halted
278  u.halted = 'extension in progress'
279  try {
280    const now = await $.clock.now()
281    const extended = extend(s, key, u, extra, now)
282    if (!extended) {
283      u.halted = halted
284      return false
285    }
286    await save($, s)
287    await arm($, s, key)
288    await showStatus($, s)
289    return true
290  } catch (err) {
291    u.halted = halted
292    throw err
293  } finally {
294    s.extending.delete(key)
295  }
296}
297
298function extensionMinutes(u: Unit): number[] {
299  const asked = u.report?.more_min
300  return [...new Set([10, 20, 30, ...(asked && asked > 0 ? [asked] : [])])]
301}
302
303function unreadReports(s: S, now: number): string[] {
304  const out: string[] = []
305  for (const [id, u] of Object.entries(s.units)) {
306    if (!u.unread || !u.report) continue
307    u.unread = false
308    out.push(`Subagent report (${u.agentType ?? 'subagent'} ${id}, ${mins(now - u.start)}/${Math.round(u.budgetMin ?? 0)} min${u.halted ? `, stopped: ${u.halted}` : ''}):\n${reportText(u.report)}\nGrant it time with ${TOOL.grant} {agent_id: "${id}", minutes}, or narrow its scope with SendMessage.`)
309  }
310  return out
311}
312
313async function saveLoop($: Ctx, s: S) {
314  try {
315    await $.state.set(UNLIMITED, s.unlimited)
316    await $.state.set(RESUMES, s.resumes)
317  } catch (err) {
318    log($, `state.set (unlimited) failed: ${errText(err)}`)
319  }
320}
321
322// Turning unlimited on lifts every stop already in force, and the abort waiting on it.
323async function liftHalts($: Ctx, s: S) {
324  const now = await $.clock.now()
325  for (const [key, u] of Object.entries(s.units)) {
326    if (!u.halted || u.budgetMin === undefined) continue
327    resume(s, key, u)
328    u.fired = crossed(now - u.start, u.budgetMin)
329    await arm($, s, key)
330  }
331  disarm(s, MAIN_ABORT)
332  await save($, s)
333  await showStatus($, s)
334}
335
336// At a main turn's end in unlimited mode: queue a resume prompt while clm work is open.
337async function autoResume($: Ctx, s: S) {
338  if (!isUnlimited(s) || s.resumes >= MAX_RESUMES) return
339  let open
340  try {
341    open = (await $.clm.issues()).filter(i => OPEN_STATUSES.has(i.status))
342  } catch (err) {
343    log($, `clm.issues failed, no auto-resume: ${errText(err)}`)
344    return
345  }
346  if (!open.length) return
347  s.resumes++
348  await saveLoop($, s)
349  const list = open.map(i => `- ${i.id} [${i.status}] ${i.title}`).join('\n')
350  const text = `Unlimited mode, auto-resume ${s.resumes}/${MAX_RESUMES}: the user is away and these clm tasks are still open. Continue them now, verify each before marking it done, and mark a task dropped if it cannot be finished without the user.\n${list}`
351  // The prompt starts its own turn once the session is idle, so awaiting it here would hold this turn's end.
352  void $.prompt.submit({ text }).catch(err => log($, `auto-resume submit failed: ${errText(err)}`))
353}
354
355export const register: Register = (on) => {
356  registerAwake(on)
357  const s: S = { units: {}, parked: {}, timers: new Map(), handedBack: new Map(), extending: new Set(), unlimited: false, envUnlimited: false, resumes: 0 }
358
359  on('session.start', async ($, e, next) => {
360    const started = await next(e)
361    try {
362      const read = await $.state.get(UNITS)
363      s.units = read.value ? { ...read.value } : {}
364      const parked = await $.state.get(PARKED)
365      s.parked = parked.value ? { ...parked.value } : {}
366      s.unlimited = (await $.state.get(UNLIMITED)).value === true
367      s.resumes = (await $.state.get(RESUMES)).value ?? 0
368    } catch (err) {
369      log($, `state.get failed: ${errText(err)}`)
370    }
371    s.envUnlimited = /^(1|true|on|yes)$/i.test((await $.env.get('CLAUDE_TIME_BUDGET_UNLIMITED')) ?? '')
372    try {
373      await $.command.register({
374        name: UNLIMITED_COMMAND,
375        description: `on|off: unattended mode. No time budget stops the work, and a turn that ends with clm tasks open is resumed automatically, up to ${MAX_RESUMES} times in a row until you type a prompt. ${UNLIMITED_ENV}=1 in the environment turns it on too.`,
376        immediate: true,
377      })
378    } catch (err) {
379      log($, `command.register(${UNLIMITED_COMMAND}) failed: ${errText(err)}`)
380    }
381    // The old prompt-scoped state used the literal key `main`; it cannot be tied to a clm issue.
382    if (s.units.main?.kind === 'main') {
383      disarm(s, 'main')
384      disarm(s, MAIN_ABORT)
385      delete s.units.main
386      await save($, s)
387    }
388    for (const t of [...TOOLS, AWAKE_DEF]) {
389      try {
390        await $.tool.register({ ...t, inputSchema: t.inputSchema as unknown as Record<string, unknown> })
391      } catch (err) {
392        log($, `tool.register(${t.name}) failed: ${errText(err)}`)
393      }
394    }
395    for (const key of Object.keys(s.units)) await arm($, s, key)
396    if (isUnlimited(s)) await liftHalts($, s)
397    s.refresh?.cancel()
398    s.refresh = $.clock.every(STATUS_REFRESH_MS, () => { void showStatus($, s) })
399    await showStatus($, s)
400    return started
401  })
402
403  on('session.end', async ($, e, next) => {
404    await closeMain($, s, await $.clock.now())
405    await save($, s)
406    return next(e)
407  })
408
409  on('turn.start', async ($, e, next) => {
410    s.runningTurn = e.turnId
411    return next(e)
412  })
413
414  on('command.run', { command: UNLIMITED_COMMAND }, async ($, e) => {
415    const arg = (e.args ?? '').trim().toLowerCase()
416    if (arg === 'on' || arg === 'off') {
417      s.unlimited = arg === 'on'
418      s.resumes = 0
419      await saveLoop($, s)
420      if (isUnlimited(s)) await liftHalts($, s)
421    } else if (arg) {
422      return { text: `Usage: /${UNLIMITED_COMMAND} on|off (or ${UNLIMITED_ENV}=1 in the environment)` }
423    }
424    const env = s.envUnlimited ? ` ${UNLIMITED_ENV} is set, so it stays on until that is unset.` : ''
425    return { text: isUnlimited(s)
426      ? `Unlimited mode is on: no time budget stops the work, and a turn ending with clm tasks open is resumed automatically (${s.resumes}/${MAX_RESUMES} in a row so far).${env}`
427      : 'Unlimited mode is off: the time budget stops the work as usual.' }
428  })
429
430  // The user's own prompt ends a run of auto-resumes; a plugin's (ours included) does not.
431  on('prompt.submit', async ($, e, next) => {
432    if (e.origin.kind !== 'plugin' && e.origin.kind !== 'unclassified' && s.resumes !== 0) {
433      s.resumes = 0
434      await saveLoop($, s)
435    }
436    return next(e)
437  })
438
439  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
440    const entry = mainEntry(s)
441    const u = entry?.[1]
442    if (!u?.halted || isUnlimited(s) || e.props.hasSurvey) return next(e)
443    const { Box, Text, Button } = $.ui.resolve(e)
444    const buttons = extensionMinutes(u).map(extra => Button({
445      key: `time-budget:+${extra}m`,
446      label: `+${extra}m`,
447      onPress: () => {
448        void (async () => {
449          if (await extendMain($, s, extra)) await $.prompt.submit({ text: 'Continue this task.', asUser: true })
450        })().catch(err => log($, `extension +${extra}m failed: ${errText(err)}`))
451      },
452    }))
453    return Box({
454      gap: 1,
455      children: [Text({ children: `Time budget stopped: ${u.halted}` }), ...buttons],
456    })
457  })
458
459  on('turn.complete', async ($, e, next) => {
460    const r = await next(e)
461    const now = await $.clock.now()
462    if (e.agentId) {
463      const u = s.units[e.agentId]
464      if (u) {
465        disarm(s, e.agentId)
466        s.handedBack.delete(e.agentId)
467        delete s.units[e.agentId]
468        await logPair($, s, u, now)
469        await save($, s)
470      }
471      return r
472    }
473    if (s.runningTurn === e.turnId) s.runningTurn = undefined
474    await syncMainTask($, s, now)
475    // An aborted turn is the user's Esc, which the loop leaves alone.
476    if (!e.isAborted) await autoResume($, s)
477    const u = mainEntry(s)?.[1]
478    if (!u) return r
479    u.lastEnd = now
480    await save($, s)
481    // The status line already shows the budget; a line under every answer becomes a focus-view notice row.
482    if (u.budgetMin === undefined || !(u.halted || u.reportDue)) return r
483    const head = `⏱ estimate ${u.estimateMin}m → ${statusText(u, now)}`
484    const line = u.report ? `${head}\n${reportText(u.report)}` : head
485    // A text other than the answer is shown beneath it; another plugin's line stays above ours.
486    return { ...r, text: r.text === e.answer ? line : `${r.text}\n${line}` }
487  })
488
489  on('agent.spawn', async ($, e, next) => {
490    const codex = e.subagentType.startsWith(CODEX)
491    const declared = budgetLine(e.prompt)
492    if (codex && declared === undefined) {
493      return { deny: `Start the prompt with the line "Budget: <N> min", N being your estimate in minutes for this ${e.subagentType} task, then call Agent again. Codex runs its own shell where no estimate gate reaches it, so the spawner's line is its budget.` }
494    }
495    const r = await next(e)
496    if (!r.agentId) return r
497    const now = await $.clock.now()
498    const u: Unit = { kind: 'sub', agentType: e.subagentType, start: now, fired: 0 }
499    if (declared !== undefined) {
500      u.estimateMin = declared
501      u.budgetMin = budgetFor(declared, calibrationFactor(await loadPairs($, s)))
502    }
503    s.units[r.agentId] = u
504    await save($, s)
505    await arm($, s, r.agentId)
506    return r
507  })
508
509  on('tool.call', async ($, e, next) => {
510    const key = e.agentId
511    const main = key === undefined
512    const now = await $.clock.now()
513    // An agent id no spawn of ours named (the engine's own forks) runs ungated.
514    if (main) await syncMainTask($, s, now)
515    let unitKey = key ?? mainEntry(s)?.[0]
516    let u = unitKey ? s.units[unitKey] : undefined
517    if (u && unitKey) await advance($, s, unitKey, now)
518    const args = e as unknown as Record<string, unknown>
519
520    if (e.tool === TOOL.estimate) {
521      const m = args.minutes
522      if (typeof m !== 'number' || !(m > 0)) return { deny: `Call ${TOOL.estimate} with minutes as a positive number, scope as one line and steps as a list.` }
523      const factor = calibrationFactor(await loadPairs($, s))
524      const budget = budgetFor(m, factor)
525      const scope = typeof args.scope === 'string' ? args.scope : ''
526      if (main) {
527        const requested = args.task
528        if (requested !== undefined && (typeof requested !== 'string' || !requested.trim()))
529          return { deny: `Call ${TOOL.estimate} with task as an existing clm issue id when resuming one.` }
530        const targetKey = typeof requested === 'string' ? requested : undefined
531        // Opening, switching or raising a task would bypass the stop, the due report, or the user's extension buttons.
532        if (u?.halted) return { deny: `Stop here: ${u.halted}. Call ${TOOL.report}, then end the turn; the user extends with a button above the prompt.` }
533        if (u?.reportDue && targetKey !== unitKey) return { deny: `Checkpoint ${u.reportDue}: call ${TOOL.report} before opening or re-estimating a task.` }
534        const prior = targetKey === undefined ? undefined : targetKey === unitKey ? u : s.parked[targetKey]
535        if (prior?.budgetMin !== undefined && budget > prior.budgetMin)
536          return { deny: `The budget of ${Math.round(prior.budgetMin)} min only grows by the user's extension button; put more_min in your ${TOOL.report} call and the button appears when it projects past the budget.` }
537        let issue
538        try {
539          issue = await $.clm.track({ title: scope, status: 'doing', ...(targetKey ? { issue: targetKey } : {}) })
540        } catch (err) {
541          return { deny: `${TOOL.estimate}: ${errText(err)}` }
542        }
543        if (unitKey && unitKey !== issue.id) await closeMain($, s, now, unitKey)
544        unitKey = issue.id
545        u = s.units[unitKey]
546        const parked = s.parked[unitKey]
547        delete s.parked[unitKey]
548        // A resumed task keeps its first start and budget; a lower estimate applies below.
549        if (!u || u.kind !== 'main') u = s.units[unitKey] = parked
550          ? { ...parked, lastEnd: undefined }
551          : { kind: 'main', start: now, fired: 0 }
552      } else {
553        unitKey = key
554        u = s.units[unitKey] ?? (s.units[unitKey] = { kind: 'sub', start: now, fired: 0 })
555        if (u.budgetMin !== undefined && budget > u.budgetMin)
556          return { deny: `Put more_min in your ${TOOL.report} call; main can grant time.` }
557      }
558      if (!u || !unitKey) return { deny: `Could not open a time-budget unit for ${TOOL.estimate}.` }
559      const unit = u
560      const first = unit.budgetMin === undefined
561      if (first) unit.estimateMin = m
562      unit.budgetMin = budget
563      unit.scope = scope || undefined
564      unit.steps = Array.isArray(args.steps) ? args.steps.filter((s): s is string => typeof s === 'string') : undefined
565      if (first) {
566        unit.fired = crossed(now - unit.start, unit.budgetMin)
567        unit.reportDue = undefined
568      } else {
569        await advance($, s, unitKey, now)
570      }
571      await save($, s)
572      await arm($, s, unitKey)
573      if (main) await showStatus($, s)
574      const at = CHECKPOINTS.slice(0, LAST - 1).map((c, i) => `${c.label} at ${mins(dueAt(i, unit.budgetMin!))}m`).join(', ')
575      return { result: `Budget ${Math.round(unit.budgetMin)} min (estimate ${m} min × calibration ${factor.toFixed(2)}, floor 15); ${mins(now - unit.start)} min elapsed. Reports are due at ${at}; the budget ends at ${Math.round(unit.budgetMin)}m. State this budget and the steps to the user in your reply as you start.` }
576    }
577
578    const denied = gate(u, e.tool, now)
579    if (denied) return denied
580
581    if (e.tool === TOOL.report) {
582      if (!u) return { deny: `Call ${TOOL.estimate} first; a report measures against its budget.` }
583      if (!isReport(args)) return { deny: `Call ${TOOL.report} with done as [{item, check}], open as [{item, next}], and more_min as a number of minutes when the open items need more time.` }
584      const report: Report = { done: args.done, open: args.open, ...(args.more_min !== undefined ? { more_min: args.more_min } : {}) }
585      u.report = report
586      u.reportAt = now
587      u.reportDue = undefined
588      if (!main) u.unread = true
589      if (unitKey) await advance($, s, unitKey, now)
590      await save($, s)
591      if (main) await showStatus($, s)
592      const projected = Math.round(projectedMin(now - u.start, report))
593      if (u.halted) {
594        const close = main
595          ? 'End the turn now with this report; use an extension button above the prompt to continue.'
596          : `End your run now: call SubagentHandback with this report; main reads it and can grant time.`
597        return { result: `Report recorded. Stop here: ${u.halted}. ${close}` }
598      }
599      return { result: `Report recorded: projected ${projected} of ${Math.round(u.budgetMin!)} min. Continue at the same depth; verification stays part of the work.` }
600    }
601
602    if (e.tool === TOOL.grant) {
603      if (e.agentId) return { deny: `Ask the main thread for time: put more_min in your ${TOOL.report} call.` }
604      const id = String(args.agent_id ?? '')
605      const extra = args.minutes
606      const sub = s.units[id]
607      if (!sub || sub.kind === 'main' || sub.budgetMin === undefined) return { deny: `Call ${TOOL.grant} with the agent_id of a running subagent that has a budget (one of: ${Object.entries(s.units).filter(([, candidate]) => candidate.kind === 'sub').map(([candidateId]) => candidateId).join(', ') || 'none'}).` }
608      if (typeof extra !== 'number' || !(extra > 0)) return { deny: `Call ${TOOL.grant} with minutes as a positive number.` }
609      const before = sub.budgetMin
610      const wasHalted = Boolean(sub.halted)
611      extend(s, id, sub, extra, now)
612      await save($, s)
613      await arm($, s, id)
614      const after = Math.round(sub.budgetMin)
615      return { result: wasHalted || sub.budgetMin > before
616        ? `Granted +${extra} min from now to ${id}; its budget is now ${after} min.`
617        : `${id} already has more than +${extra} min left; its budget stays ${after} min.` }
618    }
619
620    const r = await next(e)
621    if (!main || 'deny' in r && r.deny !== undefined) return r
622    const reports = unreadReports(s, now)
623    if (!reports.length) return r
624    await save($, s)
625    return { ...r, context: [...(r.context ?? []), ...reports] } as typeof r
626  })
627
628  // A Claude subagent that keeps stepping GRACE_MS after its stop hands back
629  // the harness's report instead of making another request. A codex agent's
630  // steps are codex-subagent's to answer, so it is left to the tool gate.
631  on('turn.step', async function* ($, e, next) {
632    const id = e.agentId
633    const u = id ? s.units[id] : undefined
634    if (!id || !u?.halted || u.agentType?.startsWith(CODEX)) return yield* next(e)
635    const now = await $.clock.now()
636    if (now - (u.haltedAt ?? now) < GRACE_MS) return yield* next(e)
637    const text = `Stopped by the time budget: ${u.halted} (${mins(now - u.start)}/${Math.round(u.budgetMin ?? 0)} min).\n${u.report ? reportText(u.report) : 'No report was filed.'}`
638    if (s.handedBack.get(id)) return yield* endTurn(e, text)
639    s.handedBack.set(id, true)
640    const toolUseId = `toolu_time_budget_${id.replace(/[^A-Za-z0-9_-]/g, '_')}_${e.index}`
641    const input = { message: text }
642    yield { kind: 'tool', index: 0, id: toolUseId, name: 'SubagentHandback' }
643    yield { kind: 'input', index: 0, json: JSON.stringify(input) }
644    yield { kind: 'stop', stopReason: 'tool_use', usage: null }
645    return { turnId: e.turnId, index: e.index, answer: '', toolUses: [{ name: 'SubagentHandback', input }], stopReason: 'tool_use', usage: null } satisfies TurnStepResult
646  })
647}
648
649async function* endTurn(e: { turnId: string; index: number }, text: string): AsyncGenerator<TurnStepChunk, TurnStepResult> {
650  yield { kind: 'text', index: 0, text }
651  yield { kind: 'stop', stopReason: 'end_turn', usage: null }
652  return { turnId: e.turnId, index: e.index, answer: text, toolUses: [], stopReason: 'end_turn', usage: null }
653}
654

1// The time budget's arithmetic and its gate, with no `$`: what register.ts
2// asks of a unit at each tool call, timer and report.
3//
4// WHY THE HARNESS MEASURES: a model's own time estimate runs 3-10x off, and a
5// mid-task "percent done" self-report is unreliable (arXiv 2609.08589). So the
6// wall clock lives here, and the estimate only seeds a budget after it is
7// scaled by the p90 of actual/estimate over past units, clamped to [0.25, 10].
8// A report lists verified-done and open items, never a percentage, which
9// relieves the late-task pull toward closing over verifying (arXiv 2609.00823).
10
11import type { Report, Unit } from '../types'
12
13// The user's notation; keep exactly these five.
14export const CHECKPOINTS = [
15  { label: '1/3', f: 1 / 3 },
16  { label: '1/2', f: 1 / 2 },
17  { label: '2/3', f: 2 / 3 },
18  { label: '5/6', f: 5 / 6 },
19  { label: '6/6', f: 1 },
20] as const
21export const LAST = CHECKPOINTS.length
22
23export const FLOOR_MIN = 15
24const MIN_HISTORY = 5
25export const HISTORY_WINDOW = 50
26const FACTOR_MIN = 0.25
27const FACTOR_MAX = 10
28// After 6/6, a main turn that is still running gets this long to write its
29// report with every tool but `report` denied, then it is aborted.
30export const GRACE_MS = 2 * 60_000
31// Unlimited mode auto-resumes a turn that ended with clm tasks open at most this many times in a row.
32export const MAX_RESUMES = 50
33
34export const TOOL = {
35  estimate: 'mcp__time-budget__estimate',
36  report: 'mcp__time-budget__report',
37  grant: 'mcp__time-budget__grant',
38} as const
39const HANDBACK = 'SubagentHandback'
40const BEFORE_ESTIMATE = new Set<string>([TOOL.estimate, 'ToolSearch', HANDBACK])
41const WHILE_REPORT_DUE = new Set<string>([TOOL.report, TOOL.estimate, 'ToolSearch', HANDBACK])
42const WHILE_HALTED = new Set<string>([TOOL.report, HANDBACK])
43
44/** One line of calibration.jsonl, as time-budget.py wrote it. */
45export type Pair = { kind: 'main' | 'sub'; agent_type: string | null; estimate_min: number; budget_min: number; actual_min: number; at: number }
46
47export function parsePairs(text: string): Pair[] {
48  const out: Pair[] = []
49  for (const line of text.split('\n')) {
50    if (!line.trim()) continue
51    try {
52      const p = JSON.parse(line) as unknown
53      if (p && typeof p === 'object' && !Array.isArray(p)) out.push(p as Pair)
54    } catch { /* a torn line is skipped, as the python reader did */ }
55  }
56  return out
57}
58
59export function calibrationFactor(pairs: readonly Partial<Pair>[]): number {
60  const ratios = pairs.slice(-HISTORY_WINDOW)
61    .filter(p => (p.estimate_min ?? 0) > 0 && (p.actual_min ?? -1) >= 0)
62    .map(p => p.actual_min! / p.estimate_min!)
63    .sort((a, b) => a - b)
64  if (ratios.length < MIN_HISTORY) return 1
65  const p90 = ratios[Math.max(0, Math.ceil(0.9 * ratios.length) - 1)]!
66  return Math.min(FACTOR_MAX, Math.max(FACTOR_MIN, p90))
67}
68
69export const budgetFor = (estimateMin: number, factor: number) => Math.max(FLOOR_MIN, estimateMin * factor)
70
71/** When checkpoint `i` (0-based) falls, in whole ms since the unit's start. Whole ms keep the timer and
72 * `crossed` in agreement: with a fractional due time, start + due rounded to the timer's ms while
73 * elapsed stayed below due, so the timer fired, found nothing crossed, and re-armed at delay 0 forever. */
74export const dueAt = (i: number, budgetMin: number) => Math.ceil(CHECKPOINTS[i]!.f * budgetMin * 60_000)
75
76/** How many checkpoints the elapsed time has passed. */
77export const crossed = (elapsedMs: number, budgetMin: number) => CHECKPOINTS.filter((_, i) => elapsedMs >= dueAt(i, budgetMin)).length
78
79/** `Budget: <N> min` as the prompt's first line: an optional spawner estimate; codex requires it because no gate reaches inside codex. */
80export const budgetLine = (prompt: string): number | undefined => {
81  const m = /^\s*Budget:\s*(\d+(?:\.\d+)?)\s*min\s*$/i.exec(prompt.split('\n', 1)[0] ?? '')
82  const n = m ? Number(m[1]) : NaN
83  return n > 0 ? n : undefined
84}
85
86/** Minutes the work will take at the pace the report shows: elapsed per done item, times the open items left. */
87export function projectedMin(elapsedMs: number, report: Report): number {
88  const elapsed = elapsedMs / 60_000
89  const rate = elapsed / Math.max(1, report.done.length)
90  return elapsed + rate * report.open.length
91}
92
93export const mins = (ms: number) => Math.round(ms / 60_000)
94
95const SHAPE = 'done (each item with the check that proved it), open (each item with its next action), and more_min when the open items need more time'
96
97export function gate(unit: Unit | undefined, tool: string, now: number): { deny: string } | undefined {
98  if (!unit) return undefined
99  if (unit.budgetMin === undefined) {
100    if (unit.kind === 'main') return undefined
101    if (BEFORE_ESTIMATE.has(tool)) return undefined
102    return { deny: `Call ${TOOL.estimate} first, with minutes (your estimate for the whole task), scope (one line) and steps (the plan). The user sees the calibrated budget; then call ${tool} again.` }
103  }
104  const elapsed = now - unit.start
105  if (unit.halted) {
106    if (WHILE_HALTED.has(tool)) return undefined
107    return { deny: `Stop here: ${unit.halted}. Call ${TOOL.report} with ${SHAPE}, then end the turn with that report; use an extension button above the prompt to continue.` }
108  }
109  if (unit.reportDue && !WHILE_REPORT_DUE.has(tool)) {
110    return { deny: `Checkpoint ${unit.reportDue} (${mins(elapsed)}/${Math.round(unit.budgetMin)} min): call ${TOOL.report} with ${SHAPE}, then call ${tool} again.` }
111  }
112  return undefined
113}
114
115/** `unlimited`: the user's unattended mode, where no budget ever stops the work. */
116export function haltReason(unit: Unit, now: number, unlimited = false): string | undefined {
117  if (unlimited || unit.budgetMin === undefined) return undefined
118  const elapsed = now - unit.start
119  if (elapsed >= unit.budgetMin * 60_000) return `the budget of ${Math.round(unit.budgetMin)} min is spent (${mins(elapsed)} min elapsed)`
120  if (unit.report) {
121    const p = projectedMin(unit.reportAt! - unit.start, unit.report)
122    if (p > unit.budgetMin) return `the last report projects ${Math.round(p)} min against a budget of ${Math.round(unit.budgetMin)} min`
123  }
124  return undefined
125}
126
127/** The line the user reads: in the status line and beneath a main answer. */
128export function statusText(unit: Unit, now: number): string {
129  if (unit.budgetMin === undefined) return 'budget: no estimate yet'
130  const elapsed = now - unit.start
131  const base = `budget ${mins(elapsed)}/${Math.round(unit.budgetMin)}m`
132  if (unit.halted) return `${base} · stopped: ${unit.halted}`
133  if (unit.reportDue) return `${base} · checkpoint ${unit.reportDue} report due`
134  const next = CHECKPOINTS[unit.fired]
135  return next ? `${base} · next ${next.label} at ${mins(dueAt(unit.fired, unit.budgetMin))}m` : base
136}
137
138export function reportText(r: Report): string {
139  const done = r.done.map(d => `- ${d.item} (checked: ${d.check})`).join('\n') || '- (none)'
140  const open = r.open.map(o => `- ${o.item} -> ${o.next}`).join('\n') || '- (none)'
141  return `Verified done:\n${done}\nOpen:\n${open}${r.more_min ? `\nAsks +${r.more_min} min` : ''}`
142}
143
144export function isReport(x: unknown): x is Report {
145  const r = x as Report
146  return !!r && Array.isArray(r.done) && Array.isArray(r.open)
147    && r.done.every(d => d && typeof d.item === 'string' && typeof d.check === 'string')
148    && r.open.every(o => o && typeof o.item === 'string' && typeof o.next === 'string')
149    && (r.more_min === undefined || (typeof r.more_min === 'number' && r.more_min >= 0))
150}
151

1import type { EngineInterface, Register } from 'claude-code'
2
3// stay_awake: the user's ask, "덮개 닫아도 절전 안 빠지는 걸 클로드가 툴로 연장할 수
4// 있어야 함". The tool only writes an epoch-seconds deadline; the root launchd
5// daemon `org.ranolp.claude-awake` (nix/darwin/default.nix) reads it every
6// minute and flips `pmset disablesleep`, so no sudo ever runs from here.
7
8export const AWAKE_TOOL = 'mcp__time-budget__stay_awake'
9export const MAX_AWAKE_MIN = 240
10
11/** Registered by register.ts's session.start alongside estimate/report/grant: one handler per event per module. */
12export const AWAKE_DEF = {
13  name: 'stay_awake',
14  description: `Keep the Mac awake with the lid closed for the next N minutes (max ${MAX_AWAKE_MIN} per call; 0 lets it sleep again). Call it again before the deadline to extend. A root daemon applies it within a minute, and only while the battery is at 20% or more or on AC power.`,
15  inputSchema: {
16    type: 'object',
17    properties: { minutes: { type: 'number', minimum: 0, description: `Minutes from now to stay awake, 0 to ${MAX_AWAKE_MIN}; 0 clears it` } },
18    required: ['minutes'],
19  },
20}
21
22/** The deadline, in epoch seconds, that `minutes` from `nowMs` writes; undefined clears it. */
23export function awakeUntil(nowMs: number, minutes: number): number | undefined {
24  const m = Math.min(Math.max(0, minutes), MAX_AWAKE_MIN)
25  if (m === 0) return undefined
26  return Math.floor(nowMs / 1000 + m * 60)
27}
28
29async function untilFile($: EngineInterface) {
30  const dir = `${(await $.env.get('HOME')) ?? '.'}/.local/state/claude-awake`
31  return { dir, path: `${dir}/until` }
32}
33
34const errText = (err: unknown) => (err instanceof Error ? err.message : String(err))
35
36export const register: Register = (on) => {
37  on('tool.call', { tool: AWAKE_TOOL }, async ($, e) => {
38    const m = (e as unknown as Record<string, unknown>).minutes
39    if (typeof m !== 'number' || !Number.isFinite(m) || m < 0) return { deny: `Call ${AWAKE_TOOL} with minutes as a number from 0 to ${MAX_AWAKE_MIN}.` }
40    const until = awakeUntil(await $.clock.now(), m)
41    const { dir, path } = await untilFile($)
42    try {
43      if (until === undefined) {
44        const rm = await $.process.run(['rm', '-f', path], { timeoutMs: 5_000 })
45        if (rm.exitCode !== 0) throw new Error(`rm -f ${path} failed (${rm.exitCode}): ${rm.stderr}`)
46        return { result: 'Stay-awake cleared; the Mac may sleep with the lid closed again within a minute.' }
47      }
48      const mk = await $.process.run(['mkdir', '-p', dir], { timeoutMs: 5_000 })
49      if (mk.exitCode !== 0) throw new Error(`mkdir -p ${dir} failed (${mk.exitCode}): ${mk.stderr}`)
50      await $.fs.write(path, `${until}\n`)
51    } catch (err) {
52      return { isError: true, text: `stay_awake could not update ${path}: ${errText(err)}` }
53    }
54    const clamped = m > MAX_AWAKE_MIN ? ` (clamped from ${m} to ${MAX_AWAKE_MIN} min)` : ''
55    return { result: `Staying awake until ${new Date(until * 1000).toISOString()} (epoch ${until})${clamped}; call again before then to extend.` }
56  })
57}
58

1export type Report = {
2  done: { item: string; check: string }[]
3  open: { item: string; next: string }[]
4  more_min?: number
5}
6
7/** One timed unit: a main-thread clm task, or one subagent run. */
8export type Unit = {
9  kind: 'main' | 'sub'
10  agentType?: string
11  /** `$.clock.now()` milliseconds when the unit began. */
12  start: number
13  estimateMin?: number
14  /** Calibrated budget; undefined until the estimate is declared, which is what the gate waits on. */
15  budgetMin?: number
16  scope?: string
17  steps?: string[]
18  /** How many of CHECKPOINTS have fired. */
19  fired: number
20  /** The label of the checkpoint whose report every other tool waits on. */
21  reportDue?: string
22  /** Why the unit stopped; every tool but the report is denied while set. */
23  halted?: string
24  haltedAt?: number
25  report?: Report
26  reportAt?: number
27  /** A subagent report main has not been shown yet. */
28  unread?: boolean
29  /** The last time a main turn of this unit ended: the unit's actual end for calibration. */
30  lastEnd?: number
31  /** Wall-clock time spent waiting for the user after a main turn ended. */
32  pausedMs?: number
33  /** The unit's calibration pair is already logged; a resumed main task logs none again. */
34  calibrated?: boolean
35}
36
37declare module 'claude-code' {
38  interface PluginState {
39    /** `parked`: closed main units whose clm task can still be resumed, keyed by issue id. */
40    /** `unlimited`: `/budget-unlimited on`; `resumes`: auto-resume prompts since the user's last own prompt. */
41    'time-budget': { units: Record<string, Unit>; parked: Record<string, Unit>; unlimited: boolean; resumes: number }
42  }
43}
44