better-compaction

1import type { Register, TurnStepInput } from 'claude-code'
2
3export const register: Register = (on, options) => {
4  const defer = options.deferCompaction !== false
5  // Never defer once the context is this full (% of the model's window): past it,
6  // waiting risks a prompt-too-long error, so the compaction runs mid-turn.
7  const deferUnderPercent = typeof options.deferUnderPercent === 'number' ? options.deferUnderPercent : 80
8  const effort = typeof options.compactionEffort === 'string' ? options.compactionEffort : 'inherit'
9  const overrideEffort = effort !== 'inherit'
10
11  let midTurn = false // the main loop has sent a request this turn
12  let lastStepFailed = false // the main loop's last request got no response (e.g. prompt too long)
13  let pending = false // an auto-compaction was deferred this turn
14  let saved: { value: string | undefined } | undefined // CLAUDE_CODE_EFFORT_LEVEL before the override
15  let turnEffort: TurnStepInput['effort'] // the conversation's own effort, last seen
16  // Each running subagent's own effort, from its steps outside a summary.
17  // ponytail: a subagent whose first step falls inside a summary gets the
18  // conversation's effort, which is wrong only if its definition sets another.
19  const agentEffort = new Map<string, TurnStepInput['effort']>()
20
21  // Every compaction path (threshold, prompt-too-long, /compact, a plugin's)
22  // runs the classic PreCompact hooks right before its summarizer request.
23  on('classic.PreCompact', async ($, e, next) => {
24    if (e.agent_id) return next(e) // a subagent's own transcript: leave it alone
25
26    if (defer && e.trigger === 'auto' && midTurn && !lastStepFailed) {
27      const { context } = await $.session.usage()
28      if ((context.percent ?? 0) < deferUnderPercent) {
29        pending = true
30        return { block: 'better-compaction: compacting when this turn ends' }
31      }
32    }
33
34    // The summarizer reads effort when it builds its request; the env var
35    // outranks the session's setting. Same model, so the cache still hits.
36    // It is process-wide; turn.step hands subagents their own effort back.
37    if (overrideEffort && !saved) {
38      saved = { value: await $.env.get('CLAUDE_CODE_EFFORT_LEVEL') }
39      await $.env.set('CLAUDE_CODE_EFFORT_LEVEL', effort)
40    }
41    return next(e)
42  })
43
44  on('classic.PostCompact', async ($, e, next) => {
45    if (saved && !e.agent_id) {
46      await $.env.set('CLAUDE_CODE_EFFORT_LEVEL', saved.value)
47      saved = undefined
48    }
49    return next(e)
50  })
51
52  // A precomputed summary is made in the background while the turn runs, where
53  // the effort override would leak into the turn's own requests. Skip it; the
54  // summary is then made when compaction happens, at the chosen effort.
55  on('session.compact', ($, e, next) => {
56    if (overrideEffort && e.trigger === 'precompute') {
57      return { skip: 'better-compaction: summaries run at their own effort, not ahead of time' }
58    }
59    return next(e)
60  })
61
62  on('turn.step', async function* ($, e, next) {
63    if (e.agentId) {
64      // A subagent stepping while a summary runs would read the override too;
65      // a step's own effort outranks the env var, so give it back its own.
66      if (!saved) agentEffort.set(e.agentId, e.effort)
67      else e = { ...e, effort: agentEffort.has(e.agentId) ? agentEffort.get(e.agentId) : turnEffort }
68      return yield* next(e)
69    }
70    if (saved) {
71      // A compaction ended without PostCompact (it failed, or another hook
72      // blocked it): put the effort back before this request goes out.
73      await $.env.set('CLAUDE_CODE_EFFORT_LEVEL', saved.value)
74      saved = undefined
75      if (turnEffort !== undefined) e = { ...e, effort: turnEffort }
76    }
77    turnEffort = e.effort
78    midTurn = true
79    const r = yield* next(e)
80    lastStepFailed = r.stopReason === null
81    return r
82  })
83
84  on('turn.complete', async ($, e, next) => {
85    const r = await next(e)
86    if (e.agentId) {
87      agentEffort.delete(e.agentId)
88      return r
89    }
90    midTurn = false
91    lastStepFailed = false
92    if (pending) {
93      pending = false
94      // $.session.compact refuses while the turn is still held; this hook
95      // returning releases it. Compacting now, not at the next prompt, keeps
96      // the prompt cache warm. Where it's refused (headless -p / SDK), the
97      // next prompt's own auto-compaction check picks it up instead.
98      $.clock.after(0, () => {
99        $.session.compact().catch((err) =>
100          $.ui.log(`better-compaction: deferred compaction left to the next prompt: ${err}`, { to: 'debug' }),
101        )
102      })
103    }
104    return r
105  })
106}
107