session-band

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Reading } from '../types'
5
6const reading = atom({ plugin: 'session-band', key: 'reading' } as const, null)
7
8const WARN_TOKENS = 100_000
9
10const short = (tokens: number): string =>
11  tokens >= 1_000_000 ? `${(tokens / 1_000_000).toFixed(1)}M` : `${Math.round(tokens / 1000)}k`
12
13const LIMIT_LABELS: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'spend' }
14
15const measure = async ($: EngineInterface, seen: { model?: string; effort?: string | null }) => {
16  const usage = await $.session.usage({ breakdown: 'summary' })
17  const { context, rateLimits } = usage
18  const previous = await read($, reading)
19
20  const next: Reading = {
21    model: seen.model ?? context.breakdown?.model ?? previous?.model ?? '',
22    effort: seen.effort === undefined ? (previous?.effort ?? null) : seen.effort,
23    tokens: context.tokens ?? null,
24    window: context.window,
25    compactAt: context.breakdown?.autoCompactThreshold ?? null,
26    limits: rateLimits.map(({ kind, percentUsed }) => ({ kind, percentUsed })),
27  }
28
29  await update($, reading, () => next)
30}
31
32export const register: Register = on => {
33  on('session.start', async ($, e, next) => {
34    const started = await next(e)
35    await measure($, {})
36
37    return started
38  })
39
40  on('turn.step', async function* ($, e, next) {
41    if (e.agentId === undefined) {
42      await measure($, { model: e.model, effort: e.effort === undefined ? null : String(e.effort) })
43    }
44
45    return yield* next(e)
46  })
47
48  on('turn.complete', async ($, e, next) => {
49    const done = await next(e)
50
51    if (e.agentId === undefined) {
52      await measure($, {})
53    }
54
55    return done
56  })
57
58  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
59    const now = await read($, reading)
60
61    if (e.props.hasSurvey || now === null) {
62      return next(e)
63    }
64
65    const { Box, Text } = $.ui.resolve(e)
66    const used = now.tokens ?? 0
67    const percent = Math.round((used / now.window) * 100)
68    const left = now.compactAt === null ? null : now.compactAt - used
69    const isClose = left !== null && left < WARN_TOKENS
70    const hot = now.limits.filter(limit => limit.percentUsed >= 50)
71
72    return (
73      <Box>
74        <Text dimColor>
75          {now.model}
76          {now.effort === null ? '' : ` · ${now.effort}`} · {short(used)}/{short(now.window)} {percent}%
77        </Text>
78        {left === null ? null : (
79          <Text dimColor={!isClose} color={isClose ? 'yellow' : undefined}>
80            {' · '}
81            {left > 0 ? `compacts in ${short(left)}` : 'compacting next'}
82          </Text>
83        )}
84        {hot.map(limit => (
85          <Text key={limit.kind} dimColor={limit.percentUsed < 80} color={limit.percentUsed >= 80 ? 'yellow' : undefined}>
86            {' · '}
87            {LIMIT_LABELS[limit.kind] ?? limit.kind} {Math.round(limit.percentUsed)}%
88          </Text>
89        ))}
90      </Box>
91    )
92  })
93}
94

1export type Reading = {
2  model: string
3  effort: string | null
4  tokens: number | null
5  window: number
6  compactAt: number | null
7  limits: { kind: string; percentUsed: number }[]
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'session-band': { reading: Reading | null }
13  }
14}
15