token-turns

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelUsage, Register, SessionUsage } from 'claude-code'
3
4import type { Base, Live, Meters, Spend, Turn } from '../types'
5
6const KEEP = 50
7
8const meters = atom({ plugin: 'token-turns', key: 'meters' } as const, null)
9const turns = atom({ plugin: 'token-turns', key: 'turns' } as const, [])
10const live = atom({ plugin: 'token-turns', key: 'live' } as const, null)
11const isHidden = atom({ plugin: 'token-turns', key: 'isHidden' } as const, false)
12const view = atom({ plugin: 'token-turns', key: 'view' } as const, 'compact')
13
14const toMeters = (usage: Pick<SessionUsage, 'context' | 'rateLimits'>): Meters => {
15  const limit = (kind: string) => {
16    const found = usage.rateLimits.find(one => one.kind === kind)
17    return found && { percent: found.percentUsed, resetsAt: found.resetsAt }
18  }
19
20  return {
21    ctxTokens: usage.context.tokens,
22    ctxWindow: usage.context.window,
23    ctxPercent: usage.context.percent,
24    five: limit('five_hour'),
25    week: limit('seven_day'),
26  }
27}
28
29const toBase = (m: Meters | null): Base | null =>
30  m && { ctxTokens: m.ctxTokens, five: m.five?.percent, week: m.week?.percent }
31
32const minus = (after?: number, before?: number) =>
33  after === undefined || before === undefined ? undefined : Math.round((after - before) * 10) / 10
34
35const withDeltas = (turn: Turn, m: Meters | null): Turn =>
36  turn.base === null || m === null
37    ? turn
38    : {
39        ...turn,
40        ctxDelta: minus(m.ctxTokens, turn.base.ctxTokens),
41        fiveDelta: minus(m.five?.percent, turn.base.five),
42        weekDelta: minus(m.week?.percent, turn.base.week),
43      }
44
45const toSpend = (u: ModelUsage): Spend => ({
46  fresh: u.input_tokens + u.cache_creation_input_tokens,
47  cacheRead: u.cache_read_input_tokens,
48  out: u.output_tokens,
49})
50
51const add = <S extends Spend>(a: S, b: Spend): S => ({
52  ...a,
53  fresh: a.fresh + b.fresh,
54  cacheRead: a.cacheRead + b.cacheRead,
55  out: a.out + b.out,
56})
57
58const delta = (s: Spend) => s.fresh + s.out
59
60const measure = async ($: EngineInterface) => {
61  try {
62    const m = toMeters(await $.session.usage())
63    await update($, meters, () => m)
64    return m
65  } catch {
66    return read($, meters)
67  }
68}
69
70// 950, 18.4k, 2.1M
71export const fmt = (n: number) => {
72  const abs = Math.abs(n)
73  if (abs < 1000) return String(Math.round(n))
74  if (abs < 1e6) return `${(n / 1e3).toFixed(abs < 1e5 ? 1 : 0).replace(/\.0$/, '')}k`
75  return `${(n / 1e6).toFixed(1).replace(/\.0$/, '')}M`
76}
77
78const signed = (n: number, f: (n: number) => string = fmt) => (n < 0 ? `−${f(-n)}` : `+${f(n)}`)
79
80const pct = (n: number) => `${n.toFixed(1).replace(/\.0$/, '')}%`
81
82const heat = (percent?: number) =>
83  percent === undefined ? 'subtle' : percent < 60 ? 'success' : percent < 85 ? 'warning' : 'error'
84
85type Seg = { text: string; color?: string; bold?: boolean; dim?: boolean }
86
87// The fill in its heat colour, the last turn's growth in red on top of it.
88const meter = (percent: number | undefined, width: number, grown?: number, window?: number): Seg[] => {
89  const filled = Math.min(width, Math.round(((percent ?? 0) / 100) * width))
90  const fresh =
91    grown && grown > 0 && window ? Math.min(filled, Math.max(1, Math.round((grown / window) * width))) : 0
92  return [
93    { text: '▰'.repeat(filled - fresh), color: heat(percent) },
94    { text: '▰'.repeat(fresh), color: 'error' },
95    { text: '▱'.repeat(width - filled), dim: true },
96  ].filter(s => s.text.length > 0)
97}
98
99const point = (n: number) => n.toFixed(1)
100
101// The grey ` +0.4` after a limit's percentage.
102const nudge = (value: number | undefined): Seg[] =>
103  value === undefined ? [] : [{ text: ` ${signed(value, point)}`, dim: true }]
104
105export const register: Register = on => {
106  on('session.start', async ($, e, next) => {
107    await $.command.register({
108      name: 'tokens',
109      description: 'Show or hide the token display; `/tokens compact` or `/tokens expanded` picks the view',
110      argumentHint: '[compact|expanded]',
111    })
112    void measure($)
113
114    return next(e)
115  })
116
117  on('command.run', { command: 'tokens' }, async ($, e) => {
118    const arg = e.args.trim()
119    if (arg === 'compact' || arg === 'expanded') {
120      await update($, view, () => arg)
121      await update($, isHidden, () => false)
122      return { text: `Tokens: ${arg} view.` }
123    }
124    if (arg !== '') return { text: 'Usage: /tokens [compact|expanded]' }
125
126    const hidden = await update($, isHidden, h => !h)
127    return { text: `Tokens ${hidden ? 'hidden' : 'shown'}.` }
128  })
129
130  on('turn.start', async ($, e, next) => {
131    if (e.agentId !== undefined) return next(e)
132
133    const done = await read($, turns)
134    const n = (done.at(-1)?.n ?? 0) + 1
135    const base = toBase(await read($, meters))
136    await update($, live, () => ({ turnId: e.turnId, n, base, fresh: 0, cacheRead: 0, out: 0 }))
137
138    return next(e)
139  })
140
141  // Every model request, the main loop's and each subagent's, adds to the turn
142  // in flight; one that lands after the turn ended goes to the last turn.
143  on('turn.step', async function* ($, e, next) {
144    const response = yield* next(e)
145    if (response.usage) {
146      const spend = toSpend(response.usage)
147      const current = await read($, live)
148      if (current) {
149        await update($, live, now => now && add(now, spend))
150      } else {
151        await update($, turns, list =>
152          list.length === 0 ? list : [...list.slice(0, -1), add(list.at(-1) as Turn, spend)],
153        )
154      }
155    }
156
157    return response
158  })
159
160  on('turn.complete', async ($, e, next) => {
161    const current = await read($, live)
162    // Any main-loop end closes the open turn, so none is left looking live.
163    if (e.agentId !== undefined || !current) return next(e)
164
165    // No request reached the steps hook (it loaded mid-turn): use the turn's own sum.
166    const spend: Spend = delta(current) === 0 && e.usage ? add(current, toSpend(e.usage)) : current
167    const after = await measure($)
168    const turn = withDeltas(
169      {
170        n: current.n,
171        base: current.base,
172        fresh: spend.fresh,
173        cacheRead: spend.cacheRead,
174        out: spend.out,
175        durationMs: e.durationMs,
176        isAborted: e.isAborted,
177      },
178      after,
179    )
180    await update($, turns, list => [...list, turn].slice(-KEEP))
181    await update($, live, () => null)
182
183    return next(e)
184  })
185
186  // The limits report a moment after the turn: settle the last turn's deltas
187  // against them until the next turn starts.
188  on('session.measure', async ($, e, next) => {
189    const m = toMeters(e)
190    await update($, meters, () => m)
191    if ((await read($, live)) === null) {
192      await update($, turns, list =>
193        list.length === 0 ? list : [...list.slice(0, -1), withDeltas(list.at(-1) as Turn, m)],
194      )
195    }
196
197    return next(e)
198  })
199  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
200    if (e.props.hasSurvey) return next(e)
201
202    if (await read($, isHidden)) return next(e)
203    const isCard = (await read($, view)) === 'expanded'
204
205    const m = await read($, meters)
206    const running = e.props.isWorking ? await read($, live) : null
207    const done = await read($, turns)
208    const last = done.at(-1)
209    if (!m && !running && !last) return next(e)
210
211    const { Box, Text } = $.ui.resolve(e)
212    const room = e.props.bodyColumns
213    const shown = running ?? last
214    // The last turn's changes ride inline on the meters; none while one runs.
215    const settled = running ? undefined : last
216
217    const segs = (list: Seg[]) =>
218      list.map(s => (
219        <Text color={s.color} bold={s.bold} dimColor={s.dim} wrap="truncate-end">
220          {s.text}
221        </Text>
222      ))
223
224    const theDelta: Seg[] = shown
225      ? [
226          { text: 'Δ ', dim: true },
227          { text: signed(delta(shown)), bold: true, color: running ? 'claude' : 'text' },
228        ]
229      : []
230    const cache: Seg[] = shown ? [{ text: `  ${fmt(shown.cacheRead)} cached`, dim: true }] : []
231
232    if (isCard) {
233      // Two borderless rows, right-aligned above the prompt: the delta with the
234      // recent turns' sparkline, then the meters with their inline changes.
235      const limit = (label: string, l: { percent: number } | undefined, d?: number): Seg[] =>
236        l ? [{ text: label, dim: true }, { text: pct(l.percent), color: heat(l.percent) }, ...nudge(d)] : []
237      const recent = done.slice(-12)
238      const peak = Math.max(1, ...recent.map(delta))
239      const spark = recent.map(t => '▁▂▃▄▅▆▇█'[Math.min(7, Math.floor((delta(t) / peak) * 7.999))]).join('')
240
241      const top: Seg[] = [
242        ...(spark.length > 1 ? [{ text: spark, color: 'claude' }, { text: '  ' }] : []),
243        ...(theDelta.length ? [...theDelta, ...cache] : [{ text: 'waiting…', dim: true }]),
244      ]
245      const bottom: Seg[] = m
246        ? [
247            { text: 'ctx ', dim: true },
248            ...meter(m.ctxPercent, 10, settled?.ctxDelta, m.ctxWindow),
249            { text: ` ${m.ctxPercent ?? 0}%` },
250            ...limit('   5h ', m.five, settled?.fiveDelta),
251            ...limit('   wk ', m.week, settled?.weekDelta),
252          ]
253        : []
254
255      return (
256        <Box flexDirection="column" alignItems="flex-end" width={room}>
257          <Box flexDirection="row">{segs(top)}</Box>
258          {bottom.length > 0 && <Box flexDirection="row">{segs(bottom)}</Box>}
259        </Box>
260      )
261    }
262
263    const gauges = (withDeltas: boolean): Seg[] => {
264      if (!m) return []
265      const d = withDeltas ? settled : undefined
266      return [
267        { text: 'ctx ', dim: true },
268        ...meter(m.ctxPercent, 8, d?.ctxDelta, m.ctxWindow),
269        { text: ` ${m.ctxPercent ?? 0}%` },
270        ...(m.five
271          ? [{ text: ' · 5h ', dim: true }, { text: pct(m.five.percent), color: heat(m.five.percent) }, ...nudge(d?.fiveDelta)]
272          : []),
273        ...(m.week
274          ? [{ text: ' · wk ', dim: true }, { text: pct(m.week.percent), color: heat(m.week.percent) }, ...nudge(d?.weekDelta)]
275          : []),
276        { text: '  │  ', dim: true },
277      ]
278    }
279
280    // Drop from the right until it fits: cache reads, the inline deltas, then the meters.
281    const width = (list: Seg[]) => list.reduce((sum, s) => sum + s.text.length, 0)
282    const line =
283      [
284        [...gauges(true), ...theDelta, ...cache],
285        [...gauges(true), ...theDelta],
286        [...gauges(false), ...theDelta],
287      ].find(l => width(l) <= room) ?? theDelta
288
289    return (
290      <Box flexDirection="row" justifyContent="flex-end" width={room}>
291        {segs(line)}
292      </Box>
293    )
294  })
295}
296

1export type Limit = { percent: number; resetsAt?: string }
2
3export type Meters = {
4  ctxTokens?: number
5  ctxWindow: number
6  ctxPercent?: number
7  five?: Limit
8  week?: Limit
9}
10
11/** The meters as a turn began, to measure the turn's deltas against. */
12export type Base = { ctxTokens?: number; five?: number; week?: number }
13
14/** Token counts summed over a turn's model requests, subagents' included. */
15export type Spend = { fresh: number; cacheRead: number; out: number }
16
17export type Live = Spend & { turnId: string; n: number; base: Base | null }
18
19export type Turn = Spend & {
20  n: number
21  base: Base | null
22  ctxDelta?: number
23  fiveDelta?: number
24  weekDelta?: number
25  durationMs: number
26  isAborted: boolean
27}
28
29declare module 'claude-code' {
30  interface PluginState {
31    'token-turns': {
32      meters: Meters | null
33      turns: Turn[]
34      live: Live | null
35      isHidden: boolean
36      view: 'compact' | 'expanded'
37    }
38  }
39}
40