spend-guard

1import type { Register } from 'claude-code'
2
3// Spend Guard (display name). Real figures from $.session.usage() / session.measure
4// (the same numbers /cost and the status line use). No estimating, no network, no files.
5// Hard stop: prompt.submit answers { drop: reason } once cost >= capUsd.
6// Override: start a prompt with "override:" -> it is let through (prefix stripped) once.
7// Limit: cost only updates after each turn, so one runaway turn can overshoot the cap.
8
9export const register: Register = (on, options) => {
10  const threshold = Number(options?.thresholdUsd ?? 2)
11  const cap = Number(options?.capUsd ?? 5)
12  let usd = 0
13  let turnStartUsd = 0
14  let lastTurnUsd = 0
15  let ctxPercent: number | null = null
16  let warned = false
17  let lockToasted = false
18
19  const fmt = (n: number) => '$' + n.toFixed(n < 10 ? 2 : 1)
20  const isLocked = () => usd >= cap
21  const OVERRIDE = /^\s*override:\s*/i
22
23  on('session.start', async ($, e, next) => {
24    const result = await next(e)
25    try {
26      const u = await $.session.usage()
27      usd = u.cost?.usd ?? 0
28      ctxPercent = u.context?.percent ?? null
29    } catch {}
30    $.ui.invalidate('ui.render')
31    return result
32  })
33
34  // The hard stop. Refresh the cost first so it is current, then hold or pass.
35  on('prompt.submit', async ($, e, next) => {
36    try {
37      const u = await $.session.usage()
38      if (typeof u.cost?.usd === 'number') usd = u.cost.usd
39    } catch {}
40
41    if (isLocked()) {
42      if (OVERRIDE.test(e.text)) {
43        const stripped = e.text.replace(OVERRIDE, '')
44        if (stripped.trim() === '') {
45          return { drop: `Spend Guard: type "override: <your prompt>" with a prompt after it.` }
46        }
47        turnStartUsd = usd
48        $.ui.toast(`Spend Guard: one more turn allowed past the ${fmt(cap)} cap`, { timeoutMs: 6000 })
49        return next({ ...e, text: stripped })
50      }
51      if (!lockToasted) {
52        lockToasted = true
53        $.ui.toast(`Spend Guard: session hit ${fmt(cap)} cap (now ${fmt(usd)}). Type "override: <prompt>" to allow one more turn, or raise the cap.`, { timeoutMs: 10000 })
54      }
55      $.ui.invalidate('ui.render')
56      return { drop: `Spend Guard: session hit ${fmt(cap)} cap (now ${fmt(usd)}). Start your prompt with "override: " to allow one more turn, or raise the cap in the config menu.` }
57    }
58
59    lockToasted = false
60    turnStartUsd = usd
61    return next(e)
62  }).catch(() => undefined)
63
64  // Fires after each turn and whenever the cost total grows.
65  on('session.measure', ($, e, next) => {
66    if (e.cost) {
67      usd = e.cost.usd
68      lastTurnUsd = Math.max(0, usd - turnStartUsd)
69    }
70    ctxPercent = e.context?.percent ?? ctxPercent
71
72    if (!warned && usd >= threshold) {
73      warned = true
74      $.ui.toast(`Spend Guard: session passed ${fmt(threshold)} (now ${fmt(usd)})`, { timeoutMs: 8000 })
75    }
76    $.ui.status(isLocked() ? `LOCKED ${fmt(usd)}/${fmt(cap)}` : `cost ${fmt(usd)}`)
77    $.ui.invalidate('ui.render')
78    return next(e)
79  })
80
81  on('ui.render', { component: 'AbovePrompt' }, ($, e, next) => {
82    if (e.props.hasSurvey || usd === 0) return next(e)
83    const { Box, Text } = $.ui.resolve(e)
84    const locked = isLocked()
85    const color = locked || usd >= threshold ? 'red' : usd >= threshold / 2 ? 'yellow' : 'green'
86    return (
87      <Box paddingX={1}>
88        <Text color={color} bold>{fmt(usd)}</Text>
89        <Text dimColor>
90          {'  this session   last turn '}{fmt(lastTurnUsd)}
91          {ctxPercent !== null ? `   context ${Math.round(ctxPercent)}%` : ''}
92          {`   warn at ${fmt(threshold)}   `}
93        </Text>
94        {locked
95          ? <Text color="red" bold>{`LOCKED  CAP ${fmt(cap)}  (override: <prompt>)`}</Text>
96          : <Text dimColor>{`CAP ${fmt(cap)}`}</Text>}
97      </Box>
98    )
99  })
100}
101