session-usage

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Gauge, Segment } from '../types'
5
6const BAR_HEIGHT = 5
7
8// The swatches and order the desktop's context menu draws; every other used category is its grey.
9const SWATCHES: [string, string][] = [
10  ['Messages', '#256dd0'],
11  ['System tools', '#e85d2e'],
12  ['MCP tools', '#1aa66f'],
13  ['Memory files', '#ea9700'],
14]
15const OTHER = '#7e7c76'
16
17const gauge = atom({ plugin: 'session-usage', key: 'gauge' } as const, null as Gauge)
18
19export const colorFor = (percent: number) =>
20  percent >= 90 ? 'error' : percent >= 80 ? 'warning' : 'success'
21
22// Orders categories as the menu lists them and folds the greys into one segment.
23export const segmentsOf = (categories: { name: string; tokens: number }[]): Segment[] => {
24  const named = SWATCHES.flatMap(([name, color]) =>
25    categories.filter(c => c.name === name && c.tokens > 0).map(c => ({ color, tokens: c.tokens })),
26  )
27  const other = categories
28    .filter(c => !SWATCHES.some(([name]) => name === c.name))
29    .reduce((sum, c) => sum + c.tokens, 0)
30
31  return other > 0 ? [...named, { color: OTHER, tokens: other }] : named
32}
33
34// The menu's bar: a rounded track with each segment at its share of the window and a 1px gap after it,
35// the last segment's end rounded where it meets the track.
36export const barSvg = (segments: Segment[], window: number, width: number) => {
37  const r = BAR_HEIGHT / 2
38  let x = 0
39  const rects = segments.map((s, i) => {
40    const w = Math.max(1, (s.tokens / window) * width - 1)
41    const rect = (width: number, rx = '') =>
42      `<rect x="${x.toFixed(1)}" width="${width.toFixed(1)}" height="${BAR_HEIGHT}"${rx} fill="${s.color}"/>`
43    const drawn =
44      i === segments.length - 1 && w > r ? rect(w - r) + rect(w, ` rx="${r}"`) : rect(w)
45    x += w + 1
46
47    return drawn
48  })
49
50  return (
51    `<svg xmlns="http://www.w3.org/2000/svg" width="${width}" height="${BAR_HEIGHT}" viewBox="0 0 ${width} ${BAR_HEIGHT}" preserveAspectRatio="none">` +
52    `<style>.track{fill:#e9e9e7}@media (prefers-color-scheme: dark){.track{fill:#3b3b39}}</style>` +
53    `<clipPath id="c"><rect width="${width}" height="${BAR_HEIGHT}" rx="${r}"/></clipPath>` +
54    `<g clip-path="url(#c)"><rect class="track" width="${width}" height="${BAR_HEIGHT}"/>${rects.join('')}</g></svg>`
55  )
56}
57
58const show = async ($: EngineInterface) => {
59  const { context } = await $.session.usage({ breakdown: 'summary' })
60  const used = (context.breakdown?.categories ?? []).filter(c => c.kind === 'used')
61
62  await update($, gauge, () =>
63    context.percent === undefined
64      ? null
65      : { percent: context.percent, window: context.window, segments: segmentsOf(used) },
66  )
67}
68
69export const register: Register = on => {
70  on('session.start', async ($, e, next) => {
71    const result = await next(e)
72    // Clears the status line an earlier version pinned.
73    $.ui.status(undefined)
74    await show($)
75
76    return result
77  })
78
79  on('session.measure', async ($, e, next) => {
80    if (e.changed.includes('context')) {
81      await show($)
82    }
83
84    return next(e)
85  })
86
87  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
88    const g = await read($, gauge)
89
90    if (e.props.hasSurvey || g === null) {
91      return next(e)
92    }
93
94    const label = `${g.percent}%`
95    const elements = $.ui.resolve(e)
96    const { Box, Text } = elements
97
98    if (!('Svg' in elements)) {
99      return <Text color={colorFor(g.percent)}>{label}</Text>
100    }
101
102    const { Svg } = elements
103    // Pixels are estimated from the band's columns; the image is capped at its slot and stretches to fill it.
104    const width = Math.max(100, (e.props.bodyColumns - label.length - 3) * 9)
105
106    return (
107      <Box flexDirection="row" alignItems="center" gap={1} paddingX={1}>
108        <Text color={colorFor(g.percent)}>{label}</Text>
109        <Box flexGrow={1}>
110          <Svg source={barSvg(g.segments, g.window, width)} alt={`Context ${label} used`} height={BAR_HEIGHT} />
111        </Box>
112      </Box>
113    )
114  })
115}
116

1export type Segment = { color: string; tokens: number }
2
3export type Gauge = { percent: number; window: number; segments: Segment[] } | null
4
5declare module 'claude-code' {
6  interface PluginState {
7    'session-usage': { gauge: Gauge }
8  }
9}
10