token-warn

1import { atom, update } from 'claude-code'
2import type { Register, TurnUsage } from 'claude-code'
3
4import type { TokenWarnTally } from '../types'
5
6const empty: TokenWarnTally = { total: 0, owner: {}, byKey: {}, flagged: [] }
7const tally = atom({ plugin: 'token-warn', key: 'tally' } as const, empty)
8
9export const parseTokens = (s: string): number[] =>
10  s.split(',').map(p => p.trim().toLowerCase()).filter(Boolean).map(p => {
11    const mult = p.endsWith('m') ? 1e6 : p.endsWith('k') ? 1e3 : 1
12    return Number(mult === 1 ? p : p.slice(0, -1)) * mult
13  }).filter(n => Number.isFinite(n) && n > 0).sort((a, b) => a - b)
14
15export const fmt = (n: number): string =>
16  n >= 1e6 ? `${+(n / 1e6).toFixed(1)}M` : n >= 1e3 ? `${Math.round(n / 1e3)}k` : String(n)
17
18// Input, output and cache writes. Cache reads are left out: they re-count the same context every turn.
19const sum = (u: TurnUsage) =>
20  u.input_tokens + u.output_tokens + u.cache_creation_input_tokens
21
22export const register: Register = (on, options) => {
23  const thresholds = parseTokens(String(options.thresholds ?? '1M,5M,10M'))
24  const limit = Number(options.agentLimit ?? 500000)
25
26  // Remember which spawned agent counts toward which subagent or Workflow run.
27  on('agent.spawn', async ($, e, next) => {
28    const r = await next(e)
29    const id = 'agentId' in r ? r.agentId : undefined
30    if (id) {
31      const own = e.workflow
32        ? { key: `wf:${e.workflow.runId}`, label: `Workflow run ${e.workflow.runId}` }
33        : { key: `agent:${id}`, label: `subagent "${e.description ?? id}"` }
34      await update($, tally, t => ({ ...t, owner: { ...t.owner, [id]: own } }))
35    }
36    return r
37  }).catch(($, e, next) => next(e)) // never block a spawn over bookkeeping
38
39  on('turn.complete', async ($, e, next) => {
40    const r = await next(e)
41    const usage = r.usage ?? e.usage
42    if (!usage) return r
43    const n = sum(usage)
44    const toasts: string[] = []
45
46    const after = await update($, tally, t => {
47      toasts.length = 0 // update() retries on a version miss; keep the last attempt's toasts
48      const total = t.total + n
49      const flagged = [...t.flagged]
50      const byKey = { ...t.byKey }
51      for (const th of thresholds) {
52        const id = `session:${th}`
53        if (total >= th && !flagged.includes(id)) {
54          flagged.push(id)
55          toasts.push(`⚠ session passed ${fmt(th)} tokens`)
56        }
57      }
58      const own = e.agentId ? t.owner[e.agentId] : undefined
59      if (own) {
60        byKey[own.key] = (byKey[own.key] ?? 0) + n
61        if (byKey[own.key]! > limit && !flagged.includes(own.key)) {
62          flagged.push(own.key)
63          toasts.push(`⚠ ${own.label} used ${fmt(byKey[own.key]!)} tokens`)
64        }
65      }
66      return { ...t, total, byKey, flagged }
67    })
68
69    $.ui.status(`tok ${fmt(after.total)}`)
70    for (const text of toasts) $.ui.toast(text, { timeoutMs: 8000 })
71    return r
72  })
73}
74

1export type TokenWarnTally = {
2  total: number
3  // agentId -> what it counts toward: `agent:<id>` or `wf:<runId>`, with a label
4  owner: Record<string, { key: string; label: string }>
5  byKey: Record<string, number>
6  flagged: string[]
7}
8
9declare module 'claude-code' {
10  interface PluginState {
11    'token-warn': { tally: TokenWarnTally }
12  }
13}
14