usage-weather

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, RenderSurface, SessionContextUsage, SessionRateLimit } from 'claude-code'
3
4import type { Step } from '../types'
5import {
6  HANDOFF_PROMPT,
7  cacheOf,
8  continuation,
9  forecastFor,
10  limitLabel,
11  orderLimits,
12  savedTurn,
13  short,
14  span,
15  ttlOf,
16  until,
17  usageColor,
18} from './lib'
19
20const STORE_HANDOFF = 'handoff'
21// The last turn, kept so a reload of the mod does not lose whether the cache is warm.
22const STORE_LAST = 'last'
23// How often the countdowns move.
24const HEARTBEAT_MS = 15_000
25
26const context = atom({ plugin: 'usage-weather', key: 'context' } as const, null)
27const limits = atom({ plugin: 'usage-weather', key: 'limits' } as const, [])
28const last = atom({ plugin: 'usage-weather', key: 'last' } as const, null)
29const beat = atom({ plugin: 'usage-weather', key: 'beat' } as const, 0)
30const owner = atom({ plugin: 'usage-weather', key: 'owner' } as const, '')
31const step = atom({ plugin: 'usage-weather', key: 'step' } as const, 'idle')
32const note = atom({ plugin: 'usage-weather', key: 'note' } as const, '')
33
34type Figures = { context: SessionContextUsage; rateLimits: readonly SessionRateLimit[] }
35
36// Takes in what the engine measured: the window's fill and the account's limits.
37async function apply($: EngineInterface, figures: Figures) {
38  const { tokens, window } = figures.context
39  if (tokens !== undefined && tokens > 0 && window > 0) {
40    const percent = Math.round(figures.context.percent ?? (tokens / window) * 100)
41    await update($, context, () => ({ tokens, window, percent }))
42  }
43  if (figures.rateLimits.length > 0) {
44    const next = figures.rateLimits.map(l => ({ kind: l.kind, percentUsed: l.percentUsed, resetsAt: l.resetsAt }))
45    await update($, limits, () => next)
46  }
47}
48
49async function fail($: EngineInterface, why: string) {
50  await update($, note, () => why)
51  await update($, step, () => 'idle')
52  void $.ui.toast(why)
53}
54
55// Summarizes the session, clears it and hands the summary to the fresh one.
56// Nothing is cleared until the summary is in hand, and if the clear or the
57// hand-over fails the summary goes to the clipboard so it is not lost.
58async function clearAndContinue($: EngineInterface, surface: RenderSurface) {
59  await update($, note, () => '')
60  await update($, step, () => 'summarizing')
61  let handoff = ''
62  try {
63    const r = await $.model.fork({ prompt: HANDOFF_PROMPT })
64    if (!r.isAnswered) {
65      await fail($, r.reason === 'nothing-to-fork' ? 'Nothing to carry over yet.' : `Could not summarize the session (${r.reason}).`)
66
67      return
68    }
69    handoff = r.text.trim()
70    if (handoff === '') {
71      await fail($, 'The summary came back empty, so nothing was cleared.')
72
73      return
74    }
75    await $.store.set(STORE_HANDOFF, { text: handoff, at: await $.clock.now() })
76    await update($, step, () => 'clearing')
77    await $.command.run({ command: 'clear' })
78    await update($, step, () => 'continuing')
79    await $.prompt.submit({ text: continuation(handoff), asUser: true })
80    await update($, step, () => 'idle')
81  } catch (error) {
82    const reason = error instanceof Error ? error.message : String(error)
83    const copied = handoff !== '' && (await $.ui.copy({ text: handoff, surface }).catch(() => undefined))?.isCopied === true
84    await fail($, copied ? `Clear & continue stopped (${reason}). The summary is on your clipboard.` : `Clear & continue stopped (${reason}).`)
85  }
86}
87
88export const register: Register = (on, options) => {
89  const ttlMs = ttlOf(options.cacheTtl)
90  const instance = Math.random().toString(36).slice(2)
91
92  on('session.start', async ($, e, next) => {
93    // The heartbeat moves the countdowns. A reload starts another and takes
94    // the ownership, so the one left from before stops at its next beat.
95    await update($, owner, () => instance)
96    const timer = $.clock.every(HEARTBEAT_MS, () => {
97      void read($, owner)
98        .then(async who => {
99          if (who !== instance) {
100            timer.cancel()
101
102            return
103          }
104          await update($, beat, n => n + 1)
105        })
106        .catch(() => timer.cancel())
107    })
108    // A reload drops a Clear & continue that was under way.
109    await update($, step, () => 'idle')
110    // The last turn outlives a reload, so the cache stays known.
111    const saved = savedTurn(await $.store.get(STORE_LAST), await $.session.id())
112    if (saved !== null) {
113      await update($, last, was => was ?? saved)
114    }
115    try {
116      await apply($, await $.session.usage())
117    } catch {
118      // No reading yet; the first measurement fills the band.
119    }
120
121    return next(e)
122  })
123
124  on('session.measure', async ($, e, next) => {
125    await apply($, e)
126
127    return next(e)
128  })
129
130  on('turn.complete', async ($, e, next) => {
131    if (e.agentId === undefined) {
132      // The next turn retires the last failure's line.
133      await update($, note, () => '')
134    }
135    if (e.agentId === undefined && e.usage !== undefined) {
136      const at = await $.clock.now()
137      const { cache_read_input_tokens: cached, cache_creation_input_tokens: wrote, input_tokens: fresh } = e.usage
138      await update($, last, () => ({ at, read: cached, wrote, fresh }))
139      await $.store.set(STORE_LAST, { sessionId: await $.session.id(), at, read: cached, wrote, fresh })
140    }
141
142    return next(e)
143  })
144
145  // /clear starts the figures over; Clear & continue's own step is left alone.
146  on('session.end', async ($, e, next) => {
147    if (e.reason === 'clear') {
148      await update($, context, () => null)
149      await update($, last, () => null)
150      await $.store.set(STORE_LAST, null)
151    }
152
153    return next(e)
154  })
155
156  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
157    if (e.props.hasSurvey) {
158      return next(e)
159    }
160    const window = await read($, context)
161    const windows = orderLimits(await read($, limits))
162    const turn = await read($, last)
163    const doing: Step = await read($, step)
164    const why = await read($, note)
165    // Read for the subscription: each beat redraws the countdowns.
166    await read($, beat)
167    if (window === null && windows.length === 0 && turn === null) {
168      return next(e)
169    }
170    const { Box, Text, Button } = $.ui.resolve(e)
171    const now = await $.clock.now()
172    const weather = window === null ? undefined : forecastFor(window.percent)
173    const cache = cacheOf(turn, now, ttlMs)
174    const isIdle = doing === 'idle'
175    const canClear = isIdle && !e.props.isWorking && (window !== null || turn !== null)
176
177    // One quiet line. Every child sits straight in the row: a fragment would be
178    // drawn as a column and break the line. Clearing takes the whole line while
179    // it asks and works.
180    return (
181      <Box flexDirection="column" paddingX={1}>
182        <Box flexDirection="row" flexWrap="wrap" columnGap={2} alignItems="center">
183          {isIdle && window !== null && weather !== undefined && (
184            <Box flexDirection="row" gap={1}>
185              <Text color={weather.color} bold>{`${weather.icon} ${window.percent}%`}</Text>
186              <Text dimColor>{`${short(window.tokens)}/${short(window.window)}`}</Text>
187            </Box>
188          )}
189          {isIdle &&
190            windows.map(w => (
191              <Box flexDirection="row" gap={1}>
192                <Text dimColor>{limitLabel(w.kind)}</Text>
193                <Text color={usageColor(w.percentUsed)}>{`${Math.round(w.percentUsed)}%`}</Text>
194                {until(w.resetsAt, now) !== '' && (
195                  <Text dimColor>{`↻${until(w.resetsAt, now)}`}</Text>
196                )}
197              </Box>
198            ))}
199          {isIdle && (
200            <Box flexDirection="row" gap={1}>
201              <Text dimColor>cache</Text>
202              {cache.state === 'warm' ? (
203                <Text color="success">{`warm ${span(cache.leftMs)}`}</Text>
204              ) : cache.state === 'cold' ? (
205                <Text color="warning">cold</Text>
206              ) : (
207                <Text dimColor>–</Text>
208              )}
209            </Box>
210          )}
211          {canClear && (
212            <Button key="clear-continue" label="Clear & continue" onPress={() => void update($, step, () => 'confirm')} />
213          )}
214          {doing === 'confirm' && <Text bold>Summarize and clear?</Text>}
215          {doing === 'confirm' && (
216            <Text dimColor>{cache.state === 'warm' ? 'cheap, the cache is warm' : 'costs more, the cache is cold'}</Text>
217          )}
218          {doing === 'confirm' && (
219            <Button
220              key="confirm-yes"
221              variant="primary"
222              label="Yes, clear"
223              onPress={press => void clearAndContinue($, press.surface)}
224            />
225          )}
226          {doing === 'confirm' && (
227            <Button key="confirm-no" label="Cancel" onPress={() => void update($, step, () => 'idle')} />
228          )}
229          {doing === 'summarizing' && <Text color="claude">Summarizing the session…</Text>}
230          {doing === 'clearing' && <Text color="claude">Clearing…</Text>}
231          {doing === 'continuing' && <Text color="claude">Continuing from the summary…</Text>}
232        </Box>
233        {why !== '' && isIdle && <Text color="error">{why}</Text>}
234      </Box>
235    )
236  })
237}
238

