cockpit

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { EngineInterface } from 'claude-code'
5
6import type { Category, Context, Limit, Tokens } from '../types'
7
8const limits = atom({ plugin: 'cockpit', key: 'limits' } as const, [])
9const costUsd = atom({ plugin: 'cockpit', key: 'costUsd' } as const, null)
10const minute = atom({ plugin: 'cockpit', key: 'minute' } as const, 0)
11const tokens = atom({ plugin: 'cockpit', key: 'tokens' } as const, null)
12const context = atom({ plugin: 'cockpit', key: 'context' } as const, null)
13
14const LABELS: Record<string, string> = {
15  five_hour: '5h',
16  seven_day: 'Week',
17  spend_limit: 'On-demand',
18}
19const ORDER = ['five_hour', 'seven_day', 'spend_limit']
20const BAR_WIDTH = 10
21const DAY = 24 * 60 * 60 * 1000
22const BORDER = '#b5664a'
23const ZERO = { read: 0, write: 0, cacheRead: 0, cacheWrite: 0 }
24
25export const label = (kind: string) => LABELS[kind] ?? kind.replace(/_/g, ' ')
26
27// The used part and the rest of a thin line bar, drawn in two colours.
28export const bar = (percent: number) => {
29  const filled = Math.min(BAR_WIDTH, Math.max(0, Math.round((percent / 100) * BAR_WIDTH)))
30
31  // Both centred on the text's midline: a thick bar over a thin track.
32  return { used: '▬'.repeat(filled), rest: '─'.repeat(BAR_WIDTH - filled) }
33}
34
35// A valid IANA name, or undefined for the system's own timezone.
36export const zone = (setting: unknown) => {
37  if (typeof setting !== 'string' || setting.trim() === '' || setting.trim().toLowerCase() === 'system') {
38    return undefined
39  }
40  try {
41    new Intl.DateTimeFormat('en-US', { timeZone: setting.trim() })
42
43    return setting.trim()
44  } catch {
45    return undefined
46  }
47}
48
49// 8 -> "8", 3_000 -> "3k", 3_240 -> "3.2k", 1_250_000 -> "1.3M"
50export const count = (n: number) => {
51  const short = (v: number, unit: string) => `${Number(v.toFixed(1))}${unit}`
52
53  return n >= 1e6 ? short(n / 1e6, 'M') : n >= 1e3 ? short(n / 1e3, 'k') : `${n}`
54}
55
56// Pads each text on the left to the widest one's length.
57export const pad = (...texts: string[]) => {
58  const width = Math.max(...texts.map(t => t.length))
59
60  return texts.map(t => t.padStart(width))
61}
62
63export const usd = (amount: number) => `$${amount.toFixed(2)}`
64
65export const color = (percent: number) => (percent >= 90 ? 'red' : percent >= 70 ? 'yellow' : 'green')
66
67export const countdown = (at: number, now: number) => {
68  const minutes = Math.max(0, Math.round((at - now) / 60000))
69  const days = Math.floor(minutes / 1440)
70  const hours = Math.floor((minutes % 1440) / 60)
71  const mins = minutes % 60
72  return days > 0 ? `${days}d ${hours}h` : hours > 0 ? `${hours}h ${mins}m` : `${mins}m`
73}
74
75// "⏳ 03:31 PM (2h 30m)"; a reset more than a day out leads with its weekday.
76export const resetLabel = (resetsAt: string | undefined, now: number, timeZone?: string) => {
77  const at = resetsAt ? Date.parse(resetsAt) : NaN
78  if (Number.isNaN(at)) {
79    return ''
80  }
81  const time = new Intl.DateTimeFormat('en-US', {
82    timeZone,
83    weekday: at - now > DAY ? 'short' : undefined,
84    hour: '2-digit',
85    minute: '2-digit',
86    hour12: true,
87  }).format(at)
88
89  return `⏳ ${time} (${countdown(at, now)})`
90}
91
92const TOP = 4
93
94// The biggest rows that occupy the window, largest first.
95export const topCategories = (rows: readonly { name: string; tokens: number; color: string; kind: string }[]): Category[] =>
96  rows
97    .filter(r => r.kind === 'used' && r.tokens > 0)
98    .sort((a, b) => b.tokens - a.tokens)
99    .slice(0, TOP)
100    .map(r => ({ name: r.name, tokens: r.tokens, color: r.color }))
101
102// The window's fill and, as /context counts it (the local estimate, no requests), its biggest rows.
103const measureContext = async ($: EngineInterface): Promise<Context> => {
104  const { context: c } = await $.session.usage({ breakdown: 'summary' })
105
106  return { window: c.window, tokens: c.tokens, percent: c.percent, top: topCategories(c.breakdown?.categories ?? []) }
107}
108
109type Counts = typeof ZERO
110
111// Sums the usage of every reply in a transcript (JSONL). A reply is written
112// once per content block, each line repeating its usage, so each message id counts once.
113export const sumTranscript = (text: string): Counts => {
114  const replies = new Map<string, Counts>()
115  for (const line of text.split('\n')) {
116    if (!line.includes('"usage"')) {
117      continue
118    }
119    try {
120      const entry = JSON.parse(line)
121      const u = entry?.message?.usage
122      if (entry?.type !== 'assistant' || !u) {
123        continue
124      }
125      replies.set(entry.message.id ?? entry.uuid, {
126        read: u.input_tokens ?? 0,
127        write: u.output_tokens ?? 0,
128        cacheRead: u.cache_read_input_tokens ?? 0,
129        cacheWrite: u.cache_creation_input_tokens ?? 0,
130      })
131    } catch {
132      // A line cut short while being written: skip it.
133    }
134  }
135  const sum = { ...ZERO }
136  for (const r of replies.values()) {
137    sum.read += r.read
138    sum.write += r.write
139    sum.cacheRead += r.cacheRead
140    sum.cacheWrite += r.cacheWrite
141  }
142
143  return sum
144}
145
146// The conversation's transcript, and its subagents', counted; null when it cannot be read.
147const countTranscript = async ($: EngineInterface, cwd: string): Promise<Counts | null> => {
148  try {
149    const id = await $.session.id()
150    const home = (await $.env.get('USERPROFILE')) ?? (await $.env.get('HOME'))
151    const config = (await $.env.get('CLAUDE_CONFIG_DIR')) ?? `${home}/.claude`
152    const dir = `${config}/projects/${cwd.replace(/[^a-zA-Z0-9]/g, '-')}`
153    const sum = sumTranscript(await $.fs.read(`${dir}/${id}.jsonl`))
154    try {
155      for (const f of await $.fs.list(`${dir}/${id}/subagents`)) {
156        if (f.name.endsWith('.jsonl')) {
157          const sub = sumTranscript(await $.fs.read(`${dir}/${id}/subagents/${f.name}`))
158          sum.read += sub.read
159          sum.write += sub.write
160          sum.cacheRead += sub.cacheRead
161          sum.cacheWrite += sub.cacheWrite
162        }
163      }
164    } catch {
165      // No subagents ran.
166    }
167
168    return sum
169  } catch {
170    return null
171  }
172}
173
174const sorted = (list: readonly Limit[]) =>
175  [...list].sort((a, b) => {
176    const ia = ORDER.indexOf(a.kind)
177    const ib = ORDER.indexOf(b.kind)
178
179    return (ia === -1 ? 99 : ia) - (ib === -1 ? 99 : ib)
180  })
181
182export const register: Register = (on, options) => {
183  const timeZone = zone(options.timezone)
184
185  on('session.start', async ($, e, next) => {
186    const result = await next(e)
187    const usage = await $.session.usage()
188    await update($, limits, () => usage.rateLimits.map(l => ({ ...l })))
189    await update($, costUsd, () => usage.cost?.usd ?? null)
190    const measured = await measureContext($)
191    await update($, context, () => measured)
192    // A reload keeps the running totals; a new session (or /clear) counts its
193    // transcript once, then each turn adds to that.
194    const current = await read($, tokens)
195    if (!(current?.startedAt === usage.startedAt && current.isFull)) {
196      const counted = await countTranscript($, e.cwd)
197      await update($, tokens, () =>
198        counted
199          ? { startedAt: usage.startedAt, isFull: true, ...counted }
200          : { startedAt: usage.startedAt, isFull: false, ...ZERO },
201      )
202    }
203    // One cheap redraw a minute keeps the reset countdowns current.
204    $.clock.every(60000, () => update($, minute, n => n + 1))
205
206    return result
207  })
208
209  on('session.measure', async ($, e, next) => {
210    if (e.changed.includes('rateLimits')) {
211      await update($, limits, () => e.rateLimits.map(l => ({ ...l })))
212    }
213    if (e.changed.includes('cost')) {
214      await update($, costUsd, () => e.cost?.usd ?? null)
215    }
216    if (e.changed.includes('context')) {
217      const measured = await measureContext($)
218      await update($, context, () => measured)
219    }
220
221    return next(e)
222  })
223
224  // Read: uncached input; Write: generated output; Cache Read / Cache Write:
225  // input served from and written to the prompt cache. Nested tasks' turns count too.
226  on('turn.complete', async ($, e, next) => {
227    const u = e.usage
228    if (u) {
229      // Self-heals from a `tokens` session.start never managed to set: starts
230      // counting from this turn rather than staying stuck at nothing forever.
231      // A failed lookup must not block the write below, or the heal never happens.
232      const startedAt = (await read($, tokens))?.startedAt ?? (await $.session.usage().catch(() => null))?.startedAt ?? Date.now()
233      await update($, tokens, t => {
234        const base = t ?? { startedAt, isFull: false, ...ZERO }
235
236        return {
237          ...base,
238          read: base.read + u.input_tokens,
239          write: base.write + u.output_tokens,
240          cacheRead: base.cacheRead + u.cache_read_input_tokens,
241          cacheWrite: base.cacheWrite + u.cache_creation_input_tokens,
242        }
243      })
244    }
245
246    return next(e)
247  })
248
249  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
250    if (e.props.hasSurvey) {
251      return next(e)
252    }
253
254    const list = sorted(await read($, limits))
255    const cost = await read($, costUsd)
256    const used = await read($, tokens)
257    const ctx = await read($, context)
258    await read($, minute)
259    const now = await $.clock.now()
260    const { Box, Text } = $.ui.resolve(e)
261
262    const total = used ? used.read + used.write + used.cacheRead + used.cacheWrite : 0
263
264    if (list.length === 0 && cost === null && total === 0 && ctx === null) {
265      return (
266        <Box>
267          <Text dimColor>Usage: waiting for the first reply to report limits…</Text>
268        </Box>
269      )
270    }
271
272    const divider = (key: string) => (
273      <Box key={key} flexDirection="column" marginX={2}>
274        <Text color={BORDER}>│</Text>
275        <Text color={BORDER}>│</Text>
276      </Box>
277    )
278
279    const sections = list.map(l => {
280      const left = Math.max(0, 100 - Math.round(l.percentUsed))
281      const name = label(l.kind)
282      const tone = color(l.percentUsed)
283      const { used, rest } = bar(l.percentUsed)
284      const reset = resetLabel(l.resetsAt, now, timeZone)
285
286      // Two rows: the bar, then the reset time under it, both left-aligned.
287      return (
288        <Box key={l.kind} flexDirection="column">
289          <Box key="bar">
290            <Text dimColor>{name} </Text>
291            <Text color={tone}>{used}</Text>
292            <Text color="#3a3a3a">
293              {rest}
294            </Text>
295            <Text color={tone} bold>
296              {' '}
297              {left}%
298            </Text>
299            <Text dimColor> left</Text>
300          </Box>
301          <Box key="reset">
302            <Text color="#777777">{reset}</Text>
303          </Box>
304        </Box>
305      )
306    })
307
308    if (used !== null) {
309      // "Read/Write 8/3.2k" over "Cache R/W 3k/5k", the pairs right-aligned.
310      const [io, cache] = pad(
311        `${count(used.read)}/${count(used.write)}`,
312        `${count(used.cacheRead)}/${count(used.cacheWrite)}`,
313      )
314
315      sections.push(
316        <Box key="breakdown" flexDirection="column">
317          <Box key="io">
318            <Text dimColor>Read/Write </Text>
319            <Text bold>{io}</Text>
320          </Box>
321          <Box key="cache">
322            <Text dimColor>Cache R/W  </Text>
323            <Text bold>{cache}</Text>
324          </Box>
325        </Box>,
326      )
327    }
328
329    if (used !== null || cost !== null) {
330      // Pad both rows to one width so the numbers line up on the right.
331      const [tokensText, costText] = pad(count(total), cost === null ? '—' : usd(cost))
332
333      sections.push(
334        <Box key="totals" flexDirection="column">
335          <Box key="tokens">
336            <Text dimColor>Tokens </Text>
337            <Text bold>{tokensText}</Text>
338          </Box>
339          <Box key="cost">
340            <Text dimColor>Cost   </Text>
341            <Text bold>{costText}</Text>
342          </Box>
343        </Box>,
344      )
345    }
346
347    const usageRow = (
348      <Box key="usage" flexDirection="row" flexWrap="wrap">
349        {sections.flatMap((section, i) => (i === 0 ? [section] : [divider(`divider-${i}`), section]))}
350      </Box>
351    )
352
353    // One bordered box: the context window on top, a rule, then the usage sections.
354    const contextRow = (() => {
355      if (ctx === null) {
356        return null
357      }
358      const percent = ctx.percent ?? (ctx.tokens === undefined ? undefined : (ctx.tokens / ctx.window) * 100)
359      const { used: filled, rest } = bar(percent ?? 0)
360      const tone = color(percent ?? 0)
361
362      return (
363        <Box key="context" flexWrap="wrap">
364          <Text dimColor>Context </Text>
365          <Text color={tone}>{filled}</Text>
366          <Text color="#3a3a3a">{rest}</Text>
367          <Text color={tone} bold>
368            {' '}
369            {percent === undefined ? '—' : `${Math.round(percent)}%`}
370          </Text>
371          <Text dimColor> used  </Text>
372          <Text bold>{ctx.tokens === undefined ? '—' : count(ctx.tokens)}</Text>
373          <Text dimColor>/{count(ctx.window)}</Text>
374          {ctx.top.length === 0 ? null : (
375            // Hungry frogs, eating what fills the context.
376            <Box key="top-divider" marginX={2}>
377              <Text>🐸🐸🐸</Text>
378            </Box>
379          )}
380          {ctx.top.map((c, i) => (
381            <Box key={`top-${c.name}`} marginLeft={i === 0 ? 0 : 3}>
382              <Text color={c.color}>■ </Text>
383              <Text color="#777777">{c.name} </Text>
384              <Text color="#aaaaaa">{count(c.tokens)}</Text>
385            </Box>
386          ))}
387        </Box>
388      )
389    })()
390
391    // The band's width less the border and padding on each side.
392    const ruleWidth = Math.max(1, e.props.bodyColumns - 4)
393
394    return (
395      <Box flexDirection="column" borderStyle="round" borderColor={BORDER} paddingX={1}>
396        {contextRow}
397        {contextRow === null ? null : <Text color={BORDER}>{'─'.repeat(ruleWidth)}</Text>}
398        {usageRow}
399      </Box>
400    )
401  })
402}
403

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2
3/**
4 * Tokens summed over the whole conversation; `startedAt` is the session they
5 * belong to, `isFull` whether the transcript's earlier replies were counted in.
6 */
7export type Tokens = {
8  startedAt: number
9  isFull: boolean
10  read: number
11  write: number
12  cacheRead: number
13  cacheWrite: number
14}
15
16/** The live context window: its size, and its fill once a reply reported one. */
17export type Context = { window: number; tokens?: number; percent?: number; top: Category[] }
18
19/** One of the biggest rows /context lists (Messages, System tools, MCP tools...). */
20export type Category = { name: string; tokens: number; color: string }
21
22declare module 'claude-code' {
23  interface PluginState {
24    'cockpit': {
25      limits: Limit[]; costUsd: number | null; minute: number; tokens: Tokens | null
26      context: Context | null
27    }
28  }
29}
30