prompt-upgrade

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, ModelCompleteResult, Register } from 'claude-code'
3
4import type { Question, Upgrade } from '../types'
5
6// /improve <prompt> opens a pane: the model asks 2-4 multiple-choice questions
7// about the prompt, the person picks, and the model writes a fuller prompt from
8// the answers, to run, edit in the prompt box, or copy. Two $.model.complete
9// calls per upgrade, only when asked: prompts sent normally are never touched.
10// $.state, $.model.complete, $.session.messages, $.prompt and $.ui only: no
11// files, network or processes.
12
13const PANE = 'prompt-upgrade'
14const view = atom({ plugin: 'prompt-upgrade', key: 'view' } as const, null)
15const MAX_QUESTIONS = 4
16const MAX_OPTIONS = 5
17const CLIP = 120
18const SIDE_BY_SIDE = 100
19
20const ASK = [
21  'You help a person turn a short request for an AI assistant into a strong prompt.',
22  'Ask 2 to 4 multiple-choice questions whose answers would change the result most: the goal, who it is for, scope, format, length, tone, constraints, budget.',
23  'Each question has 2 to 4 short options. Do not ask what the request already says. If the request is already specific, ask nothing.',
24  'Reply with JSON only, no prose: {"questions":[{"q":"...","options":["...","..."]}]}',
25].join('\n')
26
27const WRITE = [
28  "Rewrite the person's request as one clear, complete prompt for an AI assistant, using their answers.",
29  'Keep their intent and their voice. State the goal, the context, the constraints and the output they want.',
30  'Use only what the request and the answers say; leave out what is unknown rather than inventing it.',
31  'Reply with the prompt alone: no preamble, no quotes, no notes.',
32].join('\n')
33
34const clip = (text: string) => {
35  const t = text.replace(/\s+/g, ' ').trim()
36  return t.length > CLIP ? t.slice(0, CLIP - 1) + '…' : t
37}
38
39/** The questions in the model's reply: at most 4, each with 2-5 distinct options; [] when the reply has none that parse. */
40export function parseQuestions(reply: string): Question[] {
41  const start = reply.indexOf('{')
42  const end = reply.lastIndexOf('}')
43  if (start < 0 || end <= start) return []
44  let data: unknown
45  try {
46    data = JSON.parse(reply.slice(start, end + 1))
47  } catch {
48    return []
49  }
50  const list = (data as { questions?: unknown } | null)?.questions
51  if (!Array.isArray(list)) return []
52  const out: Question[] = []
53  for (const item of list as { q?: unknown; options?: unknown }[]) {
54    const text = clip(String(item?.q ?? ''))
55    const raw = Array.isArray(item?.options) ? (item.options as unknown[]) : []
56    const options = [...new Set(raw.map(o => clip(String(o))))].filter(Boolean).slice(0, MAX_OPTIONS)
57    if (text && options.length >= 2) out.push({ id: `q${out.length + 1}`, text, options })
58    if (out.length === MAX_QUESTIONS) break
59  }
60  return out
61}
62
63/** What the rewrite reads: the request, then each answered question and the extra note. */
64export function writeRequest(v: Pick<Upgrade, 'original' | 'questions' | 'answers' | 'extra'>): string {
65  const lines = v.questions.filter(q => v.answers[q.id]).map(q => `- ${q.text} ${v.answers[q.id]}`)
66  if (v.extra.trim()) lines.push(`- Also: ${v.extra.trim()}`)
67  return `Request:\n${v.original}\n\nAnswers:\n${lines.length > 0 ? lines.join('\n') : '- (none given)'}`
68}
69
70/** The reply without a fence or quotes around the whole of it. */
71export function cleanReply(text: string): string {
72  let t = text.trim()
73  const fenced = /^```[\w-]*\n([\s\S]*?)\n```$/.exec(t)
74  if (fenced?.[1] !== undefined) t = fenced[1].trim()
75  if (t.length >= 2 && ((t.startsWith('"') && t.endsWith('"')) || (t.startsWith('“') && t.endsWith('”')))) t = t.slice(1, -1).trim()
76  return t
77}
78
79const reasonOf = (r: Exclude<ModelCompleteResult, { isAnswered: true }>) =>
80  r.reason === 'api-error' ? `the API answered with an error${r.status ? ` (${r.status})` : ''}` : r.reason === 'empty-reply' ? 'the model replied with nothing' : 'it was cancelled'
81
82// a user row the person never typed: a task notification, a reminder, a command's echo (<task-notification>,
83// <system-reminder>, <local-command-stdout>: the engine's tags are hyphenated), an interrupt marker. A
84// SessionMessage carries no origin, so its text is all there is to tell by.
85const NOT_TYPED = /^(?:<[a-z]+(?:-[a-z]+)+[\s>]|\[Request interrupted)/i
86
87/** The person's last prompt in this session, skipping tool results, slash commands and rows they never typed. */
88async function lastPrompt($: EngineInterface): Promise<string> {
89  const messages = await $.session.messages()
90  if (!Array.isArray(messages)) return ''
91  for (let i = messages.length - 1; i >= 0; i--) {
92    const m = messages[i]
93    const t = m?.text.trim() ?? ''
94    if (m?.role === 'user' && t && !(m.toolResults?.length ?? 0) && !t.startsWith('/') && !t.includes('<command-') && !NOT_TYPED.test(t)) return t
95  }
96  return ''
97}
98
99// the userConfig model, set as the module registers
100let model = 'sonnet'
101// each start bumps it, and so does closing: a reply for an older run is dropped
102let run = 0
103
104function complete($: EngineInterface, prompt: string, system: string, maxTokens: number): Promise<ModelCompleteResult> {
105  return $.model
106    .complete({ model, system, prompt, maxTokens })
107    .catch((err: unknown) => ({ isAnswered: false, reason: 'api-error', status: null, error: err instanceof Error ? err.message : String(err) }) as never)
108}
109
110async function fail($: EngineInterface, mine: number, error: string) {
111  if (mine === run) await update($, view, v => (v ? { ...v, phase: 'error' as const, error } : v))
112}
113
114async function write($: EngineInterface, mine: number) {
115  const v = await read($, view)
116  if (!v || mine !== run) return
117  await update($, view, () => ({ ...v, phase: 'writing' as const, error: '' }))
118  const r = await complete($, writeRequest(v), WRITE, 2000)
119  if (mine !== run) return
120  if (!r.isAnswered) return fail($, mine, reasonOf(r))
121  const upgraded = cleanReply(r.text)
122  if (!upgraded) return fail($, mine, 'the model returned an empty prompt')
123  await update($, view, x => (x ? { ...x, phase: 'ready' as const, upgraded } : x))
124}
125
126async function ask($: EngineInterface, original: string) {
127  const mine = ++run
128  await update($, view, () => ({ phase: 'asking' as const, original, questions: [], answers: {}, extra: '', upgraded: '', error: '' }))
129  await $.ui.open({ id: PANE, title: 'Prompt upgrade', focus: true, closeOnEscape: true })
130  const r = await complete($, original, ASK, 1200)
131  if (mine !== run) return
132  if (!r.isAnswered) return fail($, mine, reasonOf(r))
133  const questions = parseQuestions(r.text)
134  // already specific, or a reply that didn't parse: straight to the rewrite
135  if (questions.length === 0) return write($, mine)
136  await update($, view, v => (v ? { ...v, phase: 'answering' as const, questions } : v))
137}
138
139async function close($: EngineInterface) {
140  run++
141  await update($, view, () => null)
142  await $.ui.close({ id: PANE })
143}
144
145/** A start that failed outside the model call (the pane, the state) shows its reason in the pane. */
146async function begin($: EngineInterface, original: string) {
147  await ask($, original).catch((err: unknown) => fail($, run, err instanceof Error ? err.message : String(err)))
148}
149
150export const register: Register = (on, options) => {
151  model = String(options?.model ?? '').trim() || 'sonnet'
152
153  on('session.start', async ($, e, next) => {
154    await $.command.register({
155      name: 'improve',
156      description: 'Prompt upgrade: answer a few questions, get a fuller prompt to run, edit or copy (no text: your last prompt)',
157      argumentHint: '<your prompt>',
158    })
159    return next(e)
160  })
161
162  on('command.run', { command: 'improve' }, async ($, e) => {
163    const typed = String(e.args ?? '').trim()
164    const original = typed || (await lastPrompt($))
165    if (!original) return { text: 'Usage: /improve <your prompt>. With nothing after it, /improve takes your last prompt.' }
166    void begin($, original)
167    return { text: `Prompt upgrade: reading ${typed ? 'your prompt' : 'your last prompt'}. Answer its questions in the pane.` }
168  }).catch(() => ({ text: 'Prompt upgrade could not start: your last prompt could not be read. Try /improve <your prompt>.' }))
169
170  // the person closed the pane: drop the run and its view
171  on('ui.close', async ($, e, next) => {
172    if (e.id === PANE) {
173      run++
174      await update($, view, () => null)
175    }
176    return next(e)
177  }).catch(($, e, next) => next(e))
178
179  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
180    const els = $.ui.resolve(e)
181    const { Box, Button, Text } = els
182    // the mobile app draws no text field: there the questions alone shape the rewrite
183    const Input = 'Input' in els ? els.Input : undefined
184    const v = await read($, view)
185    if (!v) return <Text dimColor>{'Nothing to upgrade. /improve <your prompt> starts one.'}</Text>
186    const cols = e.props.bodyColumns ?? 80
187    const set = (fn: (x: Upgrade) => Upgrade) => void update($, view, x => (x ? fn(x) : x))
188    const typed = (
189      <Box flexDirection="column">
190        <Text dimColor>You typed</Text>
191        <Text>{v.original}</Text>
192      </Box>
193    )
194    const cancel = <Button key="cancel" label="Cancel" onPress={() => void close($)} />
195
196    if (v.phase === 'asking' || v.phase === 'writing')
197      return (
198        <Box flexDirection="column" rowGap={1}>
199          {typed}
200          <Text dimColor>{v.phase === 'asking' ? 'Reading your prompt…' : 'Writing your upgraded prompt…'}</Text>
201          {cancel}
202        </Box>
203      )
204
205    if (v.phase === 'error')
206      return (
207        <Box flexDirection="column" rowGap={1}>
208          {typed}
209          <Text color="#f87171">Couldn't upgrade it: {v.error}.</Text>
210          <Box flexDirection="row" columnGap={1}>
211            <Button key="retry" label="Try again" variant="primary" onPress={() => void begin($, v.original)} />
212            <Button key="close" label="Close" onPress={() => void close($)} />
213          </Box>
214        </Box>
215      )
216
217    if (v.phase === 'answering')
218      return (
219        <Box flexDirection="column" rowGap={1}>
220          {typed}
221          {v.questions.map((q, i) => {
222            const picked = v.answers[q.id] ?? ''
223            const pick = (value: string) => () => set(x => ({ ...x, answers: { ...x.answers, [q.id]: value } }))
224            return (
225              <Box key={`row-${q.id}`} flexDirection="column">
226                <Text bold>
227                  {i + 1}. {q.text}
228                </Text>
229                <Box flexDirection="row" flexWrap="wrap" columnGap={1}>
230                  {q.options.map((o, n) => (
231                    <Button
232                      key={`${q.id}-${n}`}
233                      label={picked === o ? `● ${o}` : o}
234                      variant={picked === o ? 'primary' : 'secondary'}
235                      {...(i === 0 && n === 0 ? { autoFocus: true as const } : {})}
236                      onPress={pick(o)}
237                    />
238                  ))}
239                  <Button key={`${q.id}-skip`} label={picked ? 'Skip' : '● Skip'} plain onPress={pick('')} />
240                </Box>
241              </Box>
242            )
243          })}
244          {Input ? (
245            <Input
246              key="extra"
247              label="Anything else? "
248              placeholder="optional"
249              value={v.extra}
250              onInput={(value: string) => set(x => ({ ...x, extra: value }))}
251              onSubmit={(value: string) => set(x => ({ ...x, extra: value }))}
252            />
253          ) : null}
254          <Box flexDirection="row" columnGap={1}>
255            <Button key="write" label="Write my prompt" variant="primary" onPress={() => void write($, run)} />
256            {cancel}
257          </Box>
258        </Box>
259      )
260
261    const half = Math.max(20, Math.floor((cols - 3) / 2))
262    const upgraded = (
263      <Box flexDirection="column">
264        <Text dimColor>Upgraded</Text>
265        <Text>{v.upgraded}</Text>
266      </Box>
267    )
268    return (
269      <Box flexDirection="column" rowGap={1}>
270        {cols >= SIDE_BY_SIDE ? (
271          <Box flexDirection="row" columnGap={3}>
272            <Box width={half}>{typed}</Box>
273            <Box width={half}>{upgraded}</Box>
274          </Box>
275        ) : (
276          <Box flexDirection="column" rowGap={1}>
277            {typed}
278            {upgraded}
279          </Box>
280        )}
281        <Box flexDirection="row" columnGap={1}>
282          <Button
283            key="run"
284            label="Run"
285            variant="primary"
286            onPress={() => {
287              const text = v.upgraded
288              // the person just read it in the pane: it runs as their own prompt, not one the plugin sent
289              void close($).then(() => $.prompt.submit({ text, asUser: true }))
290            }}
291          />
292          <Button
293            key="edit"
294            label="Edit first"
295            onPress={() => {
296              const text = v.upgraded
297              void $.prompt.fill({ text }).then(() => close($))
298            }}
299          />
300          <Button
301            key="copy"
302            label="Copy"
303            onPress={press =>
304              void $.ui.copy({ text: v.upgraded, surface: press.surface }).then(r => $.ui.toast(r.isCopied ? 'Upgraded prompt copied.' : `Couldn't copy: ${r.reason}`))
305            }
306          />
307          <Button key="again" label="Start over" onPress={() => void begin($, v.original)} />
308          <Button key="close" label="Close" onPress={() => void close($)} />
309        </Box>
310      </Box>
311    )
312  })
313}
314

1/** One question the model asked about the prompt, with the answers to pick from. */
2export type Question = {
3  /** 'q1', 'q2', ...: the Select's key */
4  id: string
5  text: string
6  options: string[]
7}
8
9/** The pane's flow, from the prompt as typed to the upgraded one. */
10export type Upgrade = {
11  /** asking: the questions are being written; answering: the person picks; writing: the rewrite runs */
12  phase: 'asking' | 'answering' | 'writing' | 'ready' | 'error'
13  /** the prompt as the person typed it */
14  original: string
15  questions: Question[]
16  /** question id -> the option picked; absent or '' when skipped */
17  answers: Record<string, string>
18  /** what the person typed under "Anything else?" */
19  extra: string
20  /** the upgraded prompt, once written */
21  upgraded: string
22  /** why the last model call gave nothing */
23  error: string
24}
25
26declare module 'claude-code' {
27  interface PluginState {
28    'prompt-upgrade': {
29      /** null while the pane is closed */
30      view: Upgrade | null
31    }
32  }
33}
34