resume-at-reset

1import type { EngineInterface, Register, SessionRateLimit, Timer } from 'claude-code'
2
3const WINDOWS = new Set(['five_hour', 'seven_day'])
4// Resume a minute after the reset, so the new window is surely open.
5const GRACE_MS = 60_000
6// A resume that hits the same limit again means a stale reset time: stop retrying.
7const MAX_RESUMES = 3
8// A plan this far past due belongs to a session resumed much later: drop it.
9const STALE_MS = 60 * 60_000
10const PROMPT = 'The usage limit has reset. Continue where you left off.'
11
12// stoppedAt marks a plan made after a real stop; 0.1.0 armed on a full window alone and its plans lack it.
13type Plan = { resetsAt: string; at: number; resumes: number; isArmed: boolean; stoppedAt: number }
14
15// One module per session; a reload starts these over and session.start re-arms from the store.
16let timer: Timer | undefined
17let armedFor: string | undefined
18
19const isPlan = (value: unknown): value is Plan =>
20  typeof value === 'object' &&
21  value !== null &&
22  'resetsAt' in value &&
23  typeof value.resetsAt === 'string' &&
24  'at' in value &&
25  typeof value.at === 'number' &&
26  'resumes' in value &&
27  typeof value.resumes === 'number' &&
28  'isArmed' in value &&
29  typeof value.isArmed === 'boolean' &&
30  'stoppedAt' in value &&
31  typeof value.stoppedAt === 'number'
32
33// The latest reset among the windows that are full: the session waits for all of them.
34const blockingReset = (limits: readonly SessionRateLimit[]): string | undefined =>
35  limits
36    .flatMap(limit => (WINDOWS.has(limit.kind) && limit.percentUsed >= 100 && limit.resetsAt ? [limit.resetsAt] : []))
37    .sort((a, b) => Date.parse(a) - Date.parse(b))
38    .pop()
39
40// "at 16:31" today, "Sun Oct 25 20:01" on another day.
41const resumeLabel = (at: number, now: number): string => {
42  const when = new Date(at)
43  const time = when.toTimeString().slice(0, 5)
44  return when.toDateString() === new Date(now).toDateString() ? `at ${time}` : `${when.toDateString().slice(0, 10)} ${time}`
45}
46
47async function planKey($: EngineInterface): Promise<string> {
48  return `plan:${await $.session.id()}`
49}
50
51async function save($: EngineInterface, plan: Plan): Promise<void> {
52  await $.store.set(await planKey($), plan)
53}
54
55function disarm($: EngineInterface): void {
56  timer?.cancel()
57  timer = undefined
58  armedFor = undefined
59  $.ui.status(undefined)
60}
61
62// Requests are going through again (credits, a raised cap, or someone continued): drop the plan.
63async function cancel($: EngineInterface): Promise<void> {
64  disarm($)
65  await $.store.delete(await planKey($))
66}
67
68async function resume($: EngineInterface, plan: Plan): Promise<void> {
69  disarm($)
70  await save($, { ...plan, resumes: plan.resumes + 1, isArmed: false })
71  await $.prompt.submit({ text: PROMPT })
72}
73
74async function arm($: EngineInterface, plan: Plan): Promise<void> {
75  disarm($)
76  armedFor = plan.resetsAt
77  await save($, { ...plan, isArmed: true })
78  const now = await $.clock.now()
79  $.ui.status(`usage limit · resumes ${resumeLabel(plan.at, now)}`)
80  timer = $.clock.after(Math.max(0, plan.at - now), () => {
81    void resume($, plan)
82  })
83}
84
85// Called only after a rate-limit error stopped the main thread; a full window alone can mean usage credits are paying.
86async function planResume($: EngineInterface): Promise<void> {
87  const resetsAt = blockingReset((await $.session.usage()).rateLimits)
88  // No full window: server throttling, which clears on its own.
89  if (resetsAt === undefined || resetsAt === armedFor) {
90    return
91  }
92  const stored = await $.store.get(await planKey($))
93  const resumes = isPlan(stored) && stored.resetsAt === resetsAt ? stored.resumes : 0
94  if (resumes >= MAX_RESUMES) {
95    $.ui.status('usage limit · auto-resume gave up, continue by hand')
96    return
97  }
98  const now = await $.clock.now()
99  const at = Math.max(Date.parse(resetsAt), now) + GRACE_MS
100  await arm($, { resetsAt, at, resumes, isArmed: true, stoppedAt: now })
101}
102
103async function rearm($: EngineInterface): Promise<void> {
104  // The status line outlives a reload: clear what an older load pinned, and pin it again only for a valid plan.
105  disarm($)
106  const stored = await $.store.get(await planKey($))
107  if (stored === undefined) {
108    return
109  }
110  if (!isPlan(stored) || (await $.clock.now()) - stored.at > STALE_MS) {
111    await $.store.delete(await planKey($))
112  } else if (stored.isArmed) {
113    await arm($, stored)
114  }
115}
116
117export const register: Register = on => {
118  on('session.start', async ($, e, next) => {
119    const ran = await next(e)
120    await rearm($)
121    return ran
122  })
123
124  on('classic.StopFailure', async ($, e, next) => {
125    const ran = await next(e)
126    if (e.agent_id === undefined && e.error === 'rate_limit') {
127      await planResume($)
128    }
129    return ran
130  })
131
132  // A tool call the model made means a request just went through, often long before the turn ends.
133  on('tool.call', async ($, e, next) => {
134    if (timer && next.origin.plugin === 'engine') {
135      await cancel($)
136    }
137    return next(e)
138  })
139
140  on('turn.complete', async ($, e, next) => {
141    const ran = await next(e)
142    // A main-thread turn that went through, for a turn with no tool calls.
143    if (e.agentId === undefined && e.reason === 'answer' && timer) {
144      await cancel($)
145    }
146    return ran
147  })
148}
149