dcp

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { OutputRow } from '../types'
5import { applyDrop, estimateTokens, formatTokens, listOutputs, oneLine, parseArgs, summaryLine } from './prune'
6
7const USAGE = 'Usage: /prune drops every tool output; /prune -i picks which.'
8
9const PANE = 'prune'
10// The instructions the mod's own /compact carries, so the hook answers that
11// compaction and no other.
12const MARK = 'dcp-prune'
13
14const rows = atom({ plugin: 'dcp', key: 'rows' } as const, [])
15const pending = atom({ plugin: 'dcp', key: 'pending' } as const, null)
16
17// A command hook may not compact under the turn it holds, and a plugin's own
18// $.session.compact() skips its own hooks. So the prune rides a marked
19// /compact run from a timer, which the session.compact hook below answers.
20async function compact($: EngineInterface): Promise<void> {
21  let failure: string | undefined
22
23  try {
24    await $.command.run({ command: 'compact', args: MARK })
25  } catch (error) {
26    failure = String(error)
27  }
28
29  if ((await read($, pending)) !== null) {
30    await update($, pending, () => null)
31    $.ui.log(`Prune skipped: ${failure ?? 'the compaction did not run'}`)
32  }
33}
34
35async function prune($: EngineInterface, picked: readonly OutputRow[]): Promise<string> {
36  if (picked.length === 0) {
37    return summaryLine([])
38  }
39
40  // A prune already waiting takes these ids along: one /compact for both.
41  let isWaiting = false
42  await update($, pending, ids => {
43    isWaiting = ids !== null
44
45    return [...new Set([...(ids ?? []), ...picked.map(row => row.id)])]
46  })
47
48  if (!isWaiting) {
49    $.clock.after(0, () => void compact($))
50  }
51
52  return summaryLine(picked)
53}
54
55export const register: Register = on => {
56  let isInteractive = true
57
58  on('session.start', async ($, e, next) => {
59    isInteractive = e.isInteractive
60    await update($, pending, () => null)
61    await $.command.register({
62      name: 'prune',
63      description: 'Drop tool outputs from the context (-i to pick which)',
64      argumentHint: '[-i]',
65    })
66
67    return next(e)
68  })
69
70  on('command.run', { command: 'prune' }, async ($, e) => {
71    const mode = parseArgs(e.args)
72
73    if (mode === 'usage') {
74      return { text: USAGE }
75    }
76
77    if (!isInteractive) {
78      return { text: '/prune needs an interactive session.' }
79    }
80
81    const found = listOutputs(await $.session.messages({ as: 'api' }))
82
83    if (mode === 'all' || found.length === 0) {
84      return { text: await prune($, found) }
85    }
86
87    await update($, rows, () => found)
88    await $.ui.open({ id: PANE, title: 'Prune tool outputs', focus: true, closeOnEscape: true })
89
90    return {}
91  })
92
93  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
94    const { Box, Button, Text } = $.ui.resolve(e)
95    const list = await read($, rows)
96    const marked = list.filter(row => row.isDropped)
97    const chars = marked.reduce((sum, row) => sum + row.chars, 0)
98    const sizes = list.map(row => `${formatTokens(estimateTokens(row.chars))} tok`)
99    const sizeWidth = Math.max(0, ...sizes.map(size => size.length))
100    // mark (4), a gap, the size and the gap before it are kept at any width;
101    // the tool column and then the summary give way.
102    const room = Math.max(0, e.props.bodyColumns - 4 - 2 - sizeWidth - 2)
103    const toolWidth = Math.min(16, room, Math.max(4, ...list.map(row => row.tool.length)))
104    const summaryWidth = Math.max(0, room - toolWidth - 2)
105
106    const setAll = (isDropped: boolean) => update($, rows, all => all.map(row => ({ ...row, isDropped })))
107    const toggle = (id: string) =>
108      update($, rows, all => all.map(row => (row.id === id ? { ...row, isDropped: !row.isDropped } : row)))
109    const apply = async () => {
110      // Taking the rows empties them, so a second press finds nothing to do.
111      let taken: OutputRow[] | undefined
112      await update($, rows, all => {
113        taken = all
114
115        return []
116      })
117
118      if (taken === undefined || taken.length === 0) {
119        return
120      }
121      const picked = taken.filter(row => row.isDropped)
122      await $.ui.close({ id: PANE })
123      $.ui.toast(await prune($, picked))
124    }
125
126    return (
127      <Box flexDirection="column">
128        <Box>
129          <Button key="apply" variant="primary" hotkey="y" label="Apply" onPress={apply} />
130          <Text> </Text>
131          <Button key="all" plain hotkey="a" label="all" onPress={() => setAll(true)} />
132          <Text> </Text>
133          <Button key="none" plain hotkey="n" label="none" onPress={() => setAll(false)} />
134          <Text dimColor> · y apply · esc cancel</Text>
135        </Box>
136        <Text dimColor>
137          dropping {marked.length} of {list.length} · {formatTokens(estimateTokens(chars))} tok
138        </Text>
139        {list.map((row, index) => (
140          <Button
141            key={`row:${row.id}`}
142            plain
143            autoFocus={index === 0 ? true : undefined}
144            dimColor={!row.isDropped}
145            onPress={() => toggle(row.id)}
146          >
147            {row.isDropped ? 'drop' : 'keep'}
148            {'  '}
149            {toolWidth > 0 ? `${oneLine(row.tool, toolWidth).padEnd(toolWidth)}  ` : ''}
150            {summaryWidth > 0 ? `${oneLine(row.summary, summaryWidth).padEnd(summaryWidth)}  ` : ''}
151            <Text dimColor>{(sizes[index] ?? '').padStart(sizeWidth)}</Text>
152          </Button>
153        ))}
154      </Box>
155    )
156  })
157
158  on('session.compact', { trigger: 'manual' }, async ($, e, next) => {
159    if (e.instructions !== MARK || e.agentId !== undefined) {
160      return next(e)
161    }
162
163    const ids = await read($, pending)
164
165    if (ids === null) {
166      return { skip: 'nothing to prune' }
167    }
168
169    const messages = applyDrop(e.messages, new Set(ids))
170    await update($, pending, () => null)
171
172    return { messages }
173  }).catch(($, e, next) =>
174    // The mod's own compaction never reaches the summarizer, whatever failed.
175    e.instructions !== MARK || e.agentId !== undefined ? next(e) : { skip: 'prune failed' },
176  )
177}
178

