turn-memo

1import { atom, memberOf, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Turn } from '../types'
5
6const last = atom({ plugin: 'turn-memo', key: 'last' } as const, null)
7const pending = atom({ plugin: 'turn-memo', key: 'pending' } as const, [] as Turn[])
8
9// 8 -> "8", 3_000 -> "3k", 3_240 -> "3.2k", 1_250_000 -> "1.3M"
10export const count = (n: number) => {
11  const short = (v: number, unit: string) => `${Number(v.toFixed(1))}${unit}`
12
13  return n >= 1e6 ? short(n / 1e6, 'M') : n >= 1e3 ? short(n / 1e3, 'k') : `${n}`
14}
15
16export const usd = (amount: number) => `$${amount.toFixed(2)}`
17
18// The memo's parts, formatted, plus the plain line's length (for the
19// dashed rule above and below it, since a Box border has no top/bottom-only mode).
20export const format = (turn: Turn) => {
21  const tokens = count(turn.total)
22  const readWrite = `${count(turn.read)}/${count(turn.write)}`
23  const cacheRW = `${count(turn.cacheRead)}/${count(turn.cacheWrite)}`
24  const cost = turn.costUsd === null ? '—' : usd(turn.costUsd)
25  const width = `Tokens ${tokens} · Read/Write ${readWrite} · Cache R/W ${cacheRW} · Cost ${cost}`.length
26
27  return { tokens, readWrite, cacheRW, cost, width }
28}
29
30export const register: Register = on => {
31  // Read: uncached input; Write: generated output; Cache Read / Cache Write:
32  // input served from and written to the prompt cache. Cost is the session's
33  // running total (the API gives no per-turn price).
34  on('turn.complete', async ($, e, next) => {
35    const result = await next(e)
36    if (e.agentId || !e.usage) {
37      return result
38    }
39
40    const u = e.usage
41    const { cost } = await $.session.usage()
42    const turn: Turn = {
43      total: u.input_tokens + u.output_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens,
44      read: u.input_tokens,
45      write: u.output_tokens,
46      cacheRead: u.cache_read_input_tokens,
47      cacheWrite: u.cache_creation_input_tokens,
48      costUsd: cost?.usd ?? null,
49    }
50    // Queued, not written straight to `last`: this turn's own TurnDuration
51    // row doesn't exist yet (turnId and that row's requestId are unrelated
52    // ids, minted separately), so there is nothing to key the write by yet.
53    await update($, pending, list => [...list, turn])
54
55    return result
56  })
57
58  // The "Baked for 3s" row is appended as a notice right when the turn
59  // closes; claiming the oldest queued turn here, in append order, is what
60  // ties a turn's usage to that exact row's id, which ui.render then reads by.
61  on('session.append', { door: 'notice' }, async ($, e, next) => {
62    const result = await next(e)
63    if (e.message.name === 'turn_duration') {
64      const [turn, ...rest] = await read($, pending)
65      if (turn) {
66        await update($, pending, () => rest)
67        await update($, { plugin: 'turn-memo', key: 'last', id: e.uuid } as const, () => turn)
68      }
69    }
70
71    return result
72  })
73
74  // TurnDuration is the "Baked for 3s" line that closes a turn in the
75  // terminal transcript; a memo follows it, the cost in a bright color.
76  on('ui.render', { component: 'TurnDuration' }, async ($, e, next) => {
77    const below = await next(e)
78    const turn = await read($, memberOf(last, e))
79    if (!turn) {
80      return below
81    }
82
83    const { Box, Text } = $.ui.resolve(e)
84    const { tokens, readWrite, cacheRW, cost, width } = format(turn)
85    const dashes = '╌'.repeat(width)
86
87    return (
88      <Box flexDirection="column">
89        {below}
90        <Box key="rule-top">
91          <Text dimColor>{dashes}</Text>
92        </Box>
93        <Box>
94          <Text color="claude">Tokens </Text>
95          <Text color="white">{tokens} · </Text>
96          <Text color="claude">Read/Write </Text>
97          <Text color="white">{readWrite} · </Text>
98          <Text color="claude">Cache R/W </Text>
99          <Text color="white">{cacheRW} · </Text>
100          <Text color="claude">Cost </Text>
101          <Text color="greenBright" bold>
102            {cost}
103          </Text>
104        </Box>
105        <Box key="rule-bottom">
106          <Text dimColor>{dashes}</Text>
107        </Box>
108      </Box>
109    )
110  })
111}
112

1/** One turn's usage and the session's running cost at the moment it closed. */
2export type Turn = {
3  total: number
4  read: number
5  write: number
6  cacheRead: number
7  cacheWrite: number
8  costUsd: number | null
9}
10
11declare module 'claude-code' {
12  interface PluginState {
13    'turn-memo': { last: StateFamily<Turn | null>; pending: Turn[] }
14  }
15}
16