jr-awake

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, HookStream, ProcessSpawnChunk, ProcessSpawnResult, Register } from 'claude-code'
3
4import type { JrAwakeMode, JrAwakePower, JrAwakeSessionBeat } from '../types'
5import { BEAT_STALE_MS, formatCountdown, MODES, parsePmset, secondsToSleep, shouldHold } from './power'
6
7const view = atom({ plugin: 'jr-awake', key: 'view' } as const, null)
8// Hidden until the person asks for it; the warning before sleep shows regardless.
9const bandHidden = atom({ plugin: 'jr-awake', key: 'isBandHidden' } as const, true)
10
11const bandState = { isLoaded: false }
12
13/** Reads the saved choice once per load of the mod; outside a drawing, which cannot write state. */
14async function loadBand($: EngineInterface) {
15  if (bandState.isLoaded) return
16  bandState.isLoaded = true
17  if ((await $.store.get('isBandHidden')) === false) await update($, bandHidden, () => false)
18}
19
20/** Hides or shows the band for every session, kept across restarts. */
21async function setBand($: EngineInterface, hide: boolean, name: string, command: string) {
22  await $.store.set('isBandHidden', hide)
23  await update($, bandHidden, () => hide)
24  return { text: hide ? `${name} · hidden (${command} to bring it back)` : `${name} · shown` }
25}
26const TICK_MS = 15_000
27const MODE_KEY = 'mode'
28const BEAT_PREFIX = 'beat:'
29
30// This session: whether Claude is working, and the caffeinate child it holds, if any.
31const live: {
32  id: string
33  project: string
34  working: boolean
35  caffeinate: HookStream<ProcessSpawnChunk, ProcessSpawnResult> | null
36  isTicking: boolean
37  warnedAt: number
38} = { id: '', project: '', working: false, caffeinate: null, isTicking: false, warnedAt: 0 }
39
40async function loadMode($: EngineInterface, fallback: JrAwakeMode): Promise<JrAwakeMode> {
41  const saved = (await $.store.get(MODE_KEY)) as JrAwakeMode | undefined
42  return saved && MODES.includes(saved) ? saved : fallback
43}
44
45async function readPower($: EngineInterface): Promise<JrAwakePower | null> {
46  try {
47    const pm = await $.process.run(['pmset', '-g'])
48    if (pm.exitCode !== 0) return null
49    const idle = await $.process.run(['sh', '-c', "ioreg -c IOHIDSystem | awk '/HIDIdleTime/ {print int($NF/1000000000); exit}'"])
50    return { ...parsePmset(pm.stdout), idleSeconds: Number.parseInt(idle.stdout.trim(), 10) || 0 }
51  } catch {
52    return null
53  }
54}
55
56/** Starts or stops this session's caffeinate so it matches what the shared mode asks for. */
57async function hold($: EngineInterface, wanted: boolean, keepDisplay: boolean) {
58  if (wanted && !live.caffeinate) {
59    const stream = $.process.spawn({ argv: keepDisplay ? ['caffeinate', '-i', '-d'] : ['caffeinate', '-i'] })
60    live.caffeinate = stream
61    void (async () => {
62      try {
63        for await (const _ of stream) {
64          // caffeinate prints nothing; the loop is the child's life.
65        }
66      } catch {
67        // Not on macOS, or caffeinate missing: nothing to hold.
68      }
69      if (live.caffeinate === stream) live.caffeinate = null
70    })()
71  } else if (!wanted && live.caffeinate) {
72    const stream = live.caffeinate
73    live.caffeinate = null
74    await stream.return(undefined as never)
75  }
76}
77
78async function beat($: EngineInterface) {
79  const now = await $.clock.now()
80  const mine: JrAwakeSessionBeat = { id: live.id, project: live.project, working: live.working, holding: live.caffeinate !== null, at: now }
81  await $.store.set(BEAT_PREFIX + live.id, mine)
82  const beats: JrAwakeSessionBeat[] = []
83  for (const key of await $.store.keys()) {
84    if (!key.startsWith(BEAT_PREFIX)) continue
85    const b = (await $.store.get(key)) as JrAwakeSessionBeat | undefined
86    if (!b) continue
87    if (now - b.at > BEAT_STALE_MS * 10) await $.store.delete(key)
88    else if (now - b.at < BEAT_STALE_MS) beats.push(b)
89  }
90  return beats
91}
92
93async function tick($: EngineInterface, fallbackMode: JrAwakeMode, keepDisplay: boolean, warnAt: number) {
94  const mode = await loadMode($, fallbackMode)
95  await hold($, shouldHold(mode, live.working), keepDisplay)
96  const beats = await beat($)
97  const power = await readPower($)
98  const sleepIn = power ? secondsToSleep(power, live.caffeinate !== null) : null
99  const holders = beats.filter(b => b.holding).length
100  await update($, view, () => ({ mode, holding: live.caffeinate !== null, holders, sleepIn, power }))
101  const now = await $.clock.now()
102  if (sleepIn !== null && sleepIn <= warnAt && now - live.warnedAt > 5 * 60_000) {
103    live.warnedAt = now
104    $.ui.toast(`jr awake · your Mac sleeps in ${formatCountdown(sleepIn)} · /jrawake session keeps it awake`)
105  }
106}
107
108async function ensureStarted($: EngineInterface, fallbackMode: JrAwakeMode, keepDisplay: boolean, warnAt: number) {
109  if (!live.id) {
110    live.id = await $.session.id()
111    const root = await $.session.root()
112    live.project = root.split('/').pop() || root
113  }
114  if (!live.isTicking) {
115    live.isTicking = true
116    $.clock.every(TICK_MS, () => {
117      void tick($, fallbackMode, keepDisplay, warnAt)
118    })
119    await tick($, fallbackMode, keepDisplay, warnAt)
120  }
121}
122
123export const register: Register = (on, options) => {
124  const fallbackMode = (MODES.includes(options.mode as JrAwakeMode) ? options.mode : 'working') as JrAwakeMode
125  const keepDisplay = options.keepDisplay === true
126  const warnAt = Number(options.warnAtSeconds ?? 60)
127  const showBand = options.showBand !== false
128
129  on('session.start', async ($, e, next) => {
130    const started = await next(e)
131    await $.command.register({ name: 'jrawake', description: 'Keep the Mac awake: /jrawake working | session | off | status | show | hide (the line is hidden by default)' })
132    await loadBand($)
133    await ensureStarted($, fallbackMode, keepDisplay, warnAt)
134    return started
135  })
136
137  on('session.end', async ($, e, next) => {
138    await hold($, false, keepDisplay)
139    if (live.id) await $.store.delete(BEAT_PREFIX + live.id)
140    return next(e)
141  })
142
143  // Claude starts working: hold the Mac awake for the turn in "working" mode.
144  on('prompt.submit', async ($, e, next) => {
145    await ensureStarted($, fallbackMode, keepDisplay, warnAt)
146    if (!e.command && !e.isMeta) {
147      live.working = true
148      await tick($, fallbackMode, keepDisplay, warnAt)
149    }
150    return next(e)
151  })
152
153  on('turn.complete', async ($, e, next) => {
154    const done = await next(e)
155    if (e.agentId === undefined) {
156      live.working = false
157      await tick($, fallbackMode, keepDisplay, warnAt)
158    }
159    return done
160  })
161
162  on('command.run', { command: 'jrawake' }, async ($, e) => {
163    await ensureStarted($, fallbackMode, keepDisplay, warnAt)
164    const word = e.args.trim().toLowerCase()
165    if (word === 'hide' || word === 'show') {
166      await loadBand($)
167      const answer = await setBand($, word === 'hide', 'jr awake line', '/jrawake show')
168      return word === 'hide' ? { text: `${answer.text}; the Mac is still kept awake, and the warning before sleep still shows` } : answer
169    }
170    const arg = e.args.trim() as JrAwakeMode | 'on' | 'status' | ''
171    const mode: JrAwakeMode | null = arg === 'on' ? 'session' : MODES.includes(arg as JrAwakeMode) ? (arg as JrAwakeMode) : null
172    if (mode) {
173      // Shared by every session on the machine: they all pick it up on their next tick.
174      await $.store.set(MODE_KEY, mode)
175      await tick($, fallbackMode, keepDisplay, warnAt)
176    }
177    const v = await read($, view)
178    const sleep = v?.sleepIn == null ? 'not scheduled' : `in ${formatCountdown(v.sleepIn)}`
179    return {
180      text: `jr awake · mode ${v?.mode ?? fallbackMode} (all sessions) · ${v?.holders ?? 0} session(s) holding the Mac awake · next sleep ${sleep}${
181        v?.power ? ` · sleep after ${v.power.sleepMin || 'never'} min${v.power.onBattery ? ' on battery' : ''}` : ''
182      }`,
183    }
184  })
185
186  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
187    const below = await next(e)
188    const v = await read($, view)
189    if ((!v && !live.isTicking) || !bandState.isLoaded) {
190      $.clock.after(0, () => {
191        void loadBand($)
192        void ensureStarted($, fallbackMode, keepDisplay, warnAt)
193      })
194    }
195    if (!showBand || !v || e.props.hasSurvey) return below
196    const { Box, Text } = $.ui.resolve(e)
197    // About to sleep: an alert on the left, shown even while the line is hidden.
198    if (v.sleepIn !== null && v.sleepIn <= warnAt) {
199      return (
200        <Box flexDirection="column">
201          <Text wrap="truncate-end">
202            <Text color="#ef4565">{'▲ '}</Text>
203            <Text>{`Mac sleeps in ${formatCountdown(v.sleepIn)}`}</Text>
204            <Text dimColor>{' · /jrawake on'}</Text>
205          </Text>
206          {below}
207        </Box>
208      )
209    }
210    if (await read($, bandHidden)) return below
211    // Otherwise one quiet row on the right: the shared mode, and who holds the Mac awake.
212    return (
213      <Box flexDirection="column">
214        <Box justifyContent="flex-end">
215          <Text wrap="truncate-end">
216            <Text dimColor>{'awake · '}</Text>
217            <Text color={v.holders > 0 ? '#2cb67d' : undefined}>{v.mode}</Text>
218            {v.holders > 0 && <Text dimColor>{` · held by ${v.holders}`}</Text>}
219          </Text>
220        </Box>
221        {below}
222      </Box>
223    )
224  })
225}
226

