turn-usage

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, SessionRateLimit } from 'claude-code'
3
4import type { UsageNote } from '../types'
5
6// The "current session" figure in Claude's usage settings is the five-hour window.
7const WINDOW = 'five_hour'
8// The subscription windows; with one used up, a turn that still completes ran on extra usage.
9const SUBSCRIPTION_WINDOWS = ['five_hour', 'seven_day']
10const MAX_NOTES = 100
11
12const notes = atom({ plugin: 'turn-usage', key: 'notes' } as const, [] as UsageNote[])
13
14const WINDOW_MS = 5 * 60 * 60 * 1000
15// How long before a turn a window may have begun and still count as begun by it.
16const FRESH_WINDOW_SLACK_MS = 15 * 60 * 1000
17
18const formatPercent = (n: number) => `${Math.round(n * 10) / 10}%`
19// Escaped, so a line with two amounts never reads as inline math.
20const formatUsd = (n: number) => (n < 0.005 ? '<\\$0.01' : `\\$${n.toFixed(2)}`)
21
22const LABELS = {
23  es: { session: 'Sesión', used: 'usado', turn: 'este turno', extra: 'uso extra' },
24  en: { session: 'Session', used: 'used', turn: 'this turn', extra: 'extra usage' },
25}
26
27// Spanish when Claude Code's `language` setting says so, English otherwise.
28async function readLabels($: EngineInterface): Promise<(typeof LABELS)['en']> {
29  const { language } = await $.settings.read()
30  const isSpanish = typeof language === 'string' && /^\s*(es\b|es-|spanish|espa[nñ]ol|castellano)/i.test(language)
31  return isSpanish ? LABELS.es : LABELS.en
32}
33
34// How many points of the window a turn spent, or undefined with nothing to compare against.
35function spentThisTurn(
36  after: SessionRateLimit,
37  start: { reading: SessionRateLimit | undefined; usd: number | undefined; startedAt: number } | undefined,
38  firstStep: SessionRateLimit | undefined,
39): number | undefined {
40  const since = (from: SessionRateLimit) =>
41    // A window that reset during the turn started over from zero.
42    from.resetsAt !== after.resetsAt ? after.percentUsed : Math.max(0, after.percentUsed - from.percentUsed)
43
44  if (start?.reading !== undefined) {
45    return since(start.reading)
46  }
47  // A window that began with this turn started it at zero.
48  const windowStart = after.resetsAt === undefined ? NaN : Date.parse(after.resetsAt) - WINDOW_MS
49  if (start !== undefined && windowStart >= start.startedAt - FRESH_WINDOW_SLACK_MS) {
50    return after.percentUsed
51  }
52  // Otherwise the first request's reading is the closest to the turn's start.
53  return firstStep === undefined ? undefined : since(firstStep)
54}
55
56// Keeps a finished reply's line, for the drawing of the block that ends it.
57async function addNote($: EngineInterface, answer: string, line: string): Promise<void> {
58  await update($, notes, (list) => [...list, { answer, line }].slice(-MAX_NOTES))
59}
60
61async function readWindow($: EngineInterface): Promise<SessionRateLimit | undefined> {
62  const { rateLimits } = await $.session.usage()
63  return rateLimits.find((limit) => limit.kind === WINDOW)
64}
65
66export const register: Register = (on) => {
67  // The window's reading, the session's cost and the time when each turn began, by turn id.
68  const atStart = new Map<
69    string,
70    { reading: SessionRateLimit | undefined; usd: number | undefined; startedAt: number }
71  >()
72  // For a turn that began with no reading (a session's first): the one its first request brought.
73  const afterFirstStep = new Map<string, SessionRateLimit>()
74
75  on('turn.start', async ($, e, next) => {
76    const { rateLimits, cost } = await $.session.usage()
77    atStart.set(e.turnId, {
78      reading: rateLimits.find((limit) => limit.kind === WINDOW),
79      usd: cost?.usd,
80      startedAt: await $.clock.now(),
81    })
82    return next(e)
83  })
84
85  on('turn.step', async function* ($, e, next) {
86    const result = yield* next(e)
87    const start = atStart.get(e.turnId)
88    if (e.agentId === undefined && e.index === 0 && start !== undefined && start.reading === undefined) {
89      const reading = await readWindow($)
90      if (reading !== undefined) {
91        afterFirstStep.set(e.turnId, reading)
92      }
93    }
94    return result
95  })
96
97  on('turn.complete', async ($, e, next) => {
98    const result = await next(e)
99    const start = atStart.get(e.turnId)
100    const firstStep = afterFirstStep.get(e.turnId)
101    atStart.delete(e.turnId)
102    afterFirstStep.delete(e.turnId)
103
104    const answer = e.answer.trim()
105    if (e.agentId !== undefined || answer === '') {
106      return result
107    }
108
109    const labels = await readLabels($)
110    const { rateLimits, cost } = await $.session.usage()
111    const after = rateLimits.find((limit) => limit.kind === WINDOW)
112    const usd = cost?.usd
113    const turnUsd = usd === undefined || start?.usd === undefined ? '—' : formatUsd(Math.max(0, usd - start.usd))
114
115    if (after === undefined) {
116      // No usage window after a turn that reached the API: pay as you go, so show the cost.
117      if (e.usage === undefined || usd === undefined) {
118        return result
119      }
120      const line = `${labels.session}: ${formatUsd(usd)} · ${labels.turn}: ${turnUsd}`
121      await addNote($, answer, line)
122      return result
123    }
124
125    const isExtraUsage = rateLimits.some(
126      (limit) => SUBSCRIPTION_WINDOWS.includes(limit.kind) && limit.percentUsed >= 100,
127    )
128    if (isExtraUsage) {
129      const line = `${labels.session}: ${formatPercent(after.percentUsed)} ${labels.used} · ${labels.turn}: ${turnUsd} (${labels.extra})`
130      await addNote($, answer, line)
131      return result
132    }
133
134    const spent = spentThisTurn(after, start, firstStep)
135    let spentText = '—'
136    if (spent !== undefined) {
137      // The API reports whole points, so a turn that moved nothing spent under one.
138      spentText = spent < 1 ? '<1%' : `+${formatPercent(spent)}`
139    }
140
141    const line = `${labels.session}: ${formatPercent(after.percentUsed)} ${labels.used} · ${labels.turn}: ${spentText}`
142    await addNote($, answer, line)
143    return result
144  })
145
146  // The reply's last text block is the one the turn's answer ends with.
147  on('ui.render', { component: 'AssistantMessage' }, async ($, e, next) => {
148    const text = e.props.text.trim()
149    const list = await read($, notes)
150    const note = text === '' ? undefined : list.findLast((n) => n.answer.endsWith(text))
151    if (note === undefined) {
152      return next(e)
153    }
154
155    // Line up with the reply's text: the terminal indents it past the bullet,
156    // the desktop half a cell in from the message's edge.
157    const indent = e.surface === 'terminal' ? 2 : 0.5
158    const { Box, Markdown } = $.ui.resolve(e)
159    return (
160      <Box flexDirection="column">
161        {await next(e)}
162        <Box marginTop={1} paddingLeft={indent}>
163          <Markdown dimColor text={note.line} />
164        </Box>
165      </Box>
166    )
167  })
168}
169

1/** One finished turn: the reply's final text and the usage line drawn under it. */
2export type UsageNote = { answer: string; line: string }
3
4declare module 'claude-code' {
5  interface PluginState {
6    'turn-usage': { notes: UsageNote[] }
7  }
8}
9