usage-band

1import { atom, read, update } from 'claude-code'
2import type { Hook, Register, SessionContextUsage, SessionRateLimit } from 'claude-code'
3
4import type { UsageReading } from '../types'
5
6type Engine = Parameters<Hook<'session.measure'>>[0]
7
8const reading = atom({ plugin: 'usage-band', key: 'reading' } as const, null)
9
10// A figure turns to the warning colour at this much.
11const PLAN_WARN_PERCENT = 80
12const CONTEXT_WARN_PERCENT = 90
13
14// Re-reads the plan usage file each minute, so the plan figures stay current between turns.
15const REFRESH_MS = 60000
16
17const SESSION_WINDOW = 'five_hour'
18const WEEKLY_WINDOW = 'seven_day'
19
20// Where the desktop app records plan usage samples: { samples: [{ t, u: { fh, sd } }] }, fh the 5-hour and sd the 7-day window.
21const PLAN_USAGE_FILE = 'Claude/plan-usage-history.json'
22const MAC_APP_SUPPORT = 'Library/Application Support'
23
24type PlanUsage = { sessionPercent?: number; weeklyPercent?: number }
25
26// The token count at which auto-compaction runs, for the window it was measured against; only a new window size changes it.
27let compaction: { window: number; thresholdTokens?: number } | null = null
28
29// Whether the last measurement carried live plan figures; while it did, the plan usage file is not needed.
30let hasLivePlanUsage = false
31
32async function compactionThreshold($: Engine, context: SessionContextUsage) {
33  if (compaction?.window !== context.window) {
34    try {
35      // A summary breakdown is estimated locally, so this sends no request.
36      const { breakdown } = (await $.session.usage({ breakdown: 'summary' })).context
37      compaction = {
38        window: context.window,
39        thresholdTokens: breakdown?.isAutoCompactEnabled ? breakdown.autoCompactThreshold : undefined,
40      }
41    } catch {
42      // Without the threshold the band measures against the whole window; the next measurement tries again.
43      return undefined
44    }
45  }
46
47  return compaction.thresholdTokens
48}
49
50// The context's progress towards auto-compaction; against the whole window when there is no threshold.
51async function contextPercent($: Engine, context: SessionContextUsage) {
52  const threshold = await compactionThreshold($, context)
53  if (context.tokens === undefined || threshold === undefined) {
54    return context.percent
55  }
56
57  return Math.round((context.tokens / threshold) * 100)
58}
59
60// The plan usage file's newest sample; empty when there is no file or it cannot be read.
61async function planUsageFromFile($: Engine): Promise<PlanUsage> {
62  try {
63    const appData = await $.env.get('APPDATA')
64    const home = await $.env.get('HOME')
65    const folder = appData ?? (home === undefined ? undefined : `${home}/${MAC_APP_SUPPORT}`)
66    if (folder === undefined) {
67      return {}
68    }
69
70    const history = JSON.parse(await $.fs.read(`${folder}/${PLAN_USAGE_FILE}`)) as { samples?: { u?: { fh?: unknown; sd?: unknown } }[] }
71    const newest = history.samples?.at(-1)?.u
72
73    return {
74      sessionPercent: typeof newest?.fh === 'number' ? newest.fh : undefined,
75      weeklyPercent: typeof newest?.sd === 'number' ? newest.sd : undefined,
76    }
77  } catch {
78    return {}
79  }
80}
81
82// Live figures when the engine passes them on, which the desktop app does not always do; the plan usage file otherwise.
83async function planUsage($: Engine, rateLimits: SessionRateLimit[]): Promise<PlanUsage> {
84  const session = rateLimits.find(limit => limit.kind === SESSION_WINDOW)?.percentUsed
85  const weekly = rateLimits.find(limit => limit.kind === WEEKLY_WINDOW)?.percentUsed
86  hasLivePlanUsage = session !== undefined && weekly !== undefined
87  if (hasLivePlanUsage) {
88    return { sessionPercent: session, weeklyPercent: weekly }
89  }
90
91  const fromFile = await planUsageFromFile($)
92
93  return { sessionPercent: session ?? fromFile.sessionPercent, weeklyPercent: weekly ?? fromFile.weeklyPercent }
94}
95
96async function record($: Engine, context: SessionContextUsage, rateLimits: SessionRateLimit[]) {
97  const current: UsageReading = { ...(await planUsage($, rateLimits)), contextPercent: await contextPercent($, context) }
98  await update($, reading, () => current)
99}
100
101// Between turns only the plan usage file can move.
102async function refreshPlanUsage($: Engine) {
103  if (hasLivePlanUsage) {
104    return
105  }
106
107  const fromFile = await planUsageFromFile($)
108  await update($, reading, previous => ({ ...previous, ...fromFile }))
109}
110
111export const register: Register = on => {
112  on('session.start', async ($, e, next) => {
113    const result = await next(e)
114    $.clock.every(REFRESH_MS, () => void refreshPlanUsage($))
115
116    const usage = await $.session.usage()
117    await record($, usage.context, usage.rateLimits)
118
119    return result
120  })
121
122  on('session.measure', async ($, e, next) => {
123    await record($, e.context, e.rateLimits)
124
125    return next(e)
126  })
127
128  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
129    const current = await read($, reading)
130    if (e.props.hasSurvey || current === null) {
131      return next(e)
132    }
133
134    const { Box, Text } = $.ui.resolve(e)
135
136    return (
137      <Box>
138        <Text dimColor>
139          <Text color={warnAt(current.sessionPercent, PLAN_WARN_PERCENT)}>Session tokens {percent(current.sessionPercent)}</Text>
140          {', '}
141          <Text color={warnAt(current.weeklyPercent, PLAN_WARN_PERCENT)}>weekly tokens {percent(current.weeklyPercent)}</Text>
142          {', '}
143          <Text color={warnAt(current.contextPercent, CONTEXT_WARN_PERCENT)}>context {percent(current.contextPercent)}</Text>
144        </Text>
145      </Box>
146    )
147  })
148}
149
150function percent(value: number | undefined) {
151  return value === undefined ? '–' : `${value}%`
152}
153
154function warnAt(value: number | undefined, threshold: number) {
155  return value !== undefined && value >= threshold ? 'warning' : undefined
156}
157

1// What the band shows, each 0 to 100 and absent while unknown.
2export type UsageReading = {
3  // How much of the plan's 5-hour session allowance is used.
4  sessionPercent?: number
5  // How much of the plan's weekly allowance is used.
6  weeklyPercent?: number
7  // How far the context is towards auto-compaction, which runs at 100 (of the whole window when auto-compaction is off).
8  contextPercent?: number
9}
10
11declare module 'claude-code' {
12  interface PluginState {
13    'usage-band': { reading: UsageReading | null }
14  }
15}
16