caffeine

1import type { EngineInterface, Register } from 'claude-code'
2import { update } from 'claude-code'
3
4import type { Task } from '../types'
5import {
6  MINUTE, type Output, type Schedule, firstReminder, hasExited, holdReason, nextReminder, outputPath, reminder, schedule, squash, tailOf,
7  taskId, toastText,
8} from './reminders'
9
10const tasks = { plugin: 'caffeine', key: 'tasks' } as const
11const READ_LIMIT = 4 * 1024 * 1024
12
13/** Whether the main loop is inside a turn, and whether a reminder pass is under way; both start over on a reload. */
14const loop = { running: false, ticking: false }
15
16export const register: Register = (on, options) => {
17  const plan = schedule(options.marks, options.every)
18  let tick: { cancel: () => void } | null = null
19
20  on('session.start', async ($, e, next) => {
21    const started = await next(e)
22    loop.running = false
23    tick?.cancel()
24    tick = $.clock.every(MINUTE, () => void remind($, plan))
25    return started
26  })
27
28  on('turn.step', async function* ($, e, next) {
29    if (e.agentId === undefined) loop.running = true
30    return yield* next(e)
31  })
32
33  on('tool.call', { tool: 'Bash' }, async ($, e, next) => {
34    const ran = await next(e)
35    if (ran.deny !== undefined || ran.isError === true) return ran
36    const id = ran.result.backgroundTaskId
37    if (id === undefined) return ran
38    const now = await $.clock.now()
39    const startedAt = now - (ran.result.timedOutAfterMs ?? 0)
40    const task: Task = {
41      id,
42      command: squash(e.command),
43      description: squash(e.description ?? ''),
44      output: outputPath(ran.text ?? ''),
45      startedAt,
46      remindAt: nextReminder(startedAt, now, plan),
47      agentId: e.agentId ?? null,
48      parent: null,
49      endsWithCaller: ran.result.backgroundEndsWithFinalResponse === true,
50      held: false,
51    }
52    await update($, tasks, (all = {}) => ({ ...all, [id]: task }))
53    return { ...ran, context: [...(ran.context ?? []), firstReminder(task, now, plan)] }
54  })
55
56  on('tool.call', { tool: 'TaskStop' }, async ($, e, next) => {
57    const ran = await next(e)
58    if (ran.deny !== undefined || ran.isError === true) return ran
59    const id = e.task_id ?? e.shell_id
60    if (id !== undefined) await forget($, id)
61    return ran
62  })
63
64  on('prompt.submit', async ($, e, next) => {
65    if (e.origin.kind === 'task-notification') {
66      const id = taskId(e.text)
67      if (id !== null) await forget($, id)
68    }
69    return next(e)
70  })
71
72  on('classic.SubagentStop', async ($, e, next) => {
73    const stopped = await next(e)
74    if (stopped.block !== undefined) return stopped
75    const { value: all = {} } = await $.state.get(tasks)
76    const held: Task[] = []
77    for (const task of Object.values(all)) {
78      if (task.agentId !== e.agent_id || task.held || !task.endsWithCaller) continue
79      if (!(await outputOf($, task.output))?.done) held.push(task)
80    }
81    if (held.length === 0) return stopped
82    const ids = new Set(held.map(task => task.id))
83    await update($, tasks, (cur = {}) => Object.fromEntries(
84      Object.values(cur).map(task => [task.id, ids.has(task.id) ? { ...task, held: true } : task]),
85    ))
86    return { ...stopped, block: holdReason(held) }
87  })
88
89  on('turn.complete', async ($, e, next) => {
90    const done = await next(e)
91    if (e.agentId === undefined) {
92      loop.running = false
93      return done
94    }
95    const ended = e.agentId
96    const { value: before = {} } = await $.state.get(tasks)
97    const finished = new Set<string>()
98    for (const task of Object.values(before)) {
99      if (task.agentId === ended && (await outputOf($, task.output))?.done) finished.add(task.id)
100    }
101    await update($, tasks, (all = {}) => Object.fromEntries(
102      Object.values(all)
103        .filter(task => task.agentId !== ended || !(task.endsWithCaller || finished.has(task.id)))
104        .map(task => [task.id, task]),
105    ))
106    return done
107  })
108}
109
110/** Sends every due reminder, one pass at a time. */
111async function remind($: EngineInterface, plan: Schedule): Promise<void> {
112  if (loop.ticking) return
113  loop.ticking = true
114  try {
115    const now = await $.clock.now()
116    const { value: all = {} } = await $.state.get(tasks)
117    for (const task of Object.values(all)) {
118      if (task.remindAt > now) continue
119      const output = await outputOf($, task.output)
120      if (output?.done) {
121        await forget($, task.id)
122        continue
123      }
124      let owner = task.agentId
125      if (owner !== null && !(await toAgent($, reminder(task, now, output, plan), owner))) owner = null
126      const sent = owner === null ? { ...task, agentId: null, parent: task.parent ?? task.agentId } : task
127      if (owner === null) {
128        await toMain($, reminder(sent, now, output, plan))
129        if (!loop.running) $.ui.toast(toastText(sent, now))
130      }
131      const remindAt = nextReminder(task.startedAt, now, plan)
132      await update($, tasks, (cur = {}) => cur[task.id] === undefined ? cur : { ...cur, [task.id]: { ...sent, remindAt } })
133    }
134  } finally {
135    loop.ticking = false
136  }
137}
138
139async function forget($: EngineInterface, id: string): Promise<void> {
140  await update($, tasks, (all = {}) => {
141    const { [id]: _, ...rest } = all
142    return rest
143  })
144}
145
146/** What the output file says now: its size, when it last changed, its tail when it is small enough to read. */
147async function outputOf($: EngineInterface, path: string | null): Promise<Output | null> {
148  if (path === null) return null
149  let stat
150  try {
151    stat = await $.fs.stat(path)
152  } catch {
153    return null
154  }
155  const text = stat.size > 0 && stat.size <= READ_LIMIT ? String(await $.fs.read(path)) : ''
156  return { size: stat.size, mtimeMs: stat.mtimeMs, tail: tailOf(text), done: hasExited(text) }
157}
158
159/** Appends the text to a subagent's conversation; false when its loop no longer runs. */
160async function toAgent($: EngineInterface, text: string, agentId: string): Promise<boolean> {
161  try {
162    const stored = await $.session.append({ message: { type: 'user', content: [{ type: 'text', text }] }, agentId })
163    return stored.deny === undefined
164  } catch {
165    return false
166  }
167}
168
169/** Appends the text into the running turn, or wakes an idle session with it as a prompt. */
170async function toMain($: EngineInterface, text: string): Promise<void> {
171  if (loop.running) {
172    await $.session.append({ message: { type: 'user', content: [{ type: 'text', text }] } })
173    return
174  }
175  $.prompt.submit({ text }).catch(() => undefined)
176}
177

