context-trim

1import type { Register } from 'claude-code'
2import { parseOptions } from './config'
3import { applyDrop, digest, parseDrop } from './trim-core'
4
5const PROMPT = (lines: string) => `You prune the transcript of a coding-agent session to free context. Each line below is one message: index, role, ~tokens, preview.
6Pick the messages that no longer help the work: large tool outputs already acted on, dead ends, retried attempts, repeated listings, chatter. Keep decisions, the user's requests and constraints, file paths and facts still needed, and anything unresolved.
7When unsure, keep the message. Answer with JSON only: {"drop":[indices]}
8
9${lines}`
10
11export const register: Register = (on, options) => {
12  const opts = parseOptions(options)
13  if (!opts.enabled) return
14
15  let lastBucket = 0
16  let isRunning = false
17  let hasHookRun = false
18
19  // Answers a plugin compaction with the transcript minus what the model picked: no summary, nothing sent to the summarizer.
20  // Any failure vetoes the compaction instead of calling `next(e)`, which would summarize the whole conversation.
21  on('session.compact', { trigger: 'plugin' }, async ($, e) => {
22    hasHookRun = true
23    if (e.messages.length <= opts.keepRecent + 1) return { skip: 'context-trim: transcript too short to trim' }
24
25    const asked = await $.model.complete({ model: opts.model, prompt: PROMPT(digest(e.messages)), effort: 'low', timeoutMs: 45000 })
26    if (!asked.isAnswered) return { skip: `context-trim: judge failed (${asked.reason})` }
27
28    const drop = parseDrop(asked.text, e.messages.length)
29    if (!drop) return { skip: 'context-trim: judge answer was not a drop list' }
30
31    const trimmed = applyDrop(e.messages, drop, opts.keepRecent)
32    if (trimmed.dropped.length === 0) return { skip: 'context-trim: nothing safe to drop' }
33
34    $.ui.toast(`context-trim: dropped ${trimmed.dropped.length} messages, ~${trimmed.saved} tokens`)
35    return { messages: trimmed.kept }
36  })
37
38  on('turn.complete', async ($, e, next) => {
39    const { context } = await $.session.usage()
40    const bucket = Math.floor((context.percent ?? 0) / opts.step)
41    if (bucket < lastBucket) lastBucket = bucket
42
43    if (opts.enabled && bucket > lastBucket && !isRunning) {
44      const previous = lastBucket
45      lastBucket = bucket
46      isRunning = true
47      hasHookRun = false
48
49      // `compact` rejects while a turn runs, so it waits for this hook to return.
50      $.clock.after(0, () => {
51        $.session
52          .compact()
53          .then((r) => {
54            // The engine runs every hook but the calling plugin's own; if ours was skipped, a summary
55            // compaction just happened and the mod stops itself rather than repeat it.
56            if (r.skip === undefined && !hasHookRun) {
57              opts.enabled = false
58              $.ui.toast('context-trim: its own hook is not reached by its own compact call; the mod is off until it is split in two plugins')
59            } else if (r.skip !== undefined && r.skip.startsWith('context-trim:')) {
60              $.ui.toast(r.skip)
61            }
62          })
63          .catch(() => {
64            lastBucket = previous
65          })
66          .finally(() => {
67            isRunning = false
68          })
69      })
70    }
71
72    return next(e)
73  })
74}
75

1// The mod's `userConfig` values, parsed in one place: an unknown or mistyped value takes the default.
2export type Options = { step: number; keepRecent: number; model: string; enabled: boolean }
3
4type Raw = Readonly<Record<string, unknown>> | undefined
5
6const count = (raw: Raw, key: string, allowed: readonly number[], fallback: number): number => {
7  const n = Number(raw?.[key])
8  return allowed.includes(n) ? n : fallback
9}
10
11export const parseOptions = (raw: Raw): Options => {
12  const model = raw?.model
13  return {
14    step: count(raw, 'step', [5, 10, 15, 20, 25], 10),
15    keepRecent: count(raw, 'keepRecent', [4, 8, 12, 16], 8),
16    model: typeof model === 'string' && model.trim() ? model.trim() : 'claude-haiku-5-5',
17    enabled: raw?.enabled !== 'off',
18  }
19}
20