1import type { JrAwakeMode, JrAwakePower, JrAwakeSessionBeat } from '../types'
2
3export const MODES: JrAwakeMode[] = ['working', 'session', 'off']
4export const BEAT_STALE_MS = 60_000
5
6/**
7 * Reads `pmset -g`: the idle minutes before the system and the display sleep (0 = never),
8 * and which processes currently prevent system sleep.
9 */
10export const parsePmset = (text: string): Omit<JrAwakePower, 'idleSeconds'> => {
11  const minutes = (name: string) => {
12    const m = text.match(new RegExp(`^\\s*${name}\\s+(\\d+)`, 'm'))
13    return m ? Number(m[1]) : 0
14  }
15  const prevented = text.match(/^\s*sleep\s+\d+\s+\(sleep prevented by ([^)]*)\)/m)?.[1]
16  return {
17    sleepMin: minutes('sleep'),
18    displaySleepMin: minutes('displaysleep'),
19    preventedBy: prevented ? prevented.split(',').map(s => s.trim()).filter(Boolean) : [],
20    onBattery: /Battery Power/.test(text),
21  }
22}
23
24/**
25 * Seconds until the Mac goes to idle sleep, or null when it will not: sleep disabled, or something
26 * other than this mod's own caffeinate already prevents it.
27 */
28export const secondsToSleep = (p: JrAwakePower, ownCaffeinate: boolean) => {
29  if (p.sleepMin <= 0) return null
30  const others = p.preventedBy.filter(name => !(ownCaffeinate && name === 'caffeinate'))
31  if (others.length > 0 || ownCaffeinate) return null
32  return Math.max(0, p.sleepMin * 60 - p.idleSeconds)
33}
34
35/** Sessions whose heartbeat is fresh and that say Claude is working. */
36export const workingSessions = (beats: readonly JrAwakeSessionBeat[], now: number) =>
37  beats.filter(b => b.working && now - b.at < BEAT_STALE_MS)
38
39/** Whether this session should hold the Mac awake right now. */
40export const shouldHold = (mode: JrAwakeMode, isWorking: boolean) => mode === 'session' || (mode === 'working' && isWorking)
41
42export const formatCountdown = (seconds: number) => {
43  const m = Math.floor(seconds / 60)
44  const s = Math.round(seconds % 60)
45  return `${m}:${String(s).padStart(2, '0')}`
46}
47

1/** working: awake while Claude works in any session. session: awake while any session is open. off. */
2export type JrAwakeMode = 'working' | 'session' | 'off'
3
4export type JrAwakePower = {
5  /** Idle minutes before system sleep; 0 means never. */
6  sleepMin: number
7  displaySleepMin: number
8  /** Processes that currently prevent system sleep, as pmset names them. */
9  preventedBy: string[]
10  onBattery: boolean
11  /** Seconds since the last keyboard or pointer event. */
12  idleSeconds: number
13}
14
15/** One session's heartbeat in the shared store. */
16export type JrAwakeSessionBeat = { id: string; project: string; working: boolean; holding: boolean; at: number }
17
18export type JrAwakeView = {
19  mode: JrAwakeMode
20  holding: boolean
21  /** Sessions holding the Mac awake, this one included. */
22  holders: number
23  sleepIn: number | null
24  power: JrAwakePower | null
25}
26
27declare module 'claude-code' {
28  interface PluginState {
29    'jr-awake': { view: JrAwakeView | null; isBandHidden: boolean }
30  }
31}
32