1import type { SessionMessage } from 'claude-code'
2
3import type { OutputRow } from '../types'
4
5export const STUB = '[tool output dropped by /prune]'
6
7const NAMING_FIELDS = ['command', 'file_path', 'pattern', 'url', 'query', 'prompt']
8
9export function oneLine(text: string, width: number): string {
10  const line = Array.from(text.replace(/\s+/g, ' ').trim())
11
12  return line.length > width ? `${line.slice(0, Math.max(0, width - 1)).join('')}…` : line.join('')
13}
14
15export function summarizeInput(tool: string, input: Record<string, unknown>, width: number): string {
16  const named = NAMING_FIELDS.map(field => input[field]).find(value => typeof value === 'string')
17  const picked = named ?? Object.values(input).find(value => typeof value === 'string') ?? ''
18
19  return oneLine(String(picked), width)
20}
21
22// The conversation in Messages API form, as `$.session.messages({ as: 'api' })`
23// answers it: unlike the rows, it holds only what the next request sends.
24export type ApiRow = { role: string; content: { type: string; [field: string]: unknown }[] }
25
26function textOf(content: unknown): string {
27  if (typeof content === 'string') {
28    return content
29  }
30
31  if (!Array.isArray(content)) {
32    return ''
33  }
34
35  return content.map(block => (block?.type === 'text' && typeof block.text === 'string' ? block.text : '')).join('')
36}
37
38export function listOutputs(messages: readonly ApiRow[]): OutputRow[] {
39  const blocks = messages.flatMap(message => message.content)
40  const results = new Map(
41    blocks
42      .filter(block => block.type === 'tool_result')
43      .map(block => [String(block.tool_use_id), textOf(block.content)]),
44  )
45
46  return blocks
47    .filter(block => block.type === 'tool_use')
48    .flatMap(block => {
49      const text = results.get(String(block.id))
50
51      if (text === undefined || text === STUB) {
52        return []
53      }
54      const input = (block.input ?? {}) as Record<string, unknown>
55
56      return [
57        {
58          id: String(block.id),
59          tool: String(block.name),
60          summary: summarizeInput(String(block.name), input, 200),
61          chars: text.length,
62          isDropped: true,
63        },
64      ]
65    })
66}
67
68// Only the rows holding the results are rewritten. A call's own row goes back
69// whole: parallel calls are rows the engine merges by their handles, and a
70// rebuilt one splits the group and orphans its siblings' results.
71export function applyDrop(messages: readonly SessionMessage[], ids: ReadonlySet<string>): SessionMessage[] {
72  return messages.map(message => {
73    const results = message.toolResults ?? []
74
75    if (!results.some(result => ids.has(result.tool_use_id) && result.text !== STUB)) {
76      return message
77    }
78
79    return {
80      role: message.role,
81      text: message.text,
82      toolUses: message.toolUses,
83      toolResults: results.map(result => {
84        if (!ids.has(result.tool_use_id)) {
85          return result
86        }
87        const { result: _result, ...kept } = result
88
89        return { ...kept, text: STUB }
90      }),
91    }
92  })
93}
94
95export function estimateTokens(chars: number): number {
96  return Math.round(chars / 4)
97}
98
99export function formatTokens(tokens: number): string {
100  if (tokens < 1000) {
101    return `~${tokens}`
102  }
103
104  return `~${(tokens / 1000).toFixed(tokens < 10000 ? 1 : 0).replace(/\.0$/, '')}k`
105}
106
107export function summaryLine(rows: readonly OutputRow[]): string {
108  if (rows.length === 0) {
109    return 'No tool outputs to drop.'
110  }
111  const chars = rows.reduce((sum, row) => sum + row.chars, 0)
112  const noun = rows.length === 1 ? 'tool output' : 'tool outputs'
113
114  return `Dropping ${rows.length} ${noun} (${formatTokens(estimateTokens(chars))} tokens).`
115}
116
117export function parseArgs(args: string): 'all' | 'interactive' | 'usage' {
118  const trimmed = args.trim()
119
120  if (trimmed === '') {
121    return 'all'
122  }
123
124  return trimmed === '-i' || trimmed === '--interactive' ? 'interactive' : 'usage'
125}
126

1export type OutputRow = {
2  id: string
3  tool: string
4  summary: string
5  chars: number
6  isDropped: boolean
7}
8
9declare module 'claude-code' {
10  interface PluginState {
11    dcp: { rows: OutputRow[]; pending: string[] | null }
12  }
13}
14