cost-guard

1// Cost Guard: keeps Claude Code spend visible and catches the two habits that
2// cost the most: letting one session's context grow very large, and coming
3// back to a big session after the prompt cache has expired.
4//
5// turn.complete: after each main-loop turn, read the session's cost and
6// context from $.session.usage(), add the cost since the last reading to
7// today's total in $.store, and toast when a context or budget threshold is
8// crossed.
9// prompt.submit: when the cache has expired (no turn for 5 minutes) and the
10// context is big, hold the prompt once with an estimate of what resending it
11// costs. Pressing Enter again sends it.
12// session.start: register /cost-guard and keep the status line updated.
13
14import type { EngineInterface, Register } from 'claude-code'
15
16// Prompt cache entries expire after 5 minutes without use.
17const CACHE_TTL_MS = 5 * 60 * 1000
18const CONTEXT_STEP_TOKENS = 50_000
19const DAY_KEYS_KEPT = 31
20
21// List input price in USD per million tokens, matched on the model id.
22// Writing to the cache costs 1.25x this; reading from it 0.1x or less.
23const INPUT_PRICES: [RegExp, number][] = [
24  [/fable|mythos/, 10],
25  [/opus-5-5/, 4],
26  [/opus/, 5],
27  [/sonnet-5/, 2],
28  [/sonnet/, 3],
29  [/haiku/, 1],
30]
31
32type Reading = { tokens: number; costUsd: number | undefined; model: string }
33
34let last: Reading = { tokens: 0, costUsd: undefined, model: '' }
35let lastTurnAt = 0
36let costSeen: number | undefined
37let contextWarnedAt = 0
38let sessionWarned = false
39let dailyWarned = false
40let heldAt = 0
41
42export const register: Register = (on, options) => {
43  const contextWarnTokens = Number(options.contextWarnTokens ?? 120_000)
44  const coldGuardTokens = Number(options.coldGuardTokens ?? 60_000)
45  const sessionBudgetUsd = Number(options.sessionBudgetUsd ?? 10)
46  const dailyBudgetUsd = Number(options.dailyBudgetUsd ?? 25)
47
48  on('session.start', async ($, e, next) => {
49    const result = await next(e)
50    await $.command.register({
51      name: 'cost-guard',
52      description: 'Show session and daily spend, context size and cache state',
53    })
54    lastTurnAt = await $.clock.now()
55    await takeReading($)
56    costSeen = last.costUsd
57    $.clock.every(15_000, () => void showStatus($))
58    await showStatus($)
59
60    return result
61  })
62
63  on('turn.complete', async ($, e, next) => {
64    const result = await next(e)
65    if (e.agentId !== undefined) {
66      return result
67    }
68    lastTurnAt = await $.clock.now()
69    heldAt = 0
70    await takeReading($)
71    const today = await addToToday($)
72
73    if (last.tokens >= contextWarnTokens && last.tokens >= contextWarnedAt + CONTEXT_STEP_TOKENS) {
74      contextWarnedAt = last.tokens
75      $.ui.toast(
76        `Context is ${short(last.tokens)} tokens. Every turn re-reads it: /clear when you switch tasks, /compact to keep going.`,
77        { timeoutMs: 10_000 },
78      )
79    }
80    if (!sessionWarned && (last.costUsd ?? 0) >= sessionBudgetUsd) {
81      sessionWarned = true
82      $.ui.toast(`This session has cost ${usd(last.costUsd ?? 0)}, past its ${usd(sessionBudgetUsd)} budget.`, {
83        timeoutMs: 10_000,
84      })
85    }
86    if (!dailyWarned && today >= dailyBudgetUsd) {
87      dailyWarned = true
88      $.ui.toast(`Today's spend is ${usd(today)}, past the ${usd(dailyBudgetUsd)} daily budget.`, {
89        timeoutMs: 10_000,
90      })
91    }
92    await showStatus($)
93
94    return result
95  })
96
97  on('prompt.submit', async ($, e, next) => {
98    const isTyped = e.origin.kind === 'composer' && e.turnId === undefined
99    if (!isTyped || coldGuardTokens <= 0 || e.text.trimStart().startsWith('/')) {
100      return next(e)
101    }
102    const now = await $.clock.now()
103    const idleMs = now - lastTurnAt
104    if (idleMs < CACHE_TTL_MS || last.tokens < coldGuardTokens || heldAt > lastTurnAt) {
105      return next(e)
106    }
107    heldAt = now
108    const rewrite = rewriteCost(last.tokens, last.model)
109    // Put the prompt back so Enter sends it as typed.
110    $.clock.after(50, () => void $.prompt.fill({ text: e.text }))
111
112    return {
113      drop:
114        `Held once: the prompt cache expired ${minutes(idleMs)} ago, so this message re-writes ` +
115        `${short(last.tokens)} tokens of context (~${usd(rewrite)}). Press Enter to send anyway, ` +
116        `or /clear first if this is a new task.`,
117    }
118  })
119
120  on('command.run', { command: 'cost-guard' }, async $ => {
121    await takeReading($)
122    const now = await $.clock.now()
123    const days = await recentDays($, 7)
124    const week = days.reduce((sum, day) => sum + day.usd, 0)
125    const idleMs = now - lastTurnAt
126    const cache =
127      last.tokens === 0
128        ? 'nothing cached yet'
129        : idleMs < CACHE_TTL_MS
130        ? `warm for ${minutes(CACHE_TTL_MS - idleMs)} more`
131        : `expired ${minutes(idleMs)} ago: the next message re-writes ~${usd(rewriteCost(last.tokens, last.model))}`
132    const lines = [
133      `Session  ${last.costUsd === undefined ? 'n/a' : usd(last.costUsd)} (budget ${usd(sessionBudgetUsd)})`,
134      `Today    ${usd(days[0]?.usd ?? 0)} (budget ${usd(dailyBudgetUsd)})`,
135      `7 days   ${usd(week)}  ${days.map(day => `${day.date.slice(5)} ${usd(day.usd)}`).join(' · ')}`,
136      `Context  ${short(last.tokens)} tokens on ${last.model || 'unknown model'}`,
137      `Cache    ${cache}`,
138      '',
139      'Cheaper habits: /clear between tasks, /compact before a break, lower /effort for routine work,',
140      'and /model to a Sonnet for everyday edits. Costs use list prices; a cloud provider may bill differently.',
141    ]
142
143    return { text: lines.join('\n') }
144  })
145}
146
147async function takeReading($: EngineInterface) {
148  try {
149    const usage = await $.session.usage()
150    last = {
151      tokens: usage.context.tokens ?? 0,
152      costUsd: usage.cost?.usd,
153      model: await $.session.model(),
154    }
155  } catch {
156    // Keep the last reading.
157  }
158}
159
160// Adds the session's cost since the last reading to today's total; returns it.
161async function addToToday($: EngineInterface) {
162  const key = `day:${dateKey(await $.clock.now())}`
163  const total = Number((await $.store.get(key)) ?? 0)
164  if (last.costUsd === undefined) {
165    return total
166  }
167  // /clear starts the session's cost over; count from the new figure.
168  const delta = costSeen === undefined || last.costUsd < costSeen ? last.costUsd : last.costUsd - costSeen
169  costSeen = last.costUsd
170  if (delta <= 0) {
171    return total
172  }
173  await $.store.set(key, total + delta)
174  await pruneDays($)
175
176  return total + delta
177}
178
179async function recentDays($: EngineInterface, count: number) {
180  const now = await $.clock.now()
181  const days: { date: string; usd: number }[] = []
182  for (let i = 0; i < count; i++) {
183    const date = dateKey(now - i * 86_400_000)
184    days.push({ date, usd: Number((await $.store.get(`day:${date}`)) ?? 0) })
185  }
186
187  return days
188}
189
190async function pruneDays($: EngineInterface) {
191  const keys = (await $.store.keys()).filter(key => key.startsWith('day:')).sort()
192  for (const key of keys.slice(0, Math.max(0, keys.length - DAY_KEYS_KEPT))) {
193    await $.store.delete(key)
194  }
195}
196
197async function showStatus($: EngineInterface) {
198  const idleMs = (await $.clock.now()) - lastTurnAt
199  const today = Number((await $.store.get(`day:${dateKey(await $.clock.now())}`)) ?? 0)
200  const parts = [
201    last.costUsd === undefined ? undefined : `${usd(last.costUsd)} session`,
202    `${usd(today)} today`,
203    `ctx ${short(last.tokens)}`,
204  ]
205  if (last.tokens > 0) {
206    parts.push(idleMs < CACHE_TTL_MS ? `cache ${minutes(CACHE_TTL_MS - idleMs)}` : 'cache cold')
207  }
208  $.ui.status(parts.filter(Boolean).join(' · '))
209}
210
211export function rewriteCost(tokens: number, model: string) {
212  const price = INPUT_PRICES.find(([pattern]) => pattern.test(model))?.[1] ?? 5
213
214  return (tokens / 1_000_000) * price * 1.25
215}
216
217function dateKey(ms: number) {
218  const d = new Date(ms)
219  const pad = (n: number) => String(n).padStart(2, '0')
220
221  return `${d.getFullYear()}-${pad(d.getMonth() + 1)}-${pad(d.getDate())}`
222}
223
224function minutes(ms: number) {
225  const total = Math.max(0, Math.round(ms / 1000))
226  const m = Math.floor(total / 60)
227
228  return m >= 60 ? `${Math.floor(m / 60)}h${String(m % 60).padStart(2, '0')}m` : `${m}:${String(total % 60).padStart(2, '0')}`
229}
230
231function short(n: number) {
232  if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
233  if (n >= 1_000) return `${Math.round(n / 1_000)}k`
234
235  return String(n)
236}
237
238function usd(n: number) {
239  return `$${n.toFixed(2)}`
240}
241