optimize-prompt

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register, ToolUseSummary } from 'claude-code'
3
4const isBusy = atom({ plugin: 'optimize-prompt', key: 'isBusy' } as const, false)
5const frame = atom({ plugin: 'optimize-prompt', key: 'frame' } as const, 0)
6
7// Braille dots: one column wide in every terminal font. No Spinner element on
8// the terminal, so a $.clock.every tick advances `frame` and the band redraws.
9const SPINNER = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
10const SPINNER_MS = 100
11const spinnerGlyph = (n: number) => SPINNER[n % SPINNER.length] ?? ''
12
13// Haiku: cheapest + fastest model; rewriting a prompt needs no deep reasoning.
14const MODEL = 'haiku'
15
16const SYSTEM = `You rewrite a draft prompt the user is about to send to an AI coding assistant, in the middle of an ongoing session.
17<conversation> holds the latest messages of that session and the assistant's tool calls between them (lines starting "tool", showing files read or edited, commands run and their results). Use it only to understand what the draft refers to: resolve vague references ("it", "the button", "that error") to the concrete thing, file or feature being discussed. Name a file or command from it only where the draft refers to one.
18Rules:
19- Keep the draft's scope. Never add requirements, deliverables, sections or questions the author did not ask for.
20- Never comment on what the draft leaves unspecified. If the conversation does not resolve something, leave it as written.
21- Stay proportionate: a one-line draft becomes one to three sentences, not a spec with headings.
22- Keep the author's voice, language and every concrete detail (file paths, names, code, numbers).
23- Do not answer or carry out the prompt.
24Reply with the rewritten prompt only, no preamble or quotes.`
25
26// Recent turns ground the rewrite: without them Haiku sees a bare draft and
27// generalises "the spinner" into "the application's UI".
28// Tool calls show which files and commands "it" means. The window is counted in
29// text messages so a tool-heavy turn can't push the person's own words out.
30// ponytail: plain char cuts, not token counts; CONTEXT_TOTAL is the ceiling.
31const CONTEXT_MESSAGES = 8 // text messages; tool calls between them ride along
32const CONTEXT_CHARS = 1200 // per message text, its tail kept
33const TOOL_CHARS = 200 // per tool input and per tool result, the head kept
34const CONTEXT_TOTAL = 12000 // whole block, the newest kept
35
36const head = (s: string, n: number) => (s.length > n ? `${s.slice(0, n)}…` : s)
37const tail = (s: string, n: number) => (s.length > n ? `…${s.slice(-n)}` : s)
38
39const toolLine = (u: ToolUseSummary) =>
40  `tool ${u.tool}(${head(JSON.stringify(u.input), TOOL_CHARS)})` +
41  (u.text === undefined ? '' : ` → ${u.isError ? 'error: ' : ''}${head(u.text, TOOL_CHARS)}`)
42
43const recentContext = async ($: EngineInterface) => {
44  const messages = (await $.session.messages()).filter(m => m.text.trim() || m.toolUses.length)
45  const textAt = messages.flatMap((m, i) => (m.text.trim() ? [i] : []))
46  const rendered = messages
47    .slice(textAt.at(-CONTEXT_MESSAGES) ?? 0)
48    .map(m =>
49      [...(m.text.trim() ? [`${m.role}: ${tail(m.text, CONTEXT_CHARS)}`] : []), ...m.toolUses.map(toolLine)].join('\n'),
50    )
51    .join('\n\n')
52  return tail(rendered, CONTEXT_TOTAL)
53}
54
55export const register: Register = on => {
56  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
57    if (e.props.hasSurvey) {
58      return next(e)
59    }
60
61    const busy = await read($, isBusy)
62    const { Box, Button, Text } = $.ui.resolve(e)
63
64    const optimize = async () => {
65      if (await read($, isBusy)) return
66
67      const draft = (await $.prompt.read()).text
68      if (!draft.trim()) {
69        $.ui.toast('Type a prompt first')
70        return
71      }
72
73      await update($, isBusy, () => true)
74      const tick = $.clock.every(SPINNER_MS, () => void update($, frame, n => n + 1))
75      try {
76        const context = await recentContext($)
77        const r = await $.model.complete({
78          model: MODEL,
79          system: SYSTEM,
80          prompt: `${context && `<conversation>\n${context}\n</conversation>\n\n`}<draft>\n${draft}\n</draft>`,
81          effort: 'low',
82          maxTokens: 2048,
83          timeoutMs: 30000,
84        })
85        if (!r.isAnswered) {
86          $.ui.toast(`Optimize failed: ${r.reason}`)
87          return
88        }
89        // Don't clobber anything typed while the model was working.
90        if ((await $.prompt.read()).text !== draft) {
91          $.ui.toast('Draft changed while optimizing; left it as is')
92          return
93        }
94        await $.prompt.fill({ text: r.text.trim(), mode: 'replace' })
95      } finally {
96        tick.cancel()
97        await update($, isBusy, () => false)
98      }
99    }
100
101    return (
102      <Box>
103        {busy ? (
104          <Text color="suggestion">
105            {spinnerGlyph(await read($, frame))} Optimizing prompt
106          </Text>
107        ) : (
108          <Button key="optimize" variant="primary" onPress={optimize}>
109            ✦ Optimize prompt
110          </Button>
111        )}
112      </Box>
113    )
114  })
115}
116

1export type Busy = boolean
2
3declare module 'claude-code' {
4  interface PluginState {
5    'optimize-prompt': { isBusy: Busy; frame: number }
6  }
7}
8