compact-shaper

1import type { Register } from 'claude-code'
2
3import { pruneTranscript } from './prune'
4import type { Decision, Outcome } from './prune'
5
6// A prune that saves less than this is not worth rewriting the transcript.
7const MIN_REDUCTION = 0.1
8
9export const HANDOFF_INSTRUCTIONS =
10  'Write the summary as a handoff the rest of this session continues from without asking anything. Cover, in this order: the goal; the current state, including the branch and what is done, committed or still open; decisions with their reasons; approaches that were rejected and why; open points; the exact next step; the files to read first. Quote the requirements and corrections the user gave word for word. Leave out anything that can be re-read from the code or the repository.'
11
12let preserveRecentMessages = 6
13let keepHeadChars = 300
14
15function formatTokens(tokens: number): string {
16  return tokens >= 1000 ? `${Math.round(tokens / 1000)}k` : String(tokens)
17}
18
19function reductionOf(outcome: Outcome): number {
20  return outcome.tokensBefore === 0 ? 0 : 1 - outcome.tokensAfter / outcome.tokensBefore
21}
22
23function reportOf(outcome: Outcome, verb: string): string {
24  const cut = outcome.decisions.filter(i => i.reason === 'cut' || i.reason === 'superseded')
25  const kept = outcome.decisions.filter(i => i.reason === 'protected')
26  const line = (decision: Decision) =>
27    `- ${decision.tool} ${decision.input.slice(0, 80)} (${decision.resultChars} chars, ${decision.reason})`
28  const largest = [...cut].sort((a, b) => b.resultChars - a.resultChars).slice(0, 8)
29  const lines = [
30    `prune ${verb}: ${formatTokens(outcome.tokensBefore)} → ${formatTokens(outcome.tokensAfter)} (−${Math.round(reductionOf(outcome) * 100)}%) · ${cut.length} outputs cut`,
31  ]
32  if (largest.length > 0) {
33    lines.push('Largest cuts:', ...largest.map(line))
34  }
35  if (kept.length > 0) {
36    lines.push('Protected:', ...kept.map(line))
37  }
38
39  return lines.join('\n')
40}
41
42export const register: Register = (on, options) => {
43  preserveRecentMessages = Number(options.preserveRecentMessages ?? 6)
44  keepHeadChars = Number(options.keepHeadChars ?? 300)
45
46  on('session.start', async ($, e, next) => {
47    await $.command.register({
48      name: 'prune-preview',
49      description: 'Show what /compact prune would cut from this conversation, without changing it',
50    })
51
52    return next(e)
53  })
54
55  on('command.run', { command: 'prune-preview' }, async $ => {
56    const usage = await $.session.usage()
57    const messages = await $.session.messages()
58    const outcome = pruneTranscript(messages, usage.context.tokens ?? 0, preserveRecentMessages, keepHeadChars)
59
60    return { text: reportOf(outcome, 'preview') }
61  })
62
63  // `/compact prune` cuts re-readable tool output instead of summarizing. It
64  // rides on /compact because `$.session.compact()` skips the calling plugin's
65  // own hook and may not run from a command hook at all.
66  on('session.compact', async ($, e, next) => {
67    if (e.agentId !== undefined) {
68      return next(e)
69    }
70    if (e.trigger === 'manual' && e.instructions?.trim().toLowerCase() === 'prune') {
71      const usage = await $.session.usage()
72      const outcome = pruneTranscript(e.messages, usage.context.tokens ?? 0, preserveRecentMessages, keepHeadChars)
73      if (reductionOf(outcome) < MIN_REDUCTION) {
74        return { skip: `prune: nothing worth cutting (−${Math.round(reductionOf(outcome) * 100)}%)` }
75      }
76      $.ui.toast(reportOf(outcome, 'applied').split('\n')[0] ?? '', { timeoutMs: 15_000 })
77
78      return { messages: outcome.messages, tokensBefore: outcome.tokensBefore, tokensAfter: outcome.tokensAfter }
79    }
80    const instructions = e.instructions === undefined ? HANDOFF_INSTRUCTIONS : `${HANDOFF_INSTRUCTIONS}\n\n${e.instructions}`
81
82    return next({ ...e, instructions })
83  })
84}
85

