rate-limit-guard

1import type { CommandSpec, EngineInterface, Register, Timer } from 'claude-code'
2
3import type { ResumeMode } from '../types'
4import {
5  agentCheckpoint,
6  fiveHour,
7  formatLimits,
8  orchestratorStopMessage,
9  resumeAt,
10  resumeMessage,
11  shortHash,
12  shouldTrigger,
13  stopMessage,
14} from './logic'
15
16const phase = { plugin: 'rate-limit-guard', key: 'phase' } as const
17const resetsAt = { plugin: 'rate-limit-guard', key: 'resetsAt' } as const
18const lastDrained = { plugin: 'rate-limit-guard', key: 'lastDrained' } as const
19const modeOverride = { plugin: 'rate-limit-guard', key: 'resumeMode' } as const
20
21// idle is left out: a finished subagent or a teammate between turns has nothing to checkpoint.
22const LIVE = new Set(['pending', 'running', 'waiting'])
23const POLL_MS = 30_000
24
25type Config = { dir: string; isAutoResume: boolean; bufferMs: number; resumeMode: ResumeMode }
26
27const RESUME_MODES = new Set<string>(['continue', 'compact', 'clear'])
28
29const COMMANDS: CommandSpec[] = [
30  {
31    name: 'rate-limit-guard-rearm',
32    description: 'Re-arm the rate-limit guard for the current window (clears the drained mark)',
33  },
34  { name: 'rate-limit-guard-status', description: 'Show the rate-limit windows and the guard state' },
35  {
36    name: 'rate-limit-guard-resume-mode',
37    description: 'Set how this session resumes after the reset (clear, compact, continue, default)',
38    argumentHint: '[clear|compact|continue|default]',
39    // Usually set while a long run is mid-turn.
40    immediate: true,
41  },
42]
43// Each command also answers to /rlg-<suffix>; the API has no aliases, so the short name is a second command.
44const short = (name: string) => name.replace(/^rate-limit-guard-/, 'rlg-')
45const names = (name: string) => [name, short(name)]
46const COMPACT_INSTRUCTIONS =
47  'Keep only the overall goal and the paths of the checkpoint files; the checkpoints hold the state.'
48
49let poll: Timer | undefined
50
51// $.state is per session and /clear wipes it, so the drained mark also lives in a file:
52// without it the cleared session re-fires in the window it already drained.
53const markFile = (dir: string) => `${dir}/.drained.json`
54
55async function loadDrained($: EngineInterface, dir: string): Promise<string | null> {
56  const { value } = await $.state.get(lastDrained)
57  if (typeof value === 'string') return value
58  try {
59    const parsed = JSON.parse(await $.fs.read(markFile(dir)))
60    return typeof parsed?.resetsAt === 'string' ? parsed.resetsAt : null
61  } catch {
62    return null
63  }
64}
65
66// A failed write must not stop the guard: the in-memory mark still holds until a clear.
67async function saveDrained($: EngineInterface, dir: string, at: string | null) {
68  await $.state.set(lastDrained, at)
69  await $.fs.write(markFile(dir), JSON.stringify({ resetsAt: at })).catch(() => {})
70}
71
72// A failed id lookup must not stop the guard: fall back to a fixed hash.
73async function sessionHash($: EngineInterface): Promise<string> {
74  return shortHash(await $.session.id().catch(() => 'unknown-session'))
75}
76
77// Poll rather than one long timer: it survives clock drift and a laptop sleep.
78function watchForReset($: EngineInterface, cfg: Config) {
79  poll?.cancel()
80  if (!cfg.isAutoResume) return
81
82  poll = $.clock.every(POLL_MS, async () => {
83    const [{ value: p }, { value: at }] = await Promise.all([
84      $.state.get(phase),
85      $.state.get(resetsAt),
86    ])
87    if (p !== 'draining') return poll?.cancel()
88
89    const now = await $.clock.now()
90    if (now < resumeAt(at ?? null, now, cfg.bufferMs)) return
91
92    poll?.cancel()
93    await $.state.set(phase, 'armed')
94    $.ui.status(undefined)
95    $.ui.toast('rate-limit-guard: window reset, resuming')
96    // Taken before a clear: the new session id would hash to a different checkpoint name.
97    const hash = await sessionHash($)
98    // The session's own pick (/rate-limit-guard-resume-mode) wins over the configured default.
99    const { value: picked } = await $.state.get(modeOverride)
100    const mode = picked ?? cfg.resumeMode
101    // After a 5h wait the prompt cache is cold, so re-sending the old transcript is the expensive part.
102    // A failed clear or compact must not stop the resume: continue in the same session instead.
103    if (mode === 'clear') await $.command.run({ command: 'clear' }).catch(() => {})
104    if (mode === 'compact') {
105      await $.session.compact({ instructions: COMPACT_INSTRUCTIONS }).catch(() => {})
106    }
107    await $.prompt.submit({ text: resumeMessage(cfg.dir, hash) })
108  })
109}
110
111export const register: Register = (on, options) => {
112  const threshold = Number(options.threshold ?? 90)
113  const dir = String(options.checkpointDir ?? '.claude/checkpoints')
114  const isAutoResume = options.autoResume !== false
115  const bufferMs = Number(options.resumeBufferSeconds ?? 60) * 1000
116  const mode = String(options.resumeMode ?? 'continue')
117  const resumeMode = RESUME_MODES.has(mode) ? (mode as ResumeMode) : 'continue'
118  const cfg: Config = { dir, isAutoResume, bufferMs, resumeMode }
119
120  on('session.start', async ($, e, next) => {
121    for (const spec of COMMANDS) {
122      await $.command.register(spec)
123      await $.command.register({ ...spec, name: short(spec.name), description: `Short for /${spec.name}` })
124    }
125    // A hot reload re-runs this; a draining session keeps polling.
126    if ((await $.state.get(phase)).value === undefined) {
127      await $.state.set(phase, 'armed')
128      await $.state.set(modeOverride, null)
129      await $.state.set(resetsAt, null)
130      await $.state.set(lastDrained, await loadDrained($, dir))
131    }
132    if ((await $.state.get(phase)).value === 'draining') watchForReset($, cfg)
133
134    return next(e)
135  })
136
137  for (const command of names('rate-limit-guard-status')) on('command.run', { command }, async $ => {
138    const { rateLimits } = await $.session.usage()
139    const { value: held } = await $.state.get(phase)
140    const { value: picked } = await $.state.get(modeOverride)
141    const resume = picked ? `${picked} (this session)` : `${resumeMode} (default)`
142
143    return {
144      text: `${formatLimits(rateLimits, threshold)}\nguard: ${held ?? 'unknown'}\nresume: ${resume}`,
145    }
146  })
147
148  // Lives in $.state, so it holds for this session only: a /clear (including a clear resume) drops it.
149  for (const command of names('rate-limit-guard-resume-mode')) on('command.run', { command }, async ($, e) => {
150    const arg = e.args.trim().toLowerCase()
151    if (arg === '') {
152      const { value: picked } = await $.state.get(modeOverride)
153      return {
154        text: picked
155          ? `resume mode: ${picked} (this session; default ${resumeMode})`
156          : `resume mode: ${resumeMode} (default)`,
157      }
158    }
159    if (arg === 'default') {
160      await $.state.set(modeOverride, null)
161      return { text: `resume mode back to the default: ${resumeMode}` }
162    }
163    if (!RESUME_MODES.has(arg)) {
164      return { text: `unknown resume mode "${arg}": use clear, compact, continue or default` }
165    }
166    await $.state.set(modeOverride, arg as ResumeMode)
167    return { text: `resume mode for this session: ${arg}` }
168  })
169
170  for (const command of names('rate-limit-guard-rearm')) on('command.run', { command }, async $ => {
171    poll?.cancel()
172    await $.state.set(phase, 'armed')
173    await $.state.set(resetsAt, null)
174    await saveDrained($, dir, null)
175    $.ui.status(undefined)
176
177    return { text: 'rate-limit-guard re-armed: it can fire again in the current window.' }
178  })
179
180  on('session.measure', async ($, e, next) => {
181    const limit = fiveHour(e.rateLimits)
182    const held = await $.state.get(phase)
183    const drained = await loadDrained($, dir)
184
185    if (
186      held.value !== 'draining' &&
187      limit !== undefined &&
188      shouldTrigger(limit, threshold, drained)
189    ) {
190      // ifVersion: two measures in a burst must not both trigger.
191      const claim = await $.state.set(phase, 'draining', { ifVersion: held.version })
192
193      if (claim.isSet) {
194        await $.state.set(resetsAt, limit.resetsAt ?? null)
195        await saveDrained($, dir, limit.resetsAt ?? null)
196        $.ui.status(`rate-limit-guard: ${limit.percentUsed}% — checkpointing`)
197        $.ui.toast(`5-hour window at ${limit.percentUsed}%: asking agents to checkpoint and stop`)
198
199        // A failed listing must not stop the orchestrator from being told.
200        const agents = (await $.agent.list().catch(() => [])).filter(a => LIVE.has(a.status))
201        const hash = await sessionHash($)
202        await Promise.allSettled(
203          agents.map(a =>
204            $.session.send({
205              to: { agentId: a.id },
206              text: stopMessage(
207                limit.percentUsed,
208                agentCheckpoint(dir, a.name ?? (a.description || a.type), a.id, hash),
209              ),
210            }),
211          ),
212        )
213        void $.prompt.submit({ text: orchestratorStopMessage(limit.percentUsed, dir, hash) })
214        watchForReset($, cfg)
215      }
216    }
217
218    return next(e)
219  })
220
221  // The orchestrator may be mid-turn for a long while; keep the instruction in its system prompt too.
222  on('prompt.compose', async ($, e, next) => {
223    const out = await next(e)
224    if ((await $.state.get(phase)).value !== 'draining') return out
225
226    return {
227      sections: [
228        ...out.sections,
229        {
230          id: 'rate-limit-guard:draining',
231          scope: 'session',
232          text:
233            `The 5-hour usage limit is nearly exhausted. Start no new work. Finish the current unit, ` +
234            `write checkpoints to ${dir}/, and stop. Work resumes automatically after the reset.`,
235        },
236      ],
237    }
238  })
239}
240

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2
3export const FIVE_HOURS_MS = 5 * 60 * 60 * 1000
4
5export const fiveHour = (limits: readonly Limit[]): Limit | undefined =>
6  limits.find(l => l.kind === 'five_hour')
7
8/** Trigger once per window: at/over threshold, and not the window already drained. */
9export const shouldTrigger = (
10  limit: Limit | undefined,
11  threshold: number,
12  lastDrained: string | null,
13): boolean =>
14  limit !== undefined &&
15  limit.percentUsed >= threshold &&
16  (limit.resetsAt === undefined || limit.resetsAt !== lastDrained)
17
18/** When to resume: the reset time plus a buffer; five hours from `now` if the reset is unknown. */
19export const resumeAt = (resetsAt: string | null, now: number, bufferMs: number): number => {
20  const at = resetsAt === null ? NaN : Date.parse(resetsAt)
21
22  return (Number.isNaN(at) ? now + FIVE_HOURS_MS : at) + bufferMs
23}
24
25const slug = (s: string): string => s.replace(/[^A-Za-z0-9_.-]+/g, '-').replace(/^[-.]+|[-.]+$/g, '') || 'agent'
26
27/** First 8 hex chars of a SHA-256: enough to keep sessions' checkpoints apart. */
28export const shortHash = async (text: string): Promise<string> => {
29  const bytes = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(text))
30
31  return Array.from(new Uint8Array(bytes))
32    .slice(0, 4)
33    .map(b => b.toString(16).padStart(2, '0'))
34    .join('')
35}
36
37/** <dir>/<name>-<agentId>-<sessionHash>.md, one file per agent per session. */
38export const agentCheckpoint = (dir: string, name: string, agentId: string, hash: string): string =>
39  `${dir}/${slug(name)}-${slug(agentId)}-${hash}.md`
40
41export const orchestratorCheckpoint = (dir: string, hash: string): string =>
42  `${dir}/orchestrator-${hash}.md`
43
44export const stopMessage = (percent: number, path: string): string =>
45  `RATE-LIMIT GUARD: the 5-hour usage window is at ${percent}%. Stop at the next safe point ` +
46  `(finish the current atomic unit of work; do not start new subtasks or long tool calls). ` +
47  `Write a checkpoint to exactly ${path} (use that filename, it is unique to you) covering: ` +
48  `what is done, what is in progress, exact next steps, files touched, and any state needed to ` +
49  `resume cold. Then reply "CHECKPOINTED" and stop. Work resumes automatically after the window resets.`
50
51export const orchestratorStopMessage = (percent: number, dir: string, hash: string): string =>
52  `RATE-LIMIT GUARD: the 5-hour usage window is at ${percent}%. Every running subagent has been told to ` +
53  `checkpoint and stop, each to its own ${dir}/<name>-<agentId>-${hash}.md. You must: (1) not start or ` +
54  `dispatch any new work; (2) collect each subagent's "CHECKPOINTED" reply (wait for in-flight ones; ` +
55  `do not wait on any that stall); (3) write your own orchestrator checkpoint to exactly ` +
56  `${orchestratorCheckpoint(dir, hash)} (overall goal, plan, which subagents finished / were stopped / ` +
57  `never started, their checkpoint files, and the exact next steps); (4) reply "CHECKPOINTED" and stop. ` +
58  `Work resumes automatically after the window resets.`
59
60export const resumeMessage = (dir: string, hash: string): string =>
61  `RATE-LIMIT GUARD: the usage window has reset. Read ${orchestratorCheckpoint(dir, hash)} and the ` +
62  `per-agent checkpoints matching ${dir}/*-${hash}.md (ignore files with a different hash: they belong ` +
63  `to other sessions), then continue the work: re-dispatch subagents from their checkpoints and carry on ` +
64  `from where you stopped. Rename each checkpoint to *.consumed.md once resumed.`
65
66const WINDOW_NAMES: Record<string, string> = { five_hour: '5-hour', seven_day: 'weekly' }
67
68/** One line per window: "5-hour: 35% (resets 2026-10-09 15:00 UTC)". */
69export const formatLimits = (limits: readonly Limit[], threshold: number): string => {
70  if (limits.length === 0) return 'No rate-limit windows reported (not a subscription, or no response yet).'
71
72  return limits
73    .map(l => {
74      const reset = l.resetsAt === undefined ? '' : ` (resets ${l.resetsAt.replace('T', ' ').slice(0, 16)} UTC)`
75      const guard = l.kind === 'five_hour' ? ` [guard at ${threshold}%]` : ''
76
77      return `${WINDOW_NAMES[l.kind] ?? l.kind}: ${l.percentUsed}%${reset}${guard}`
78    })
79    .join('\n')
80}
81

1export type Phase = 'armed' | 'draining'
2export type ResumeMode = 'continue' | 'compact' | 'clear'
3
4declare module 'claude-code' {
5  interface PluginState {
6    'rate-limit-guard': {
7      /** armed: watching; draining: told everyone to checkpoint, waiting for the reset. */
8      phase: Phase
9      /** ISO reset time of the five_hour window that was drained, or null. */
10      resetsAt: string | null
11      /** resetsAt of the last window we drained, so one window triggers once. */
12      lastDrained: string | null
13      /** This session's pick from /rate-limit-guard-resume-mode, or null for the configured default. */
14      resumeMode: ResumeMode | null
15    }
16  }
17}
18