statusline

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, RenderChildren, Timer } from 'claude-code'
3
4import type { StatuslineDiff, StatuslineLimit, StatuslineSnap } from '../types'
5
6const EMPTY: StatuslineSnap = { isFast: false, isDirty: false, limits: [], minute: 0 }
7const NO_DIFF: StatuslineDiff = { added: 0, removed: 0 }
8const snap = atom({ plugin: 'statusline', key: 'snap' } as const, EMPTY)
9const diff = atom({ plugin: 'statusline', key: 'diff' } as const, NO_DIFF)
10
11const HOT = 50
12const BAR_CAP = 300_000
13const CTX_WARN = 150_000
14const CTX_HOT = 300_000
15const POLL_MS = 5_000
16const MAX_SESSIONS = 50
17const STORE_VERSION = 1
18const SESSIONS_KEY = 'sessions'
19const DIFF_PREFIX = 'diff:'
20
21const BLUE = '#89b4fa'
22const SURFACE0 = '#313244'
23const MANTLE = '#181825'
24const GREEN = '#a6e3a1'
25const YELLOW = '#f9e2af'
26const RED = '#f38ba8'
27
28const CAP_LEFT = ''
29const CAP_RIGHT = ''
30const GIT_ICON = ''
31
32const settled = <T,>(p: Promise<T>) => p.then(v => v, () => undefined)
33
34export const human = (n: number) => {
35  if (n >= 1_000_000) return `${Math.floor(n / 1_000_000)}M`
36  if (n >= 1_000) return `${Math.floor(n / 1_000)}k`
37  return `${n}`
38}
39
40export const countdown = (at: number, now: number) => {
41  const secs = Math.floor((at - now) / 1000)
42  if (secs <= 0) return undefined
43  const mins = Math.floor(secs / 60)
44  const hours = Math.floor(mins / 60)
45  return hours >= 48 ? `${Math.floor(hours / 24)}d` : hours >= 1 ? `${hours}h` : `${mins}m`
46}
47
48export const heatOf = (tokens: number) => (tokens >= CTX_HOT ? RED : tokens >= CTX_WARN ? YELLOW : BLUE)
49
50export const patchLinesOf = (result: unknown): StatuslineDiff => {
51  const hunks: unknown = typeof result === 'object' && result !== null ? Reflect.get(result, 'structuredPatch') : undefined
52  const lines = Array.isArray(hunks) ? hunks.flatMap((h: unknown) => (typeof h === 'object' && h !== null ? Reflect.get(h, 'lines') : [])) : []
53  if (lines.length === 0 && typeof result === 'object' && result !== null) {
54    const content: unknown = Reflect.get(result, 'content')
55    const original: unknown = Reflect.get(result, 'originalFile')
56    const count = (s: unknown) => (typeof s === 'string' && s !== '' ? s.split('\n').length : 0)
57    return { added: count(content), removed: count(original) }
58  }
59  return {
60    added: lines.filter((l: unknown) => typeof l === 'string' && l.startsWith('+')).length,
61    removed: lines.filter((l: unknown) => typeof l === 'string' && l.startsWith('-')).length,
62  }
63}
64
65export const storedDiffOf = (raw: unknown): StatuslineDiff | undefined => {
66  if (typeof raw !== 'object' || raw === null || Reflect.get(raw, 'version') !== STORE_VERSION) return undefined
67  const added: unknown = Reflect.get(raw, 'added')
68  const removed: unknown = Reflect.get(raw, 'removed')
69  return typeof added === 'number' && typeof removed === 'number' ? { added, removed } : undefined
70}
71
72export const keptSessionsOf = (raw: unknown, id: string) =>
73  [...(Array.isArray(raw) ? raw.filter((s): s is string => typeof s === 'string' && s !== id) : []), id].slice(-MAX_SESSIONS)
74
75const persist = async ($: EngineInterface, id: string, d: StatuslineDiff) => {
76  await $.store.set(DIFF_PREFIX + id, { version: STORE_VERSION, ...d })
77  const kept = keptSessionsOf(await $.store.get(SESSIONS_KEY), id)
78  await $.store.set(SESSIONS_KEY, kept)
79  for (const key of await $.store.keys()) {
80    if (key.startsWith(DIFF_PREFIX) && !kept.includes(key.slice(DIFF_PREFIX.length))) await $.store.delete(key)
81  }
82}
83
84const rebuilt = async ($: EngineInterface): Promise<StatuslineDiff> => {
85  const messages = await settled($.session.messages())
86  return (Array.isArray(messages) ? messages : [])
87    .flatMap(m => m.toolUses)
88    .filter(u => (u.tool === 'Edit' || u.tool === 'Write') && u.isError === undefined && u.result !== undefined)
89    .map(u => patchLinesOf(u.result))
90    .reduce((a, p) => ({ added: a.added + p.added, removed: a.removed + p.removed }), NO_DIFF)
91}
92
93export const fieldOf = (o: unknown, key: string): unknown => (typeof o === 'object' && o !== null ? Reflect.get(o, key) : undefined)
94
95const git = async ($: EngineInterface, cwd: string, args: string[]) => {
96  const ran = await settled($.process.run(['git', '--no-optional-locks', '-C', cwd, ...args]))
97  return ran?.exitCode === 0 ? ran.stdout.trim() : undefined
98}
99
100const branchOf = async ($: EngineInterface, cwd: string) =>
101  (await git($, cwd, ['symbolic-ref', '--quiet', '--short', 'HEAD'])) ?? (await git($, cwd, ['rev-parse', '--short', 'HEAD']))
102
103let stepEffort: string | undefined
104let poll: Timer | undefined
105let refreshing: Promise<void> | undefined
106let isStale = false
107let saving: Promise<unknown> = Promise.resolve()
108
109function refreshSoon($: EngineInterface) {
110  isStale = true
111  refreshing ??= (async () => {
112    try {
113      while (isStale) {
114        isStale = false
115        await settled(refresh($))
116      }
117    } finally {
118      refreshing = undefined
119    }
120  })()
121  return refreshing
122}
123
124function save($: EngineInterface) {
125  saving = saving.then(async () => {
126    const id = await settled($.session.id())
127    if (id !== undefined) await settled(persist($, id, await read($, diff)))
128  })
129  return saving
130}
131
132async function refresh($: EngineInterface) {
133  const cwd = await $.session.cwd()
134  const [model, usage, effortSetting, modelSettings, fast, branch, status, now] = await Promise.all([
135    settled($.session.model()),
136    settled($.session.usage()),
137    settled($.settings.read({ key: 'effortLevel' })),
138    settled($.settings.read({ key: 'modelSettings' })),
139    settled($.settings.read({ key: 'fastMode' })),
140    branchOf($, cwd),
141    git($, cwd, ['status', '--porcelain']),
142    $.clock.now(),
143  ])
144  const limits: StatuslineLimit[] = (usage?.rateLimits ?? []).map(l => ({
145    kind: l.kind,
146    percent: Math.round(l.percentUsed),
147    ...(l.resetsAt === undefined ? {} : { resetsAt: Date.parse(l.resetsAt) }),
148  }))
149  const perModel = model === undefined ? undefined : fieldOf(fieldOf(modelSettings, model), 'effortLevel')
150  const effort = [perModel, effortSetting, stepEffort].find(v => typeof v === 'string')
151  const next: StatuslineSnap = {
152    ...(model === undefined ? {} : { model }),
153    ...(typeof effort === 'string' ? { effort } : {}),
154    isFast: fast === true,
155    ...(branch === undefined ? {} : { branch }),
156    isDirty: (status ?? '') !== '',
157    ...(usage?.cost === undefined ? {} : { cost: usage.cost.usd }),
158    ...(usage?.context.tokens === undefined ? {} : { tokens: usage.context.tokens }),
159    ...(usage?.context.percent === undefined ? {} : { percent: Math.round(usage.context.percent) }),
160    ...(usage === undefined ? {} : { window: usage.context.window }),
161    limits,
162    minute: Math.floor(now / 60_000),
163  }
164  const prev = await read($, snap)
165  if (JSON.stringify(prev) !== JSON.stringify(next)) await update($, snap, () => next)
166}
167
168export const register: Register = on => {
169  on('session.start', async ($, e, next) => {
170    const started = await next(e)
171    const id = await settled($.session.id())
172    if (id !== undefined) {
173      const before = await read($, diff)
174      const stored = storedDiffOf(await settled($.store.get(DIFF_PREFIX + id)))
175      const base = stored ?? (await rebuilt($))
176      await update($, diff, d => ({ added: base.added + d.added - before.added, removed: base.removed + d.removed - before.removed }))
177      if (stored === undefined) await save($)
178    }
179    await refreshSoon($)
180    poll?.cancel()
181    poll = $.clock.every(POLL_MS, () => void refreshSoon($))
182    return started
183  })
184
185  on('turn.step', async function* ($, e, next) {
186    const done = yield* next(e)
187    if (e.agentId === undefined && e.effort !== undefined) {
188      stepEffort = String(e.effort)
189      await refreshSoon($)
190    }
191    return done
192  })
193
194  on('session.end', async ($, e, next) => {
195    if (e.reason === 'clear' || e.reason === 'resume') {
196      stepEffort = undefined
197      await update($, diff, () => NO_DIFF)
198    }
199    return next(e)
200  })
201
202  on('tool.call', async ($, e, next) => {
203    const done = await next(e)
204    if ((e.tool === 'Edit' || e.tool === 'Write') && done.result !== undefined && done.isError === undefined) {
205      const { added, removed } = patchLinesOf(done.result)
206      if (added > 0 || removed > 0) {
207        await update($, diff, d => ({ added: d.added + added, removed: d.removed + removed }))
208        await save($)
209      }
210    }
211    if (done.isReadOnly !== true) await refreshSoon($)
212    return done
213  })
214
215  on('turn.complete', async ($, e, next) => {
216    const done = await next(e)
217    await refreshSoon($)
218    return done
219  })
220
221  on('session.append', { door: 'command' }, async ($, e, next) => {
222    const kept = await next(e)
223    await refreshSoon($)
224    return kept
225  })
226
227  on('ui.render', { component: 'PromptHint' }, async ($, e, next) => {
228    const engine = await next(e)
229    const s = await read($, snap)
230    const d = await read($, diff)
231    const now = await $.clock.now()
232    const { Box, Text } = $.ui.resolve(e)
233
234    type Seg = { color: string; body: RenderChildren }
235    const pills = (dir: 'r' | 'l', segs: Seg[]) => {
236      const shown = segs.filter(seg => seg.body !== undefined)
237      const last = shown[shown.length - 1]
238      if (last === undefined) return []
239      const out = shown.flatMap((seg, i) => {
240        const prev = shown[i - 1]
241        const text = seg.color === BLUE ? MANTLE : BLUE
242        const joint =
243          prev === undefined
244            ? [<Text key={`c${i}`} color={seg.color}>{CAP_LEFT}</Text>]
245            : dir === 'r'
246              ? [<Text key={`c${i}`} color={prev.color} backgroundColor={seg.color}>{CAP_RIGHT}</Text>]
247              : [
248                  <Text key={`s${i}`} backgroundColor={prev.color}> </Text>,
249                  <Text key={`c${i}`} color={seg.color} backgroundColor={prev.color}>{CAP_LEFT}</Text>,
250                ]
251        return [
252          ...joint,
253          <Text key={`b${i}`} backgroundColor={seg.color} color={text}>
254            {prev !== undefined && dir === 'r' ? ' ' : ''}
255            {seg.body}
256          </Text>,
257        ]
258      })
259      return [...out, <Text key="end" color={last.color}>{CAP_RIGHT}</Text>]
260    }
261
262    const model =
263      s.model === undefined ? undefined : (
264        <Text>
265          <Text bold>{s.model}</Text>
266          {s.effort === undefined ? '' : ` ${s.effort}`}
267          {s.isFast ? ' ⚡fast' : ''}
268        </Text>
269      )
270    const branch = s.branch === undefined ? undefined : `${GIT_ICON} ${s.branch}${s.isDirty ? '*' : ''}`
271
272    let ctx: RenderChildren
273    if (s.percent !== undefined) {
274      const window = s.window ?? 0
275      const barPct = window > 0 ? Math.floor((s.percent * window) / BAR_CAP) : s.percent
276      const filled = Math.min(8, Math.ceil((barPct * 8) / 100))
277      ctx = (
278        <Text color={heatOf(Math.floor((s.percent * window) / 100))}>
279          {s.tokens === undefined ? '' : `${human(s.tokens)} `}
280          {'▰'.repeat(filled)}
281          <Text dimColor>{'▱'.repeat(8 - filled)}</Text>
282          {' '}
283          <Text bold>{`${s.percent}%`}</Text>
284        </Text>
285      )
286    }
287
288    const limit = (label: string, l: StatuslineLimit) => {
289      const left = l.resetsAt === undefined ? undefined : countdown(l.resetsAt, now)
290      return (
291        <Text key={label}>
292          {`${label} `}
293          <Text bold>{`${l.percent}%`}</Text>
294          {left === undefined ? '' : ` ↻${left}`}
295        </Text>
296      )
297    }
298    const five = s.limits.find(l => l.kind === 'five_hour')
299    const seven = s.limits.find(l => l.kind === 'seven_day')
300    const lims = [
301      ...(five === undefined ? [] : [limit('5h', five)]),
302      ...(seven === undefined || seven.percent < HOT ? [] : [limit('7d', seven)]),
303    ]
304    const lim = lims.length === 0 ? undefined : lims.flatMap((l, i) => (i === 0 ? [l] : [<Text key={`dot${i}`}> · </Text>, l]))
305
306    return (
307      <Box flexDirection="column">
308        <Box flexDirection="row" marginLeft={-2} marginRight={-2}>
309          <Box flexDirection="row" flexShrink={0}>
310            {pills('r', [
311              { color: BLUE, body: model },
312              { color: SURFACE0, body: branch },
313            ])}
314            {d.added > 0 || d.removed > 0 ? (
315              <Text>
316                {' '}
317                <Text color={GREEN}>{`+${d.added}`}</Text>/<Text color={RED}>{`-${d.removed}`}</Text>
318              </Text>
319            ) : null}
320          </Box>
321          <Box flexGrow={1} minWidth={1} />
322          <Box flexDirection="row" flexShrink={0}>
323            {s.cost === undefined ? null : <Text color={GREEN} bold>{`$${Math.round(s.cost * 100) / 100} `}</Text>}
324            {pills('l', [
325              { color: SURFACE0, body: ctx },
326              { color: BLUE, body: lim },
327            ])}
328          </Box>
329        </Box>
330        {engine}
331      </Box>
332    )
333  })
334}
335

1export type StatuslineLimit = { kind: string; percent: number; resetsAt?: number }
2
3export type StatuslineDiff = { added: number; removed: number }
4
5export type StatuslineSnap = {
6  model?: string
7  effort?: string
8  isFast: boolean
9  branch?: string
10  isDirty: boolean
11  cost?: number
12  tokens?: number
13  percent?: number
14  window?: number
15  limits: StatuslineLimit[]
16  minute: number
17}
18
19declare module 'claude-code' {
20  interface PluginState {
21    statusline: { snap: StatuslineSnap; diff: StatuslineDiff }
22  }
23}
24