prompt-polish

1import { atom, read, update } from 'claude-code'
2import type { ModelTextBlock, Register, SessionMessage } from 'claude-code'
3
4import type { Polished } from '../types'
5import { GUIDE, SYSTEM } from './guide'
6
7// Start a prompt with this to send it exactly as typed.
8const RAW_PREFIX = '::'
9const ENABLED_KEY = 'enabled'
10const MODEL_KEY = 'model'
11const MIN_WORDS = 4
12// Characters kept of each turn shown to the polishing model.
13const TURN_CHARS = 1500
14// The conversation window only grows, so each request reuses the cached turns
15// of the last one; past MAX_TURNS it restarts from the last KEEP_TURNS (one miss).
16const MAX_TURNS = 16
17const KEEP_TURNS = 8
18const NOTE = "The user's prompt was automatically rewritten"
19
20const last = atom({ plugin: 'prompt-polish', key: 'last' } as const, null)
21const isHidden = atom({ plugin: 'prompt-polish', key: 'isHidden' } as const, false)
22const windowStart = atom({ plugin: 'prompt-polish', key: 'windowStart' } as const, 0)
23
24type Turn = { role: 'user' | 'assistant'; text: string }
25
26// One entry per speaker change; a turn's text is the same on every read once
27// the turn is over, which keeps the cached prefix byte-identical.
28const toTurns = (rows: readonly SessionMessage[]): Turn[] => {
29  const turns: Turn[] = []
30  for (const row of rows) {
31    let text = row.text.replace(/<system-reminder>[\s\S]*?<\/system-reminder>/g, '')
32    const note = text.indexOf(NOTE)
33    if (note >= 0) text = text.slice(0, note)
34    text = text.trim()
35    if (!text) continue
36    const prev = turns[turns.length - 1]
37    if (prev?.role === row.role) prev.text += `\n\n${text}`
38    else turns.push({ role: row.role, text })
39  }
40
41  // The person's words from the start, Claude's from the end, where it sums up.
42  return turns.map(t => ({
43    ...t,
44    text:
45      t.text.length <= TURN_CHARS
46        ? t.text
47        : t.role === 'user'
48          ? `${t.text.slice(0, TURN_CHARS)} […]`
49          : `[…] ${t.text.slice(-TURN_CHARS)}`,
50  }))
51}
52
53const MODELS = ['sonnet', 'haiku', 'opus']
54
55const label = (model: string) => model.charAt(0).toUpperCase() + model.slice(1)
56
57export const register: Register = (on, options) => {
58  const installed = typeof options.model === 'string' ? options.model : 'sonnet'
59  // A model picked with /polish <model> wins over the install option. Plugin
60  // options aren't /config rows, so $.config.set can't change that one.
61  const pickModel = (saved: unknown) =>
62    typeof saved === 'string' && MODELS.includes(saved) ? saved : installed
63
64  on('session.start', async ($, e, next) => {
65    await $.command.register({
66      name: 'polish',
67      description: 'Turn prompt polishing on or off, or pick the model that polishes',
68      argumentHint: 'on | off | sonnet | haiku | opus',
69    })
70
71    return next(e)
72  })
73
74  on('command.run', { command: 'polish' }, async ($, e) => {
75    const arg = e.args.trim().toLowerCase()
76
77    if (MODELS.includes(arg)) {
78      await $.store.set(MODEL_KEY, arg)
79      $.ui.toast(`Prompts are now polished by ${label(arg)}`)
80
81      return { text: `Prompts are now polished by ${label(arg)}.` }
82    }
83
84    const isEnabled =
85      arg === 'on' ? true : arg === 'off' ? false : (await $.store.get(ENABLED_KEY)) === false
86    await $.store.set(ENABLED_KEY, isEnabled)
87    if (!isEnabled) await update($, last, () => null)
88    $.ui.toast(`Prompt polish ${isEnabled ? 'on' : 'off'}`)
89    const model = pickModel(await $.store.get(MODEL_KEY))
90
91    return { text: `Prompt polish is ${isEnabled ? 'on' : 'off'} (model: ${model}).` }
92  })
93
94  on('prompt.submit', async ($, e, next) => {
95    const isUsers = ['composer', 'bridge', 'sdk'].includes(e.origin.kind)
96    if (e.command || e.isMeta || !isUsers) return next(e)
97
98    // A new prompt of the person's: the band shows this one's rewrite or nothing.
99    await update($, last, () => null)
100
101    if (e.text.startsWith(RAW_PREFIX)) {
102      return next({ ...e, text: e.text.slice(RAW_PREFIX.length).trimStart() })
103    }
104
105    const isEnabled = (await $.store.get(ENABLED_KEY)) !== false
106    if (!isEnabled || e.text.trim().split(/\s+/).length < MIN_WORDS) return next(e)
107
108    const model = pickModel(await $.store.get(MODEL_KEY))
109    $.ui.status(`Polishing prompt with ${label(model)}…`)
110    let text = ''
111    let cachedTokens = 0
112    let failure = 'empty reply'
113    try {
114      let turns: Turn[] = []
115      try {
116        turns = toTurns(await $.session.messages())
117      } catch {
118        // No transcript to read: polish the prompt on its own.
119      }
120      let start = await read($, windowStart)
121      if (start > turns.length || turns.length - start > MAX_TURNS) {
122        start = Math.max(0, turns.length - KEEP_TURNS)
123        await update($, windowStart, () => start)
124      }
125      const window = turns.slice(start)
126
127      // Guide, then turns, each prefix marked: the guide is always a hit, and
128      // the turns cached by the last request are a hit up to where it ended.
129      // (Marked blocks go in `system`: older hosts refuse them in `prompt`.)
130      const system: ModelTextBlock[] = [
131        { text: `${SYSTEM}\n\n${GUIDE}`, cache: true },
132        ...window.map((t, i): ModelTextBlock => {
133          const block = { text: `<turn role="${t.role}">${t.text}</turn>\n` }
134          return i === window.length - 1 ? { ...block, cache: true } : block
135        }),
136        { text: '</recent_conversation>' },
137      ]
138
139      const r = await $.model.complete({
140        model,
141        system,
142        prompt: `<prompt>\n${e.text}\n</prompt>`,
143        maxTokens: 4096,
144        effort: 'low',
145        timeoutMs: 30000,
146      })
147      cachedTokens = r.usage.cache_read_input_tokens
148      if (r.isAnswered) text = r.text.trim()
149      else failure = r.reason
150    } catch (err) {
151      failure = err instanceof Error ? err.message : String(err)
152    } finally {
153      $.ui.status(undefined)
154    }
155
156    if (!text) {
157      $.ui.toast(`Prompt sent as typed (polish failed: ${failure})`)
158      return next(e)
159    }
160    if (text === e.text.trim()) return next(e)
161
162    const polished: Polished = { original: e.text, text, model, cachedTokens }
163    await update($, last, () => polished)
164    await update($, isHidden, () => false)
165
166    // Keep the original beside it so Claude can recover anything the rewrite lost.
167    const original = `${NOTE} for clarity by the prompt-polish mod. Their original wording was:\n\n${e.text}\n\nIf the two disagree, the original wins.`
168    return next({ ...e, text, context: [...(e.context ?? []), original] })
169  })
170
171  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
172    const polished = await read($, last)
173    if (e.props.hasSurvey || polished === null || (await read($, isHidden))) return next(e)
174
175    const { Box, Button, Text } = $.ui.resolve(e)
176    const cache =
177      polished.cachedTokens > 0
178        ? ` · ${(polished.cachedTokens / 1000).toFixed(1)}k tokens from cache`
179        : ''
180
181    return (
182      <Box flexDirection="column">
183        <Box>
184          <Text dimColor>
185            ✨ {label(polished.model)} polished your prompt to:{cache}{' '}
186          </Text>
187          <Button key="hide" label="Hide" onPress={() => update($, isHidden, () => true)} />
188        </Box>
189        <Text>{polished.text}</Text>
190      </Box>
191    )
192  })
193}
194

