caffeinate

1import type { EngineInterface, HookStream, ProcessSpawnChunk, ProcessSpawnResult, Register } from 'claude-code'
2import { atom, read, update } from 'claude-code'
3import { DEFAULTS, FLAGS, LID, MODES, NOTHING_PENDING, PROGRAMS, SCHEDULED, commandFor, holdReason, normalize } from './settings.js'
4import type { Pending, Settings } from './settings.js'
5
6const PANE = 'caffeinate'
7const STORE_KEY = 'settings'
8// Another session may change the settings; each session picks that up within this long
9const REFRESH_MS = 15_000
10
11// The lid helper (lid/): a root LaunchDaemon, installed once with an admin prompt, that keeps the Mac awake with the
12// lid closed while a process holds it. 'busy' is an install or uninstall waiting on that prompt.
13type LidState = 'unknown' | 'unsupported' | 'missing' | 'outdated' | 'installed' | 'busy'
14
15// What the pane draws: the settings and what this session is doing about them
16type View = {
17  settings: Settings
18  lid: LidState
19  lidMessage: string | undefined
20  holding: string | undefined
21  reason: string | undefined
22  problem: string | undefined
23}
24const view = atom(
25  { plugin: 'caffeinate', key: 'view' } as const,
26  { settings: DEFAULTS, lid: 'unknown', lidMessage: undefined, holding: undefined, reason: undefined, problem: undefined } as View,
27)
28
29// Module variables: a reload kills the child anyway, so these start over with it
30let settings: Settings = DEFAULTS
31let turnId: string | undefined
32// What the last turn left running, from its Stop. An interrupted turn raises no Stop, so this stays as last reported
33let pending: Pending = NOTHING_PENDING
34let reason: string | undefined
35let child: { key: string; label: string; stream: HookStream<ProcessSpawnChunk, ProcessSpawnResult> } | undefined
36let problem: string | undefined
37let lid: LidState = 'unknown'
38// The last lid helper error, shown in the pane
39let lidMessage: string | undefined
40
41async function loadSettings($: EngineInterface): Promise<void> {
42  settings = normalize(await $.store.get(STORE_KEY))
43}
44
45async function saveSettings($: EngineInterface, change: Partial<Settings>): Promise<void> {
46  // Re-read first: the store is shared by every session on the machine
47  settings = normalize({ ...normalize(await $.store.get(STORE_KEY)), ...change })
48  await $.store.set(STORE_KEY, settings)
49  await sync($)
50}
51
52// Starts, swaps or stops the child so it matches the settings, whether a turn is running and what it left pending
53async function sync($: EngineInterface): Promise<void> {
54  reason = holdReason(settings, turnId !== undefined, pending)
55  const want = reason ? commandFor(settings, $.plugin.root + '/lid/hold.sh') : undefined
56  const key = want && 'argv' in want ? JSON.stringify(want.argv) : undefined
57  if (child && child.key !== key) {
58    const old = child
59    child = undefined
60    void old.stream.return({ code: null, signal: null })
61  }
62  if (want && 'error' in want) problem = want.error
63  else if (!want) problem = undefined
64  if (want && 'argv' in want && !child) start($, want.argv)
65  await show($)
66}
67
68function start($: EngineInterface, argv: string[]): void {
69  const stream = $.process.spawn({ argv })
70  // Through the lid helper's hold script, name the command it runs rather than the script
71  const shown = argv[0] === '/bin/sh' && argv[1] === $.plugin.root + '/lid/hold.sh' ? argv.slice(2) : argv
72  const mine = { key: JSON.stringify(argv), label: shown.join(' ') + (shown === argv ? '' : ' (lid closed too)'), stream }
73  child = mine
74  problem = undefined
75  void (async () => {
76    let said = ''
77    let ended = ''
78    try {
79      for (;;) {
80        const r = await stream.next()
81        if (r.done) {
82          ended = r.value.signal ? 'was killed by ' + r.value.signal : 'exited with code ' + r.value.code
83          break
84        }
85        if (r.value.stream === 'stderr') said = (said + r.value.text).slice(-200)
86      }
87    } catch (err) {
88      ended = 'could not start: ' + (err instanceof Error ? err.message : String(err))
89    }
90    // Stopped on purpose (sync replaced it, or the module unloaded): nothing to report
91    if (child !== mine) return
92    child = undefined
93    problem = shown[0] + ' ' + ended + (said.trim() ? ': ' + said.trim().split('\n').pop() : '')
94    await show($)
95  })()
96}
97
98async function show($: EngineInterface): Promise<void> {
99  // Never blocks the normal hold: a lid helper that is not ready only adds a note
100  const lidNote = !settings.lid ? '' : lid === 'missing' ? ' (lid helper not installed: /caffeinate)' : lid === 'outdated' ? ' (lid helper needs an update: /caffeinate)' : ''
101  if (problem) $.ui.status('☕ not keeping awake: ' + problem)
102  else if (child) $.ui.status('☕ keeping awake: ' + reason + lidNote)
103  else $.ui.status(undefined)
104  const snapshot: View = { settings, lid, lidMessage, holding: child?.label, reason, problem }
105  await update($, view, () => snapshot)
106}
107
108// Asks the helper's check script, which needs no privileges, whether the helper is installed and current
109async function checkLid($: EngineInterface): Promise<void> {
110  if (lid === 'busy') return
111  try {
112    const { exitCode } = await $.process.run(['/bin/sh', $.plugin.root + '/lid/check.sh'])
113    lid = exitCode === 0 ? 'installed' : exitCode === 1 ? 'missing' : exitCode === 2 ? 'outdated' : 'unsupported'
114  } catch {
115    lid = 'unsupported'
116  }
117  await show($)
118}
119
120// Runs a lid script as root through the macOS admin prompt. Returns an error message, or undefined on success.
121async function runAdmin($: EngineInterface, script: string, prompt: string): Promise<string | undefined> {
122  try {
123    const user = (await $.process.run(['id', '-un'])).stdout.trim()
124    const r = await $.process.run(
125      [
126        '/usr/bin/osascript',
127        '-e', 'on run argv',
128        '-e', 'do shell script "/bin/sh " & quoted form of (item 1 of argv) & " " & quoted form of (item 2 of argv) with prompt "' + prompt + '" with administrator privileges',
129        '-e', 'end run',
130        $.plugin.root + '/lid/' + script,
131        user,
132      ],
133      { timeoutMs: 600_000 },
134    )
135    if (r.exitCode === 0) return undefined
136    if (r.stderr.includes('-128')) return 'cancelled'
137    return r.stderr.trim().split('\n').pop() || 'exited with code ' + r.exitCode
138  } catch (err) {
139    return err instanceof Error ? err.message : String(err)
140  }
141}
142
143// Installs or updates the helper. Only ever called for an explicit action in the pane: it raises the admin prompt.
144async function installLid($: EngineInterface): Promise<boolean> {
145  const before = lid
146  lid = 'busy'
147  lidMessage = undefined
148  await show($)
149  const failed = await runAdmin($, 'install.sh', 'caffeinate wants to install a helper that keeps your Mac awake with the lid closed.')
150  lid = before
151  if (failed) lidMessage = 'Could not install the lid helper: ' + failed
152  await checkLid($)
153  return !failed && lid === 'installed'
154}
155
156async function uninstallLid($: EngineInterface): Promise<void> {
157  const before = lid
158  lid = 'busy'
159  lidMessage = undefined
160  await show($)
161  const failed = await runAdmin($, 'uninstall.sh', 'caffeinate wants to remove the helper that keeps your Mac awake with the lid closed.')
162  lid = before
163  if (failed) lidMessage = 'Could not uninstall the lid helper: ' + failed
164  await checkLid($)
165  if (!failed) await saveSettings($, { lid: false })
166}
167
168// Turning it on asks for admin only when the helper is missing or outdated. Turning it off touches nothing, so the
169// idle helper stays and turning it on again never asks.
170async function setLid($: EngineInterface, on: boolean): Promise<void> {
171  if (!on) {
172    lidMessage = undefined
173    await saveSettings($, { lid: false })
174    return
175  }
176  if (lid === 'busy') return
177  if (lid === 'unknown') await checkLid($)
178  lidMessage = undefined
179  if (lid === 'unsupported') {
180    lidMessage = 'Only on macOS'
181    await show($)
182    return
183  }
184  if ((lid === 'missing' || lid === 'outdated') && !(await installLid($))) return
185  if (lid === 'installed') await saveSettings($, { lid: true })
186}
187
188// The install button shows when the helper is outdated, or wanted and not there
189const needsInstall = (v: View): boolean => v.lid === 'outdated' || (v.settings.lid && v.lid === 'missing')
190
191export const register: Register = (on) => {
192  on('session.start', async ($, e, next) => {
193    const started = await next(e)
194    await loadSettings($)
195    // Only when asked for: the check is cheap, but a setting that is off should cost nothing
196    if (settings.lid) await checkLid($)
197    await sync($)
198    $.clock.every(REFRESH_MS, async () => {
199      const fresh = normalize(await $.store.get(STORE_KEY))
200      // Also retries a child that died, e.g. a custom command that failed to start
201      if (JSON.stringify(fresh) !== JSON.stringify(settings) || (!child && holdReason(fresh, turnId !== undefined, pending))) {
202        settings = fresh
203        await sync($)
204      }
205    })
206    try {
207      await $.command.register({ name: 'caffeinate', description: 'Keep the computer awake while Claude works: settings', immediate: true })
208    } catch (err) {
209      $.ui.log('caffeinate: could not register /caffeinate: ' + String(err))
210    }
211    return started
212  })
213
214  // Subagents raise no turn.start, and their turn.complete carries agentId: only the main loop's turn counts
215  on('turn.start', async ($, e, next) => {
216    turnId = e.turnId
217    settings = normalize(await $.store.get(STORE_KEY))
218    await sync($)
219    return next(e)
220  })
221
222  // Fires at the end of a main-loop turn, just before turn.complete (not on an interrupt): what is left running or scheduled
223  on('classic.Stop', async ($, e, next) => {
224    pending = { background: e.background_tasks?.length ?? 0, scheduled: e.session_crons?.length ?? 0 }
225    return next(e)
226  })
227
228  on('turn.complete', async ($, e, next) => {
229    if (!e.agentId && e.turnId === turnId) {
230      turnId = undefined
231      await sync($)
232    }
233    return next(e)
234  })
235
236  on('command.run', { command: 'caffeinate' }, async ($) => {
237    await checkLid($)
238    await $.ui.open({ id: PANE, title: 'caffeinate', focus: true, closeOnEscape: true })
239    return {}
240  })
241
242  on('ui.render', { component: 'Pane' }, async ($, e, next) => {
243    if (e.requestId !== PANE) return next(e)
244    // Select and Input exist only in the terminal and the desktop app
245    if (e.surface !== 'terminal' && e.surface !== 'desktop') return next(e)
246    const { Box, Text, Select, Input, Button } = $.ui.resolve(e)
247    const v = await read($, view)
248    const s = v.settings
249    const state = v.problem
250      ? Text({ color: 'warning', children: ['Not keeping awake: ' + v.problem] })
251      : v.holding
252        ? Text({ color: 'success', children: ['Keeping awake (' + v.reason + '): ' + v.holding] })
253        : Text({ dimColor: true, children: [s.mode === 'off' ? 'Off' : 'Idle: starts with the next turn'] })
254    const wantsInstall = needsInstall(v)
255    const save = (change: Partial<Settings>) => void saveSettings($, change)
256    return Box({
257      flexDirection: 'column',
258      gap: 1,
259      children: [
260        state,
261        Select({ key: 'mode', label: 'Keep awake', options: MODES, value: s.mode, onSelect: (value) => save({ mode: value as Settings['mode'] }) }),
262        s.mode === 'turn'
263          ? Text({ dimColor: true, children: ['Includes background work a turn leaves running (monitors, background shells and agents).'] })
264          : undefined,
265        s.mode === 'turn'
266          ? Select({
267              key: 'scheduled',
268              label: 'Scheduled wake-ups (/loop, ScheduleWakeup, cron)',
269              options: SCHEDULED,
270              value: s.scheduled ? 'on' : 'off',
271              onSelect: (value) => save({ scheduled: value === 'on' }),
272            })
273          : undefined,
274        Select({ key: 'program', label: 'Using', options: PROGRAMS, value: s.program, onSelect: (value) => save({ program: value as Settings['program'] }) }),
275        s.program === 'caffeinate'
276          ? Select({ key: 'flags', label: 'Flags', options: FLAGS, value: s.flags, onSelect: (value) => save({ flags: value }) })
277          : Input({
278              key: 'custom',
279              label: 'Command',
280              placeholder: 'systemd-inhibit --what=idle sleep infinity',
281              value: s.custom,
282              submitLabel: 'Save',
283              onSubmit: (value) => save({ custom: value.trim() }),
284            }),
285        s.program === 'custom'
286          ? Text({ dimColor: true, children: ['Runs while awake is wanted and is stopped after. Quotes work; no shell (no pipes or &&).'] })
287          : undefined,
288        Select({ key: 'lid', label: 'Lid closed', options: LID, value: s.lid ? 'on' : 'off', onSelect: (value) => void setLid($, value === 'on') }),
289        v.lidMessage ? Text({ color: 'warning', children: [v.lidMessage] }) : undefined,
290        v.lid === 'busy' ? Text({ dimColor: true, children: ['Waiting for the lid helper: answer the admin prompt.'] }) : undefined,
291        s.lid
292          ? Text({ dimColor: true, children: ["The lid can close while awake is held. If it's ever stuck awake: sudo pmset -a disablesleep 0"] })
293          : undefined,
294        wantsInstall
295          ? Button({ key: 'lid-install', label: 'Install/Update lid helper', onPress: () => void installLid($) })
296          : undefined,
297        v.lid === 'installed'
298          ? Button({ key: 'lid-uninstall', label: 'Uninstall lid helper', onPress: () => void uninstallLid($) })
299          : undefined,
300        Text({ dimColor: true, children: ['Saved for every session on this machine. Esc closes.'] }),
301        Button({ key: 'close', label: 'Close', role: 'dismiss', onPress: () => void $.ui.close({ id: PANE }) }),
302      ],
303    })
304  })
305}
306

