usage-aware

1import type { EngineInterface, Register, SessionCost, SessionRateLimit } from 'claude-code'
2
3import type { Band, Resume } from '../types'
4
5// Thresholds on the fullest window (five_hour or seven_day), percent used.
6const LOW = 75 // past this: say so, prefer no new fan-out
7const STOP = 90 // past this: new agents are blocked until the operator says yes
8const CUT_OFF = 98 // a turn that dies past this was cut off by the limit
9const SLACK_MS = 90_000 // wait this long past resetsAt before resuming
10const TICK_MS = 60_000
11
12const TOOL = 'mcp__usage-aware__usage'
13const GATED = new Set(['Agent', 'Workflow'])
14const RANK: Record<Band, number> = { ok: 0, low: 1, stop: 2 }
15const LABEL: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'spend' }
16
17const bandRef = { plugin: 'usage-aware', key: 'band' } as const
18const noticeRef = { plugin: 'usage-aware', key: 'pendingNotice' } as const
19const resumeRef = { plugin: 'usage-aware', key: 'resume' } as const
20const overrideRef = { plugin: 'usage-aware', key: 'override' } as const
21
22type Input = { action?: string; prompt?: string }
23
24const label = (w: SessionRateLimit) => LABEL[w.kind] ?? w.kind
25
26const fullest = (limits: readonly SessionRateLimit[]) =>
27  limits.reduce<SessionRateLimit | undefined>(
28    (top, w) => (top === undefined || w.percentUsed > top.percentUsed ? w : top),
29    undefined,
30  )
31
32const bandOf = (pct: number): Band => (pct >= STOP ? 'stop' : pct >= LOW ? 'low' : 'ok')
33
34export function until(iso: string | undefined, now: number): string {
35  if (iso === undefined) return 'unknown'
36  const ms = Date.parse(iso) - now
37  if (!(ms > 0)) return 'now'
38  const mins = Math.floor(ms / 60_000)
39  const d = Math.floor(mins / 1440)
40  const h = Math.floor((mins % 1440) / 60)
41  const m = mins % 60
42  if (d > 0) return `${d}d${h}h`
43  if (h > 0) return `${h}h${String(m).padStart(2, '0')}m`
44  return `${m}m`
45}
46
47export function statusLine(limits: readonly SessionRateLimit[], cost: SessionCost | undefined, now: number) {
48  const parts = limits.map(w => `${label(w)} ${w.percentUsed}% ↻${until(w.resetsAt, now)}`)
49  if (cost !== undefined) parts.push(`$${cost.usd.toFixed(2)}`)
50  return parts.length > 0 ? parts.join(' · ') : undefined
51}
52
53export function figures(limits: readonly SessionRateLimit[], cost: SessionCost | undefined, now: number) {
54  const lines = limits.map(
55    w => `${label(w)} window: ${w.percentUsed}% used, resets in ${until(w.resetsAt, now)} (${w.resetsAt ?? 'reset time unknown'})`,
56  )
57  if (cost !== undefined) lines.push(`this session so far: $${cost.usd.toFixed(2)} (subagents included)`)
58  return lines.join('\n')
59}
60
61const GUIDANCE: Record<Band, string> = {
62  ok: '',
63  low:
64    'Usage is getting low. Finish what is running, avoid new fan-out (parallel agents, workflows, big reviews) unless the operator agrees, and checkpoint long jobs as you go.',
65  stop:
66    'Usage is nearly out. New agents are blocked. Finish the current step, write a checkpoint (progress file, ledger), then call the usage tool with action resume_after_reset and a continuation prompt so the session picks itself up after the reset. If the operator may want to spend the rest anyway, ask them through AskUserQuestion with the figures in the question; on a yes, call the usage tool with action allow_agents.',
67}
68
69async function overrideActive($: EngineInterface, now: number) {
70  const { value } = await $.state.get(overrideRef)
71  return value != null && value.until > now
72}
73
74// Moves the band when a window moved; a rise leaves a notice for the main loop's next tool result.
75async function reband($: EngineInterface, limits: readonly SessionRateLimit[], cost: SessionCost | undefined, now: number) {
76  const top = fullest(limits)
77  if (top === undefined) return
78  const band = bandOf(top.percentUsed)
79  const { value: was = 'ok' } = await $.state.get(bandRef)
80  if (band === was) return
81  await $.state.set(bandRef, band)
82  if (RANK[band] < RANK[was]) {
83    await $.state.set(noticeRef, null)
84    if (band === 'ok') await $.state.set(overrideRef, null)
85    return
86  }
87  await $.state.set(noticeRef, `[usage-aware] ${label(top)} window crossed ${band === 'stop' ? STOP : LOW}%.\n${figures(limits, cost, now)}\n${GUIDANCE[band]}`)
88  $.ui.toast(`usage-aware: ${label(top)} at ${top.percentUsed}%, resets in ${until(top.resetsAt, now)}`)
89}
90
91async function arm($: EngineInterface, prompt: string, limits: readonly SessionRateLimit[]): Promise<Resume | undefined> {
92  const top = fullest(limits.filter(w => w.resetsAt !== undefined))
93  if (top?.resetsAt === undefined) return undefined
94  const resume: Resume = { at: Date.parse(top.resetsAt) + SLACK_MS, prompt, window: label(top) }
95  await $.state.set(resumeRef, resume)
96  return resume
97}
98
99async function tick($: EngineInterface) {
100  const { value: resume } = await $.state.get(resumeRef)
101  if (resume == null || (await $.clock.now()) < resume.at) return
102  await $.state.set(resumeRef, null)
103  await $.state.set(overrideRef, null)
104  $.ui.toast('usage-aware: the window reset, resuming')
105  await $.prompt.submit({ text: `[usage-aware] The ${resume.window} usage window has reset. ${resume.prompt}` })
106}
107
108async function serve($: EngineInterface, input: Input) {
109  const usage = await $.session.usage()
110  const now = await $.clock.now()
111  const action = input.action ?? 'read'
112
113  if (action === 'resume_after_reset') {
114    if (!input.prompt) return { deny: 'usage-aware: resume_after_reset needs a prompt: what the session should do once the window resets.' }
115    const resume = await arm($, input.prompt, usage.rateLimits)
116    if (resume === undefined) return { deny: 'usage-aware: no rate-limit window with a reset time has been reported yet, so there is nothing to wait for.' }
117    return { result: `Armed. When the ${resume.window} window resets (${new Date(resume.at).toISOString()}, in ${until(new Date(resume.at).toISOString(), now)}), this session submits:\n${resume.prompt}\nThis only fires while the session stays open.` }
118  }
119  if (action === 'cancel_resume') {
120    await $.state.set(resumeRef, null)
121    return { result: 'The armed resume is cancelled.' }
122  }
123  if (action === 'allow_agents') {
124    const top = fullest(usage.rateLimits.filter(w => w.resetsAt !== undefined))
125    const untilAt = top?.resetsAt !== undefined ? Date.parse(top.resetsAt) : now + 5 * 3_600_000
126    await $.state.set(overrideRef, { until: untilAt })
127    return { result: `New agents are allowed again until ${new Date(untilAt).toISOString()}. Only use this after the operator said yes.` }
128  }
129
130  const { value: resume } = await $.state.get(resumeRef)
131  const { value: band = 'ok' } = await $.state.get(bandRef)
132  if (usage.rateLimits.length === 0)
133    return { result: `No rate-limit reading yet (none until the first response of this session, and none off a subscription).${usage.cost ? `\nThis session so far: $${usage.cost.usd.toFixed(2)}` : ''}` }
134  return {
135    result: [
136      figures(usage.rateLimits, usage.cost, now),
137      `band: ${band} (low from ${LOW}%, agents blocked from ${STOP}%)${(await overrideActive($, now)) ? ', operator override active' : ''}`,
138      resume ? `armed resume: after the ${resume.window} reset, ${new Date(resume.at).toISOString()}` : 'no resume armed',
139    ].join('\n'),
140  }
141}
142
143export const register: Register = on => {
144  on('session.start', async ($, e, next) => {
145    await $.tool.register({
146      name: 'usage',
147      description:
148        "Read this session's plan usage: percent used of the 5-hour and 7-day windows, when each resets, and what the session has cost. " +
149        'Call it whenever the user asks how much usage is left, and before starting expensive fan-out. ' +
150        'Other actions: resume_after_reset (with prompt) makes this session submit that prompt once the fullest window resets, so a job the limit stops picks itself up; ' +
151        'cancel_resume drops it; allow_agents lifts the new-agent block for the current window, only after the operator said yes.',
152      inputSchema: {
153        type: 'object',
154        properties: {
155          action: { type: 'string', enum: ['read', 'resume_after_reset', 'cancel_resume', 'allow_agents'] },
156          prompt: { type: 'string', description: 'For resume_after_reset: the continuation prompt, naming the checkpoint to resume from.' },
157        },
158      },
159      isDeferred: false,
160    })
161    $.clock.every(TICK_MS, () => void tick($))
162    const usage = await $.session.usage()
163    $.ui.status(statusLine(usage.rateLimits, usage.cost, await $.clock.now()))
164    return next(e)
165  })
166
167  on('session.measure', async ($, e, next) => {
168    const now = await $.clock.now()
169    $.ui.status(statusLine(e.rateLimits, e.cost, now))
170    if (e.changed.includes('rateLimits')) await reband($, e.rateLimits, e.cost, now)
171    return next(e)
172  })
173
174  // Every prompt carries the figures, so the model steers without having to ask.
175  on('prompt.submit', async ($, e, next) => {
176    if (e.command !== undefined) return next(e)
177    const usage = await $.session.usage()
178    const top = fullest(usage.rateLimits)
179    if (top === undefined) return next(e)
180    const now = await $.clock.now()
181    const line = `[usage-aware] ${statusLine(usage.rateLimits, usage.cost, now)}. ${GUIDANCE[bandOf(top.percentUsed)]}`.trim()
182    await $.state.set(noticeRef, null)
183    return next({ ...e, context: [...(e.context ?? []), line] })
184  })
185
186  on('tool.call', async ($, e, next) => {
187    if (e.tool === TOOL) return serve($, e as unknown as Input)
188
189    if (GATED.has(e.tool)) {
190      const usage = await $.session.usage()
191      const top = fullest(usage.rateLimits)
192      const now = await $.clock.now()
193      if (top !== undefined && top.percentUsed >= STOP && !(await overrideActive($, now)))
194        return {
195          deny: `usage-aware: the ${label(top)} window is at ${top.percentUsed}% (resets in ${until(top.resetsAt, now)}), so no new agents. ${GUIDANCE.stop}`,
196        }
197    }
198
199    const ran = await next(e)
200    if (e.agentId !== undefined || ran.deny !== undefined || ran.isError === true) return ran
201    const { value: notice } = await $.state.get(noticeRef)
202    if (!notice) return ran
203    await $.state.set(noticeRef, null)
204    return { ...ran, context: [...(ran.context ?? []), notice] }
205  }).catch(($, e, next) => (next.called || e.tool !== TOOL ? next(e) : { deny: 'usage-aware: the usage tool failed.' }))
206
207  // A main-loop turn that dies at the limit arms its own resume, unless one is armed already.
208  on('turn.complete', async ($, e, next) => {
209    if (e.agentId !== undefined || e.reason === 'aborted') return next(e)
210    const usage = await $.session.usage()
211    const top = fullest(usage.rateLimits)
212    const looksCut = e.reason === 'error' || /usage limit|hit your limit|limit reached/i.test(e.answer)
213    if (top === undefined || top.percentUsed < CUT_OFF || !looksCut) return next(e)
214    const { value: armed } = await $.state.get(resumeRef)
215    if (armed == null) {
216      const resume = await arm(
217        $,
218        'The limit cut off the last turn mid-task. Re-read where it got to (the transcript tail, any progress or checkpoint files it wrote) and carry on with the task; if it was already finished, say so and stop.',
219        usage.rateLimits,
220      )
221      if (resume !== undefined) {
222        const now = await $.clock.now()
223        void $.ui.notify(`Cut off at the ${resume.window} limit; resuming in ${until(new Date(resume.at).toISOString(), now)}.`)
224      }
225    }
226    return next(e)
227  })
228}
229

1export type Band = 'ok' | 'low' | 'stop'
2
3export type Resume = { at: number; prompt: string; window: string }
4
5declare module 'claude-code' {
6  interface PluginState {
7    'usage-aware': {
8      band: Band
9      pendingNotice: string | null
10      resume: Resume | null
11      override: { until: number } | null
12    }
13  }
14}
15