trim-startup-context

1import type { EngineInterface, Register } from 'claude-code'
2
3import { addDeny, candidates, removeDeny, usedIn, type Candidate, type McpTool } from './lib'
4
5const SESSIONS = 30 // transcripts scanned, newest first
6const MIN_SESSIONS = 5 // fewer than this is too little history to call a server unused
7
8type Scan = { tools: McpTool[]; all: number; found: Candidate[]; scanned: number; skipped: number }
9
10// MCP servers loaded now that the last sessions never called.
11async function scan($: EngineInterface): Promise<Scan | undefined> {
12  const { context } = await $.session.usage({ breakdown: 'summary' })
13  const tools: McpTool[] = context.breakdown?.mcpTools ?? []
14  const all = tools.filter(t => t.isLoaded).reduce((n, t) => n + t.tokens, 0)
15
16  // Transcripts live under the active config folder, which is not always ~/.claude.
17  const config = await $.env.get('CLAUDE_CONFIG_DIR')
18  const home = (await $.env.get('USERPROFILE')) ?? (await $.env.get('HOME'))
19  const base = config ?? (home ? `${home}/.claude` : undefined)
20  if (!base) return undefined
21  const root = `${base}/projects`
22  const files: { path: string; mtimeMs: number }[] = []
23  for (const dir of (await $.fs.list(root)).filter(d => d.kind === 'dir')) {
24    for (const f of await $.fs.list(`${root}/${dir.name}`)) {
25      if (f.kind === 'file' && f.name.endsWith('.jsonl')) {
26        files.push({ path: `${root}/${dir.name}/${f.name}`, mtimeMs: f.mtimeMs })
27      }
28    }
29  }
30  files.sort((a, b) => b.mtimeMs - a.mtimeMs)
31  const used = new Set<string>()
32  let scanned = 0
33  let skipped = 0
34  for (const f of files.slice(0, SESSIONS)) {
35    try {
36      for (const s of usedIn((await $.fs.read(f.path)) as string).servers) used.add(s)
37      scanned += 1
38    } catch {
39      skipped += 1 // over the 4 MiB read limit, or unreadable
40    }
41  }
42  return { tools, all, found: candidates(tools, used), scanned, skipped }
43}
44
45async function files($: EngineInterface) {
46  const cwd = await $.session.cwd()
47  const local = `${cwd}/.claude/settings.local.json`
48  return {
49    cwd,
50    local,
51    rulesKey: `trim-context:${cwd}`,
52    offKey: `trim-context:off:${cwd}`, // set by undo, so the automatic pass leaves this project alone
53    read: async (): Promise<string> => ((await $.fs.exists(local)) ? ((await $.fs.read(local)) as string) : ''),
54  }
55}
56
57// Denies the found servers in this project and remembers the rules it added.
58async function apply($: EngineInterface, found: readonly Candidate[]): Promise<{ added: string[]; local: string }> {
59  const f = await files($)
60  const { text, added } = addDeny(await f.read(), found.map(c => c.rule))
61  if (added.length > 0) {
62    await $.fs.write(f.local, text)
63    await $.store.set(f.rulesKey, [...(((await $.store.get(f.rulesKey)) as string[] | undefined) ?? []), ...added])
64  }
65  return { added, local: f.local }
66}
67
68// Automatic pass: deny servers that are loaded and unused across enough history, unless undo opted this project out.
69async function auto($: EngineInterface) {
70  try {
71    const f = await files($)
72    if (await $.store.get(f.offKey)) return
73    const s = await scan($)
74    if (!s || s.scanned < MIN_SESSIONS || s.found.length === 0) return
75    const { added } = await apply($, s.found)
76    if (added.length > 0) {
77      const tokens = s.found.filter(c => added.includes(c.rule)).reduce((n, c) => n + c.tokens, 0)
78      $.ui.toast(`trim-context: denied ${added.length} unused MCP server(s), ${tokens.toLocaleString()} tokens, from next session. /trim-context undo reverts.`)
79    }
80  } catch {}
81}
82
83export const register: Register = on => {
84  let checked = false // the automatic pass has run once this session
85
86  on('session.start', async ($, e, next) => {
87    await $.command.register({
88      name: 'trim-context',
89      description: 'Unused MCP servers are denied automatically per project; this shows them, apply forces it, undo reverts and opts out',
90      argumentHint: '[apply|undo]',
91    })
92    // The session may not be bound yet; if no MCP schemas are visible, the first turn tries again.
93    try {
94      const { context } = await $.session.usage({ breakdown: 'summary' })
95      if ((context.breakdown?.mcpTools ?? []).length > 0) {
96        checked = true
97        await auto($)
98      }
99    } catch {}
100    return next(e)
101  })
102
103  on('turn.start', async ($, e, next) => {
104    if (!checked && e.agentId === undefined) {
105      checked = true
106      await auto($)
107    }
108    return next(e)
109  })
110
111  on('command.run', { command: 'trim-context' }, async ($, e) => {
112    const mode = e.args.trim()
113    const f = await files($)
114
115    if (mode === 'undo') {
116      const rules = ((await $.store.get(f.rulesKey)) as string[] | undefined) ?? []
117      await $.store.set(f.offKey, true)
118      if (rules.length === 0) return { text: 'Nothing to undo: this mod has added no rules here. Automatic trimming is now off for this project.' }
119      await $.fs.write(f.local, removeDeny(await f.read(), rules))
120      await $.store.delete(f.rulesKey)
121      return { text: `Removed ${rules.length} deny rule(s) from ${f.local}. Automatic trimming is off for this project; /trim-context apply turns it back on. Restart the session to load the servers again.` }
122    }
123    if (mode !== '' && mode !== 'apply') return { text: 'Usage: /trim-context [apply|undo]' }
124
125    const s = await scan($)
126    if (!s) return { text: 'Cannot find the home folder to read past sessions.' }
127    const saved = s.found.reduce((n, c) => n + c.tokens, 0)
128    const lines = [
129      `MCP schemas in the window: ${s.all.toLocaleString()} tokens. Scanned ${s.scanned} past session(s)${s.skipped ? `, ${s.skipped} unreadable` : ''}.`,
130      s.found.length === 0
131        ? 'Every connected MCP server was used recently. Nothing to trim.'
132        : `Unused MCP servers (${saved.toLocaleString()} tokens):\n` +
133          s.found.map(c => `  ${c.rule}  ${c.tools} tools, ${c.tokens.toLocaleString()} tokens`).join('\n'),
134    ]
135
136    if (mode !== 'apply' || s.found.length === 0) {
137      if (s.found.length > 0) lines.push('These are denied automatically at session start; /trim-context apply does it now.')
138      return { text: lines.join('\n') }
139    }
140    if (s.scanned < MIN_SESSIONS) {
141      lines.push(`Not applied: only ${s.scanned} session(s) of history, need ${MIN_SESSIONS} to call a server unused.`)
142      return { text: lines.join('\n') }
143    }
144
145    await $.store.delete(f.offKey)
146    const { added, local } = await apply($, s.found)
147    lines.push(`Applied: denied ${added.length} server(s) in ${local}. Takes effect next session. /trim-context undo reverts.`)
148    return { text: lines.join('\n') }
149  })
150}
151

