context-budget

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, PluginOptions, Register } from 'claude-code'
3
4import type { LastActivity, LogEntry } from '../types'
5import { MESSAGES, detectLocale } from './i18n'
6import type { Locale, Messages } from './i18n'
7
8type Engine = EngineInterface
9type Model = 'haiku' | 'sonnet' | 'opus'
10type Effort = 'low' | 'medium' | 'high' | 'xhigh' | 'max'
11type Language = 'auto' | Locale
12
13type Config = {
14  language: Language
15  firstThreshold: number
16  thresholdStep: number
17  idleMs: number
18  idleMinTokens: number
19  agentModels: Readonly<Record<string, Model>>
20  effortFloor: Effort | undefined
21}
22
23const LANGUAGES: readonly Language[] = ['auto', 'en', 'zh-TW']
24const MODELS: readonly Model[] = ['haiku', 'sonnet', 'opus']
25const EFFORTS: readonly Effort[] = ['low', 'medium', 'high', 'xhigh', 'max']
26// [subagent type, userConfig key, default model]
27const AGENT_MODEL_OPTIONS = [
28  ['Explore', 'explore_model', 'haiku'],
29  ['claude-code-guide', 'guide_model', 'sonnet'],
30  ['Plan', 'plan_model', 'opus'],
31] as const
32// A drop below 60% of the previous reading means the session was compacted or cleared.
33const RESET_RATIO = 0.6
34const HISTORY_LENGTH = 12
35const LOG_LIMIT = 200
36const SESSION_KEYS_LIMIT = 100
37const SEEN_AGENTS_LIMIT = 200
38const seenAgents = new Set<string>()
39const BROWSER_TOOL = /playwright|chrome-devtools/i
40const BARS = '▁▂▃▄▅▆▇█'
41
42const reading = atom({ plugin: 'context-budget', key: 'reading' } as const, null)
43const history = atom({ plugin: 'context-budget', key: 'history' } as const, [])
44const alert = atom({ plugin: 'context-budget', key: 'alert' } as const, null)
45const nextThreshold = atom({ plugin: 'context-budget', key: 'nextThreshold' } as const, null)
46const stashedDraft = atom({ plugin: 'context-budget', key: 'stashedDraft' } as const, null)
47const localeChoice = atom({ plugin: 'context-budget', key: 'locale' } as const, null)
48
49function numberOption(value: unknown, fallback: number, min: number): number {
50  return typeof value === 'number' && Number.isFinite(value) && value >= min ? value : fallback
51}
52
53function choiceOption<T extends string>(value: unknown, choices: readonly T[]): T | undefined {
54  return choices.find(choice => choice === value)
55}
56
57function readConfig(options: PluginOptions): Config {
58  const agentModels: Record<string, Model> = {}
59  for (const [subagentType, key, fallback] of AGENT_MODEL_OPTIONS) {
60    const model = options[key] === 'inherit' ? undefined : (choiceOption(options[key], MODELS) ?? fallback)
61    if (model !== undefined) {
62      agentModels[subagentType] = model
63    }
64  }
65  const floor = options.subagent_effort_floor
66  return {
67    language: choiceOption(options.language, LANGUAGES) ?? 'auto',
68    firstThreshold: numberOption(options.first_threshold, 500_000, 50_000),
69    thresholdStep: numberOption(options.threshold_step, 150_000, 10_000),
70    idleMs: numberOption(options.idle_minutes, 60, 1) * 60_000,
71    idleMinTokens: numberOption(options.idle_min_tokens, 200_000, 0),
72    agentModels,
73    effortFloor: floor === 'off' ? undefined : (choiceOption(floor, EFFORTS) ?? 'high'),
74  }
75}
76
77let config = readConfig({})
78
79async function strings($: Engine): Promise<Messages> {
80  const cached = await read($, localeChoice)
81  if (cached !== null && cached.language === config.language) {
82    return MESSAGES[cached.locale]
83  }
84  const locale = config.language === 'auto' ? detectLocale((await $.settings.read()).language) : config.language
85  await update($, localeChoice, () => ({ language: config.language, locale }))
86  return MESSAGES[locale]
87}
88
89function firstThreshold(stored: number | null): number {
90  return stored ?? config.firstThreshold
91}
92
93// The top-level effortLevel setting does not reach 5.5-generation subagents, so they run at the model's default (medium).
94function raisedEffort(effort: string | number | undefined): Effort | undefined {
95  const floor = config.effortFloor
96  const level = EFFORTS.findIndex(candidate => candidate === effort)
97  if (floor === undefined || level === -1) {
98    return undefined
99  }
100  return level < EFFORTS.indexOf(floor) ? floor : undefined
101}
102
103const k = (tokens: number) => `${Math.round(tokens / 1000)}k`
104
105function spark(values: readonly number[], window: number): string {
106  return values
107    .map(value => BARS[Math.min(BARS.length - 1, Math.floor((value / window) * BARS.length))])
108    .join('')
109}
110
111async function log($: Engine, kind: string, detail: string) {
112  const entries = ((await $.store.get('log')) as LogEntry[] | undefined) ?? []
113  const at = await $.clock.now()
114  await $.store.set('log', [...entries, { at, kind, detail }].slice(-LOG_LIMIT))
115}
116
117async function pruneStore($: Engine) {
118  const keys = await $.store.keys()
119  const sessionKeys = keys.filter(key => key.startsWith('last:') || key.startsWith('ack:'))
120  const excess = sessionKeys.length - SESSION_KEYS_LIMIT
121  for (const key of sessionKeys.slice(0, Math.max(0, excess))) {
122    await $.store.delete(key)
123  }
124}
125
126async function instructions($: Engine): Promise<string> {
127  const t = await strings($)
128  const messages = await $.session.messages()
129  const usedBrowser = messages.some(message =>
130    message.toolUses.some(use => BROWSER_TOOL.test(use.tool)),
131  )
132  return usedBrowser ? `${t.baseInstructions} ${t.browserInstructions}` : t.baseInstructions
133}
134
135async function showStatus($: Engine) {
136  const current = await read($, reading)
137  if (current === null) {
138    $.ui.status(undefined)
139    return
140  }
141  const values = await read($, history)
142  $.ui.status(`ctx ${k(current.tokens)}/${k(current.window)} ${spark(values, current.window)}`)
143}
144
145async function resetBudget($: Engine) {
146  await update($, nextThreshold, () => null)
147  await update($, alert, () => null)
148}
149
150// A draft in the prompt (such as the message the idle check put back) is stashed and restored after the manual compact.
151async function fillCompact($: Engine, text: string): Promise<boolean> {
152  const draft = (await $.prompt.read()).text
153  const isDraft = draft.trim() !== '' && !draft.trimStart().startsWith('/compact')
154  const filled = await $.prompt.fill({ text: `/compact ${text}` })
155  if (filled.isFilled && isDraft) {
156    await update($, stashedDraft, () => draft)
157    $.ui.toast((await strings($)).draftStashed)
158  }
159  return filled.isFilled
160}
161
162async function restoreDraft($: Engine) {
163  const draft = await read($, stashedDraft)
164  if (draft === null) {
165    return
166  }
167  await update($, stashedDraft, () => null)
168  if ((await $.prompt.read()).text.trim() === '') {
169    await $.prompt.fill({ text: draft })
170  }
171}
172
173async function fillTailored($: Engine): Promise<boolean> {
174  const t = await strings($)
175  $.ui.toast(t.tailoring)
176  const reply = await $.model.fork({ prompt: t.forkPrompt })
177  if (!reply.isAnswered) {
178    $.ui.toast(t.tailorFailed(reply.reason))
179    await log($, 'fork-failed', reply.reason)
180    return fillCompact($, await instructions($))
181  }
182  await log($, 'fork', 'tailored instructions filled')
183  return fillCompact($, reply.text.replace(/\s+/g, ' ').trim())
184}
185
186// usage().context.tokens is the input of the previous response and changes only on the next turn,
187// so the reading after a compact comes from its tokensAfter.
188async function recordCompaction($: Engine, tokensAfter: number | undefined) {
189  const id = await $.session.id()
190  if (tokensAfter === undefined) {
191    // With no new size, drop the old reading so the idle check never uses the pre-compact size.
192    await $.store.delete(`last:${id}`)
193    await update($, reading, () => null)
194    $.ui.status((await strings($)).compactedStatus)
195    return
196  }
197  const now = await $.clock.now()
198  const last: LastActivity = { at: now, tokens: tokensAfter }
199  await $.store.set(`last:${id}`, last)
200  const window = (await $.session.usage()).context.window
201  await update($, reading, () => ({ tokens: tokensAfter, window, at: now }))
202  await update($, history, values => [...values, tokensAfter].slice(-HISTORY_LENGTH))
203  await showStatus($)
204}
205
206async function takeReading($: Engine) {
207  const usage = await $.session.usage()
208  const tokens = usage.context.tokens
209  if (tokens === undefined) {
210    return
211  }
212  const now = await $.clock.now()
213  const previous = await read($, reading)
214  const last: LastActivity = { at: now, tokens }
215  await $.store.set(`last:${await $.session.id()}`, last)
216  await update($, reading, () => ({ tokens, window: usage.context.window, at: now }))
217  await update($, history, values => [...values, tokens].slice(-HISTORY_LENGTH))
218
219  if (previous !== null && tokens < previous.tokens * RESET_RATIO) {
220    await resetBudget($)
221  }
222
223  const threshold = firstThreshold(await read($, nextThreshold))
224  if (tokens >= threshold) {
225    let next = threshold
226    while (next <= tokens) {
227      next += config.thresholdStep
228    }
229    await update($, nextThreshold, () => next)
230    await update($, alert, () => ({ kind: 'threshold' as const, tokens }))
231    $.ui.toast((await strings($)).thresholdToast(k(tokens)))
232    await log($, 'threshold', `${k(tokens)}, next ${k(next)}`)
233  }
234  await showStatus($)
235}
236
237export const register: Register = (on, options) => {
238  config = readConfig(options)
239
240  on('session.start', async ($, e, next) => {
241    // Pick up a language setting that changed since the last session.
242    await update($, localeChoice, () => null)
243    await $.command.register({
244      name: 'context-budget',
245      description: (await strings($)).commandDescription,
246      argumentHint: '[compact|tailor]',
247    })
248    await pruneStore($)
249    await showStatus($)
250    return next(e)
251  })
252
253  on('turn.complete', async ($, e, next) => {
254    const result = await next(e)
255    if (e.agentId === undefined) {
256      await takeReading($)
257    }
258    return result
259  })
260
261  on('session.compact', async ($, e, next) => {
262    const result = await next(e)
263    // precompute only prepares a summary and keeps the conversation; a subagent's compact leaves the main session alone.
264    const isSkipped = 'skip' in result && typeof result.skip === 'string'
265    if (isSkipped || e.trigger === 'precompute' || e.agentId !== undefined) {
266      return result
267    }
268    const tokensAfter = 'tokensAfter' in result ? result.tokensAfter : undefined
269    await resetBudget($)
270    await recordCompaction($, tokensAfter)
271    await log($, 'compacted', `${e.trigger}, ${tokensAfter === undefined ? 'size unknown' : k(tokensAfter)}`)
272    if (e.trigger === 'manual') {
273      await restoreDraft($)
274    }
275    return result
276  })
277
278  on('prompt.submit', async ($, e, next) => {
279    // Sending anything other than /compact abandons the filled command, so the stashed draft is dropped.
280    if (e.origin.kind === 'composer' && !e.text.trimStart().startsWith('/compact')) {
281      await update($, stashedDraft, () => null)
282    }
283    const isTypedWhileIdle = e.origin.kind === 'composer' && e.turnId === undefined
284    const isCommand = e.text.trimStart().startsWith('/')
285    const hasAttachments = (e.attachments?.length ?? 0) > 0
286    if (!isTypedWhileIdle || isCommand || hasAttachments) {
287      return next(e)
288    }
289
290    const id = await $.session.id()
291    const last = (await $.store.get(`last:${id}`)) as LastActivity | undefined
292    if (last === undefined) {
293      return next(e)
294    }
295    const idleMs = (await $.clock.now()) - last.at
296    // usage() still reports the old size after a compact; `last` is updated by the compact.
297    const tokens = last.tokens
298    const isAcknowledged = (await $.store.get(`ack:${id}`)) === last.at
299    if (idleMs < config.idleMs || tokens < config.idleMinTokens || isAcknowledged) {
300      return next(e)
301    }
302
303    const minutes = Math.round(idleMs / 60_000)
304    await $.store.set(`ack:${id}`, last.at)
305    await update($, alert, () => ({ kind: 'idle' as const, tokens, minutes }))
306    await $.prompt.fill({ text: e.text })
307    await log($, 'idle-drop', `${minutes}m, ${k(tokens)}`)
308    return { drop: (await strings($)).idleDrop(minutes, k(tokens), k(tokens * 2)) }
309  })
310
311  on('agent.spawn', async ($, e, next) => {
312    const model = config.agentModels[e.subagentType]
313    if (e.fork || e.model !== undefined || model === undefined) {
314      return next(e)
315    }
316    await log($, 'agent-model', `${e.subagentType} -> ${model}`)
317    return next({ ...e, model })
318  })
319
320  on('turn.step', async function* ($, e, next) {
321    if (e.agentId === undefined) {
322      return yield* next(e)
323    }
324    const raised = raisedEffort(e.effort)
325    // Log only each subagent's first step, to show which effort it started with.
326    if (!seenAgents.has(e.agentId)) {
327      if (seenAgents.size >= SEEN_AGENTS_LIMIT) {
328        seenAgents.clear()
329      }
330      seenAgents.add(e.agentId)
331      const after = raised === undefined ? '' : ` -> ${raised}`
332      await log($, 'agent-effort', `${e.effort ?? 'unset'}${after} (${e.model})`)
333    }
334    return yield* next(raised === undefined ? e : { ...e, effort: raised })
335  })
336
337  on('command.run', { command: 'context-budget' }, async ($, e) => {
338    const t = await strings($)
339    const action = e.args.trim()
340    if (action === 'compact' || action === 'tailor') {
341      const isFilled = action === 'compact' ? await fillCompact($, await instructions($)) : await fillTailored($)
342      await log($, 'command', action)
343      return { text: isFilled ? t.filled : t.fillFailed }
344    }
345    if (action !== '') {
346      return { text: t.usage }
347    }
348    const current = await read($, reading)
349    const threshold = firstThreshold(await read($, nextThreshold))
350    const entries = (((await $.store.get('log')) as LogEntry[] | undefined) ?? []).slice(-15)
351    const now = await $.clock.now()
352    const lines = [
353      current === null ? t.noReading : t.context(k(current.tokens), k(current.window)),
354      t.nextThreshold(k(threshold)),
355      t.manualCompact,
356      t.recentEvents,
357      ...(entries.length === 0
358        ? [t.none]
359        : entries.map(entry => t.logLine(Math.round((now - entry.at) / 60_000), entry.kind, entry.detail))),
360    ]
361    return { text: lines.join('\n') }
362  })
363
364  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
365    const active = await read($, alert)
366    if (active === null || e.props.hasSurvey) {
367      return next(e)
368    }
369    const t = await strings($)
370    const { Box, Button, Text } = $.ui.resolve(e)
371    const dismiss = () => update($, alert, () => null)
372    const isIdle = active.kind === 'idle'
373    const message = isIdle ? t.idleBand(active.minutes, k(active.tokens)) : t.thresholdBand(k(active.tokens))
374    const buttons = [
375      <Button
376        key="compact"
377        label={t.compactLabel}
378        variant="primary"
379        onPress={async () => {
380          await fillCompact($, await instructions($))
381        }}
382      />,
383      <Button key="tailor" label={t.tailorLabel} onPress={() => fillTailored($)} />,
384      ...(isIdle
385        ? [
386            <Button
387              key="clear"
388              label={t.clearLabel}
389              onPress={async () => {
390                await $.command.run({ command: 'clear' })
391                await dismiss()
392              }}
393            />,
394          ]
395        : []),
396      <Button key="dismiss" label={t.dismissLabel} role="dismiss" onPress={dismiss} />,
397    ]
398
399    return (
400      <Box flexDirection="column">
401        <Text color="yellow" wrap="truncate-end">
402          {message}
403        </Text>
404        <Box gap={1}>{buttons}</Box>
405      </Box>
406    )
407  })
408}
409