1import type { Task } from '../types'
2
3export const MINUTE = 60_000
4export const MARKS = [2, 5, 10, 30, 60, 120]
5export const EVERY = 120
6const TAIL_LINES = 5
7const LINE_WIDTH = 160
8const COMMAND_WIDTH = 200
9
10export type Schedule = { marks: number[]; every: number }
11
12/** What an output file says: `done` when the engine wrote its exit line, so the command has ended. */
13export type Output = { size: number; mtimeMs: number; tail: string[]; done: boolean }
14
15/** The schedule from the options: `marks` as "2 5 10 30 60 120" in minutes of run time, `every` in minutes; the defaults where a value does not parse. */
16export function schedule(marks: unknown, every: unknown): Schedule {
17  const parsed = [...new Set(String(marks ?? '').split(/[\s,]+/).map(Number).filter(n => Number.isFinite(n) && n > 0))].sort((a, b) => a - b)
18  const period = Number(every)
19  return { marks: parsed.length > 0 ? parsed : MARKS, every: Number.isFinite(period) && period > 0 ? period : EVERY }
20}
21
22/** The first reminder time strictly after `now` for a task started at `startedAt`: the next mark, then every period past the last. */
23export function nextReminder(startedAt: number, now: number, plan: Schedule): number {
24  const elapsed = now - startedAt
25  for (const mark of plan.marks) if (mark * MINUTE > elapsed) return startedAt + mark * MINUTE
26  const last = plan.marks[plan.marks.length - 1] ?? 0
27  const periods = Math.floor((elapsed - last * MINUTE) / (plan.every * MINUTE)) + 1
28  return startedAt + (last + periods * plan.every) * MINUTE
29}
30
31/** The note in the tool result that moved the shell to the background. */
32export function firstReminder(task: Task, now: number, plan: Schedule): string {
33  const ahead = plan.marks.filter(mark => task.startedAt + mark * MINUTE > now)
34  const when = ahead.length > 0
35    ? `at ${list(ahead)} minutes of run time, then every ${plural(plan.every, 'minute')}`
36    : `every ${plural(plan.every, 'minute')}`
37  return `Background task ${task.id} is yours: \`${task.command}\`. Nobody else watches it. ` +
38    `You will be reminded ${when}, with the tail of its output, until it ends or you stop it. ` +
39    `When it is no longer needed: TaskStop ${task.id}.`
40}
41
42/** A reminder on the schedule, with what the output file says. */
43export function reminder(task: Task, now: number, output: Output | null, plan: Schedule): string {
44  const by = task.parent === null ? '' : ` It was started by subagent ${task.parent}, which has ended.`
45  const next = elapsedText(nextReminder(task.startedAt, now, plan) - now)
46  return `Reminder: background task ${task.id} (\`${task.command}\`) has run for ${elapsedText(now - task.startedAt)}.${by} ` +
47    `${outputText(output, now)}\n` +
48    `Is it still needed? If not, stop it: TaskStop ${task.id}. If it is, carry on; the next reminder comes in ${next}.`
49}
50
51/** The reason a subagent's end is held for the background commands that the engine kills when its turn ends. */
52export function holdReason(tasks: Task[]): string {
53  const listed = tasks.map(task => `- ${task.id}: \`${task.command}\`${task.output === null ? '' : `, output in ${task.output}`}`).join('\n')
54  return `These background commands you started are still running:\n${listed}\n` +
55    'You are a foreground subagent: they will be killed when you end your turn, and the one who called you gets no result. ' +
56    'Wait for them now in the foreground (a Bash until-loop on the output file ending with "[exited with code", ' +
57    'with a timeout), read the result and finish your report. ' +
58    'A dev server cannot outlive you: if that is what a command is, say so in the report and end the turn again.'
59}
60
61export function toastText(task: Task, now: number): string {
62  return `caffeine: ${task.id} has run for ${elapsedText(now - task.startedAt)}; the agent is reminded`
63}
64
65/** The output file's path from the backgrounding text, or null when it names none. */
66export function outputPath(text: string): string | null {
67  return /Output is being written to: (\S+?\.output)/.exec(text)?.[1] ?? null
68}
69
70/** The task id a completion notification names, or null for any other prompt. */
71export function taskId(text: string): string | null {
72  return /<task-id>([^<\s]+)<\/task-id>/.exec(text)?.[1] ?? null
73}
74
75/** Whether an output ends with the exit line the engine writes when the command ends. */
76export function hasExited(text: string): boolean {
77  return text.slice(-200).includes('[exited with code')
78}
79
80/** The last non-empty lines of an output, each cut to a line's width. */
81export function tailOf(text: string): string[] {
82  const lines = text.split('\n').map(line => line.trimEnd()).filter(line => line !== '')
83  return lines.slice(-TAIL_LINES).map(line => line.length > LINE_WIDTH ? `${line.slice(0, LINE_WIDTH - 1)}…` : line)
84}
85
86/** One line of a command or description, whitespace collapsed, cut to a width. */
87export function squash(text: string): string {
88  const line = text.replace(/\s+/g, ' ').trim()
89  return line.length > COMMAND_WIDTH ? `${line.slice(0, COMMAND_WIDTH - 1)}…` : line
90}
91
92export function elapsedText(ms: number): string {
93  const minutes = Math.round(ms / MINUTE)
94  if (minutes < 1) return 'under a minute'
95  if (minutes < 60) return plural(minutes, 'minute')
96  const hours = Math.floor(minutes / 60)
97  const rest = minutes % 60
98  return rest === 0 ? plural(hours, 'hour') : `${plural(hours, 'hour')} ${plural(rest, 'minute')}`
99}
100
101function outputText(output: Output | null, now: number): string {
102  if (output === null) return 'No output file was found.'
103  if (output.size === 0) return 'Its output file is still empty.'
104  const changed = `Its output last changed ${elapsedText(now - output.mtimeMs)} ago, ${sizeText(output.size)} in all`
105  if (output.tail.length === 0) return `${changed}.`
106  return `${changed}; its last lines:\n${output.tail.map(line => `  ${line}`).join('\n')}`
107}
108
109function sizeText(bytes: number): string {
110  if (bytes < 1024) return `${bytes} bytes`
111  if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} kB`
112  return `${(bytes / (1024 * 1024)).toFixed(1)} MB`
113}
114
115function plural(n: number, unit: string): string {
116  return `${n} ${unit}${n === 1 ? '' : 's'}`
117}
118
119function list(numbers: number[]): string {
120  if (numbers.length === 1) return String(numbers[0])
121  return `${numbers.slice(0, -1).join(', ')} and ${numbers[numbers.length - 1]}`
122}
123

1/** A shell the model left in the background, tracked until it ends or is stopped. */
2export type Task = {
3  id: string
4  command: string
5  description: string
6  /** The file its output goes to; null when the tool result named none. */
7  output: string | null
8  /** When the command started, in milliseconds. */
9  startedAt: number
10  /** When the next reminder is due, in milliseconds. */
11  remindAt: number
12  /** The subagent whose loop runs it; null for the main conversation. */
13  agentId: string | null
14  /** The subagent that started it, once its loop has ended and the task passed to the main conversation. */
15  parent: string | null
16  /** True when the engine ends the shell with its caller's final response, so it cannot outlive the subagent. */
17  endsWithCaller: boolean
18  /** True once the subagent's end was held for this task, which is done at most once. */
19  held: boolean
20}
21
22export type Tasks = Record<string, Task>
23
24declare module 'claude-code' {
25  interface PluginState {
26    'caffeine': {
27      tasks: Tasks
28    }
29  }
30}
31