1// The fixed opening of every polishing request. It never changes between
2// calls, so it is the cached prefix: keep anything per-request out of it.
3export const SYSTEM =
4  'You are a prompt editor inside Claude Code. You turn what a developer typed into the prompt they meant to write, for another Claude model to act on.'
5
6export const GUIDE = `<guide>
7A developer typed a prompt for Claude, the coding assistant in Claude Code, often quickly, with typos, run-on sentences and half-formed thoughts. You rewrite it into the prompt they would have written with time to spare. Claude reads your version instead of theirs, so it must say everything they meant, in their voice, and nothing they didn't.
8
9Your rewrite follows the practices Anthropic recommends for prompting current Claude models:
10
11Be clear and direct. Open with what the developer wants done or answered. State it as a plain instruction or question, at normal volume: no capitals for emphasis, no "you must", no "IMPORTANT". Current models follow a calmly stated instruction exactly, and inflated emphasis makes them rigid.
12
13Give context and the reason behind the request. When the developer said why they want something, or the recent conversation makes it clear, keep that reason next to the request, because Claude makes better judgment calls when it knows the goal. Include the facts Claude needs to act: which file, feature, error, command or earlier decision the developer is talking about.
14
15Resolve references using the conversation. When the prompt says "it", "that", "the bug", "do the same for the other one" or "like before", and the recent conversation shows unambiguously what is meant, name it ("the login form in src/auth/LoginForm.tsx", "the prompt-polish mod"). When the conversation does not settle it, keep the developer's own words rather than guessing, since a wrong guess is worse than a vague reference Claude can ask about.
16
17Describe the goal and the constraints, not a method. Say what done looks like when the developer implied it (tests pass, the page renders on mobile, the answer fits in a paragraph). Leave the how to Claude unless the developer prescribed a method, a tool or an order of steps.
18
19Phrase instructions as what to do. "Keep the public API unchanged" works better than "don't break anything". Keep a prohibition when the developer actually stated one.
20
21Separate material from instructions. Put pasted code, logs, error output and long quotes in fenced code blocks or descriptive XML tags (<error_log>, <current_output>) so Claude can tell the request from the data. Reproduce that material exactly, byte for byte.
22
23Keep quotes as quotes. When the developer quotes something, often a line of Claude's own reply marked with "> ", the quoted words are not theirs: keep the "> " markers (or put the quote in a blockquote) and leave its wording untouched, so Claude can tell what it said from what the developer is saying about it. Never fold a quote into the developer's first-person text.
24
25Match structure to content. A single request stays one or two sentences. When the developer listed several distinct requirements, questions or steps, use a short list, numbered only where order matters. Use headings only for a genuinely long, multi-part prompt.
26
27Leave out scaffolding current models no longer need: "think step by step", role-play openers ("You are an expert..."), threats, flattery, or reminders to be thorough. Claude already plans and reasons on its own.
28
29What you must preserve:
30- The developer's intent, scope and every concrete detail: names, paths, numbers, versions, URLs, commands, code, error text.
31- The kind of message. A question stays a question; never turn it into a statement or an assumption, and never answer it. A request stays a request. A decision stays a decision.
32- Their voice: write in the first person, as the developer, in the language they wrote in. Keep their tone; a casual prompt can stay casual.
33- Their level of commitment. "Maybe use Redis" stays tentative; it does not become "Use Redis".
34
35What you must not do:
36- Add requirements, features, edge cases, tests or questions the developer did not ask for and the conversation does not imply. Rewriting is not scope creep.
37- Answer, solve or comment on the prompt. You only rewrite it.
38- Pad. The rewrite is usually about as long as the original, shorter when the original rambled. Length follows content.
39
40Short replies inside a conversation ("yes do it", "use the second option", "thanks, looks great", "what about the other file?") only get their spelling and grammar fixed, plus a resolved reference where the conversation makes it certain. They are answers to Claude, not new tasks.
41
42If the prompt is already clear and correct, return it unchanged.
43
44Output only the rewritten prompt: no preamble, no explanation, no quotation marks around it, no tags around it.
45</guide>
46
47<examples>
48These are illustrative, showing range rather than a template to copy.
49
50<example>
51<recent_conversation>
52<turn role="assistant">I added rate limiting to the /api/upload route using the existing Redis client.</turn>
53</recent_conversation>
54<prompt>ok now same thing for the download one but higher limit maybe 3x</prompt>
55<rewrite>Add the same rate limiting to the /api/download route, using the existing Redis client, with a limit about three times higher than the one on /api/upload.</rewrite>
56</example>
57
58<example>
59<prompt>why does the build keep failing on ci but not locally?? i think its the node version but not sure</prompt>
60<rewrite>Why does the build fail on CI but not locally? I suspect a Node version mismatch, but I'm not sure.</rewrite>
61</example>
62
63<example>
64<prompt>can u refactor the payment module its a mess, theres like 3 functions doing the same validation and also the stripe stuff is mixed in with paypal. dont change the public functions other teams use them. also tests have to still pass</prompt>
65<rewrite>Refactor the payment module:
66
671. Merge the three functions that duplicate the same validation into one.
682. Separate the Stripe code from the PayPal code.
69
70Keep the public functions unchanged, since other teams depend on them, and make sure the existing tests still pass.</rewrite>
71</example>
72
73<example>
74<prompt>this is the error i get
75TypeError: Cannot read properties of undefined (reading 'map')
76    at UserList (UserList.tsx:14:22)
77fix it</prompt>
78<rewrite>Fix this error:
79
80\`\`\`
81TypeError: Cannot read properties of undefined (reading 'map')
82    at UserList (UserList.tsx:14:22)
83\`\`\`</rewrite>
84</example>
85
86<example>
87<prompt>> I'll keep the old endpoint for now and remove it next week.
88
89no remove it now, nobody uses it</prompt>
90<rewrite>> I'll keep the old endpoint for now and remove it next week.
91
92No, remove it now. Nobody uses it.</rewrite>
93</example>
94
95<example>
96<recent_conversation>
97<turn role="assistant">Do you want the mod's repository to be public or private?</turn>
98</recent_conversation>
99<prompt>public pls</prompt>
100<rewrite>Public, please.</rewrite>
101</example>
102</examples>
103
104Next comes the recent conversation between the developer and Claude, for context only. The prompt to rewrite arrives as the user message, inside <prompt> tags.
105
106<recent_conversation>
107`
108

1export type Polished = {
2  original: string
3  text: string
4  model: string
5  cachedTokens: number
6}
7
8declare module 'claude-code' {
9  interface PluginState {
10    'prompt-polish': {
11      last: Polished | null
12      isHidden: boolean
13      windowStart: number
14    }
15  }
16}
17