token-limit

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionContextUsage, SessionRateLimit } from 'claude-code'
3
4import type { Engine, Gauge } from '../types'
5
6const LABELS: Record<string, string> = {
7  five_hour: '5h',
8  seven_day: '7d',
9  spend_limit: 'spend',
10}
11
12const EFFORTS: Record<string, string> = {
13  medium: 'med',
14}
15
16const SEPARATOR = ' │ '
17const BAR_CELLS = 4
18
19const gauges = atom({ plugin: 'token-limit', key: 'gauges' } as const, [])
20const engine = atom({ plugin: 'token-limit', key: 'engine' } as const, {})
21
22// claude-haiku-4-5-20251001 → haiku-4.5; an alias (opus) is left as it is.
23export const shortModel = (id: string): string =>
24  id
25    .replace(/\[.*\]$/, '')
26    .replace(/^claude-/, '')
27    .replace(/-\d{8}$/, '')
28    .replace(/-(\d+)-(\d+)$/, '-$1.$2')
29
30export const labelOf = ({ model, effort }: Engine): string | undefined =>
31  model === undefined ? undefined : effort === undefined ? model : `${model} ${effort}`
32
33const compact = (n: number): string =>
34  n >= 1_000_000
35    ? `${(n / 1_000_000).toFixed(1).replace(/\.0$/, '')}M`
36    : n >= 1000
37      ? `${Math.round(n / 1000)}k`
38      : `${n}`
39
40const clock = (iso: string): string | undefined => {
41  const at = new Date(iso)
42
43  if (Number.isNaN(at.getTime())) return undefined
44
45  const hours = `${at.getHours()}`.padStart(2, '0')
46  const minutes = `${at.getMinutes()}`.padStart(2, '0')
47
48  return `${hours}:${minutes}`
49}
50
51const contextGauge = (context: SessionContextUsage): Gauge => {
52  const window = compact(context.window)
53
54  if (context.tokens === undefined) return { label: 'ctx', detail: `-/${window}` }
55
56  return {
57    label: 'ctx',
58    percent: context.percent ?? Math.round((context.tokens / context.window) * 100),
59    detail: `${compact(context.tokens)}/${window}`,
60  }
61}
62
63const limitGauge = (limit: SessionRateLimit): Gauge => {
64  const resets = limit.resetsAt === undefined ? undefined : clock(limit.resetsAt)
65  // The weekly window resets days away: a bare HH:MM would mislead.
66  const hasNearReset = resets !== undefined && limit.kind === 'five_hour'
67
68  return {
69    label: LABELS[limit.kind] ?? limit.kind,
70    percent: limit.percentUsed,
71    detail: hasNearReset ? `~${resets}` : undefined,
72  }
73}
74
75export const toGauges = (
76  context: SessionContextUsage,
77  rateLimits: readonly SessionRateLimit[],
78): Gauge[] => [contextGauge(context), ...rateLimits.map(limitGauge)]
79
80export const colorOf = (percent: number): string =>
81  percent >= 80 ? 'error' : percent >= 50 ? 'warning' : 'success'
82
83const bar = (percent: number): { filled: string; empty: string } => {
84  const cells = Math.min(BAR_CELLS, Math.max(0, Math.round((percent / 100) * BAR_CELLS)))
85
86  return { filled: '█'.repeat(cells), empty: '░'.repeat(BAR_CELLS - cells) }
87}
88
89const liveGauges = async ($: EngineInterface): Promise<Gauge[]> => {
90  const { context, rateLimits } = await $.session.usage()
91
92  return toGauges(context, rateLimits)
93}
94
95export const register: Register = on => {
96  on('session.start', async ($, e, next) => {
97    const started = await next(e)
98    const live = await liveGauges($)
99    const model = shortModel(await $.session.model())
100    await update($, gauges, () => live)
101    await update($, engine, known => ({ ...known, model }))
102
103    return started
104  })
105
106  on('session.measure', async ($, e, next) => {
107    await update($, gauges, () => toGauges(e.context, e.rateLimits))
108
109    return next(e)
110  })
111
112  on('turn.step', async function* ($, e, next) {
113    // A subagent's request names its own model and effort, not the session's.
114    if (e.agentId === undefined) {
115      const effort = e.effort === undefined ? undefined : (EFFORTS[e.effort] ?? `${e.effort}`)
116      await update($, engine, () => ({ model: shortModel(e.model), effort }))
117    }
118
119    return yield* next(e)
120  })
121
122  on('ui.render', { component: 'SessionMode' }, async ($, e) => {
123    const { Box, Text } = $.ui.resolve(e)
124    const measured = await read($, gauges)
125    // A /clear empties both values and raises no session.start: read them live until the next measurement.
126    const isCleared = measured.length === 0
127    const shown = isCleared ? await liveGauges($) : measured
128    const known = isCleared ? { model: shortModel(await $.session.model()) } : await read($, engine)
129    // The context window is the model's, so its gauge carries the model's name.
130    const contextLabel = labelOf(known)
131
132    return (
133      <Box>
134        {shown.map((gauge, index) => {
135          const { filled, empty } = bar(gauge.percent ?? 0)
136          const label = index === 0 ? (contextLabel ?? gauge.label) : gauge.label
137
138          return (
139            <Text key={gauge.label} wrap="truncate">
140              {index > 0 ? <Text dimColor>{SEPARATOR}</Text> : null}
141              <Text bold>{label} </Text>
142              {gauge.percent === undefined ? null : (
143                <Text color={colorOf(gauge.percent)}>
144                  {filled}
145                  <Text dimColor>{empty}</Text> {gauge.percent}%
146                </Text>
147              )}
148              {gauge.detail === undefined ? null : (
149                <Text dimColor>{gauge.percent === undefined ? gauge.detail : ` ${gauge.detail}`}</Text>
150              )}
151            </Text>
152          )
153        })}
154        {e.props.modes.length === 0 ? null : (
155          <Text dimColor>{`${SEPARATOR}${e.props.modes.join(' & ')}`}</Text>
156        )}
157      </Box>
158    )
159  })
160}
161

1export type Gauge = {
2  label: string
3  percent?: number
4  detail?: string
5}
6
7export type Engine = {
8  model?: string
9  effort?: string
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    'token-limit': { gauges: Gauge[]; engine: Engine }
15  }
16}
17