dymaxion

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionRateLimit } from 'claude-code'
3import type { Pause, Quota } from '../types'
4
5const PERSON_ORIGINS: ReadonlySet<string> = new Set(['composer', 'bridge', 'slack-ping'])
6const USAGE = 'Usage: /pause [compact | status | until-reset [<prompt>]], or /pause limit [<percent> [resume] [compact] | off]'
7const WINDOW_NAMES: Readonly<Record<string, string>> = { five_hour: '5-hour', seven_day: 'weekly' }
8const RESUME_PROMPT = 'The quota window has reset. Continue where you left off.'
9// The reset time is the server's; the margin keeps the resumed turn past it.
10const RESUME_MARGIN_MS = 60_000
11const RESUME_POLL_MS = 60_000
12const LIMIT_MODES: ReadonlySet<string> = new Set(['resume', 'compact'])
13
14// Held by the host, so the pause outlives a reload of this module.
15const pause = atom({ plugin: 'dymaxion', key: 'state' } as const, {
16  isPaused: false,
17  runningTurnId: null,
18  abortedTurnId: null,
19  nextTurnOrigin: null,
20  windowsAtLimit: [],
21  sessionQuota: null,
22  resume: null,
23})
24
25// A value an earlier version of the mod wrote lacks the fields added since.
26function current(s: Pause): Pause {
27  return { ...pause.initial, ...s }
28}
29
30async function load($: EngineInterface) {
31  return current(await read($, pause))
32}
33
34// Holds until the person's prompt; a pause that resumes on its own sets resume after it.
35function paused(s: Pause): Pause {
36  return { ...s, isPaused: true, nextTurnOrigin: null, resume: null }
37}
38
39// windowsAtLimit is emptied so a window still at the limit pauses the session again.
40function resumed(s: Pause): Pause {
41  return { ...s, isPaused: false, resume: null, windowsAtLimit: [] }
42}
43
44function atLimit(windows: readonly SessionRateLimit[], quota: Quota) {
45  return quota.percent > 0 ? windows.filter(w => w.percentUsed >= quota.percent) : []
46}
47
48// When every one of `windows` has reset; null when one does not say.
49function resetOf(windows: readonly SessionRateLimit[]) {
50  const times = windows.map(w => (w.resetsAt === undefined ? NaN : Date.parse(w.resetsAt)))
51  return times.some(Number.isNaN) ? null : Math.max(...times)
52}
53
54function quotaText(q: Quota) {
55  if (q.percent <= 0) return 'does not pause for quota'
56  const actions = [
57    `pauses at ${q.percent}% of a quota window`,
58    ...(q.isCompacting ? ['compacts'] : []),
59    ...(q.isResuming ? ['resumes when it resets'] : []),
60  ]
61  return actions.length < 3 ? actions.join(' and ') : `${actions.slice(0, -1).join(', ')}, and ${actions.at(-1)}`
62}
63
64function isRunning(s: Pause) {
65  return s.runningTurnId !== null && s.runningTurnId !== s.abortedTurnId
66}
67
68function statusOf(s: Pause, setting: Quota) {
69  if (s.isPaused) {
70    if (isRunning(s)) return 'pausing'
71    return s.resume === null ? 'paused' : `paused until ${new Date(s.resume.at).toISOString()}`
72  }
73  const quota = s.sessionQuota ?? setting
74  return quota.percent > 0 && quota.isResuming ? `pauses at ${quota.percent}%, resumes at reset` : undefined
75}
76
77// Also answers the state the write replaced, so of two callers racing to make
78// the same change, each can tell whether it was the one that made it.
79async function write($: EngineInterface, setting: Quota, change: (s: Pause) => Pause) {
80  let before = pause.initial
81  const after = await update($, pause, s => change((before = current(s))))
82  const status = statusOf(after, setting)
83  if (status !== statusOf(before, setting)) $.ui.status(status)
84
85  return { before, after }
86}
87
88// Parallel tool calls of one response each stop the turn; only the first aborts it.
89async function stop($: EngineInterface, setting: Quota, turnId: string) {
90  const { before } = await write($, setting, s => ({ ...s, abortedTurnId: turnId }))
91  if (before.abortedTurnId !== turnId) await $.turn.abort({ turnId })
92}
93
94// Queued until the session is idle: it runs once the turn stops, while the prompt cache is warm.
95function compactWhenIdle($: EngineInterface) {
96  $.clock.after(0, () => void $.command.run({ command: 'compact' }))
97}
98
99async function resumeIfDue($: EngineInterface, setting: Quota) {
100  const now = await $.clock.now()
101  const s = await load($)
102  if (s.resume === null || now < s.resume.at) return
103  const { before } = await write($, setting, t => (t.resume === null ? t : resumed(t)))
104  if (before.resume !== null) await $.prompt.submit({ text: before.resume.prompt })
105}
106
107// Waits for the most-used window, the one holding the session back; of several
108// used as much, for the last to reset.
109async function runUntilReset($: EngineInterface, setting: Quota, prompt: string) {
110  const { rateLimits } = await $.session.usage()
111  const most = Math.max(...rateLimits.map(w => w.percentUsed))
112  const windows = rateLimits.filter(w => w.percentUsed === most)
113  const reset = windows.length === 0 ? null : resetOf(windows)
114  if (reset === null || reset <= (await $.clock.now())) {
115    return { text: 'No quota reading gives a reset still to come, so there is nothing to wait for.' }
116  }
117  const at = reset + RESUME_MARGIN_MS
118  await write($, setting, s => ({ ...paused(s), resume: { at, prompt: prompt === '' ? RESUME_PROMPT : prompt } }))
119  const names = windows.map(w => WINDOW_NAMES[w.kind] ?? w.kind).join(' and ')
120
121  return { text: `Paused until ${new Date(at).toISOString()}, after the ${names} window resets.` }
122}
123
124async function runLimit($: EngineInterface, args: readonly string[], setting: Quota) {
125  const [value, ...modes] = args
126  if (value === undefined) return { text: `This session ${quotaText((await load($)).sessionQuota ?? setting)}.` }
127  const percent = value === 'off' ? 0 : Number(value)
128  const isValid =
129    percent >= 0 &&
130    percent <= 100 &&
131    (modes.length === 0 || percent > 0) &&
132    modes.every(m => LIMIT_MODES.has(m)) &&
133    new Set(modes).size === modes.length
134  if (!isValid) return { text: USAGE }
135  const isResuming = modes.includes('resume')
136  const quota = { percent, isResuming, isCompacting: modes.includes('compact') }
137  await write($, setting, s => ({ ...s, sessionQuota: quota, resume: isResuming ? s.resume : null }))
138
139  return { text: `This session ${quotaText(quota)}.` }
140}
141
142async function runStatus($: EngineInterface, setting: Quota) {
143  const s = await load($)
144  const { rateLimits } = await $.session.usage()
145  const pauseLine = !s.isPaused
146    ? 'Not paused.'
147    : isRunning(s)
148      ? 'Pausing: Claude stops before its next step.'
149      : s.resume === null
150        ? 'Paused until you send a prompt.'
151        : `Paused until ${new Date(s.resume.at).toISOString()}, or until you send a prompt.`
152  const source = s.sessionQuota === null ? 'from the settings' : 'set by /pause limit'
153  const windowLines = rateLimits.map(w => {
154    const resets = w.resetsAt === undefined ? '' : `, resets ${w.resetsAt}`
155    const held = s.windowsAtLimit.includes(w.kind) ? '; paused for it once, again only after it resets' : ''
156    return `${WINDOW_NAMES[w.kind] ?? w.kind} limit: ${w.percentUsed}% used${resets}${held}`
157  })
158
159  return {
160    text: [pauseLine, `Quota: this session ${quotaText(s.sessionQuota ?? setting)} (${source}).`, ...windowLines].join(
161      '\n',
162    ),
163  }
164}
165
166export const register: Register = (on, options) => {
167  const setting: Quota = {
168    percent: typeof options.quotaLimit === 'number' ? options.quotaLimit : 0,
169    isResuming: options.quotaResume === true,
170    isCompacting: options.quotaCompact === true,
171  }
172
173  on('session.start', async ($, e, next) => {
174    await $.command.register({
175      name: 'pause',
176      description: 'Stop once the running tool finishes and stay paused until you send a prompt',
177      argumentHint: '[compact | status | until-reset [prompt] | limit <percent> [resume] [compact] | limit off]',
178      immediate: true,
179    })
180    // After a reload, show the pause the host still holds.
181    $.ui.status(statusOf(await load($), setting))
182    // Polled rather than one long wait, so a machine that slept past a reset
183    // still resumes soon after it wakes.
184    $.clock.every(RESUME_POLL_MS, () => void resumeIfDue($, setting))
185
186    return next(e)
187  })
188
189  on('command.run', { command: 'pause' }, async ($, e) => {
190    const [verb = '', ...rest] = e.args.trim().split(/\s+/)
191    if (verb === 'limit') return runLimit($, rest, setting)
192    if (verb === 'status' && rest.length === 0) return runStatus($, setting)
193    if (verb === 'until-reset') return runUntilReset($, setting, e.args.trim().slice(verb.length).trim())
194    if (rest.length > 0 || (verb !== '' && verb !== 'compact')) return { text: USAGE }
195    const isCompacting = verb === 'compact'
196
197    const { after } = await write($, setting, paused)
198    if (isCompacting) compactWhenIdle($)
199    if (!isRunning(after)) {
200      return { text: isCompacting ? 'Paused; compacting now.' : 'Paused until you send a prompt.' }
201    }
202
203    return {
204      text: isCompacting
205        ? 'Claude will stop before its next step, then compact.'
206        : 'Claude will stop before its next step.',
207    }
208  })
209
210  // A window pauses the session when it reaches the limit, which is off unless
211  // set. Once the session resumes, the same window pauses it again only after
212  // it has reset. A pause holding for the person's prompt keeps holding; one
213  // that resumes on its own waits for this window too. Only a new pause compacts.
214  on('session.measure', async ($, e, next) => {
215    const now = await $.clock.now()
216    const { before, after } = await write($, setting, s => {
217      const quota = s.sessionQuota ?? setting
218      const windows = atLimit(e.rateLimits, quota)
219      const t = { ...s, windowsAtLimit: windows.map(w => w.kind) }
220      if (!windows.some(w => !s.windowsAtLimit.includes(w.kind))) return t
221      const reset = resetOf(windows)
222      const at = reset !== null && reset > now ? reset + RESUME_MARGIN_MS : null
223      if (s.isPaused) {
224        return s.resume === null || at === null ? t : { ...t, resume: { ...s.resume, at: Math.max(s.resume.at, at) } }
225      }
226
227      return { ...paused(t), resume: quota.isResuming && at !== null ? { at, prompt: RESUME_PROMPT } : null }
228    })
229    if (!before.isPaused && after.isPaused && (after.sessionQuota ?? setting).isCompacting) compactWhenIdle($)
230    const reached = atLimit(e.rateLimits, before.sessionQuota ?? setting).find(
231      w => !before.windowsAtLimit.includes(w.kind),
232    )
233    if (reached !== undefined) {
234      const name = WINDOW_NAMES[reached.kind] ?? reached.kind
235      const then =
236        after.resume === null ? 'Send a prompt to go on.' : `Resuming at ${new Date(after.resume.at).toISOString()}.`
237      $.ui.toast(`Pausing: the ${name} limit is at ${reached.percentUsed}%. ${then}`)
238    }
239
240    return next(e)
241  })
242
243  on('prompt.submit', async ($, e, next) => {
244    if (e.turnId === undefined) await write($, setting, s => ({ ...s, nextTurnOrigin: e.origin.kind }))
245
246    return next(e)
247  })
248
249  on('turn.start', async ($, e, next) => {
250    const started = await next(e)
251    if (e.agentId !== undefined) return started
252
253    await write($, setting, s => {
254      const t = { ...s, runningTurnId: started.turnId, nextTurnOrigin: null }
255      return s.isPaused && PERSON_ORIGINS.has(s.nextTurnOrigin ?? '') ? { ...t, isPaused: false, resume: null } : t
256    })
257
258    return started
259  })
260
261  on('turn.complete', async ($, e, next) => {
262    if (e.agentId === undefined) {
263      await write($, setting, s => (s.runningTurnId === e.turnId ? { ...s, runningTurnId: null } : s))
264    }
265
266    return next(e)
267  })
268
269  // The two boundaries of the main loop: the next model request, and the next
270  // tool of a response whose earlier tools ran one after another.
271  on('turn.step', async function* ($, e, next) {
272    if (e.agentId !== undefined || !(await load($)).isPaused) return yield* next(e)
273    await stop($, setting, e.turnId)
274
275    return { turnId: e.turnId, index: e.index, answer: '', toolUses: [], stopReason: null, usage: null }
276  })
277
278  on('tool.call', async ($, e, next) => {
279    if (e.agentId !== undefined) return next(e)
280    const s = await load($)
281    if (!s.isPaused || s.runningTurnId === null) return next(e)
282    await stop($, setting, s.runningTurnId)
283
284    return { deny: 'Paused.' }
285  })
286}
287

1/** When the session pauses for quota: at `percent` of a window (0: never), whether it compacts, and whether it resumes at the reset. */
2export type Quota = { percent: number; isResuming: boolean; isCompacting: boolean }
3
4export type Pause = {
5  isPaused: boolean
6  /** The main turn that started and has not completed. */
7  runningTurnId: string | null
8  /** The main turn the mod aborted; its turn.complete may come later or not at all. */
9  abortedTurnId: string | null
10  /** The origin of the prompt the next main turn starts from. */
11  nextTurnOrigin: string | null
12  /** The rate-limit windows at or past the limit at the last measurement. */
13  windowsAtLimit: string[]
14  /** The quota /pause limit set for this session; null follows the settings. */
15  sessionQuota: Quota | null
16  /** When a pause resumes on its own (epoch milliseconds) and the prompt it sends; null holds until the person's prompt. */
17  resume: { at: number; prompt: string } | null
18}
19
20declare module 'claude-code' {
21  interface PluginState {
22    dymaxion: { state: Pause }
23  }
24}
25