prompt-optimizer

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4import type { Busy, Suggestion } from '../types'
5
6const MODEL = 'claude-opus-5-5'
7// Writes the first suggestion guess while Opus works
8const FAST_MODEL = 'claude-haiku-5-5'
9// The conversation the suggester and the optimizer read, from its end
10const CONVERSATION_CHARS = 12000
11// Saved (answer end, next prompt) pairs: how many to keep, and how many to show the model
12const HABITS_KEY = 'habits'
13const MAX_HABITS = 40
14const HABIT_EXAMPLES = 10
15
16const isBusy = atom({ plugin: 'prompt-optimizer', key: 'isBusy' } as const, false as Busy)
17const suggestion = atom(
18  { plugin: 'prompt-optimizer', key: 'suggestion' } as const,
19  { status: 'none' } as Suggestion,
20)
21
22const MAX_NEXT_WORDS = 20
23const MAX_SUGGESTIONS = 3
24
25const NEXT_RULES = `You write the user's next prompt in a chat with an AI coding agent.
26Read the workspace, the user's recent prompts, and the agent's last answer. Write up to ${MAX_SUGGESTIONS} different prompts the user is likely to send next, most likely first, one per line:
271. The answer asks a question or offers a choice: answer it; the first line picks the agent's recommended option, another line picks the alternative.
282. The answer names a next step: ask for that step.
293. The answer changed something it did not check: ask to test it where it runs.
304. Otherwise: ask for the next step toward the goal of the recent prompts.
31Rules for each line:
32- Write it the way this user writes next prompts in the examples: their words, their length, and the kind of step they take after an answer like this one.
33- ${MAX_NEXT_WORDS} words or fewer, an imperative the user sends as is, with no edits.
34- Stay on the task of the workspace and the recent prompts; never start an unrelated task.
35- Never state a result, observation, or fact the user has not said.
36- Name the exact file, command, function, or item; never only "it" or "this".
37- No numbering, bullets, placeholders, quotes, greetings, or "please".
38Output only the lines. Fewer lines are better than a weak one.`
39
40const RULES = `You rewrite a prompt draft so another AI agent gets the real outcome the writer wants, not only the task they named.
41
42Before you write, think silently about what a top professional in this field would make true here that the writer did not ask for. Use it to raise the bar inside the writer's scope. Do not widen the scope.
43
44Output exactly one paragraph that starts with "Goal: " and describes the outcome: what is true for the user when the work is done. Output plain text with nothing before or after it (no code fence, no other sections, no notes, no questions).
45
46Goal rules:
47- Exactly one goal. Fold every part of the draft into one end result for the user.
48- 1-2 sentences, 50 words or fewer in all. Each sentence is outcome-first (the subject is the user or the finished thing), specific, measurable, and atomic.
49- Name no method (audit, refactor, review). Limits are fine.
50- Keep every standard, book, framework, file, path, name, and number the draft mentions.
51- Give a time frame or end point when the draft has one.
52- Replace every vague word with an exact value: "about 30 miles" becomes "within 30 miles", "good" or "easy" becomes what the user does and how fast or how often.
53- Make every claim checkable: "complete" or "current" names the source of truth it is checked against (for example "every job posted on careers.example.com today").
54- Resolve contradictions: when two parts can conflict (a location filter and remote roles), state which one wins.
55- Cut every clause that only restates the outcome or explains why it helps ("so the user can...", "letting the user...", "within one scan").
56- Positive framing. No "think carefully" or "double-check".
57
58Intent rules:
59- Read the context first: the workspace, the writer's recent prompts, and the agent's last answer. Find what the draft points at ("this", "that", "it", "the bug", "those 2") and what the writer is trying to finish.
60- Name those things exactly in the Goal: the file, feature, item, or number from the context.
61- Aim the Goal at the result the writer needs next in that work, not at a general reading of the draft.
62- Use the context only to sharpen the draft; never add a task the draft does not ask for. Same length as without context.
63
64When a fact is missing, choose the most likely value and keep going. Never ask questions.`
65
66// The main loop's recent prompts, newest last, and the id of the newest suggestion request.
67const recentPrompts: string[] = []
68// The main loop's last answer, so the optimizer knows what a draft points at
69let lastAnswer = ''
70let suggestRun = 0
71// Suggestion calls this module instance has running; a reload starts a new instance with none.
72let callsInFlight = 0
73
74async function isDesktop($: EngineInterface) {
75  return (await $.session.surfaces()).includes('desktop')
76}
77
78// Lines of the model's reply that can be sent as is: no list marks or quotes, 1 to MAX_NEXT_WORDS words, no repeats
79export function parseSuggestions(reply: string): string[] {
80  const lines = reply
81    .split('\n')
82    .map(l => l.trim().replace(/^(\d+[.)]|[-*•])\s+/, '').replace(/^["'`]+|["'`]+$/g, '').trim())
83    .filter(l => l !== '' && l !== 'NONE')
84    .filter(l => l.split(/\s+/).length <= MAX_NEXT_WORDS)
85  return [...new Set(lines)].slice(0, MAX_SUGGESTIONS)
86}
87
88// A reply the agent would write, not a prompt the user would send: the fast model sometimes answers instead
89export function soundsLikeAgent(reply: string) {
90  return /^\s*(I|I'm|I'll|I've|Sure|Here|Okay|OK)\b|\?\s*$/m.test(reply)
91}
92
93async function workspace($: EngineInterface) {
94  const cwd = await $.session.cwd().catch(() => '')
95  const repo = await $.session.repo().catch(() => null)
96  return [cwd && `Folder: ${cwd}`, repo && `Repository: ${JSON.stringify(repo)}`].filter(Boolean).join('\n')
97}
98
99// The conversation so far, oldest first, cut to its last CONVERSATION_CHARS; falls back to the recent prompts
100async function conversation($: EngineInterface) {
101  const rows = await $.session.messages().catch(() => [])
102  const text = rows
103    .filter(m => m.text.trim() !== '')
104    .map(m => `${m.role === 'user' ? 'User' : 'Agent'}: ${m.text.trim().slice(0, 3000)}`)
105    .join('\n\n')
106  return text !== '' ? text.slice(-CONVERSATION_CHARS) : recentPrompts.map(p => `User: ${p}`).join('\n\n')
107}
108
109// Pairs of (end of the agent's answer, the prompt the user sent next), kept on this computer across sessions
110type HabitPair = { after: string; next: string }
111
112async function habits($: EngineInterface): Promise<HabitPair[]> {
113  const saved = await $.store.get(HABITS_KEY).catch(() => undefined)
114  return Array.isArray(saved) ? (saved as HabitPair[]) : []
115}
116
117async function learnHabit($: EngineInterface, next: string) {
118  if (lastAnswer.trim() === '' || next.trim() === '') return
119  const pairs = [...(await habits($)), { after: lastAnswer.slice(-400), next: next.slice(0, 300) }]
120  await $.store.set(HABITS_KEY, pairs.slice(-MAX_HABITS)).catch(() => {})
121}
122
123function habitExamples(pairs: HabitPair[]) {
124  return pairs
125    .slice(-HABIT_EXAMPLES)
126    .map(p => `<example>\n<answer_end>${p.after}</answer_end>\n<user_sent>${p.next}</user_sent>\n</example>`)
127    .join('\n')
128}
129
130// Everything a suggestion is written from: workspace, the user's habits, the conversation, and the last answer
131async function suggestContext($: EngineInterface, answer: string) {
132  return [
133    `<workspace>\n${await workspace($)}\n</workspace>`,
134    `<how_this_user_writes_next_prompts>\n${habitExamples(await habits($))}\n</how_this_user_writes_next_prompts>`,
135    `<conversation>\n${await conversation($)}\n</conversation>`,
136    `<answer>\n${answer.slice(-2500)}\n</answer>`,
137    `You are not the agent and you do not answer the user. Write the prompts the user sends the agent next: up to ${MAX_SUGGESTIONS} lines in the user's voice. Output only those lines.`,
138  ].join('\n\n')
139}
140
141async function suggestNext($: EngineInterface, answer: string) {
142  const run = ++suggestRun
143  // Counted until the state leaves 'thinking', so a redraw in between does not start a second call
144  callsInFlight++
145  let isFinal = false
146  const ask = (model: string, effort: 'low' | 'medium', context: string) =>
147    $.model.complete({
148      model,
149      effort,
150      system: [{ text: NEXT_RULES, cache: true }],
151      prompt: context,
152      // Room for thinking; the lines themselves are short
153      maxTokens: 1024,
154      timeoutMs: 30000,
155    })
156  try {
157    await update($, suggestion, (): Suggestion => ({ status: 'thinking', prompt: '', answer }))
158    const context = await suggestContext($, answer)
159    // A fast first guess shows in seconds; Opus replaces it when it is done
160    const fast = ask(FAST_MODEL, 'low', context).then(async r => {
161      const texts = r.isAnswered && !soundsLikeAgent(r.text) ? parseSuggestions(r.text) : []
162      if (run === suggestRun && !isFinal && texts.length > 0) {
163        await update($, suggestion, (): Suggestion => ({ status: 'ready', texts }))
164      }
165    })
166    const r = await ask(MODEL, 'medium', context)
167    isFinal = true
168    if (run !== suggestRun) return
169    // A line over the limit is dropped, not cut: a cut-off prompt reads as a broken one
170    const texts = r.isAnswered ? parseSuggestions(r.text) : []
171    if (texts.length > 0) await update($, suggestion, (): Suggestion => ({ status: 'ready', texts }))
172    else {
173      // Opus gave nothing usable; keep the fast guess when it has one
174      await fast.catch(() => {})
175      const now = await read($, suggestion)
176      if (now.status !== 'ready') await update($, suggestion, (): Suggestion => ({ status: 'none' }))
177    }
178  } finally {
179    callsInFlight--
180  }
181}
182
183async function useSuggestion($: EngineInterface, index: number) {
184  const now = await read($, suggestion)
185  if (now.status !== 'ready') return
186  const text = now.texts[index]
187  if (text === undefined) return
188  suggestRun++
189  await update($, suggestion, () => ({ status: 'none' }))
190  if (!(await isDesktop($))) {
191    await $.prompt.fill({ text })
192    return
193  }
194  // The desktop box takes no text from a mod, and a prompt a mod sends carries a "plugin sent" label,
195  // so the click copies it: Ctrl+V and Enter send it as the user's own
196  if ((await $.ui.copy({ text, surface: 'desktop' })).isCopied) $.ui.toast('Copied. Press Ctrl+V, then Enter.')
197  else await $.prompt.submit({ text })
198}
199
200async function dismissSuggestion($: EngineInterface) {
201  suggestRun++
202  await update($, suggestion, () => ({ status: 'none' }))
203}
204
205// The draft before the last rewrite, and the rewrite itself, so a second press undoes it.
206let original: string | null = null
207let rewritten: string | null = null
208
209const MAX_GOAL_WORDS = 50
210// Words that leave a Goal unmeasurable; one found sends the Goal back once for an exact value.
211const VAGUE =
212  /\b(about|around|roughly|approximately|near|nearby|modern|good|better|nice|easy|easily|quick|quickly|fast|faster|soon|cheap|some|several|few|various|etc|kinda|pretty|or so)\b|\w+-ish\b/gi
213
214function goalProblems(goal: string): string[] {
215  const problems: string[] = []
216  const words = goal.split(/\s+/).filter(Boolean).length
217  if (words > MAX_GOAL_WORDS) problems.push(`It has ${words} words; the limit is ${MAX_GOAL_WORDS}.`)
218  const vague = [...new Set((goal.match(VAGUE) ?? []).map(w => w.toLowerCase()))]
219  if (vague.length > 0) problems.push(`It uses vague words: ${vague.map(w => `"${w}"`).join(', ')}. Replace each with an exact value or cut it.`)
220  return problems
221}
222
223function cleanGoal(text: string): string {
224  return text.trim().replace(/^```\w*\n?|\n?```$/g, '').trim()
225}
226
227async function optimize($: EngineInterface, draft: string): Promise<string | null> {
228  await update($, isBusy, () => true)
229  $.ui.status('Optimizing prompt...')
230  try {
231    const ask = async (prompt: string) =>
232      $.model.complete({
233        model: MODEL,
234        effort: 'medium',
235        system: [{ text: RULES, cache: true }],
236        prompt,
237        maxTokens: 2048,
238        timeoutMs: 60000,
239      })
240    const context = `<workspace>\n${await workspace($)}\n</workspace>\n\n<conversation>\n${await conversation($)}\n</conversation>`
241    const first = `${context}\n\nRewrite this prompt draft:\n\n<draft>\n${draft}\n</draft>`
242    const r = await ask(first)
243    if (!r.isAnswered) {
244      $.ui.toast(`Prompt optimizer failed: ${r.reason}`)
245      return null
246    }
247    const goal = cleanGoal(r.text)
248    const problems = goalProblems(goal)
249    if (problems.length === 0) return goal
250
251    // One retry: shorten and make exact. Keep the first Goal if the retry fails.
252    const retry = await ask(
253      `${first}\n\nYour Goal was:\n\n<goal>\n${goal}\n</goal>\n\n${problems.join(' ')} Rewrite the Goal to fix this. Keep every name, number, and source.`,
254    )
255    return retry.isAnswered ? cleanGoal(retry.text) : goal
256  } finally {
257    $.ui.status(undefined)
258    await update($, isBusy, () => false)
259  }
260}
261
262async function press($: EngineInterface) {
263  if (await read($, isBusy)) return
264  const { text } = await $.prompt.read()
265  const draft = text.trim()
266
267  if (rewritten !== null && original !== null && draft === rewritten.trim()) {
268    await $.prompt.fill({ text: original })
269    rewritten = null
270    $.ui.toast('Original prompt restored')
271    return
272  }
273  if (draft === '') {
274    $.ui.toast('Type a prompt first, then click Optimize prompt')
275    return
276  }
277
278  const better = await optimize($, draft)
279  if (better === null) return
280  const { isFilled } = await $.prompt.fill({ text: better })
281  if (!isFilled) {
282    $.ui.toast('Could not write to the prompt box')
283    return
284  }
285  original = text
286  rewritten = better
287}
288
289export const register: Register = on => {
290  on('session.start', async ($, e, next) => {
291    // A reload drops a suggestion call in flight; start it again from what it was written from.
292    const left = await read($, suggestion)
293    if (left.status === 'thinking') {
294      void suggestNext($, left.answer)
295    }
296    await $.command.register({
297      name: 'optimize',
298      description: 'Rewrite a prompt with Opus 5.5 and put it in the prompt box',
299      argumentHint: '<prompt draft>',
300    })
301    return next(e)
302  })
303
304  on('turn.start', async ($, e, next) => {
305    if (e.agentId === undefined) {
306      await learnHabit($, e.text)
307      recentPrompts.push(e.text.slice(-600))
308      if (recentPrompts.length > 3) recentPrompts.shift()
309    }
310    return next(e)
311  })
312
313  on('prompt.submit', async ($, e, next) => {
314    await dismissSuggestion($)
315    return next(e)
316  })
317
318  on('turn.complete', async ($, e, next) => {
319    const done = await next(e)
320    if (e.agentId === undefined && e.reason === 'answer' && e.answer.trim() !== '') {
321      // Not awaited: the suggestion fills the band later and never holds up the turn.
322      lastAnswer = e.answer
323      void suggestNext($, e.answer)
324    }
325    return done
326  })
327
328  on('command.run', { command: 'optimize' }, async ($, e) => {
329    const draft = e.args.trim()
330    if (draft === '') return { text: 'Usage: /optimize <prompt draft>. Or type the draft and click Optimize prompt.' }
331    const better = await optimize($, draft)
332    if (better === null) return { text: 'Prompt optimizer failed.' }
333    await $.prompt.fill({ text: better })
334    original = draft
335    rewritten = better
336    return { text: 'Optimized prompt is in the prompt box. Click Optimize prompt to restore the original.' }
337  })
338
339  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
340    if (e.props.hasSurvey) return next(e)
341    const { Box, Button, Text } = $.ui.resolve(e)
342    const busy = await read($, isBusy)
343    const next_ = await read($, suggestion)
344
345    return (
346      <Box flexDirection="column">
347        {next_.status === 'thinking' && <Text dimColor>Next prompt: thinking...</Text>}
348        {next_.status === 'ready' && (
349          <Box flexDirection="row" gap={1}>
350            <Box flexDirection="column">
351              {next_.texts.map((text, i) => (
352                <Button key={`use-suggestion-${i}`} variant={i === 0 ? 'primary' : undefined} onPress={() => useSuggestion($, i)}>
353                  {text}
354                </Button>
355              ))}
356            </Box>
357            <Button key="dismiss-suggestion" label="x" plain dimColor onPress={() => dismissSuggestion($)} />
358          </Box>
359        )}
360        <Button
361          key="optimize"
362          label={busy ? 'Optimizing prompt...' : 'Optimize prompt'}
363          plain
364          dimColor
365          onPress={() => press($)}
366        />
367      </Box>
368    )
369  })
370}
371

1export type Busy = boolean
2
3// The next prompt offered after a turn: being written (with what it is written from, so a reload can restart it), ready, or none.
4export type Suggestion =
5  | { status: 'none' }
6  | { status: 'thinking'; prompt: string; answer: string }
7  | { status: 'ready'; texts: string[] }
8
9declare module 'claude-code' {
10  interface PluginState {
11    'prompt-optimizer': { isBusy: Busy; suggestion: Suggestion }
12  }
13}
14