limit-guard

1// burnrate: one source for the full mod and the two standalone ones.
2// scripts/build.py writes plugins/<name>/hooks/register.tsx from this file,
3// flipping FEATURES and the plugin name; edit here, then run the script.
4
5import { atom, read, update } from 'claude-code'
6import type { Engine, Register, SessionRateLimit, SessionUsage } from 'claude-code'
7
8import type { Category, Snapshot, Tokens } from '../types'
9
10const FEATURES = { band: false, guard: true }
11
12// ---- band ----
13
14const ZERO: Tokens = { input: 0, cacheRead: 0, cacheWrite: 0, output: 0 }
15
16const snap = atom({ plugin: 'limit-guard', key: 'snap' } as const, null)
17const lastTurn = atom({ plugin: 'limit-guard', key: 'lastTurn' } as const, null)
18const total = atom({ plugin: 'limit-guard', key: 'total' } as const, ZERO)
19const categories = atom({ plugin: 'limit-guard', key: 'categories' } as const, [])
20const isExpanded = atom({ plugin: 'limit-guard', key: 'isExpanded' } as const, false)
21
22// One accent for normal fill, two alarm tones; category hues for the context bar.
23const ACCENT = '#d97757'
24const WARN = '#d4a72c'
25const DANGER = '#e5534b'
26const TRACK = '#8a8a8a'
27const HUES = ['#d97757', '#6c9ef8', '#57ab5a', '#b083f0', '#e0a84f', '#4fb3c4', '#c96198', '#8b949e']
28
29const LABEL: Record<string, string> = { five_hour: '5h', seven_day: '7d', spend_limit: 'spend' }
30
31const fmt = (n: number) =>
32  n >= 1_000_000
33    ? `${(n / 1_000_000).toFixed(1)}M`
34    : n >= 10_000
35      ? `${Math.round(n / 1000)}k`
36      : n >= 1000
37        ? `${(n / 1000).toFixed(1)}k`
38        : `${n}`
39
40const tone = (pct: number) => (pct >= 80 ? DANGER : pct >= 50 ? WARN : ACCENT)
41
42const untilReset = (resetsAt: string | undefined, now: number) => {
43  if (!resetsAt) return ''
44  const ms = Date.parse(resetsAt) - now
45  if (!(ms > 0)) return 'now'
46  const h = Math.floor(ms / 3_600_000)
47  const m = Math.floor((ms % 3_600_000) / 60_000)
48  return h >= 24 ? `${Math.floor(h / 24)}d ${h % 24}h` : h > 0 ? `${h}h ${m}m` : `${m}m`
49}
50
51const svgMeter = (pct: number, w: number, h: number, fill: string) => {
52  const r = h / 2
53  const filled = pct > 0 ? Math.max(h, (Math.min(pct, 100) / 100) * w) : 0
54  return (
55    `<svg xmlns="http://www.w3.org/2000/svg" width="${w}" height="${h}" viewBox="0 0 ${w} ${h}">` +
56    `<rect width="${w}" height="${h}" rx="${r}" fill="${TRACK}" fill-opacity="0.22"/>` +
57    (filled ? `<rect width="${filled.toFixed(1)}" height="${h}" rx="${r}" fill="${fill}"/>` : '') +
58    `</svg>`
59  )
60}
61
62const svgStack = (parts: { tokens: number; color: string }[], max: number, w: number, h: number) => {
63  let x = 0
64  const rects = parts
65    .map(p => {
66      const pw = (p.tokens / max) * w
67      const rect = pw >= 1 ? `<rect x="${x.toFixed(1)}" width="${Math.max(1, pw - 1.5).toFixed(1)}" height="${h}" fill="${p.color}"/>` : ''
68      x += pw
69      return rect
70    })
71    .join('')
72  return (
73    `<svg xmlns="http://www.w3.org/2000/svg" width="${w}" height="${h}" viewBox="0 0 ${w} ${h}">` +
74    `<defs><clipPath id="c"><rect width="${w}" height="${h}" rx="${h / 2}"/></clipPath></defs>` +
75    `<g clip-path="url(#c)"><rect width="${w}" height="${h}" fill="${TRACK}" fill-opacity="0.22"/>${rects}</g>` +
76    `</svg>`
77  )
78}
79
80const toSnapshot = (u: Pick<SessionUsage, 'rateLimits' | 'context' | 'cost'>): Snapshot => ({
81  limits: u.rateLimits,
82  ctxTokens: u.context.tokens,
83  ctxWindow: u.context.window,
84  ctxPercent: u.context.percent,
85  usd: u.cost?.usd,
86})
87
88async function refreshBreakdown($: Engine) {
89  try {
90    const u = await $.session.usage({ breakdown: 'summary' })
91    const rows: Category[] = (u.context.breakdown?.categories ?? [])
92      .filter(c => c.kind === 'used' && c.tokens > 0)
93      .map(c => ({ name: c.name, tokens: c.tokens }))
94      .sort((a, b) => b.tokens - a.tokens)
95    await update($, categories, () => rows)
96    await update($, snap, () => toSnapshot(u))
97  } catch {
98    // no session bound yet: keep the last figures
99  }
100}
101
102// ---- guard ----
103
104// The 5h window share at which the brake engages, unless `/guard <n>` stored another.
105const DEFAULT_THRESHOLD = 85
106// The weekly window brakes on its own near the very end.
107const WEEKLY_THRESHOLD = 95
108// Below the brake, a softer "be concise" nudge starts this many points earlier.
109const SOFT_MARGIN = 15
110// A turn that writes at least this much to the cache, more than it reads, lost its cache.
111const CACHE_MISS_TOKENS = 20_000
112
113const SUBAGENT_TOOLS = new Set(['Agent', 'Task', 'mcp__swarm__swarm_dispatch'])
114
115// Shell commands that walk a whole tree.
116const WIDE_SHELL =
117  /\bfind\s+(\.|\/|~|[A-Za-z]:[\\/]?)(\s|$)|\bgrep\s+-[A-Za-z]*[rR]|\brg\s+--files\b|\bls\s+-[A-Za-z]*R|\btree\b|\bdu\s+-[A-Za-z]*a|Get-ChildItem\b[^|;]*-Recurse|\bdir\s+\/s\b/i
118
119const ROOTISH = /^([A-Za-z]:[\\/]?|\/|~|\.|\.\/)$/
120
121// Module state: starts over on a reload; the engine re-measures right away.
122let limits: SessionRateLimit[] = []
123let threshold = DEFAULT_THRESHOLD
124let isOff = false
125let turns = 0
126let lastModel: string | undefined
127let lastTurnAt = 0
128
129const window = (kind: string) => limits.find(l => l.kind === kind)
130const fiveHour = () => window('five_hour')?.percentUsed ?? 0
131const weekly = () => window('seven_day')?.percentUsed ?? 0
132const isBraked = () => !isOff && (fiveHour() >= threshold || weekly() >= WEEKLY_THRESHOLD)
133const isSoft = () => !isOff && !isBraked() && fiveHour() >= threshold - SOFT_MARGIN
134
135const why = () => (weekly() >= WEEKLY_THRESHOLD ? `недельный лимит на ${weekly()}%` : `лимит 5ч на ${fiveHour()}%`)
136
137const deny = (what: string) => ({
138  deny:
139    `burnrate: ${why()} — ${what} отключены, чтобы не сжечь остаток. ` +
140    'Сделай задачу сам и точечно (конкретный путь, конкретный файл, Read с offset/limit). ' +
141    'Пользователь может снять блок командой /guard off.',
142})
143
144function showStatus($: Engine) {
145  $.ui.status(isBraked() ? `🛑 burnrate: ${why()}` : isSoft() ? `burnrate: ${fiveHour()}%, коротко` : undefined)
146}
147
148// ---- hooks ----
149
150async function guardStart($: Engine) {
151  const stored = await $.store.get('threshold')
152  if (typeof stored === 'number') threshold = stored
153  $.command.register({ name: 'guard', description: 'burnrate guard: status, on, off, or a 5h threshold (e.g. /guard 90)', argumentHint: '[on|off|<percent>]' })
154  limits = (await $.session.usage()).rateLimits
155  showStatus($)
156}
157
158function guardMeasure($: Engine, rateLimits: SessionRateLimit[]) {
159  const wasBraked = isBraked()
160  limits = rateLimits
161  if (!wasBraked && isBraked()) {
162    $.ui.toast(`🛑 burnrate: ${why()}. Субагенты и широкие сканы заблокированы, ответы — коротко. /guard off чтобы снять.`)
163  }
164  showStatus($)
165}
166
167async function bandTurn($: Engine, t: Tokens) {
168  await update($, lastTurn, () => t)
169  await update($, total, prev => ({
170    input: prev.input + t.input,
171    cacheRead: prev.cacheRead + t.cacheRead,
172    cacheWrite: prev.cacheWrite + t.cacheWrite,
173    output: prev.output + t.output,
174  }))
175}
176
177async function guardTurn($: Engine, t: Tokens, model: string) {
178  const now = await $.clock.now()
179  const gapMin = turns > 0 ? Math.round((now - lastTurnAt) / 60_000) : 0
180  const isMiss = turns > 0 && t.cacheWrite >= CACHE_MISS_TOKENS && t.cacheWrite > t.cacheRead
181
182  if (isMiss) {
183    const cause =
184      lastModel && lastModel !== model
185        ? `сменилась модель (${lastModel} → ${model})`
186        : gapMin >= 60
187          ? `кэш протух: пауза ${gapMin} мин`
188          : gapMin >= 5
189            ? `пауза ${gapMin} мин — возможно, кэш на 5 мин истёк`
190            : 'поменялся системный промпт или инструменты (CLAUDE.md, MCP, скиллы, плагины) либо был /compact'
191    // Writing costs 1.25–2× input, reading 0.1×: a rewrite is roughly 12–20× a cached read.
192    $.ui.toast(`⚠ Кэш сброшен: записано ${fmt(t.cacheWrite)}, прочитано ${fmt(t.cacheRead)} — ход ~в 12–20 раз дороже. Причина: ${cause}.`)
193  }
194
195  turns += 1
196  lastModel = model
197  lastTurnAt = now
198}
199
200export const register: Register = on => {
201  on('session.start', async ($, e, next) => {
202    const result = await next(e)
203    if (FEATURES.band) await refreshBreakdown($)
204    if (FEATURES.guard) await guardStart($)
205    return result
206  })
207
208  on('session.measure', async ($, e, next) => {
209    if (FEATURES.band) {
210      await update($, snap, () => toSnapshot(e))
211      if (e.changed.includes('context')) await refreshBreakdown($)
212    }
213    if (FEATURES.guard) guardMeasure($, e.rateLimits)
214    return next(e)
215  })
216
217  on('turn.complete', async ($, e, next) => {
218    const u = e.usage
219    if (u && !e.agentId) {
220      const t: Tokens = {
221        input: u.input_tokens,
222        cacheRead: u.cache_read_input_tokens,
223        cacheWrite: u.cache_creation_input_tokens,
224        output: u.output_tokens,
225      }
226      if (FEATURES.band) await bandTurn($, t)
227      if (FEATURES.guard) await guardTurn($, t, u.model)
228    }
229    return next(e)
230  })
231
232  on('command.run', async ($, e, next) => {
233    if (!FEATURES.guard || e.command !== 'guard') return next(e)
234    const arg = e.args.trim().toLowerCase()
235    if (arg === 'off') isOff = true
236    else if (arg === 'on') isOff = false
237    else if (/^\d{1,3}$/.test(arg)) {
238      threshold = Math.min(100, Math.max(1, Number(arg)))
239      await $.store.set('threshold', threshold)
240    } else if (arg !== '') {
241      return { text: 'Использование: /guard, /guard on, /guard off, /guard <процент>' }
242    }
243    showStatus($)
244    const state = isOff ? 'выключен на эту сессию' : isBraked() ? `ТОРМОЗ (${why()})` : 'следит'
245    return { text: `burnrate guard: ${state}. Порог 5ч: ${threshold}% (мягкий режим с ${threshold - SOFT_MARGIN}%). Сейчас 5ч ${fiveHour()}%, 7д ${weekly()}%.` }
246  })
247
248  on('prompt.submit', async ($, e, next) => {
249    if (!FEATURES.guard) return next(e)
250    if (isBraked()) {
251      const reset = untilReset(window('five_hour')?.resetsAt, await $.clock.now())
252      return next({
253        ...e,
254        context: [
255          ...(e.context ?? []),
256          `[burnrate] ${why()}${reset && reset !== 'now' ? `, сброс через ${reset}` : ''}. Отвечай максимально коротко: без вступлений, пересказов и итогов — только суть и код. ` +
257            'Не запускай субагентов и не сканируй проект целиком; читай только нужные куски файлов.',
258        ],
259      })
260    }
261    if (isSoft()) {
262      return next({
263        ...e,
264        context: [...(e.context ?? []), `[burnrate] Лимит 5ч на ${fiveHour()}%. Будь лаконичен, не трать токены на лишнее.`],
265      })
266    }
267    return next(e)
268  })
269
270  on('tool.call', ($, e, next) => {
271    if (!FEATURES.guard || !isBraked()) return next(e)
272    if (SUBAGENT_TOOLS.has(e.tool)) return deny('субагенты')
273    if (e.tool === 'Glob' && (!e.path || ROOTISH.test(e.path)) && /^\*\*/.test(e.pattern)) {
274      return deny('широкие сканы (Glob по всему проекту)')
275    }
276    if (e.tool === 'Grep' && (!e.path || ROOTISH.test(e.path)) && !e.glob && !e.type) {
277      return deny('широкие сканы (Grep без пути и фильтра)')
278    }
279    if ((e.tool === 'Bash' || e.tool === 'PowerShell') && WIDE_SHELL.test(e.command)) {
280      return deny('рекурсивные обходы файлов в шелле')
281    }
282    return next(e)
283  })
284
285  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
286    const s = await read($, snap)
287    if (!FEATURES.band || e.props.hasSurvey || s === null) {
288      return next(e)
289    }
290
291    const els = $.ui.resolve(e)
292    const { Box, Button, Text } = els
293    const Svg = 'Svg' in els ? els.Svg : null
294    const now = await $.clock.now()
295    const expanded = await read($, isExpanded)
296    const turn = await read($, lastTurn)
297    const sum = await read($, total)
298    const cats = await read($, categories)
299
300    // A meter: an SVG pill where the surface draws SVG, a thin text bar otherwise.
301    const meter = (id: string, pct: number, color: string, width: number) => {
302      if (Svg) {
303        return <Svg key={`m-${id}`} source={svgMeter(pct, width, 6, color)} alt={`${Math.round(pct)}%`} width={width} height={6} />
304      }
305      const cells = Math.round(width / 8)
306      const lit = Math.max(0, Math.min(cells, Math.round((pct / 100) * cells)))
307      return (
308        <Box key={`m-${id}`}>
309          <Text color={color}>{'━'.repeat(lit)}</Text>
310          <Text dimColor>{'━'.repeat(cells - lit)}</Text>
311        </Box>
312      )
313    }
314
315    const stat = (id: string, label: string, pct: number, value: string, hint: string, color: string) => (
316      <Box key={id} alignItems="center" gap={1}>
317        <Text dimColor>{label}</Text>
318        {meter(id, pct, color, 72)}
319        <Text bold color={pct >= 50 ? color : undefined}>
320          {value}
321        </Text>
322        {hint ? <Text dimColor>{hint}</Text> : null}
323      </Box>
324    )
325
326    const limitStats = s.limits.map(l =>
327      stat(`lim-${l.kind}`, LABEL[l.kind] ?? l.kind, l.percentUsed, `${l.percentUsed}%`, untilReset(l.resetsAt, now), tone(l.percentUsed)),
328    )
329
330    const ctxPct = s.ctxPercent ?? 0
331    const ctxStat = stat(
332      'ctx',
333      'ctx',
334      ctxPct,
335      `${ctxPct}%`,
336      s.ctxTokens !== undefined ? `${fmt(s.ctxTokens)}/${fmt(s.ctxWindow)}` : `—/${fmt(s.ctxWindow)}`,
337      ctxPct >= 80 ? DANGER : ctxPct >= 60 ? WARN : '#6c9ef8',
338    )
339
340    const allIn = sum.input + sum.cacheRead + sum.cacheWrite
341    const hit = allIn > 0 ? Math.round((sum.cacheRead / allIn) * 100) : null
342
343    const header = (
344      <Box key="head" alignItems="center" gap={3} flexWrap="wrap">
345        {limitStats.length > 0 ? limitStats : <Text dimColor>limits appear after the first reply</Text>}
346        {ctxStat}
347        {hit !== null ? (
348          <Box key="hit" gap={1}>
349            <Text dimColor>cache</Text>
350            <Text bold>{hit}%</Text>
351          </Box>
352        ) : null}
353        {s.usd !== undefined ? <Text dimColor>${s.usd.toFixed(2)}</Text> : null}
354        <Button key="toggle" label={expanded ? '▴' : '▾'} onPress={() => update($, isExpanded, v => !v)} />
355      </Box>
356    )
357
358    if (!expanded) {
359      return <Box flexDirection="column">{header}</Box>
360    }
361
362    const used = cats.reduce((a, c) => a + c.tokens, 0)
363    const parts = cats.map((c, i) => ({ ...c, color: HUES[i % HUES.length] as string }))
364
365    const composition =
366      parts.length > 0 ? (
367        <Box key="comp" flexDirection="column" gap={0} marginTop={1}>
368          {Svg ? (
369            <Svg key="stack" source={svgStack(parts, Math.max(used, 1), 480, 8)} alt="context composition" width={480} height={8} />
370          ) : (
371            <Box key="stack">
372              {parts.map(p => (
373                <Text key={`seg-${p.name}`} color={p.color}>
374                  {'█'.repeat(Math.max(1, Math.round((p.tokens / Math.max(used, 1)) * 60)))}
375                </Text>
376              ))}
377            </Box>
378          )}
379          <Box key="legend" gap={2} flexWrap="wrap">
380            {parts.map(p => (
381              <Box key={`lg-${p.name}`} gap={1}>
382                <Text color={p.color}>●</Text>
383                <Text dimColor>{p.name}</Text>
384                <Text>{fmt(p.tokens)}</Text>
385              </Box>
386            ))}
387          </Box>
388        </Box>
389      ) : null
390
391    const turnRow = turn ? (
392      <Box key="turn" gap={2} flexWrap="wrap">
393        <Text dimColor>last turn</Text>
394        <Box key="t-read" gap={1}>
395          <Text dimColor>cache read</Text>
396          <Text>{fmt(turn.cacheRead)}</Text>
397        </Box>
398        <Box key="t-write" gap={1}>
399          <Text dimColor>cache write</Text>
400          <Text color={turn.cacheWrite > turn.cacheRead ? WARN : undefined}>{fmt(turn.cacheWrite)}</Text>
401        </Box>
402        <Box key="t-in" gap={1}>
403          <Text dimColor>uncached</Text>
404          <Text>{fmt(turn.input)}</Text>
405        </Box>
406        <Box key="t-out" gap={1}>
407          <Text dimColor>output</Text>
408          <Text>{fmt(turn.output)}</Text>
409        </Box>
410      </Box>
411    ) : null
412
413    return (
414      <Box flexDirection="column">
415        {header}
416        {composition}
417        {turnRow}
418      </Box>
419    )
420  })
421}
422

1export type Limit = { kind: string; percentUsed: number; resetsAt?: string }
2export type Tokens = { input: number; cacheRead: number; cacheWrite: number; output: number }
3export type Category = { name: string; tokens: number }
4export type Snapshot = {
5  limits: Limit[]
6  ctxTokens?: number
7  ctxWindow: number
8  ctxPercent?: number
9  usd?: number
10}
11
12declare module 'claude-code' {
13  interface PluginState {
14    'limit-guard': {
15      snap: Snapshot | null
16      lastTurn: Tokens | null
17      total: Tokens
18      categories: Category[]
19      isExpanded: boolean
20    }
21  }
22}
23