1// Settings and the command line they make. Plain functions, no `$`, so the tests can call them directly.
2
3export type Mode = 'turn' | 'session' | 'off'
4export type Program = 'caffeinate' | 'custom'
5
6export type Settings = {
7  mode: Mode
8  // Turn mode: also stay awake while a scheduled wake-up (CronCreate, ScheduleWakeup, /loop) is pending
9  scheduled: boolean
10  program: Program
11  flags: string
12  custom: string
13  // macOS: also keep the Mac awake with the lid closed, through the lid helper (see lid/)
14  lid: boolean
15}
16
17export const DEFAULTS: Settings = { mode: 'turn', scheduled: true, program: 'caffeinate', flags: '-i', custom: '', lid: false }
18
19// What the session still has going after a turn ends, as the last Stop reported it
20export type Pending = { background: number; scheduled: number }
21
22export const NOTHING_PENDING: Pending = { background: 0, scheduled: 0 }
23
24export const MODES: readonly { value: Mode; label: string }[] = [
25  { value: 'turn', label: 'While Claude is working on a turn' },
26  { value: 'session', label: 'For the whole session' },
27  { value: 'off', label: 'Off' },
28]
29
30export const SCHEDULED: readonly { value: 'off' | 'on'; label: string }[] = [
31  { value: 'on', label: 'Stay awake for them' },
32  { value: 'off', label: 'Let it sleep until then' },
33]
34
35export const LID: readonly { value: 'off' | 'on'; label: string }[] = [
36  { value: 'off', label: 'Sleep as usual' },
37  { value: 'on', label: 'Stay awake with the lid closed (macOS; one-time admin install)' },
38]
39
40export const PROGRAMS: readonly { value: Program; label: string }[] = [
41  { value: 'caffeinate', label: 'caffeinate (macOS)' },
42  { value: 'custom', label: 'Custom command' },
43]
44
45export const FLAGS: readonly { value: string; label: string }[] = [
46  { value: '-i', label: '-i    no idle sleep (display may still sleep)' },
47  { value: '-di', label: '-di   no idle sleep, display stays on' },
48  { value: '-s', label: '-s    no system sleep (on AC power only)' },
49  { value: '-ims', label: '-ims  no idle, disk or system sleep' },
50]
51
52// Accepts whatever the store holds and keeps only valid fields, so an old or hand-edited value can't break the mod.
53export function normalize(raw: unknown): Settings {
54  const v = (raw && typeof raw === 'object' ? raw : {}) as Record<string, unknown>
55  const pick = <T extends string>(x: unknown, allowed: readonly { value: T }[], d: T): T =>
56    allowed.some((o) => o.value === x) ? (x as T) : d
57  return {
58    mode: pick(v.mode, MODES, DEFAULTS.mode),
59    scheduled: typeof v.scheduled === 'boolean' ? v.scheduled : DEFAULTS.scheduled,
60    program: pick(v.program, PROGRAMS, DEFAULTS.program),
61    flags: pick(v.flags, FLAGS, DEFAULTS.flags),
62    custom: typeof v.custom === 'string' ? v.custom : DEFAULTS.custom,
63    lid: typeof v.lid === 'boolean' ? v.lid : DEFAULTS.lid,
64  }
65}
66
67// Splits a command line the way a shell would for plain words, quotes and backslashes, without running a shell:
68// no variables, globs, pipes or `&&`. Throws on an unclosed quote.
69export function splitCommand(line: string): string[] {
70  const out: string[] = []
71  let word = ''
72  let inWord = false
73  let quote: '"' | "'" | undefined
74  for (let i = 0; i < line.length; i++) {
75    const c = line[i]!
76    if (quote === "'") {
77      if (c === "'") quote = undefined
78      else word += c
79    } else if (quote === '"') {
80      if (c === '"') quote = undefined
81      else if (c === '\\' && i + 1 < line.length && '"\\$`'.includes(line[i + 1]!)) word += line[++i]
82      else word += c
83    } else if (c === "'" || c === '"') {
84      quote = c
85      inWord = true
86    } else if (c === '\\' && i + 1 < line.length) {
87      word += line[++i]
88      inWord = true
89    } else if (c === ' ' || c === '\t' || c === '\n') {
90      if (inWord) out.push(word)
91      word = ''
92      inWord = false
93    } else {
94      word += c
95      inWord = true
96    }
97  }
98  if (quote) throw new Error('unclosed ' + quote + ' quote')
99  if (inWord) out.push(word)
100  return out
101}
102
103// The command that keeps the machine awake, or a reason there is none. With the lid setting on and a hold script,
104// the command runs through it (it registers the process with the lid helper, then execs the command).
105export function commandFor(s: Settings, holdScript?: string): { argv: string[] } | { error: string } {
106  const wrap = (argv: string[]) => ({ argv: s.lid && holdScript ? ['/bin/sh', holdScript, ...argv] : argv })
107  if (s.program === 'caffeinate') return wrap(['caffeinate', s.flags])
108  let argv: string[]
109  try {
110    argv = splitCommand(s.custom)
111  } catch (err) {
112    return { error: 'custom command: ' + (err instanceof Error ? err.message : String(err)) }
113  }
114  return argv.length ? wrap(argv) : { error: 'no custom command set' }
115}
116
117// Why the machine should stay awake right now, or undefined when it may sleep
118export function holdReason(s: Settings, isTurnRunning: boolean, pending: Pending): string | undefined {
119  if (s.mode === 'session') return 'for the session'
120  if (s.mode === 'off') return undefined
121  if (isTurnRunning) return 'Claude is working'
122  if (pending.background > 0) return pending.background === 1 ? '1 background task running' : pending.background + ' background tasks running'
123  if (s.scheduled && pending.scheduled > 0) return pending.scheduled === 1 ? 'a scheduled wake-up is pending' : pending.scheduled + ' scheduled wake-ups are pending'
124  return undefined
125}
126

1// The state contract: self-contained, so the settings shape is spelled out here (hooks/settings.ts mirrors it)
2declare module 'claude-code' {
3  interface PluginState {
4    caffeinate: {
5      view: {
6        settings: { mode: 'turn' | 'session' | 'off'; scheduled: boolean; program: 'caffeinate' | 'custom'; flags: string; custom: string; lid: boolean }
7        lid: 'unknown' | 'unsupported' | 'missing' | 'outdated' | 'installed' | 'busy'
8        lidMessage: string | undefined
9        holding: string | undefined
10        reason: string | undefined
11        problem: string | undefined
12      }
13    }
14  }
15}
16