token-saver

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { TokenTally } from '../types'
5
6const EMPTY: TokenTally = {
7  turns: 0,
8  input: 0,
9  output: 0,
10  cacheRead: 0,
11  cacheWrite: 0,
12  subagentTokens: 0,
13  routed: 0,
14  dedupedReads: 0,
15}
16const tally = atom({ plugin: 'token-saver', key: 'tally' } as const, EMPTY)
17
18const LEAN_RULES = [
19  'Token discipline (accuracy first, then brevity):',
20  '- Answer directly. No preamble, no restating the request, no closing recap of what you just did unless asked.',
21  '- Locate before you read: Grep/Glob first, then Read only the relevant range (offset/limit) of large files.',
22  '- Do not re-read a file that is already in your context and has not changed.',
23  '- Put independent tool calls in one message so they run in parallel.',
24  '- Hand broad, read-heavy searches to an Explore subagent so raw output stays out of this context.',
25  '- Prefer Edit over rewriting whole files; show only changed code in replies.',
26  '- Never trade correctness for brevity: verify with tools instead of guessing.',
27].join('\n')
28
29// Built-in agent types whose work is search and lookup: cheap models do it as well.
30const LOOKUP_AGENTS = new Set(['Explore', 'claude-code-guide', 'statusline-setup'])
31const ROUTINE_AGENTS = new Set(['general-purpose', 'Plan'])
32
33const TRIVIAL = /^(rename|fix (the )?typo|format|lint|add (a )?comment|what (is|does)|explain|list|show|summari[sz]e|commit|bump|translate)\b/i
34const HARD = /\b(architect|design|race condition|deadlock|memory leak|security|vulnerab|refactor (the )?(whole|entire)|migrat|concurren|why (does|is) .* (fail|crash|slow))/i
35
36function tier(model: string): 'haiku' | 'sonnet' | 'opus' | 'other' {
37  const m = model.toLowerCase()
38  if (m.includes('haiku')) return 'haiku'
39  if (m.includes('sonnet')) return 'sonnet'
40  if (m.includes('opus') || m.includes('fable')) return 'opus'
41  return 'other'
42}
43
44function k(n: number): string {
45  return n >= 1_000_000 ? `${(n / 1_000_000).toFixed(2)}M` : n >= 1000 ? `${(n / 1000).toFixed(1)}k` : String(n)
46}
47
48async function refreshHud($: EngineInterface) {
49  const u = await $.session.usage()
50  const parts: string[] = []
51  if (u.context.percent !== undefined) parts.push(`ctx ${u.context.percent}%`)
52  if (u.cost) parts.push(`$${u.cost.usd.toFixed(2)}`)
53  for (const r of u.rateLimits) {
54    const label = r.kind === 'five_hour' ? '5h' : r.kind === 'seven_day' ? '7d' : r.kind
55    parts.push(`${label} ${r.percentUsed}%`)
56  }
57  $.ui.status(parts.length > 0 ? parts.join(' · ') : undefined)
58  return u
59}
60
61export const register: Register = (on, options) => {
62  const routing = String(options.routing ?? 'balanced')
63  const isLean = options.leanRules !== false
64  const isDedupe = options.dedupeReads !== false
65  const isAdvisor = options.advisor !== false
66
67  // file reads per context (main loop or subagent): path+range -> mtime
68  const seen = new Map<string, number>()
69  // a read refused once is let through when asked again
70  const refused = new Set<string>()
71  let warnedContext = false
72  let warnedLimit = false
73  let lastAdvice = -Infinity
74
75  on('session.start', async ($, e, next) => {
76    await $.command.register({
77      name: 'tokens',
78      description: 'Token report: spend, cache hit rate, subagent share, what token-saver saved, and what to change',
79    })
80    await refreshHud($)
81
82    return next(e)
83  })
84
85  on('session.compact', async ($, e, next) => {
86    // Compaction drops file contents from context: they may be read again.
87    seen.clear()
88    refused.clear()
89    warnedContext = false
90
91    return next(e)
92   }).catch(($, e, next) => next(e))
93
94  if (isLean) {
95    on('prompt.compose', async ($, e, next) => {
96      const composed = await next(e)
97
98      return { sections: [...composed.sections, { id: 'token-saver:lean', text: LEAN_RULES, scope: 'session' as const }] }
99    }).catch(($, e, next) => next(e))
100  }
101
102  if (routing !== 'off') {
103    on('agent.spawn', async ($, e, next) => {
104      const isBuiltin = e.provider.plugin === 'engine'
105      const isFixed = e.model !== undefined || e.fork || e.resume !== undefined || e.workflow !== undefined || e.isTeammate
106      if (!isBuiltin || isFixed) return next(e)
107
108      const parent = tier(e.parentModel)
109      let model: string | undefined
110      if (LOOKUP_AGENTS.has(e.subagentType) && parent !== 'haiku') model = 'haiku'
111      else if (routing === 'aggressive' && ROUTINE_AGENTS.has(e.subagentType) && parent === 'opus') model = 'sonnet'
112      if (model === undefined) return next(e)
113
114      const started = await next({ ...e, model })
115      if (started.deny === undefined) await update($, tally, t => ({ ...t, routed: t.routed + 1 }))
116
117      return started
118    }).catch(($, e, next) => next(e))
119  }
120
121  if (isDedupe) {
122    on('tool.call', { tool: 'Read' }, async ($, e, next) => {
123      const key = `${e.agentId ?? 'main'}|${e.file_path}|${e.offset ?? ''}|${e.limit ?? ''}|${e.pages ?? ''}`
124      const stat = await $.fs.stat(e.file_path).catch(() => null)
125      const mtime = stat?.mtimeMs
126      if (mtime !== undefined && seen.get(key) === mtime && !refused.has(key)) {
127        refused.add(key)
128        await update($, tally, t => ({ ...t, dedupedReads: t.dedupedReads + 1 }))
129
130        return {
131          deny: `token-saver: ${e.file_path} was already read with this range in this context and is unchanged on disk. Use the copy you have. If it is no longer in your context, repeat this exact call and it will go through.`,
132        }
133      }
134      refused.delete(key)
135      const ran = await next(e)
136      if (mtime !== undefined && ran.deny === undefined && ran.isError !== true) seen.set(key, mtime)
137
138      return ran
139    }).catch(($, e, next) => next(e))
140  }
141
142  on('turn.complete', async ($, e, next) => {
143    const u = e.usage
144    if (u) {
145      const total = u.input_tokens + u.output_tokens + u.cache_read_input_tokens + u.cache_creation_input_tokens
146      await update($, tally, t =>
147        e.agentId !== undefined
148          ? { ...t, subagentTokens: t.subagentTokens + total }
149          : {
150              ...t,
151              turns: t.turns + 1,
152              input: t.input + u.input_tokens,
153              output: t.output + u.output_tokens,
154              cacheRead: t.cacheRead + u.cache_read_input_tokens,
155              cacheWrite: t.cacheWrite + u.cache_creation_input_tokens,
156            },
157      )
158    }
159    if (e.agentId === undefined) {
160      const usage = await refreshHud($)
161      const pct = usage.context.percent ?? 0
162      if (pct >= 70 && !warnedContext) {
163        warnedContext = true
164        $.ui.toast(`Context ${pct}% full: every turn now resends it. /compact keeps the thread, /clear starts a new task cheaply.`)
165      }
166      const limit = usage.rateLimits.find(r => r.percentUsed >= 80)
167      if (limit && !warnedLimit) {
168        warnedLimit = true
169        $.ui.toast(`${limit.kind} usage at ${limit.percentUsed}%. Consider /model sonnet for routine work until it resets.`)
170      }
171    }
172
173    return next(e)
174  })
175
176  if (isAdvisor) {
177    on('prompt.submit', async ($, e, next) => {
178      if (e.command || e.isMeta || !['composer', 'bridge', 'sdk'].includes(e.origin.kind)) return next(e)
179      const now = await $.clock.now()
180      if (now - lastAdvice < 15 * 60_000) return next(e)
181
182      const model = tier(await $.session.model())
183      const text = e.text.trim()
184      if (model === 'opus' && text.length < 160 && TRIVIAL.test(text)) {
185        lastAdvice = now
186        $.ui.toast('Simple ask on an Opus-class model. /model sonnet (or haiku) does this for a fraction of the tokens.')
187      } else if (model === 'haiku' && HARD.test(text)) {
188        lastAdvice = now
189        $.ui.toast('This looks like hard reasoning. /model opus (or sonnet) is likely more accurate here.')
190      }
191
192      return next(e)
193    }).catch(($, e, next) => next(e))
194  }
195
196  on('command.run', { command: 'tokens' }, async $ => {
197    const t = await read($, tally)
198    const u = await refreshHud($)
199    const model = await $.session.model()
200    const prompt = t.input + t.cacheRead + t.cacheWrite
201    const hit = prompt > 0 ? Math.round((t.cacheRead / prompt) * 100) : 0
202    const main = prompt + t.output
203    const sub = main + t.subagentTokens > 0 ? Math.round((t.subagentTokens / (main + t.subagentTokens)) * 100) : 0
204    const tips: string[] = []
205    if ((u.context.percent ?? 0) >= 60) tips.push('Context is past 60%: /compact, or /clear before the next unrelated task.')
206    if (prompt > 50_000 && hit < 60) tips.push('Cache hit rate is low: avoid switching model mid-thread and long idle gaps (the cache lapses after ~5 min).')
207    if (tier(model) === 'opus' && t.routed === 0 && t.turns > 5) tips.push('No subagent was routed yet: ask Claude to delegate searches to Explore subagents.')
208    if (t.turns > 0 && t.output / t.turns > 3000) tips.push('Replies are long on average: ask for diffs or bullet answers.')
209
210    const lines = [
211      `Model ${model} · ${t.turns} turns · ${u.cost ? `$${u.cost.usd.toFixed(2)}` : 'cost n/a'} · context ${u.context.percent ?? '?'}%`,
212      `Main loop: ${k(t.input)} fresh in, ${k(t.cacheRead)} cache read, ${k(t.cacheWrite)} cache write, ${k(t.output)} out (cache hit ${hit}%)`,
213      `Subagents: ${k(t.subagentTokens)} tokens (${sub}% of all)`,
214      `token-saver: ${t.routed} subagent${t.routed === 1 ? '' : 's'} routed to a cheaper model, ${t.dedupedReads} duplicate read${t.dedupedReads === 1 ? '' : 's'} stopped`,
215      ...u.rateLimits.map(r => `Limit ${r.kind}: ${r.percentUsed}%${r.resetsAt ? `, resets ${r.resetsAt}` : ''}`),
216      ...(tips.length > 0 ? ['', ...tips.map(x => `• ${x}`)] : []),
217    ]
218
219    return { text: lines.join('\n') }
220  })
221}
222

1export type TokenTally = {
2  turns: number
3  input: number
4  output: number
5  cacheRead: number
6  cacheWrite: number
7  subagentTokens: number
8  routed: number
9  dedupedReads: number
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    'token-saver': { tally: TokenTally }
15  }
16}
17