1import type { Limit, Turn } from '../types'
2
3// Context forecast bands, by percent of the window used. Theme colors, so the
4// band follows light and dark.
5const FORECAST = [
6  { upTo: 25, icon: '☀', color: 'success' },
7  { upTo: 50, icon: '☁', color: 'suggestion' },
8  { upTo: 75, icon: '☂', color: 'warning' },
9  { upTo: 90, icon: '☇', color: 'error' },
10] as const
11const COMPACT = { icon: '↯', color: 'error' } as const
12
13export const forecastFor = (percent: number) => FORECAST.find(band => percent < band.upTo) ?? COMPACT
14
15export const short = (n: number) => {
16  if (n >= 1_000_000) {
17    return `${(n / 1_000_000).toFixed(n % 1_000_000 === 0 ? 0 : 1)}M`
18  }
19  if (n >= 1_000) {
20    return `${(n / 1_000).toFixed(n % 1_000 === 0 ? 0 : 1)}k`
21  }
22
23  return String(n)
24}
25
26// Green while there is room, amber as it fills, red near the limit.
27export const usageColor = (percent: number) => (percent >= 85 ? 'error' : percent >= 60 ? 'warning' : 'success')
28
29// A span as the band writes it: <1m, 45m, 2h13m, 4d 3h. Minutes round up, so
30// a countdown never reads zero while time is left.
31export const span = (ms: number) => {
32  if (ms <= 0) {
33    return ''
34  }
35  if (ms < 60_000) {
36    return '<1m'
37  }
38  const minutes = Math.ceil(ms / 60_000)
39  if (minutes < 60) {
40    return `${minutes}m`
41  }
42  const hours = Math.floor(minutes / 60)
43  if (hours < 24) {
44    return minutes % 60 === 0 ? `${hours}h` : `${hours}h${String(minutes % 60).padStart(2, '0')}m`
45  }
46  const days = Math.floor(hours / 24)
47
48  return hours % 24 === 0 ? `${days}d` : `${days}d ${hours % 24}h`
49}
50
51// Time to a window's reset; empty when it is unknown or already past.
52export const until = (resetsAt: string | undefined, now: number) => {
53  const at = resetsAt === undefined ? Number.NaN : Date.parse(resetsAt)
54
55  return Number.isNaN(at) ? '' : span(at - now)
56}
57
58const LABELS: Record<string, string> = { five_hour: '5h', seven_day: 'week' }
59export const limitLabel = (kind: string) => LABELS[kind] ?? kind.replace(/_/g, ' ')
60
61// The 5-hour window first, then the week, then whatever else the account has.
62export const orderLimits = (limits: readonly Limit[]) => {
63  const rank = (kind: string) => (kind === 'five_hour' ? 0 : kind === 'seven_day' ? 1 : 2)
64
65  return [...limits].sort((a, b) => rank(a.kind) - rank(b.kind))
66}
67
68export const TTL = { '5m': 5 * 60_000, '1h': 60 * 60_000 } as const
69export const ttlOf = (option: unknown) => (option === '5m' ? TTL['5m'] : TTL['1h'])
70
71// A turn as it is kept between loads, tied to the session it belongs to.
72export type SavedTurn = Turn & { sessionId: string }
73
74// What the store holds, read back as a turn of this session, or null when it
75// is another session's or not a turn at all.
76export const savedTurn = (value: unknown, sessionId: string): Turn | null => {
77  const v = value as Partial<SavedTurn> | null
78  if (v === null || typeof v !== 'object' || v.sessionId !== sessionId) {
79    return null
80  }
81  const { at, read, wrote, fresh } = v
82
83  return typeof at === 'number' && typeof read === 'number' && typeof wrote === 'number' && typeof fresh === 'number'
84    ? { at, read, wrote, fresh }
85    : null
86}
87
88export type Cache = { state: 'none' | 'cold' | 'warm'; leftMs: number; hit: number }
89
90// Whether the prompt cache still holds the session: warm from the last
91// main-thread turn until the cache lifetime has passed. `hit` is how much of
92// that turn's prompt the cache served.
93export const cacheOf = (last: Turn | null, now: number, ttlMs: number): Cache => {
94  if (last === null || last.read + last.wrote === 0) {
95    return { state: 'none', leftMs: 0, hit: 0 }
96  }
97  const hit = Math.round((last.read / (last.read + last.wrote + last.fresh)) * 100)
98  const leftMs = last.at + ttlMs - now
99
100  return leftMs > 0 ? { state: 'warm', leftMs, hit } : { state: 'cold', leftMs: 0, hit }
101}
102
103// What the fork writes before the conversation is cleared. It answers over
104// the session's own transcript, so it needs no copy of it here.
105export const HANDOFF_PROMPT = [
106  'The person is about to clear this conversation and carry on in a fresh one. Write the handoff that fresh conversation will start from.',
107  'Be specific and compact (under 600 words), in the language we have been speaking. Do not call tools. Write only these sections:',
108  'Goal: what we are trying to achieve overall.',
109  'Done: what is finished, with the file paths and commands that matter.',
110  'State: where the work stands now: what is in progress, half done or running.',
111  'Decisions: choices made and why, constraints, and preferences the person stated.',
112  'Next: the very next steps, in order.',
113  'Open: questions or approvals still needed from the person.',
114].join('\n')
115
116export const continuation = (handoff: string) =>
117  [
118    '[Clear & continue] The previous conversation was cleared to free context. Its handoff is below; the files on disk are as that conversation left them.',
119    'Continue from "Next". If a step needs my approval or a decision, ask before doing it.',
120    '<handoff>',
121    handoff,
122    '</handoff>',
123  ].join('\n')
124

1// The context window as one turn left it.
2export type Reading = { tokens: number; window: number; percent: number }
3
4// A rate-limit window: `five_hour`, `seven_day`, or a gateway's `spend_limit`.
5export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
6
7// The last main-thread turn: when it ended and how its prompt cache was used.
8export type Turn = { at: number; read: number; wrote: number; fresh: number }
9
10// Where Clear & continue stands.
11export type Step = 'idle' | 'confirm' | 'summarizing' | 'clearing' | 'continuing'
12
13declare module 'claude-code' {
14  interface PluginState {
15    'usage-weather': {
16      // The window's fill as the last turn left it; null before one.
17      context: Reading | null
18      limits: Limit[]
19      last: Turn | null
20      // Ticks every few seconds so countdowns move; only its writes matter.
21      beat: number
22      // The module instance that owns the heartbeat; a reload hands it on.
23      owner: string
24      step: Step
25      // Why the last Clear & continue stopped; empty otherwise.
26      note: string
27    }
28  }
29}
30