caffeinate

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Caffeine, CaffeineLease, CaffeineMode } from '../types'
5import { IDLE, describe, isLow, parsePower, wantAwake } from './hold'
6import type { Power } from './hold'
7
8/** Each caffeinate lapses on its own after this long, so a crashed session never keeps the Mac up. */
9const LEASE_S = 600
10/** Renewed at half its life while still wanted. */
11const RENEW_MS = 5 * 60_000
12const TICK_MS = 15_000
13const POWER_EVERY_MS = 60_000
14const MODES: readonly string[] = ['auto', 'on', 'off']
15
16const caffeine = atom({ plugin: 'caffeinate', key: 'caffeine' } as const, IDLE)
17const lease = atom({ plugin: 'caffeinate', key: 'lease' } as const, null)
18
19// What the session is doing, as the turn and stop hooks see it. A reload
20// forgets these and finds the session idle until its next turn; the lease is
21// the host's, so the caffeinate it names is still renewed or ended.
22let turnSince: number | null = null
23let background = 0
24let backgroundAt = 0
25let isEnding = false
26let power: Power = { source: 'unknown', battery: null }
27let powerAt = 0
28let chain: Promise<void> = Promise.resolve()
29
30function loopOf(e: object): string | undefined {
31  const { agentId } = e as { agentId?: unknown }
32
33  return typeof agentId === 'string' ? agentId : undefined
34}
35
36async function startLease($: EngineInterface): Promise<CaffeineLease> {
37  // Detached, its output thrown away, so the call returns at once with its pid.
38  const { stdout } = await $.process.run([
39    '/bin/sh',
40    '-c',
41    `/usr/bin/caffeinate -i -t ${LEASE_S} </dev/null >/dev/null 2>&1 & echo $!`,
42  ])
43  const pid = Number.parseInt(stdout.trim(), 10)
44
45  return Number.isInteger(pid) && pid > 0 ? { pid, startedAt: await $.clock.now() } : null
46}
47
48async function endLease($: EngineInterface, held: CaffeineLease): Promise<void> {
49  if (!held) return
50  // Only while that pid is still a caffeinate: once one lapses its pid can be reused.
51  await $.process.run([
52    '/bin/sh',
53    '-c',
54    'ps -p "$1" -o comm= | grep -q caffeinate && kill "$1"; true',
55    'caffeinate-end',
56    String(held.pid),
57  ])
58}
59
60async function readPower($: EngineInterface, now: number): Promise<Power> {
61  if (now - powerAt < POWER_EVERY_MS) return power
62  const { stdout } = await $.process.run(['/usr/bin/pmset', '-g', 'batt'])
63  power = parsePower(stdout)
64  powerAt = now
65
66  return power
67}
68
69async function sync($: EngineInterface): Promise<void> {
70  const now = await $.clock.now()
71  const was = await read($, caffeine)
72  const why = isEnding ? null : wantAwake({ mode: was.mode, turnSince, background, backgroundAt, now })
73  const supply = why === null ? power : await readPower($, now)
74  const isHeldBack = why !== null && isLow(supply)
75  const isWanted = why !== null && !isHeldBack
76
77  let held = await read($, lease)
78  if (isWanted && (held === null || now - held.startedAt >= RENEW_MS)) {
79    // The new one first, so the Mac is never left unheld between the two.
80    const fresh = await startLease($)
81    await endLease($, held)
82    held = fresh
83    await update($, lease, () => fresh)
84  } else if (!isWanted && held !== null) {
85    await endLease($, held)
86    held = null
87    await update($, lease, () => null)
88  }
89
90  const isAwake = isWanted && held !== null
91  const since = isAwake ? (was.isAwake ? was.since : now) : 0
92  const next: Caffeine = {
93    mode: was.mode,
94    isAwake,
95    reason: isAwake ? why : null,
96    since,
97    minutes: isAwake ? Math.floor((now - since) / 60_000) : 0,
98    source: supply.source,
99    battery: supply.battery,
100    isHeldBack,
101  }
102  if (isHeldBack && !was.isHeldBack) {
103    $.ui.toast(`☕ Battery at ${supply.battery}%: letting the Mac sleep until it is on charge`)
104  }
105  // Written only when what a reader draws changed: once a minute while held, never while idle.
106  if (JSON.stringify(next) !== JSON.stringify(was)) await update($, caffeine, held => ({ ...next, mode: held.mode }))
107}
108
109// Changes run one at a time, so a tick and a turn's end landing together
110// cannot both start a caffeinate.
111function refresh($: EngineInterface, change?: (held: Caffeine) => Caffeine): Promise<void> {
112  chain = chain
113    .then(async () => {
114      if (change) await update($, caffeine, change)
115      await sync($)
116    })
117    .catch(error => $.ui.log(`caffeinate: ${String(error)}`, { to: 'debug' }))
118
119  return chain
120}
121
122export const register: Register = on => {
123  on('session.start', async ($, e, next) => {
124    const started = await next(e)
125    isEnding = false
126    await $.command.register({
127      name: 'caffeinate',
128      description: 'Keep the Mac awake while Claude works: on, off, auto, or show what it is doing',
129      argumentHint: '[on|off|auto]',
130    })
131    $.clock.every(TICK_MS, () => void refresh($))
132    void refresh($)
133
134    return started
135  })
136
137  on('turn.start', async ($, e, next) => {
138    if (loopOf(e) === undefined && turnSince === null) turnSince = await $.clock.now()
139
140    return next(e)
141  })
142
143  // Each turn's end says what background work is still in flight; that work
144  // wakes the session with a turn of its own when it finishes, which recounts.
145  on('classic.Stop', async ($, e, next) => {
146    const answered = await next(e)
147    if (e.agent_id === undefined) {
148      background = e.background_tasks?.length ?? 0
149      backgroundAt = await $.clock.now()
150      await refresh($)
151    }
152
153    return answered
154  }).catch(($, e, next) => next(e))
155
156  on('turn.complete', async ($, e, next) => {
157    const done = await next(e)
158    if (loopOf(e) !== undefined) return done
159    turnSince = null
160    await refresh($)
161
162    return done
163  })
164
165  // A /clear starts the watch over and keeps the mode; leaving the session lets go now
166  // rather than when the lease lapses.
167  on('session.end', async ($, e, next) => {
168    const ended = await next(e)
169    turnSince = null
170    background = 0
171    isEnding = e.reason !== 'clear'
172    await refresh($)
173
174    return ended
175  })
176
177  on('command.run', { command: 'caffeinate' }, async ($, e) => {
178    const asked = e.args.trim().toLowerCase()
179    if (asked !== '' && !MODES.includes(asked)) return { text: 'Usage: /caffeinate [on|off|auto]' }
180    const mode = asked as CaffeineMode
181    await refresh($, asked === '' ? undefined : held => ({ ...held, mode }))
182
183    return { text: describe(await read($, caffeine)) }
184  })
185}
186

