session-autocompact

1import type { EngineInterface, Register } from 'claude-code'
2
3// Auto-compact for one session while the global `autoCompactEnabled` stays off. A mod
4// cannot compact mid-turn ($.session.compact rejects while a turn runs), so the check
5// runs when a main-loop turn ends. A turn that grows from below the threshold past the
6// whole window still hits the context limit; the compaction then runs when it ends.
7// The setting is keyed by session id in $.store so it survives `claude -r`; /clear and
8// --fork-session start a new id, which starts off.
9
10export type Threshold = { percent: number } | { tokens: number }
11
12const COMMAND = 'session-autocompact'
13const DEFAULT: Threshold = { percent: 80 }
14
15let compacting = false
16let turnRunning = false
17
18export function parse(arg: string): Threshold | 'off' | 'status' | undefined {
19  const a = arg.trim().toLowerCase()
20  if (a === '' || a === 'status') return 'status'
21  if (a === 'off') return 'off'
22  if (a === 'on') return DEFAULT
23  const m = /^(\d+(?:\.\d+)?)\s*(%|k|m)?$/.exec(a)
24  if (!m) return undefined
25  const n = Number(m[1])
26  const unit = m[2]
27  if (unit === '%' || (unit === undefined && n <= 100)) return n > 0 && n < 100 ? { percent: n } : undefined
28  const tokens = Math.round(n * (unit === 'k' ? 1e3 : unit === 'm' ? 1e6 : 1))
29  return tokens >= 1000 ? { tokens } : undefined
30}
31
32export function fmtTokens(n: number): string {
33  return n >= 1e6 ? `${+(n / 1e6).toFixed(2)}M` : `${Math.round(n / 1e3)}k`
34}
35
36export const label = (t: Threshold) => ('percent' in t ? `${t.percent}%` : fmtTokens(t.tokens))
37
38export function isOver(t: Threshold, tokens: number, window: number): boolean {
39  return 'percent' in t ? (tokens / window) * 100 >= t.percent : tokens >= t.tokens
40}
41
42const storeKey = async ($: EngineInterface) => `on:${await $.session.id()}`
43
44async function load($: EngineInterface): Promise<Threshold | undefined> {
45  const v = (await $.store.get(await storeKey($))) as Threshold | undefined
46  return v && ('percent' in v || 'tokens' in v) ? v : undefined
47}
48
49function showStatus($: EngineInterface, t: Threshold | undefined): void {
50  $.ui.status(t ? `on (≥${label(t)})` : undefined)
51}
52
53// One attempt per turn end: a failed compaction on a near-full window is expensive, and the
54// next turn end tries again anyway.
55function schedule($: EngineInterface): void {
56  $.clock.after(200, () => void check($))
57}
58
59async function check($: EngineInterface): Promise<void> {
60  if (compacting || turnRunning) return
61  const t = await load($)
62  if (!t) return
63  const { context } = await $.session.usage()
64  if (context.tokens === undefined || !isOver(t, context.tokens, context.window)) return
65  const pct = Math.round((context.tokens / context.window) * 100)
66  $.ui.log(`context at ${fmtTokens(context.tokens)} (${pct}%), past ${label(t)}; compacting`)
67  compacting = true
68  try {
69    const r = await $.session.compact()
70    if (r.skip !== undefined) $.ui.log(`compaction skipped: ${r.skip}`)
71  } catch (err) {
72    $.ui.log(`compaction failed: ${err instanceof Error ? err.message : String(err)}`)
73  } finally {
74    compacting = false
75  }
76}
77
78async function run($: EngineInterface, args: string): Promise<string> {
79  const parsed = parse(args)
80  const usage = `/${COMMAND} on (≥${label(DEFAULT)}) | off | <percent> (70, 70%) | <tokens> (400k, 1m)`
81  if (parsed === undefined) return `Couldn't parse '${args.trim()}'. Usage: ${usage}`
82  const key = await storeKey($)
83  if (parsed === 'off') {
84    await $.store.delete(key)
85    showStatus($, undefined)
86    return `off for this session.`
87  }
88  const { context } = await $.session.usage()
89  const now = context.tokens === undefined ? '' : ` (now ${fmtTokens(context.tokens)}, ${Math.round((context.tokens / context.window) * 100)}%)`
90  if (parsed === 'status') {
91    const t = await load($)
92    return t
93      ? `on, compacts when a turn ends with context ≥ ${label(t)}${now}.`
94      : `off for this session${now}. Usage: ${usage}`
95  }
96  await $.store.set(key, parsed)
97  showStatus($, parsed)
98  const isPast = context.tokens !== undefined && isOver(parsed, context.tokens, context.window)
99  if (isPast) schedule($)
100  return `on, compacts when a turn ends with context ≥ ${label(parsed)}${now}.${isPast ? ' Already past it: compacting now.' : ''}`
101}
102
103export const register: Register = on => {
104  on('session.start', async ($, e, next) => {
105    await $.command.register({
106      name: COMMAND,
107      description: 'Auto-compact this session only, when a turn ends past a context threshold',
108      argumentHint: '[on|off|70%|400k]',
109    })
110    showStatus($, await load($))
111    return next(e)
112  })
113
114  on('command.run', { command: COMMAND }, async ($, e) => ({ text: await run($, e.args) }))
115
116  on('turn.start', async ($, e, next) => {
117    turnRunning = true
118    return next(e)
119  })
120
121  on('turn.complete', async ($, e, next) => {
122    const done = await next(e)
123    if (e.agentId !== undefined) return done
124    turnRunning = false
125    showStatus($, await load($))
126    if (!e.isAborted) schedule($)
127    return done
128  })
129}
130