savvy-progress

1import { atom, read, update } from 'claude-code'
2import type { Register } from 'claude-code'
3
4import type { Crab, Progress } from '../types'
5
6// Edit these to restyle the mod. Each helper crab gets a look by agent type.
7const LOOKS: Record<string, string> = {
8  'general-purpose': '(o_o)',
9  Explore: '(o.O)',
10  Plan: '(-_-)',
11  designer: '(^_~)',
12  developer: '(>_<)',
13}
14const DEFAULT_LOOK = '(o_o)'
15const BAR_WIDTH = 24
16
17// Rough USD per million tokens: input, output, cache read, cache write.
18const PRICES: Array<[RegExp, [number, number, number, number]]> = [
19  [/opus|fable/i, [15, 75, 1.5, 18.75]],
20  [/sonnet/i, [3, 15, 0.3, 3.75]],
21  [/haiku/i, [1, 5, 0.1, 1.25]],
22]
23
24const progress = atom({ plugin: 'savvy-progress', key: 'progress' } as const, {
25  steps: 0,
26  created: 0,
27  done: 0,
28} satisfies Progress)
29const crabs = atom({ plugin: 'savvy-progress', key: 'crabs' } as const, [] as Crab[])
30
31const costOf = (
32  model: string,
33  u: {
34    input_tokens: number
35    output_tokens: number
36    cache_read_input_tokens?: number
37    cache_creation_input_tokens?: number
38  },
39): number => {
40  const row = PRICES.find(([re]) => re.test(model))?.[1] ?? [3, 15, 0.3, 3.75]
41  return (
42    (u.input_tokens * row[0] +
43      u.output_tokens * row[1] +
44      (u.cache_read_input_tokens ?? 0) * row[2] +
45      (u.cache_creation_input_tokens ?? 0) * row[3]) /
46    1_000_000
47  )
48}
49
50const bar = (done: number, total: number): string => {
51  const filled = total === 0 ? 0 : Math.round((done / total) * BAR_WIDTH)
52  return '#'.repeat(filled) + '-'.repeat(BAR_WIDTH - filled)
53}
54
55export const register: Register = on => {
56  on('prompt.submit', async ($, e, next) => {
57    await update($, progress, () => ({ steps: 0, created: 0, done: 0 }))
58    await update($, crabs, () => [])
59
60    return next(e)
61  })
62
63  // Watch: every main-loop tool call is a step; task tools say how far along we are.
64  on('tool.call', async ($, e, next) => {
65    if (e.agentId === undefined) {
66      const input = e as unknown as Record<string, unknown>
67      const tool = String(e.tool)
68
69      await update($, progress, p => {
70        const out: Progress = { ...p, steps: p.steps + 1 }
71        if (tool === 'TaskCreate') out.created += 1
72        if (tool === 'TaskUpdate' && input.status === 'completed') out.done += 1
73        if (tool === 'TodoWrite' && Array.isArray(input.todos)) {
74          const todos = input.todos as Array<{ status?: string }>
75          out.created = todos.length
76          out.done = todos.filter(t => t.status === 'completed').length
77        }
78
79        return out
80      })
81    }
82
83    return next(e)
84  })
85
86  // Watch: every subagent turn adds its model, tokens and cost to its crab.
87  on('turn.complete', async ($, e, next) => {
88    if (e.agentId !== undefined && e.usage !== undefined) {
89      const id = e.agentId
90      const usage = e.usage
91      const info = (await $.agent.list()).find(a => a.id === id)
92      const tokens = usage.input_tokens + usage.output_tokens
93      const cost = costOf(usage.model, usage)
94
95      await update($, crabs, list => {
96        const known = list.find(c => c.id === id)
97        if (known === undefined) {
98          const fresh: Crab = {
99            id,
100            label: info?.description ?? info?.name ?? id.slice(0, 8),
101            type: info?.type ?? 'general-purpose',
102            model: usage.model,
103            tokens,
104            cost,
105            isDone: false,
106          }
107
108          return [...list, fresh]
109        }
110
111        return list.map(c =>
112          c.id === id ? { ...c, model: usage.model, tokens: c.tokens + tokens, cost: c.cost + cost } : c,
113        )
114      })
115    }
116
117    return next(e)
118  })
119
120  // Change: draw the band above the prompt.
121  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
122    const p = await read($, progress)
123    const list = await read($, crabs)
124    const isQuiet = e.props.hasSurvey || (!e.props.isWorking && list.length === 0)
125
126    if (isQuiet) {
127      return next(e)
128    }
129
130    const { Box, Text } = $.ui.resolve(e)
131    const total = p.created
132    const line =
133      total > 0
134        ? `[${bar(p.done, total)}] ${p.done}/${total} tasks, ${p.steps} steps`
135        : `[${'~'.repeat(BAR_WIDTH)}] ${p.steps} steps`
136    const spend = list.reduce((sum, c) => sum + c.cost, 0)
137
138    return (
139      <Box flexDirection="column">
140        <Text color="claude">{line}</Text>
141        {list.map(c => (
142          <Text dimColor={c.isDone}>
143            {LOOKS[c.type] ?? DEFAULT_LOOK} {c.label} | {c.model} | {c.tokens.toLocaleString()} tok | ~${c.cost.toFixed(3)}
144          </Text>
145        ))}
146        {list.length > 1 && <Text dimColor>helpers total ~${spend.toFixed(3)} (estimate)</Text>}
147      </Box>
148    )
149  })
150}
151

1export type Crab = {
2  id: string
3  label: string
4  type: string
5  model: string
6  tokens: number
7  cost: number
8  isDone: boolean
9}
10
11export type Progress = { steps: number; created: number; done: number }
12
13declare module 'claude-code' {
14  interface PluginState {
15    'savvy-progress': {
16      progress: Progress
17      crabs: Crab[]
18    }
19  }
20}
21