1import type { Caffeine, CaffeineMode, CaffeineReason, CaffeineSource } from '../types'
2
3/** A turn counts as long once it has run this long; shorter ones never hold the Mac. */
4export const GRACE_MS = 60_000
5/** On battery at or below this charge the hold is let go, and taken again on charge. */
6export const LOW_BATTERY = 20
7/**
8 * Background work reports back with a turn of its own when it ends, and each
9 * turn's end recounts it. A count no turn has refreshed in this long is stale
10 * (work stopped with no turn after it), and the hold it asked for lapses.
11 */
12export const BACKGROUND_MAX_MS = 12 * 60 * 60_000
13
14export const IDLE: Caffeine = {
15  mode: 'auto',
16  isAwake: false,
17  reason: null,
18  since: 0,
19  minutes: 0,
20  source: 'unknown',
21  battery: null,
22  isHeldBack: false,
23}
24
25export type Power = { source: CaffeineSource; battery: number | null }
26
27/** What the session is doing, as the turn and stop hooks have seen it. */
28export type Watch = {
29  mode: CaffeineMode
30  turnSince: number | null
31  background: number
32  backgroundAt: number
33  now: number
34}
35
36/** Reads `pmset -g batt`: the power source, and the charge where there is a battery. */
37export function parsePower(out: string): Power {
38  const source = out.includes("'AC Power'") ? 'ac' : out.includes("'Battery Power'") ? 'battery' : 'unknown'
39  const charge = /(\d{1,3})%;/.exec(out)
40
41  return { source, battery: charge ? Number(charge[1]) : null }
42}
43
44/** Why the Mac should be held awake now, or null to let it sleep as it normally would. */
45export function wantAwake(watch: Watch): CaffeineReason | null {
46  if (watch.mode === 'off') return null
47  if (watch.mode === 'on') return 'manual'
48  if (watch.turnSince !== null && watch.now - watch.turnSince >= GRACE_MS) return 'turn'
49  if (watch.background > 0 && watch.now - watch.backgroundAt < BACKGROUND_MAX_MS) return 'background'
50
51  return null
52}
53
54export function isLow(power: Power): boolean {
55  return power.source === 'battery' && power.battery !== null && power.battery <= LOW_BATTERY
56}
57
58/** `7` → `7m`, `72` → `1h12m`. */
59export function elapsed(minutes: number): string {
60  if (minutes < 1) return '<1m'
61  if (minutes < 60) return `${minutes}m`
62
63  return `${Math.floor(minutes / 60)}h${String(minutes % 60).padStart(2, '0')}m`
64}
65
66const FOR: Record<CaffeineReason, string> = {
67  turn: 'while Claude works on this turn',
68  background: 'while background work from this session finishes',
69  manual: 'until you run /caffeinate auto or /caffeinate off',
70}
71
72const LID = 'The display can still sleep and lock; closing the lid still puts the Mac to sleep.'
73
74/** What `/caffeinate` answers. */
75export function describe(caffeine: Caffeine): string {
76  if (caffeine.mode === 'off') return 'Off for this session: the Mac sleeps as it normally would. /caffeinate auto turns it back on.'
77  if (caffeine.isHeldBack) return `On battery at ${caffeine.battery}%: letting the Mac sleep until it is on charge.`
78  if (caffeine.isAwake && caffeine.reason) {
79    return `Keeping the Mac awake ${FOR[caffeine.reason]} (${elapsed(caffeine.minutes)} so far). ${LID}`
80  }
81
82  return `Standing by: the Mac is held awake once a turn has run a minute, and while background work runs. ${LID}`
83}
84

