usage-meters

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { ContextReading, TallyReading } from '../types'
5
6const REFRESH_INTERVAL_MS = 30_000
7const GROUP_GAP = 4
8
9const context = atom({ plugin: 'usage-meters', key: 'context' } as const, null as ContextReading | null)
10const tally = atom({ plugin: 'usage-meters', key: 'tally' } as const, { kind: 'pending' } as TallyReading)
11
12// Module state restarts on every reload, which is harmless for a re-entrancy guard.
13let isTallyRunning = false
14
15/** Compact token count: `950`, `134k`, `2.98M`. */
16export function formatTokens(count: number): string {
17  if (count < 1_000) return String(count)
18  if (count < 1_000_000) return `${Math.round(count / 1_000)}k`
19  return `${(count / 1_000_000).toFixed(2)}M`
20}
21
22/** The context's token count alone, or a dash until the live window's first response reports one. */
23export function formatContextTokens(reading: ContextReading | null): string {
24  return reading?.tokens === undefined ? '–' : formatTokens(reading.tokens)
25}
26
27/** Reads the script's JSON report; anything not in the shape it documents is a failure, never a number. */
28export function parseTallyReport(stdout: string): TallyReading {
29  const report = JSON.parse(stdout) as { transcriptWritten?: unknown; today?: unknown; lastHour?: unknown }
30  if (report.transcriptWritten === false) return { kind: 'awaitingTranscript' }
31  if (typeof report.today !== 'number' || typeof report.lastHour !== 'number') {
32    return { kind: 'failed', reason: `unexpected report: ${stdout.trim()}` }
33  }
34  return { kind: 'counted', today: report.today, lastHour: report.lastHour }
35}
36
37/** One value of the NCR row: a count, a dash before the transcript exists, an ellipsis while counting. */
38export function formatTallyValue(reading: TallyReading, pick: (counted: { today: number; lastHour: number }) => number): string {
39  switch (reading.kind) {
40    case 'counted':
41      return formatTokens(pick(reading))
42    case 'awaitingTranscript':
43      return '–'
44    case 'pending':
45    case 'failed':
46      return '…'
47  }
48}
49
50async function countTokens($: EngineInterface): Promise<TallyReading> {
51  try {
52    const sessionId = await $.session.id()
53    const script = `${$.plugin.root}/bin/tally_tokens.py`
54    const result = await $.process.run(['/usr/bin/python3', '-I', script, sessionId], { timeoutMs: 10_000 })
55    if (result.exitCode !== 0) {
56      return { kind: 'failed', reason: result.stderr.trim().split('\n').pop() ?? `exit ${result.exitCode}` }
57    }
58    return parseTallyReport(result.stdout)
59  } catch (error) {
60    return { kind: 'failed', reason: error instanceof Error ? error.message : String(error) }
61  }
62}
63
64async function refreshTally($: EngineInterface): Promise<void> {
65  // The timer and per-response measurements can overlap; one run at a time is enough.
66  if (isTallyRunning) return
67  isTallyRunning = true
68  try {
69    const reading = await countTokens($)
70    await update($, tally, () => reading)
71  } finally {
72    isTallyRunning = false
73  }
74}
75
76export const register: Register = on => {
77  on('session.start', async ($, e, next) => {
78    // An earlier version pinned these figures in the status line; the band holds them now.
79    $.ui.status(undefined)
80    const usage = await $.session.usage()
81    await update($, context, () => usage.context)
82    // The timer moves the "last 60m" window forward while no responses arrive.
83    $.clock.every(REFRESH_INTERVAL_MS, () => void refreshTally($))
84    void refreshTally($)
85    return next(e)
86  })
87
88  // Fires when the context fill, a rate-limit window or the session's cost moves; cost grows with
89  // every response, subagents' included, which makes it the per-request refresh trigger.
90  on('session.measure', async ($, e, next) => {
91    await update($, context, () => e.context)
92    if (e.changed.includes('cost')) void refreshTally($)
93    return next(e)
94  })
95
96  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
97    if (e.props.hasSurvey) return next(e)
98    const { Box, Text } = $.ui.resolve(e)
99    const fill = await read($, context)
100    const tokens = await read($, tally)
101
102    // Styled like the rate-limit rows of the status line: bold label, plain value.
103    return (
104      <Box flexDirection="row" gap={GROUP_GAP}>
105        <Box key="context" flexDirection="row" gap={1}>
106          <Text bold>Context</Text>
107          <Text>{formatContextTokens(fill)}</Text>
108        </Box>
109        {tokens.kind === 'failed' ? (
110          <Box key="ncr-failed" flexDirection="row" gap={1}>
111            <Text bold>NCR tok</Text>
112            <Text color="error">{`error: ${tokens.reason}`}</Text>
113          </Box>
114        ) : (
115          <Box key="ncr" flexDirection="row" gap={GROUP_GAP}>
116            <Box flexDirection="row" gap={1}>
117              <Text bold>NCR tok</Text>
118              <Text>{formatTallyValue(tokens, counted => counted.today)}</Text>
119            </Box>
120            <Box flexDirection="row" gap={1}>
121              <Text bold>Last hour</Text>
122              <Text>{formatTallyValue(tokens, counted => counted.lastHour)}</Text>
123            </Box>
124          </Box>
125        )}
126      </Box>
127    )
128  })
129}
130

1/** The live context window's fill; `tokens`/`percent` absent until the window's first response. */
2export type ContextReading = {
3  window: number
4  tokens?: number
5  percent?: number
6}
7
8/** This session's (and its subagents') non-cache-read tokens, or why they could not be counted. */
9export type TallyReading =
10  | { kind: 'pending' }
11  | { kind: 'awaitingTranscript' }
12  | { kind: 'counted'; today: number; lastHour: number }
13  | { kind: 'failed'; reason: string }
14
15declare module 'claude-code' {
16  interface PluginState {
17    'usage-meters': {
18      /** The context fill from the latest usage measurement; null before the first. */
19      context: ContextReading | null
20      /** The latest token tally. */
21      tally: TallyReading
22    }
23  }
24}
25