1export type McpTool = { name: string; serverName: string; tokens: number; isLoaded: boolean }
2
3export type Candidate = { server: string; rule: string; tools: number; tokens: number }
4
5// The wire name an MCP server goes by in `mcp__<server>__<tool>`.
6export const wireName = (server: string): string => server.replace(/[^a-zA-Z0-9_-]/g, '_')
7
8// Servers and skills the assistant actually called in one transcript's text.
9export function usedIn(transcript: string): { servers: Set<string>; skills: Set<string> } {
10  const servers = new Set<string>()
11  const skills = new Set<string>()
12  for (const line of transcript.split('\n')) {
13    if (!line.includes('"tool_use"')) continue
14    let row: any
15    try {
16      row = JSON.parse(line)
17    } catch {
18      continue
19    }
20    const blocks = row?.message?.content
21    if (row?.type !== 'assistant' || !Array.isArray(blocks)) continue
22    for (const b of blocks) {
23      if (b?.type !== 'tool_use') continue
24      const m = /^mcp__(.+?)__/.exec(b.name)
25      if (m) servers.add(m[1])
26      if (b.name === 'Skill' && typeof b.input?.skill === 'string') skills.add(b.input.skill)
27    }
28  }
29  return { servers, skills }
30}
31
32// MCP servers loaded this session and never called: biggest schema cost first.
33export function candidates(tools: readonly McpTool[], used: ReadonlySet<string>): Candidate[] {
34  const by = new Map<string, Candidate>()
35  for (const t of tools) {
36    const wire = wireName(t.serverName)
37    if (used.has(wire)) continue
38    const c = by.get(wire) ?? { server: t.serverName, rule: `mcp__${wire}`, tools: 0, tokens: 0 }
39    c.tools += 1
40    if (t.isLoaded) c.tokens += t.tokens
41    by.set(wire, c)
42  }
43  // A server whose schemas cost nothing now (deferred) saves nothing to deny.
44  return [...by.values()].filter(c => c.tokens > 0).sort((a, b) => b.tokens - a.tokens)
45}
46
47type Json = Record<string, any>
48const parse = (text: string): Json => {
49  if (!text.trim()) return {}
50  const v = JSON.parse(text)
51  if (v === null || typeof v !== 'object' || Array.isArray(v)) throw new Error('settings is not a JSON object')
52  return v
53}
54
55// Adds deny rules to a settings.local.json text; `added` is only the new ones.
56export function addDeny(text: string, rules: readonly string[]): { text: string; added: string[] } {
57  const s = parse(text)
58  const deny: string[] = Array.isArray(s.permissions?.deny) ? s.permissions.deny : []
59  const added = rules.filter(r => !deny.includes(r))
60  s.permissions = { ...s.permissions, deny: [...deny, ...added] }
61  return { text: JSON.stringify(s, null, 2) + '\n', added }
62}
63
64export function removeDeny(text: string, rules: readonly string[]): string {
65  const s = parse(text)
66  if (Array.isArray(s.permissions?.deny)) {
67    s.permissions.deny = s.permissions.deny.filter((r: string) => !rules.includes(r))
68    if (s.permissions.deny.length === 0) delete s.permissions.deny
69    if (Object.keys(s.permissions).length === 0) delete s.permissions
70  }
71  return JSON.stringify(s, null, 2) + '\n'
72}
73