answer-limit

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { LastTurn, Limit, Limits } from '../types'
5
6const DEFAULTS: Limits = { lines: { value: 3, isOn: false }, tools: { value: 3, isOn: false } }
7const STORE_KEY = 'limits'
8
9const limitsAtom = atom({ plugin: 'answer-limit', key: 'limits' } as const, DEFAULTS)
10const lastAtom = atom({ plugin: 'answer-limit', key: 'last' } as const, null as LastTurn | null)
11
12type Which = keyof Limits
13
14function countLines(text: string): number {
15  const trimmed = text.trim()
16  return trimmed === '' ? 0 : trimmed.split('\n').length
17}
18
19function instruction(limits: Limits): string {
20  const rules: string[] = []
21  if (limits.tools.isOn) {
22    rules.push(
23      `- At most ${limits.tools.value} tool calls this turn, subagents included; the harness denies any call beyond that. ` +
24        'Spend them only on what the conclusion needs, nothing outside the request unless it fails. ' +
25        'When the budget runs out, answer with what you have and say in one line what is still unknown.',
26    )
27  }
28  if (limits.lines.isOn) {
29    rules.push(
30      `- Final answer at most ${limits.lines.value} lines: the conclusion only (for a test: pass, or fail with the reason), ` +
31        'plus ids needed to look it up again. No step-by-step narration, no tables unless they fit the limit. ' +
32        'If a decision is needed from the user, ask it in one line.',
33    )
34  }
35  return 'Answer limits set by the user for every turn until changed:\n' + rules.join('\n')
36}
37
38async function setLimit($: EngineInterface, which: Which, fn: (l: Limit) => Limit) {
39  await update($, limitsAtom, limits => ({ ...limits, [which]: fn(limits[which]) }))
40  await $.store.set(STORE_KEY, await read($, limitsAtom))
41}
42
43export const register: Register = on => {
44  let toolCalls = 0
45
46  on('session.start', async ($, e, next) => {
47    const saved = (await $.store.get(STORE_KEY)) as Limits | undefined
48    if (saved !== undefined) {
49      await update($, limitsAtom, () => saved)
50    }
51    return next(e)
52  })
53
54  on('prompt.submit', ($, e, next) => {
55    toolCalls = 0
56    return next(e)
57  })
58
59  on('tool.call', async ($, e, next) => {
60    const { tools } = await read($, limitsAtom)
61    if (tools.isOn && toolCalls >= tools.value) {
62      return {
63        deny:
64          `answer-limit: tool budget of ${tools.value} calls for this turn is used up. ` +
65          'Answer now with what you have; say in one line what is still unknown.',
66      }
67    }
68    toolCalls += 1
69    return next(e)
70  })
71
72  on('prompt.compose', async ($, e, next) => {
73    const composed = await next(e)
74    const limits = await read($, limitsAtom)
75    if (!limits.lines.isOn && !limits.tools.isOn) {
76      return composed
77    }
78    const section = { id: 'answer-limit:limits', text: instruction(limits), scope: 'session' } as const
79    return { sections: [...composed.sections, section] }
80  })
81
82  on('turn.complete', async ($, e, next) => {
83    if (e.agentId === undefined) {
84      const last: LastTurn = { tools: toolCalls, lines: countLines(e.answer) }
85      await update($, lastAtom, () => last)
86    }
87    return next(e)
88  })
89
90  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
91    if (e.props.hasSurvey) {
92      return next(e)
93    }
94    // The band is shared by every plugin: draw the tree of the plugins beneath too, not over it
95    const below = await next(e)
96    const limits = await read($, limitsAtom)
97    const last = await read($, lastAtom)
98    const { Box, Button, Text } = $.ui.resolve(e)
99
100    const controls = (which: Which, label: string, keys: [string, string, string]) => {
101      const limit = limits[which]
102      return [
103        <Text key={`${which}-label`}>
104          {label}: {limit.isOn ? String(limit.value) : 'off'}{' '}
105        </Text>,
106        <Button
107          key={`${which}-dec`}
108          label="-"
109          hotkey={keys[0]}
110          onPress={() => setLimit($, which, l => ({ value: Math.max(1, l.value - 1), isOn: true }))}
111        />,
112        <Text key={`${which}-sp1`}> </Text>,
113        <Button
114          key={`${which}-inc`}
115          label="+"
116          hotkey={keys[1]}
117          onPress={() => setLimit($, which, l => ({ value: l.value + 1, isOn: true }))}
118        />,
119        <Text key={`${which}-sp2`}> </Text>,
120        <Button
121          key={`${which}-toggle`}
122          label={limit.isOn ? 'off' : 'on'}
123          hotkey={keys[2]}
124          onPress={() => setLimit($, which, l => ({ ...l, isOn: !l.isOn }))}
125        />,
126      ]
127    }
128
129    return (
130      <Box flexDirection="column">
131        <Box>
132        <Text dimColor>Answer limit: </Text>
133        {controls('lines', 'Lines', ['a', 's', 'd'])}
134        <Text dimColor> · </Text>
135        {controls('tools', 'Tool calls', ['j', 'k', 'l'])}
136        {last !== null ? <Text dimColor>  last: {last.lines} lines, {last.tools} tools</Text> : null}
137        </Box>
138        {below}
139      </Box>
140    )
141  })
142}
143

1/** value: the cap when on; isOn false keeps value for when it is turned back on */
2export type Limit = { value: number; isOn: boolean }
3export type Limits = { lines: Limit; tools: Limit }
4export type LastTurn = { tools: number; lines: number }
5
6declare module 'claude-code' {
7  interface PluginState {
8    'answer-limit': { limits: Limits; last: LastTurn | null }
9  }
10}
11