1import type { SessionMessage, ToolResultSummary, ToolUseSummary } from 'claude-code'
2
3export type Call = {
4  id: string
5  toolUseId: string
6  tool: string
7  input: Readonly<Record<string, unknown>>
8  callIndex: number
9  resultText: string
10  isError: boolean
11  isPinned: boolean
12}
13
14export type Reason = 'pinned' | 'protected' | 'superseded' | 'short' | 'cut'
15
16export type Decision = {
17  id: string
18  tool: string
19  input: string
20  resultChars: number
21  reason: Reason
22}
23
24export type Outcome = {
25  messages: SessionMessage[]
26  tokensBefore: number
27  tokensAfter: number
28  decisions: Decision[]
29}
30
31const WRITE_TOOLS = new Set(['Edit', 'Write', 'MultiEdit', 'NotebookEdit'])
32const WRITE_FIELDS = new Set(['content', 'new_string', 'old_string', 'new_source', 'edits'])
33
34// Getting these back costs more than a re-read: a subagent has to run again.
35const AGENT_TOOLS = new Set(['Agent', 'Task'])
36
37// Snapshots of state that has moved on since, and CI logs that expire.
38const UNREPEATABLE_COMMAND = /\b(git (status|stash list)|gh run view|gh pr checks)\b/
39
40export function isPinned(index: number, total: number, preserveRecentMessages: number): boolean {
41  return index === 0 || index >= total - preserveRecentMessages
42}
43
44export function collectCalls(messages: readonly SessionMessage[], preserveRecentMessages: number): Call[] {
45  const results = new Map<string, { index: number; result: ToolResultSummary }>()
46  messages.forEach((message, index) => {
47    for (const result of message.toolResults ?? []) {
48      results.set(result.tool_use_id, { index, result })
49    }
50  })
51  const calls: Call[] = []
52  messages.forEach((message, callIndex) => {
53    for (const use of message.toolUses) {
54      const found = results.get(use.tool_use_id)
55      const resultText = found?.result.text ?? use.text
56      if (resultText === undefined) {
57        continue
58      }
59      const resultIndex = found?.index ?? callIndex
60      calls.push({
61        id: `t${calls.length + 1}`,
62        toolUseId: use.tool_use_id,
63        tool: use.tool,
64        input: use.input,
65        callIndex,
66        resultText,
67        isError: found?.result.isError ?? use.isError === true,
68        isPinned:
69          isPinned(callIndex, messages.length, preserveRecentMessages) ||
70          isPinned(resultIndex, messages.length, preserveRecentMessages),
71      })
72    }
73  })
74
75  return calls
76}
77
78function pathOf(call: Call): string | undefined {
79  const path = call.input.file_path ?? call.input.notebook_path
80
81  return typeof path === 'string' ? path : undefined
82}
83
84function isFullRead(call: Call): boolean {
85  return call.tool === 'Read' && call.input.offset === undefined && call.input.limit === undefined
86}
87
88// A Read result is outdated once the same file was written, and redundant once
89// it was read again in full. The file on disk stays the source of truth.
90export function supersededReads(calls: readonly Call[]): Set<string> {
91  const superseded = new Set<string>()
92  calls.forEach((call, index) => {
93    const path = pathOf(call)
94    if (call.tool !== 'Read' || call.isPinned || path === undefined) {
95      return
96    }
97    const isReplaced = calls
98      .slice(index + 1)
99      .some(i => pathOf(i) === path && (WRITE_TOOLS.has(i.tool) || isFullRead(i)))
100    if (isReplaced) {
101      superseded.add(call.id)
102    }
103  })
104
105  return superseded
106}
107
108function shortened(text: string, headChars: number, note: string): string {
109  if (text.length <= headChars + 200) {
110    return text
111  }
112
113  return `${text.slice(0, headChars)}\n[pruned ${text.length - headChars} chars; ${note}]`
114}
115
116export function truncateResult(text: string, headChars: number, tool: string): string {
117  return shortened(text, headChars, `run ${tool} again if the content is needed`)
118}
119
120function shrinkValue(value: unknown, headChars: number): unknown {
121  if (typeof value === 'string') {
122    return shortened(value, headChars, 'the file on disk holds the current content')
123  }
124  if (Array.isArray(value)) {
125    return value.map(i => shrinkValue(i, headChars))
126  }
127  if (value !== null && typeof value === 'object') {
128    return Object.fromEntries(Object.entries(value).map(([key, inner]) => [key, shrinkValue(inner, headChars)]))
129  }
130
131  return value
132}
133
134// What an edit wrote is on disk; the transcript only needs to show that it
135// happened and roughly what it touched.
136export function shrinkWriteInput(
137  tool: string,
138  input: Readonly<Record<string, unknown>>,
139  headChars: number,
140): Record<string, unknown> | null {
141  if (!WRITE_TOOLS.has(tool)) {
142    return null
143  }
144  const shrunk = Object.fromEntries(
145    Object.entries(input).map(([key, value]) => [key, WRITE_FIELDS.has(key) ? shrinkValue(value, headChars) : value]),
146  )
147
148  return JSON.stringify(shrunk) === JSON.stringify(input) ? null : shrunk
149}
150
151function rebuiltUse(use: ToolUseSummary, text: string | undefined, input: Record<string, unknown> | null): ToolUseSummary {
152  const copy: ToolUseSummary = { tool_use_id: use.tool_use_id, tool: use.tool, input: input ?? use.input }
153  if (text !== undefined) {
154    copy.text = text
155  }
156  if (use.isError === true) {
157    copy.isError = true
158  }
159
160  return copy
161}
162
163/**
164 * Rebuilds the transcript with the dropped results cut to their head and the
165 * write inputs of unpinned calls shrunk. Messages left untouched are returned
166 * as the same objects, so they keep the engine's handle.
167 */
168export function applyDrops(
169  messages: readonly SessionMessage[],
170  calls: readonly Call[],
171  dropped: ReadonlySet<string>,
172  headChars: number,
173): SessionMessage[] {
174  const byToolUseId = new Map(calls.map(i => [i.toolUseId, i]))
175  const isDropped = (toolUseId: string) => {
176    const call = byToolUseId.get(toolUseId)
177
178    return call !== undefined && dropped.has(call.id)
179  }
180  const shrinkable = (use: ToolUseSummary) => {
181    const call = byToolUseId.get(use.tool_use_id)
182
183    return call !== undefined && !call.isPinned ? shrinkWriteInput(use.tool, use.input, headChars) : null
184  }
185
186  return messages.map(message => {
187    const uses = message.toolUses.map(use => {
188      const input = shrinkable(use)
189      const text = isDropped(use.tool_use_id) && use.text !== undefined ? truncateResult(use.text, headChars, use.tool) : use.text
190      if (input === null && text === use.text) {
191        return use
192      }
193
194      return rebuiltUse(use, text, input)
195    })
196    const results = (message.toolResults ?? []).map(result => {
197      if (!isDropped(result.tool_use_id)) {
198        return result
199      }
200      const tool = byToolUseId.get(result.tool_use_id)?.tool ?? 'the tool'
201      const text = truncateResult(result.text, headChars, tool)
202
203      return text === result.text ? result : { tool_use_id: result.tool_use_id, text, isError: result.isError }
204    })
205    const isTouched =
206      uses.some((use, index) => use !== message.toolUses[index]) ||
207      results.some((result, index) => result !== message.toolResults?.[index])
208    if (!isTouched) {
209      return message
210    }
211    const rebuilt: SessionMessage = { role: message.role, text: message.text, toolUses: uses }
212    if (message.toolResults !== undefined) {
213      rebuilt.toolResults = results
214    }
215
216    return rebuilt
217  })
218}
219
220export function charsOf(messages: readonly SessionMessage[]): number {
221  const hasResultRows = messages.some(i => (i.toolResults ?? []).length > 0)
222  let total = 0
223  for (const message of messages) {
224    total += message.text.length
225    for (const use of message.toolUses) {
226      total += JSON.stringify(use.input).length
227      if (!hasResultRows) {
228        total += use.text?.length ?? 0
229      }
230    }
231    for (const result of message.toolResults ?? []) {
232      total += result.text.length
233    }
234  }
235
236  return total
237}
238
239// Output that cannot simply be produced again stays: errors are usually what
240// the work is about, and the rest is listed above.
241export function isProtected(call: Call): boolean {
242  if (call.isError || AGENT_TOOLS.has(call.tool)) {
243    return true
244  }
245
246  return call.tool === 'Bash' && typeof call.input.command === 'string' && UNREPEATABLE_COMMAND.test(call.input.command)
247}
248
249/**
250 * Cuts every tool output that can be read again: outside the first and the
251 * newest messages, not protected, and long enough for a cut to matter. Read
252 * results a later write or full read replaced count as superseded. Text
253 * messages are never touched.
254 */
255export function pruneTranscript(
256  messages: readonly SessionMessage[],
257  tokensBefore: number,
258  preserveRecentMessages: number,
259  keepHeadChars: number,
260): Outcome {
261  const calls = collectCalls(messages, preserveRecentMessages)
262  const superseded = supersededReads(calls)
263  const reasonOf = (call: Call): Reason => {
264    if (call.isPinned) {
265      return 'pinned'
266    }
267    if (superseded.has(call.id)) {
268      return 'superseded'
269    }
270    if (isProtected(call)) {
271      return 'protected'
272    }
273
274    return call.resultText.length > keepHeadChars + 200 ? 'cut' : 'short'
275  }
276  const decisions = calls.map(call => ({
277    id: call.id,
278    tool: call.tool,
279    input: JSON.stringify(call.input).slice(0, 160),
280    resultChars: call.resultText.length,
281    reason: reasonOf(call),
282  }))
283  const dropped = new Set(decisions.filter(i => i.reason === 'cut' || i.reason === 'superseded').map(i => i.id))
284  const pruned = applyDrops(messages, calls, dropped, keepHeadChars)
285  const charsBefore = Math.max(1, charsOf(messages))
286
287  return {
288    messages: pruned,
289    tokensBefore,
290    tokensAfter: Math.round((tokensBefore * charsOf(pruned)) / charsBefore),
291    decisions,
292  }
293}
294