cost-pane

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface as Engine, Register, TurnUsage } from 'claude-code'
3
4import type { Totals } from '../types'
5
6const PANE = 'cost-pane'
7const KEEP_DAYS = 30
8const ZERO: Totals = { usd: 0, input: 0, output: 0, cacheRead: 0, cacheWrite: 0, turns: 0 }
9
10const sidA = atom({ plugin: 'cost-pane', key: 'sid' } as const, '')
11const lastUsdA = atom({ plugin: 'cost-pane', key: 'lastUsd' } as const, 0)
12const sessionA = atom({ plugin: 'cost-pane', key: 'session' } as const, ZERO)
13// this session's spend per day; the store holds the same under s:<sid>:<day>
14const mineA = atom({ plugin: 'cost-pane', key: 'mine' } as const, {})
15// every session's spend per day, summed from the store
16const daysA = atom({ plugin: 'cost-pane', key: 'days' } as const, {})
17
18// ponytail: fixed UTC+7 (Vietnam has no DST); use the host zone if you travel
19export const dayOf = (ms: number) => new Date(ms + 7 * 3600e3).toISOString().slice(0, 10)
20
21export const add = (t: Totals, usd: number, u?: TurnUsage): Totals => ({
22  usd: t.usd + usd,
23  input: t.input + (u?.input_tokens ?? 0),
24  output: t.output + (u?.output_tokens ?? 0),
25  cacheRead: t.cacheRead + (u?.cache_read_input_tokens ?? 0),
26  cacheWrite: t.cacheWrite + (u?.cache_creation_input_tokens ?? 0),
27  turns: t.turns + (u ? 1 : 0),
28})
29
30export const sum = (a: Totals, b: Totals): Totals => ({
31  usd: a.usd + b.usd,
32  input: a.input + b.input,
33  output: a.output + b.output,
34  cacheRead: a.cacheRead + b.cacheRead,
35  cacheWrite: a.cacheWrite + b.cacheWrite,
36  turns: a.turns + b.turns,
37})
38
39const tok = (n: number) =>
40  n >= 1e6 ? `${(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${Math.round(n / 1e3)}k` : `${n}`
41
42const money = (usd: number) => `$${usd.toFixed(2)}`
43
44// share of input tokens served from the prompt cache
45export const hitRate = (t: Totals) => {
46  const all = t.input + t.cacheRead + t.cacheWrite
47  return all ? Math.round((t.cacheRead / all) * 100) : 0
48}
49
50// one block per value, scaled to the largest; zero draws the lowest block
51export const spark = (values: number[]) => {
52  const max = Math.max(...values)
53  return values.map(v => '▁▂▃▄▅▆▇█'[max > 0 ? Math.round((v / max) * 7) : 0]).join('')
54}
55
56export const bar = (value: number, max: number, width: number) =>
57  max > 0 ? '█'.repeat(Math.max(value > 0 ? 1 : 0, Math.round((value / max) * width))) : ''
58
59const usdNow = async ($: Engine) => (await $.session.usage()).cost?.usd ?? 0
60
61// ponytail: reads every key each refresh; fine for ~30 days of sessions, index by day if it gets slow
62async function refresh($: Engine) {
63  const cutoff = dayOf((await $.clock.now()) - KEEP_DAYS * 86400e3)
64  const days: Record<string, Totals> = {}
65  for (const key of await $.store.keys()) {
66    const day = key.slice(-10)
67    if (!key.startsWith('s:')) continue
68    if (day < cutoff) {
69      await $.store.delete(key)
70      continue
71    }
72    days[day] = sum(days[day] ?? ZERO, ((await $.store.get(key)) as Totals) ?? ZERO)
73  }
74  await update($, daysA, () => days)
75  const session = await read($, sessionA)
76  const now = await $.clock.now()
77  const week = [6, 5, 4, 3, 2, 1, 0].map(i => days[dayOf(now - i * 86400e3)]?.usd ?? 0)
78  $.ui.status(
79    `💸 ${money(session.usd)} session · ${money(week[6] ?? 0)} today · 7d ${spark(week)} ${money(week.reduce((a, b) => a + b, 0))}`,
80  )
81}
82
83export const register: Register = on => {
84  on('session.start', async ($, e, next) => {
85    await $.command.register({ name: 'cost-pane', description: 'Show token and USD spend in a pane' })
86
87    const sid = await $.session.id()
88    if ((await read($, sidA)) !== sid) {
89      // new session, /clear or a resume: keep what the store already has for this id
90      const mine: Record<string, Totals> = {}
91      for (const key of await $.store.keys())
92        if (key.startsWith(`s:${sid}:`)) mine[key.slice(-10)] = (await $.store.get(key)) as Totals
93      await update($, sidA, () => sid)
94      await update($, mineA, () => mine)
95      await update($, sessionA, () => Object.values(mine).reduce(sum, ZERO))
96      const usd = await usdNow($)
97      await update($, lastUsdA, () => usd)
98    }
99
100    await refresh($)
101    $.clock.every(60_000, () => refresh($)) // picks up other sessions' spend
102
103    return next(e)
104  })
105
106  on('command.run', { command: 'cost-pane' }, async $ => {
107    await $.ui.open({ id: PANE, title: 'Cost' })
108
109    return { text: 'Cost pane opened.' }
110  })
111
112  on('turn.complete', async ($, e, next) => {
113    const done = await next(e)
114
115    const usd = await usdNow($)
116    let delta = 0
117    await update($, lastUsdA, last => {
118      delta = Math.max(0, usd - last) // a ledger that restarted counts from zero
119      return Math.max(usd, last)
120    })
121    const day = dayOf(await $.clock.now())
122    const sid = await read($, sidA)
123    await update($, sessionA, t => add(t, delta, e.usage))
124    await update($, mineA, m => ({ ...m, [day]: add(m[day] ?? ZERO, delta, e.usage) }))
125    await $.store.set(`s:${sid}:${day}`, (await read($, mineA))[day])
126    await refresh($)
127
128    return done
129  })
130
131  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
132    const { Box, Text } = $.ui.resolve(e)
133    const session = await read($, sessionA)
134    const days = await read($, daysA)
135    const now = await $.clock.now()
136    const today = dayOf(now)
137    const week = [...Array(7).keys()].map(i => {
138      const d = dayOf(now - i * 86400e3)
139      return [d, days[d] ?? ZERO] as const
140    })
141    const weekTotal = week.map(([, t]) => t).reduce(sum, ZERO)
142    const month = Object.values(days).reduce(sum, ZERO)
143    const max = Math.max(...week.map(([, t]) => t.usd))
144    const cols = e.viewport?.columns ?? 60
145    const barWidth = Math.max(4, Math.min(24, cols - 22))
146
147    const Row = ({ title, t }: { title: string; t: Totals }) => (
148      <Text>
149        <Text dimColor>{title.padEnd(8)}</Text>
150        <Text bold color="green">{money(t.usd).padStart(8)}</Text>
151        <Text dimColor>
152          {'  '}↑{tok(t.input + t.cacheRead + t.cacheWrite)} ↓{tok(t.output)} · {t.turns}t · cache {hitRate(t)}%
153        </Text>
154      </Text>
155    )
156
157    return (
158      <Box flexDirection="column">
159        <Row title="session" t={session} />
160        <Row title="today" t={days[today] ?? ZERO} />
161        <Box flexDirection="column" marginTop={1}>
162          {week.map(([d, t]) => (
163            <Box flexDirection="row" gap={1}>
164              <Text bold={d === today} dimColor={d !== today}>{d.slice(5)}</Text>
165              <Box width={barWidth}>
166                <Text color={d === today ? 'green' : 'cyan'}>{bar(t.usd, max, barWidth)}</Text>
167              </Box>
168              <Text bold={d === today}>{money(t.usd).padStart(8)}</Text>
169            </Box>
170          ))}
171        </Box>
172        <Box flexDirection="row" gap={3}>
173          <Text>
174            <Text dimColor>7d </Text>
175            <Text bold>{money(weekTotal.usd)}</Text>
176          </Text>
177          <Text>
178            <Text dimColor>{KEEP_DAYS}d </Text>
179            <Text bold>{money(month.usd)}</Text>
180          </Text>
181          <Text>
182            <Text dimColor>avg/day </Text>
183            <Text bold>{money(weekTotal.usd / 7)}</Text>
184          </Text>
185        </Box>
186      </Box>
187    )
188  })
189}
190

1export type Totals = {
2  usd: number
3  input: number
4  output: number
5  cacheRead: number
6  cacheWrite: number
7  turns: number
8}
9
10declare module 'claude-code' {
11  interface PluginState {
12    'cost-pane': {
13      sid: string
14      lastUsd: number
15      session: Totals
16      mine: Record<string, Totals>
17      days: Record<string, Totals>
18    }
19  }
20}
21