1export type Locale = 'en' | 'zh-TW'
2
3export type Messages = {
4  commandDescription: string
5  usage: string
6  filled: string
7  fillFailed: string
8  noReading: string
9  context: (tokens: string, window: string) => string
10  nextThreshold: (tokens: string) => string
11  manualCompact: string
12  recentEvents: string
13  none: string
14  logLine: (minutes: number, kind: string, detail: string) => string
15  thresholdToast: (tokens: string) => string
16  draftStashed: string
17  tailoring: string
18  tailorFailed: (reason: string) => string
19  compactedStatus: string
20  idleDrop: (minutes: number, tokens: string, rewrite: string) => string
21  idleBand: (minutes: number, tokens: string) => string
22  thresholdBand: (tokens: string) => string
23  compactLabel: string
24  tailorLabel: string
25  clearLabel: string
26  dismissLabel: string
27  baseInstructions: string
28  browserInstructions: string
29  forkPrompt: string
30}
31
32// Matches the free-form `language` setting ("Chinese (Traditional, Taiwan)", "zh-TW", "中文", ...).
33const CHINESE = /chinese|mandarin|cantonese|taiwan|hong kong|\bzh\b|中文|漢語|汉语|國語|国语|繁體|简体/i
34
35export function detectLocale(language: unknown): Locale {
36  return typeof language === 'string' && CHINESE.test(language) ? 'zh-TW' : 'en'
37}
38
39export const MESSAGES: Readonly<Record<Locale, Messages>> = {
40  en: {
41    commandDescription: 'Show context-budget readings and log; add compact or tailor to put a /compact command in the prompt',
42    usage: 'Usage: /context-budget [compact|tailor]',
43    filled: '/compact is in the prompt. Review it and press Enter to send.',
44    fillFailed: 'Could not fill the prompt (a dialog may be open). Try again in a moment.',
45    noReading: 'context: no reading yet (waiting for the first turn to finish)',
46    context: (tokens, window) => `context: ${tokens} / ${window}`,
47    nextThreshold: tokens => `next reminder at: ${tokens}`,
48    manualCompact: 'manual compact: /context-budget compact (default instructions) | /context-budget tailor (tailored instructions)',
49    recentEvents: 'recent events:',
50    none: '(none)',
51    logLine: (minutes, kind, detail) => `- ${minutes} min ago ${kind} ${detail}`,
52    thresholdToast: tokens => `Context is at ${tokens}. Use the buttons above the prompt to compact.`,
53    draftStashed: 'Your draft is set aside and goes back in the prompt after the compact.',
54    tailoring: 'Writing compact instructions for this conversation…',
55    tailorFailed: reason => `Could not write tailored instructions (${reason}), so the default ones are in the prompt.`,
56    compactedStatus: 'ctx compacted, updates next turn',
57    idleDrop: (minutes, tokens, rewrite) =>
58      `Idle ${minutes} min with about ${tokens} of context. The prompt cache has expired, so sending now writes the whole context again (about ${rewrite} of input). Your message is back in the prompt: press Enter to send it anyway, or use Compact, Tailored or Clear above the prompt first.`,
59    idleBand: (minutes, tokens) => `☂ Idle ${minutes} min, context ${tokens}: your message is back in the prompt, press Enter to send it`,
60    thresholdBand: tokens => `↯ context ${tokens}, consider compacting`,
61    compactLabel: 'Compact',
62    tailorLabel: 'Tailored',
63    clearLabel: 'Clear',
64    dismissLabel: 'Dismiss',
65    baseInstructions:
66      'Keep: the current task goal, what is done and what is left, important decisions and why, relevant file paths and commands. Drop: details of errors already fixed, long tool output.',
67    browserInstructions:
68      'Browser: drop snapshots, console logs, screenshots and DOM details. Keep only selectors that work, verification results and code logic still to handle.',
69    forkPrompt:
70      'Write keep-instructions for the next /compact: English, one line, at most 60 words, no preamble. List what the rest of the work has to keep (task goal, open items, decisions, file paths) and the noise that can go. Output only the instructions.',
71  },
72  'zh-TW': {
73    commandDescription: '顯示 context-budget 的讀數與紀錄；加 compact／tailor 把 /compact 指令帶進輸入框',
74    usage: '用法：/context-budget [compact|tailor]',
75    filled: '已把 /compact 指令放進輸入框，確認後按 Enter 送出',
76    fillFailed: '輸入框目前無法填入（可能有對話框開著），稍後再試',
77    noReading: 'context：尚無讀數（等第一個回合結束）',
78    context: (tokens, window) => `context：${tokens} / ${window}`,
79    nextThreshold: tokens => `下一個提醒門檻：${tokens}`,
80    manualCompact: '手動 compact：/context-budget compact（預設指令）｜/context-budget tailor（量身指令）',
81    recentEvents: '最近的判定：',
82    none: '（無）',
83    logLine: (minutes, kind, detail) => `- ${minutes} 分鐘前 ${kind} ${detail}`,
84    thresholdToast: tokens => `context 已到 ${tokens}，可以用輸入框上方的按鈕 compact`,
85    draftStashed: '原本的草稿先收起來，compact 完成後會放回輸入框',
86    tailoring: '正在依目前對話產生 compact 指令…',
87    tailorFailed: reason => `產生失敗（${reason}），改填通用指令`,
88    compactedStatus: 'ctx compact 完成，下一回合更新',
89    idleDrop: (minutes, tokens, rewrite) =>
90      `閒置 ${minutes} 分鐘、context 約 ${tokens}：快取已過期，直接送出會把整段重寫一次（約等於 ${rewrite} 的 input）。原文已放回輸入框，按 Enter 照常送出；或先按輸入框上方的 Compact／量身指令／Clear。`,
91    idleBand: (minutes, tokens) => `☂ 閒置 ${minutes} 分、context ${tokens}：原文已放回輸入框，按 Enter 照常送出`,
92    thresholdBand: tokens => `↯ context ${tokens}，建議 compact`,
93    compactLabel: 'Compact',
94    tailorLabel: '量身指令',
95    clearLabel: 'Clear',
96    dismissLabel: '略過',
97    baseInstructions:
98      '保留：目前任務目標、已完成與待辦事項、做過的決定與理由、相關檔案路徑與指令。可丟棄：已解決錯誤的細節、冗長的工具輸出。',
99    browserInstructions:
100      '瀏覽器：丟棄 snapshot、console log、截圖與 DOM 細節，只保留可用的 selector、驗證結論與待處理的程式邏輯。',
101    forkPrompt:
102      '請替接下來的 /compact 寫一段保留指示：繁體中文、單行、120 字以內、不要前言。列出接下來工作一定要保留的重點（任務目標、待辦、做過的決定、檔案路徑），以及可以丟棄的雜訊。只輸出指示本身。',
103  },
104}
105

1/** The context reading taken when the main loop's last turn ended. */
2export type Reading = { tokens: number; window: number; at: number }
3
4/** The reminder the band above the prompt shows; null hides it. */
5export type Alert =
6  | { kind: 'threshold'; tokens: number }
7  | { kind: 'idle'; tokens: number; minutes: number }
8
9/** Last activity per session in $.store, under `last:<sessionId>`. */
10export type LastActivity = { at: number; tokens: number }
11
12/** The event log in $.store, under `log`. */
13export type LogEntry = { at: number; kind: string; detail: string }
14
15declare module 'claude-code' {
16  interface PluginState {
17    'context-budget': {
18      reading: Reading | null
19      history: number[]
20      alert: Alert | null
21      /** Token count that triggers the next reminder; null means the configured first threshold. */
22      nextThreshold: number | null
23      /** The prompt draft set aside when `/compact` was filled in, restored after the manual compact; null when there is none. */
24      stashedDraft: string | null
25      /** The interface language resolved for the `language` option it was resolved under; null until first use. */
26      locale: { language: 'auto' | 'en' | 'zh-TW'; locale: 'en' | 'zh-TW' } | null
27    }
28  }
29}
30