1// `$`-free core of context-trim: what the model sees, how its answer is read, and which messages may go.
2
3export type Msg = {
4  role: 'user' | 'assistant'
5  text: string
6  toolUses: { tool_use_id: string; input?: unknown }[]
7  toolResults?: { tool_use_id: string; text: string }[]
8}
9
10export type Trimmed<T> = { kept: T[]; dropped: number[]; saved: number }
11
12/** Rough token estimate (4 characters a token) over everything a message carries. */
13export function tokensOf(m: Msg): number {
14  const chars =
15    m.text.length +
16    m.toolUses.reduce((n, u) => n + JSON.stringify(u.input ?? '').length, 0) +
17    (m.toolResults ?? []).reduce((n, r) => n + r.text.length, 0)
18  return Math.ceil(chars / 4)
19}
20
21/** One line per message, `index<TAB>role<TAB>~tokens<TAB>preview`: all the model needs to pick what to drop. */
22export function digest(messages: readonly Msg[], previewChars = 200): string {
23  return messages
24    .map((m, i) => {
25      const body = m.text || m.toolResults?.map((r) => r.text).join(' ') || m.toolUses.map((u) => JSON.stringify(u.input ?? '')).join(' ')
26      const preview = body.replace(/\s+/g, ' ').trim().slice(0, previewChars)
27      return `${i}\t${m.role}\t~${tokensOf(m)}t\t${preview}`
28    })
29    .join('\n')
30}
31
32/** Reads `{ "drop": [index, ...] }` from the model's answer, fenced or not; null when it is not that. */
33export function parseDrop(raw: string, count: number): number[] | null {
34  const body = raw.match(/```(?:json)?\s*([\s\S]*?)```/)?.[1] ?? raw
35  const start = body.indexOf('{')
36  const end = body.lastIndexOf('}')
37  if (start < 0 || end < start) return null
38  try {
39    const drop = (JSON.parse(body.slice(start, end + 1)) as { drop?: unknown }).drop
40    if (!Array.isArray(drop)) return null
41    const valid = drop.filter((i): i is number => Number.isInteger(i) && i >= 0 && i < count)
42    return [...new Set(valid)].sort((a, b) => a - b)
43  } catch {
44    return null
45  }
46}
47
48/** Indices joined by a tool_use / tool_result pair, grouped so a pair is dropped or kept whole. */
49function pairGroups(messages: readonly Msg[]): number[][] {
50  const parent = messages.map((_, i) => i)
51  const find = (i: number): number => (parent[i] === i ? i : (parent[i] = find(parent[i] as number)))
52  const useAt = new Map<string, number>()
53  messages.forEach((m, i) => m.toolUses.forEach((u) => useAt.set(u.tool_use_id, i)))
54  messages.forEach((m, i) =>
55    m.toolResults?.forEach((r) => {
56      const at = useAt.get(r.tool_use_id)
57      if (at !== undefined) parent[find(i)] = find(at)
58    }),
59  )
60  const groups = new Map<number, number[]>()
61  messages.forEach((_, i) => groups.set(find(i), [...(groups.get(find(i)) ?? []), i]))
62  return [...groups.values()]
63}
64
65/**
66 * Drops the requested messages, with two guards: the first message and the last `keepRecent` never go,
67 * and a tool_use goes with its tool_result or not at all (a group touching a protected message stays whole).
68 */
69export function applyDrop<T extends Msg>(messages: readonly T[], drop: readonly number[], keepRecent: number): Trimmed<T> {
70  const asked = new Set(drop)
71  const protectedFrom = messages.length - keepRecent
72  const gone = new Set<number>()
73  for (const group of pairGroups(messages)) {
74    const wanted = group.some((i) => asked.has(i))
75    const safe = group.every((i) => i > 0 && i < protectedFrom)
76    if (wanted && safe) group.forEach((i) => gone.add(i))
77  }
78  const dropped = [...gone].sort((a, b) => a - b)
79  return {
80    kept: messages.filter((_, i) => !gone.has(i)),
81    dropped,
82    saved: dropped.reduce((n, i) => n + tokensOf(messages[i] as Msg), 0),
83  }
84}
85