tool-output-diet

1import type { EngineInterface, Register } from 'claude-code'
2
3import { condense } from './condense'
4
5// Output is left alone unless this many characters (about 75 tokens) go.
6const MIN_SAVED_CHARS = 300
7// Rough size of a token for English text and logs.
8const CHARS_PER_TOKEN = 4
9
10type ShellResult = { stdout?: unknown; stderr?: unknown; isImage?: unknown }
11
12const isEnabled = async ($: EngineInterface) => (await $.store.get('enabled')) !== false
13
14const addSaved = async ($: EngineInterface, chars: number) => {
15  const total = Number(await $.store.get('savedChars')) || 0
16  const calls = Number(await $.store.get('calls')) || 0
17
18  await $.store.set('savedChars', total + chars)
19  await $.store.set('calls', calls + 1)
20}
21
22export const register: Register = on => {
23  on('session.start', async ($, e, next) => {
24    await $.command.register({
25      name: 'diet',
26      description: 'Tool-output diet: on, off, or how much it has saved',
27      argumentHint: '[on|off]',
28      immediate: true,
29    })
30
31    return next(e)
32  })
33
34  on('command.run', { command: 'diet' }, async ($, e) => {
35    const arg = e.args.trim().toLowerCase()
36
37    if (arg === 'on' || arg === 'off') {
38      await $.store.set('enabled', arg === 'on')
39
40      return { text: `Tool-output diet: ${arg}.` }
41    }
42    const chars = Number(await $.store.get('savedChars')) || 0
43    const calls = Number(await $.store.get('calls')) || 0
44    const state = (await isEnabled($)) ? 'on' : 'off'
45
46    return {
47      text: `Tool-output diet is ${state}. It has trimmed ${calls} command outputs, about ${Math.round(chars / CHARS_PER_TOKEN)} tokens in all. Use /diet on or /diet off.`,
48    }
49  })
50
51  // Only the output of a command that succeeded. A failed command's output is
52  // left exactly as it came: that is where the exact text matters most.
53  on('tool.call', { tool: ['Bash', 'PowerShell'] }, async ($, e, next) => {
54    const ran = await next(e)
55
56    if (ran.deny !== undefined || ran.isError === true || !(await isEnabled($))) {
57      return ran
58    }
59    const result = ran.result as ShellResult | null
60
61    if (typeof result !== 'object' || result === null || typeof result.stdout !== 'string' || result.isImage === true) {
62      return ran
63    }
64    const stdout = condense(result.stdout)
65    const stderr = typeof result.stderr === 'string' ? condense(result.stderr) : undefined
66    const before = result.stdout.length + (typeof result.stderr === 'string' ? result.stderr.length : 0)
67    const after = stdout.text.length + (stderr?.text.length ?? 0)
68
69    if (before - after < MIN_SAVED_CHARS) {
70      return ran
71    }
72    await addSaved($, before - after)
73
74    // A result of its own, without core's `ref` and `text`, so core maps the
75    // new one for the model. Nothing is added to it: what was cut was noise.
76    return {
77      result: {
78        ...result,
79        stdout: stdout.text,
80        ...(stderr === undefined ? {} : { stderr: stderr.text }),
81      } as typeof ran.result,
82      context: ran.context,
83    }
84  })
85}
86

1// What a terminal would not have shown, or showed once: colour codes, progress
2// frames drawn over each other, and the same line over and over. Everything
3// else in the output is returned as it was.
4
5// Colour and cursor sequences (CSI), and title or link sequences (OSC).
6const ANSI = /\u001b\[[0-?]*[ -/]*[@-~]|\u001b\][^\u0007\u001b]*(?:\u0007|\u001b\\)/g
7
8// Runs of fewer identical lines than this are left alone.
9const MIN_REPEATS = 4
10
11export type Condensed = {
12  text: string
13  /** Colour or cursor sequences removed. */
14  ansi: number
15  /** Progress frames that a carriage return drew over. */
16  frames: number
17  /** Lines dropped from runs of identical ones. */
18  repeats: number
19}
20
21export const condense = (input: string): Condensed => {
22  const ansi = (input.match(ANSI) ?? []).length
23  const lines = input.replace(ANSI, '').replaceAll('\r\n', '\n').split('\n')
24  let frames = 0
25
26  // "10%\r20%\r100%" is one line on a terminal: the last frame.
27  const drawn = lines.map(line => {
28    if (!line.includes('\r')) {
29      return line
30    }
31    const parts = line.split('\r').filter(part => part !== '')
32
33    frames += Math.max(0, parts.length - 1)
34
35    return parts[parts.length - 1] ?? ''
36  })
37
38  let repeats = 0
39  const kept: string[] = []
40
41  for (let at = 0; at < drawn.length; ) {
42    const line = drawn[at] ?? ''
43    let end = at + 1
44
45    while (end < drawn.length && drawn[end] === line) {
46      end += 1
47    }
48    const run = end - at
49
50    if (line.trim() !== '' && run >= MIN_REPEATS) {
51      kept.push(line, `[the line above is repeated ${run - 1} more times]`)
52      repeats += run - 1
53    } else {
54      kept.push(...drawn.slice(at, end))
55    }
56    at = end
57  }
58
59  return { text: kept.join('\n'), ansi, frames, repeats }
60}
61