1/** How the session keeps the Mac awake: by what Claude is doing, always, or never. */
2export type CaffeineMode = 'auto' | 'on' | 'off'
3
4/**
5 * What the Mac is held awake for: a turn that has run past a minute, background
6 * work the session left running when its turn ended, or `/caffeinate on`.
7 */
8export type CaffeineReason = 'turn' | 'background' | 'manual'
9
10/** Where the Mac draws its power from, as `pmset -g batt` says. */
11export type CaffeineSource = 'ac' | 'battery' | 'unknown'
12
13/** The hold as other mods draw it: Sprout's coffee cup reads this. */
14export type Caffeine = {
15  mode: CaffeineMode
16  /** A caffeinate assertion is held now. */
17  isAwake: boolean
18  reason: CaffeineReason | null
19  /** When the current hold began, in ms since the epoch; 0 while none is held. */
20  since: number
21  /** Whole minutes held, so a reader redraws once a minute and no more often. */
22  minutes: number
23  source: CaffeineSource
24  /** Battery charge in percent; null on a Mac with no battery. */
25  battery: number | null
26  /** The hold was let go because the battery is low; it is taken again on charge. */
27  isHeldBack: boolean
28}
29
30/** The caffeinate process this session started, kept by the host so a reload can still end it. */
31export type CaffeineLease = { pid: number; startedAt: number } | null
32
33declare module 'claude-code' {
34  interface PluginState {
35    caffeinate: { caffeine: Caffeine; lease: CaffeineLease }
36  }
37}
38