prompt-upgrade

1import { atom, read, update } from 'claude-code'
2import type { Elements, EngineInterface, Register, RenderElement, UiPressArgument } from 'claude-code'
3
4import type { PromptUpgradeAnswer, PromptUpgradeEntry, PromptUpgradeFlow, PromptUpgradeStep } from '../types'
5import {
6  QUESTIONS_SYSTEM,
7  REWRITE_SYSTEM,
8  bigNumber,
9  clip,
10  decide,
11  packRows,
12  parseQuestions,
13  parseRewrite,
14  questionsPrompt,
15  rewritePrompt,
16  wordCount,
17  wrapText,
18} from './logic'
19
20type Engine = EngineInterface
21
22const PANE = 'prompt-upgrade'
23const TITLE = 'Prompt Upgrade'
24const HISTORY_SIZE = 20
25const QUESTIONS_MODEL = 'haiku'
26const REWRITE_MODEL = 'sonnet'
27
28const GOLD = '#F2C46D'
29const HAIR = '#262B33'
30const MUTED = '#7D8590'
31const TEXT = '#E6EDF3'
32const GREEN = '#3DD68C'
33const AMBER = '#E3B341'
34const INK = '#11151B'
35
36const flowAtom = atom({ plugin: 'prompt-upgrade', key: 'flow' } as const, null)
37const viewAtom = atom({ plugin: 'prompt-upgrade', key: 'view' } as const, 'flow')
38const autoAtom = atom({ plugin: 'prompt-upgrade', key: 'auto' } as const, false)
39const historyAtom = atom({ plugin: 'prompt-upgrade', key: 'history' } as const, [])
40
41const NARROW = 'Prompt Upgrade needs a wider window for its pane, so your prompt is back in the box.'
42
43function newId(): string {
44  return `${Date.now().toString(36)}${Math.random().toString(36).slice(2, 8)}`
45}
46
47function reason(err: unknown): string {
48  return err instanceof Error && err.message !== '' ? err.message : String(err)
49}
50
51function newFlow(original: string, step: PromptUpgradeStep): PromptUpgradeFlow {
52  return {
53    id: newId(),
54    original,
55    step,
56    questions: [],
57    answers: [],
58    otherOpen: null,
59    upgraded: '',
60    added: [],
61    isEditing: false,
62    isCopied: false,
63    failedDuring: null,
64    callId: '',
65  }
66}
67
68// ------------------------------------------------------------------ the flow
69
70/** Opens the pane; true once a surface draws it. */
71async function showPane($: Engine): Promise<boolean> {
72  const opened = await $.ui.open({ id: PANE, title: TITLE, focus: true })
73  return opened.isPlaced
74}
75
76async function closePane($: Engine): Promise<void> {
77  await $.ui.close({ id: PANE })
78}
79
80/** Starts an upgrade of `text`. False when the pane could not be placed (the flow is then undone). */
81async function startFlow($: Engine, text: string): Promise<boolean> {
82  const flow = newFlow(text, 'reading')
83  await update($, flowAtom, () => flow)
84  await update($, viewAtom, () => 'flow' as const)
85  const isPlaced = await showPane($)
86  if (!isPlaced) {
87    await update($, flowAtom, () => null)
88    await closePane($)
89    return false
90  }
91  void askQuestions($, flow.id)
92  return true
93}
94
95async function failStep($: Engine, callId: string, during: 'questions' | 'rewrite'): Promise<void> {
96  await update($, flowAtom, f => (f !== null && f.callId === callId ? { ...f, step: 'error', failedDuring: during } : f))
97}
98
99async function askQuestions($: Engine, flowId: string): Promise<void> {
100  const callId = newId()
101  const flow = await update($, flowAtom, f =>
102    f !== null && f.id === flowId ? { ...f, step: 'reading', callId, failedDuring: null } : f,
103  )
104  if (flow === null || flow.callId !== callId) return
105  try {
106    const reply = await $.model.complete({
107      model: QUESTIONS_MODEL,
108      system: QUESTIONS_SYSTEM,
109      prompt: questionsPrompt(flow.original),
110      maxTokens: 800,
111      effort: 'low',
112      timeoutMs: 45_000,
113    })
114    if (!reply.isAnswered) throw new Error(reply.reason)
115    const questions = parseQuestions(reply.text)
116    if (questions.length === 0) {
117      const now = await read($, flowAtom)
118      if (now !== null && now.callId === callId) await writeRewrite($, flowId)
119      return
120    }
121    await update($, flowAtom, f =>
122      f !== null && f.callId === callId
123        ? { ...f, step: 'questions', questions, answers: questions.map(() => null), otherOpen: null }
124        : f,
125    )
126  } catch {
127    await failStep($, callId, 'questions')
128  }
129}
130
131async function writeRewrite($: Engine, flowId: string): Promise<void> {
132  const callId = newId()
133  const flow = await update($, flowAtom, f =>
134    f !== null && f.id === flowId ? { ...f, step: 'rewriting', callId, failedDuring: null, otherOpen: null } : f,
135  )
136  if (flow === null || flow.callId !== callId) return
137  try {
138    const reply = await $.model.complete({
139      model: REWRITE_MODEL,
140      system: REWRITE_SYSTEM,
141      prompt: rewritePrompt(flow.original, flow.questions, flow.answers),
142      maxTokens: 3000,
143      timeoutMs: 120_000,
144    })
145    if (!reply.isAnswered) throw new Error(reply.reason)
146    const { prompt, added } = parseRewrite(reply.text)
147    const done = await update($, flowAtom, f =>
148      f !== null && f.callId === callId ? { ...f, step: 'review', upgraded: prompt, added } : f,
149    )
150    if (done !== null && done.callId === callId) await remember($, done)
151  } catch {
152    await failStep($, callId, 'rewrite')
153  }
154}
155
156/** Keeps the upgrade among the last twenty, in the session's state and the mod's own store. */
157async function remember($: Engine, flow: PromptUpgradeFlow): Promise<void> {
158  const at = await $.clock.now()
159  const entry: PromptUpgradeEntry = { id: flow.id, at, original: flow.original, upgraded: flow.upgraded, added: flow.added }
160  const list = await update($, historyAtom, old => [entry, ...old.filter(one => one.id !== flow.id)].slice(0, HISTORY_SIZE))
161  await $.store.set('history', list)
162}
163
164async function answer($: Engine, index: number, given: PromptUpgradeAnswer): Promise<void> {
165  await update($, flowAtom, f => {
166    if (f === null || f.step !== 'questions') return f
167    const answers = f.answers.slice()
168    answers[index] = given
169    return { ...f, answers, otherOpen: null }
170  })
171}
172
173async function runUpgrade($: Engine, text: string): Promise<void> {
174  const flow = await read($, flowAtom)
175  if (flow === null || text.trim() === '') return
176  if (text !== flow.upgraded) await remember($, { ...flow, upgraded: text })
177  await update($, flowAtom, () => null)
178  await closePane($)
179  // Sent as the mod's own prompt (origin `plugin`), which the interceptor never takes.
180  await $.prompt.submit({ text, asUser: true })
181}
182
183async function sendOriginal($: Engine): Promise<void> {
184  const flow = await read($, flowAtom)
185  if (flow === null) return
186  await update($, flowAtom, () => null)
187  await closePane($)
188  await $.prompt.submit({ text: flow.original, asUser: true })
189}
190
191/** Closes the pane and puts the original prompt back in the box. */
192async function cancel($: Engine): Promise<void> {
193  const flow = await read($, flowAtom)
194  await update($, flowAtom, () => null)
195  await update($, viewAtom, () => 'flow' as const)
196  await closePane($)
197  if (flow !== null && flow.original !== '') await $.prompt.fill({ text: flow.original })
198}
199
200async function retry($: Engine): Promise<void> {
201  const flow = await read($, flowAtom)
202  if (flow === null) return
203  if (flow.failedDuring === 'rewrite') await writeRewrite($, flow.id)
204  else await askQuestions($, flow.id)
205}
206
207async function openEntry($: Engine, entry: PromptUpgradeEntry): Promise<void> {
208  const flow: PromptUpgradeFlow = {
209    ...newFlow(entry.original, 'review'),
210    id: entry.id,
211    upgraded: entry.upgraded,
212    added: entry.added,
213  }
214  await update($, flowAtom, () => flow)
215  await update($, viewAtom, () => 'flow' as const)
216}
217
218/** Runs a press handler, turning a failure into one plain line. */
219function act($: Engine, work: () => Promise<void>): () => void {
220  return () => {
221    work().catch(err => $.ui.toast(`Prompt Upgrade: that did not work (${reason(err)}). Your prompt is safe.`))
222  }
223}
224
225/** After a reload, picks the flow back up where it was. */
226async function resume($: Engine): Promise<void> {
227  const flow = await read($, flowAtom)
228  const view = await read($, viewAtom)
229  if (flow === null && view !== 'history') return
230  const isOpen = (await $.ui.panes()).some(pane => pane.id === PANE)
231  if (!isOpen) await $.ui.open({ id: PANE, title: TITLE })
232  if (flow === null) return
233  if (flow.step === 'reading') await askQuestions($, flow.id)
234  else if (flow.step === 'rewriting') await writeRewrite($, flow.id)
235}
236
237// ------------------------------------------------------------------ drawing
238
239type Els = Elements['terminal']
240
241type ChipKind = 'gold' | 'slate' | 'green' | 'quiet'
242type Chip = { key: string; label: string; kind: ChipKind; onPress: (e: UiPressArgument) => void }
243
244function chipWidth(chip: Chip): number {
245  return chip.kind === 'quiet' ? chip.label.length : chip.label.length + 2
246}
247
248function drawChip(els: Els, chip: Chip, room: number): RenderElement {
249  const { Button, Text } = els
250  const label = clip(chip.label, Math.max(1, chip.kind === 'quiet' ? room : room - 2))
251  if (chip.kind === 'quiet') {
252    return (
253      <Button key={chip.key} plain onPress={chip.onPress}>
254        <Text color={MUTED}>{label}</Text>
255      </Button>
256    )
257  }
258  const back = chip.kind === 'gold' ? GOLD : chip.kind === 'green' ? GREEN : HAIR
259  const fore = chip.kind === 'slate' ? TEXT : INK
260  return (
261    <Button key={chip.key} plain onPress={chip.onPress}>
262      <Text backgroundColor={back} color={fore} bold={chip.kind !== 'slate'}>
263        {` ${label} `}
264      </Text>
265    </Button>
266  )
267}
268
269function chipRows(els: Els, chips: readonly Chip[], width: number, indent = 0): RenderElement {
270  const { Box } = els
271  const rows = packRows(chips, chipWidth, width - indent)
272  return (
273    <Box flexDirection="column" paddingLeft={indent}>
274      {rows.map(row => (
275        <Box flexDirection="row" gap={1}>
276          {row.map(chip => drawChip(els, chip, width - indent))}
277        </Box>
278      ))}
279    </Box>
280  )
281}
282
283function header(els: Els, width: number, isAuto: boolean): RenderElement {
284  const { Box, Text } = els
285  const title = 'PROMPT UPGRADE'
286  const status = isAuto ? 'AUTO ON' : 'AUTO OFF'
287  const showStatus = width >= title.length + status.length + 4
288  const titleText = clip(title, width)
289  const gold = Math.min(titleText.length, width)
290  return (
291    <Box flexDirection="column">
292      <Box flexDirection="row" justifyContent="space-between" width={width}>
293        <Text bold color={TEXT}>
294          {titleText}
295        </Text>
296        {showStatus && <Text color={isAuto ? GREEN : MUTED}>{status}</Text>}
297      </Box>
298      <Text wrap="truncate">
299        <Text color={GOLD}>{'━'.repeat(gold)}</Text>
300        <Text color={HAIR}>{'─'.repeat(Math.max(0, width - gold))}</Text>
301      </Text>
302    </Box>
303  )
304}
305
306function heading(els: Els, label: string, width: number): RenderElement {
307  const { Text } = els
308  const text = clip(label.toUpperCase(), width)
309  const rest = width - text.length - 1
310  return (
311    <Text wrap="truncate">
312      <Text color={MUTED}>{text}</Text>
313      {rest > 0 && <Text color={HAIR}>{` ${'─'.repeat(rest)}`}</Text>}
314    </Text>
315  )
316}
317
318const STEPS_LONG = ['Your prompt', 'Questions', 'Rewrite', 'Review']
319const STEPS_SHORT = ['Prompt', 'Ask', 'Rewrite', 'Review']
320
321function stepIndex(flow: PromptUpgradeFlow): number {
322  switch (flow.step) {
323    case 'compose':
324    case 'reading':
325      return 0
326    case 'questions':
327      return 1
328    case 'rewriting':
329      return 2
330    case 'review':
331      return 3
332    case 'error':
333      return flow.failedDuring === 'rewrite' ? 2 : 1
334  }
335}
336
337function tracker(els: Els, current: number, width: number): RenderElement {
338  const { Box, Text } = els
339  const fits = (labels: string[], sep: number) =>
340    labels.reduce((sum, label) => sum + label.length + 2, 0) + sep * (labels.length - 1) <= width
341  const labels = fits(STEPS_LONG, 5) ? STEPS_LONG : fits(STEPS_SHORT, 3) ? STEPS_SHORT : null
342  if (labels === null) {
343    return (
344      <Text wrap="truncate" color={GOLD} bold>
345        {clip(`Step ${current + 1} of 4  ${STEPS_LONG[current]}`, width)}
346      </Text>
347    )
348  }
349  const sep = labels === STEPS_LONG ? '  ›  ' : ' › '
350  return (
351    <Box flexDirection="row">
352      {labels.map((label, i) => (
353        <Text wrap="truncate">
354          {i > 0 && <Text color={HAIR}>{sep}</Text>}
355          {i < current && <Text color={GREEN}>{`✓ `}</Text>}
356          {i < current && <Text color={TEXT}>{label}</Text>}
357          {i === current && (
358            <Text color={GOLD} bold>
359              {`● ${label}`}
360            </Text>
361          )}
362          {i > current && <Text color={MUTED}>{`○ ${label}`}</Text>}
363        </Text>
364      ))}
365    </Box>
366  )
367}
368
369function card(
370  els: Els,
371  opts: { label: string; text: string; width: number; color?: string; maxLines?: number; key: string },
372): RenderElement {
373  const { Box, Text } = els
374  const inner = Math.max(1, opts.width - 4)
375  let lines = wrapText(opts.text, inner)
376  let more = 0
377  if (opts.maxLines !== undefined && lines.length > opts.maxLines) {
378    more = lines.length - opts.maxLines
379    lines = lines.slice(0, opts.maxLines)
380  }
381  return (
382    <Box
383      key={opts.key}
384      flexDirection="column"
385      borderStyle="round"
386      borderColor={opts.color ?? HAIR}
387      paddingX={1}
388      width={opts.width}
389    >
390      <Text color={opts.color === GOLD ? GOLD : MUTED} wrap="truncate">
391        {clip(opts.label.toUpperCase(), inner)}
392      </Text>
393      {lines.map(line => (
394        <Text color={TEXT} wrap="truncate">
395          {line === '' ? ' ' : line}
396        </Text>
397      ))}
398      {more > 0 && (
399        <Text color={MUTED} wrap="truncate">
400          {clip(`… ${more} more line${more === 1 ? '' : 's'}`, inner)}
401        </Text>
402      )}
403    </Box>
404  )
405}
406
407function note(els: Els, text: string, width: number, color = MUTED): RenderElement {
408  const { Box, Text } = els
409  return (
410    <Box flexDirection="column">
411      {wrapText(text, width).map(line => (
412        <Text color={color} wrap="truncate">
413          {line}
414        </Text>
415      ))}
416    </Box>
417  )
418}
419
420function tile(els: Els, n: number, label: string, color: string): { width: number; node: RenderElement } {
421  const { Box, Text } = els
422  const rows = bigNumber(n)
423  const inner = Math.max(rows[0].length, label.length)
424  const pad = (s: string) => {
425    const left = Math.floor((inner - s.length) / 2)
426    return `${' '.repeat(left)}${s}${' '.repeat(inner - s.length - left)}`
427  }
428  return {
429    width: inner + 4,
430    node: (
431      <Box flexDirection="column" borderStyle="round" borderColor={HAIR} paddingX={1} width={inner + 4}>
432        {rows.map(row => (
433          <Text color={color}>{pad(row)}</Text>
434        ))}
435        <Text color={MUTED}>{pad(label)}</Text>
436      </Box>
437    ),
438  }
439}
440
441function stats(els: Els, flow: PromptUpgradeFlow, width: number): RenderElement {
442  const { Box, Text } = els
443  const before = wordCount(flow.original)
444  const after = wordCount(flow.upgraded)
445  const added = flow.added.length
446  const tiles = [
447    tile(els, before, 'WORDS BEFORE', TEXT),
448    tile(els, after, 'WORDS AFTER', GOLD),
449    tile(els, added, 'DETAILS ADDED', GREEN),
450  ]
451  const total = tiles.reduce((sum, one) => sum + one.width, 0) + 2 * (tiles.length - 1)
452  if (total <= width) {
453    return (
454      <Box flexDirection="row" gap={2}>
455        {tiles.map(one => one.node)}
456      </Box>
457    )
458  }
459  return (
460    <Text wrap="truncate" color={TEXT}>
461      {clip(`${before} words before  ·  ${after} words after  ·  ${added} details added`, width)}
462    </Text>
463  )
464}
465
466function drawCompose($: Engine, els: Els, width: number, hasInput: boolean): RenderElement[] {
467  const { Input } = els
468  const out: RenderElement[] = [heading(els, 'Your prompt', width)]
469  if (hasInput) {
470    out.push(
471      <Input
472        key="compose"
473        placeholder="Type a rough prompt and press Enter"
474        submitLabel="upgrade"
475        autoFocus
476        onSubmit={value => {
477          const text = value.trim()
478          if (text === '') return
479          void startFlowFromPane($, text)
480        }}
481      />,
482    )
483  } else {
484    out.push(note(els, 'This surface has no text field. Type /better followed by your prompt instead.', width))
485  }
486  out.push(chipRows(els, [{ key: 'cancel', label: 'Cancel', kind: 'quiet', onPress: act($, () => cancel($)) }], width))
487  return out
488}
489
490async function startFlowFromPane($: Engine, text: string): Promise<void> {
491  try {
492    const flow = newFlow(text, 'reading')
493    await update($, flowAtom, () => flow)
494    await askQuestions($, flow.id)
495  } catch (err) {
496    $.ui.toast(`Prompt Upgrade: that did not work (${reason(err)}). Your prompt is safe.`)
497  }
498}
499
500function drawWaiting($: Engine, els: Els, flow: PromptUpgradeFlow, width: number): RenderElement[] {
501  const answered = flow.answers.filter(a => a !== null && a.kind !== 'skip').length
502  const line =
503    flow.step === 'reading'
504      ? 'Reading your prompt and choosing the questions that would help most.'
505      : answered === 0
506        ? 'Writing the upgrade from your prompt.'
507        : `Writing the upgrade from your prompt and ${answered} answer${answered === 1 ? '' : 's'}.`
508  return [
509    heading(els, 'Your prompt', width),
510    card(els, { key: 'original', label: 'Your prompt', text: flow.original, width, maxLines: 10 }),
511    note(els, line, width, TEXT),
512    chipRows(els, [{ key: 'cancel', label: 'Cancel', kind: 'quiet', onPress: act($, () => cancel($)) }], width),
513  ]
514}
515
516function drawQuestions($: Engine, els: Els, flow: PromptUpgradeFlow, width: number, hasInput: boolean): RenderElement[] {
517  const { Box, Text, Input } = els
518  const out: RenderElement[] = [heading(els, 'Questions', width)]
519  const indent = 4
520  flow.questions.forEach((q, i) => {
521    const given = flow.answers[i] ?? null
522    const isAnswered = given !== null
523    const lines = wrapText(q.question, Math.max(1, width - indent))
524    const chips: Chip[] = q.options.map((option, j) => ({
525      key: `q${i}-o${j}`,
526      label: option,
527      kind: given?.kind === 'pick' && given.text === option ? 'green' : 'slate',
528      onPress: act($, () => answer($, i, { kind: 'pick', text: option })),
529    }))
530    chips.push({
531      key: `q${i}-other`,
532      label: given?.kind === 'other' ? `Other: ${given.text}` : 'Other (type it)',
533      kind: given?.kind === 'other' ? 'green' : 'slate',
534      onPress: act($, async () => {
535        await update($, flowAtom, f => (f === null ? f : { ...f, otherOpen: i }))
536      }),
537    })
538    chips.push({
539      key: `q${i}-skip`,
540      label: given?.kind === 'skip' ? 'Skipped' : 'Skip',
541      kind: 'quiet',
542      onPress: act($, () => answer($, i, { kind: 'skip', text: '' })),
543    })
544    out.push(
545      <Box flexDirection="column" marginTop={i === 0 ? 0 : 1}>
546        {lines.map((line, k) => (
547          <Box flexDirection="row">
548            {k === 0 ? (
549              <Text backgroundColor={isAnswered ? GREEN : GOLD} color={INK} bold>
550                {isAnswered ? ' ✓ ' : ` ${i + 1} `}
551              </Text>
552            ) : (
553              <Text>{'   '}</Text>
554            )}
555            <Text color={TEXT} wrap="truncate">
556              {` ${line}`}
557            </Text>
558          </Box>
559        ))}
560        {chipRows(els, chips, width, indent)}
561        {hasInput && flow.otherOpen === i && (
562          <Box paddingLeft={indent}>
563            <Input
564              key={`q${i}-other-text`}
565              placeholder="Type your answer and press Enter"
566              submitLabel="save"
567              autoFocus
568              onSubmit={value => {
569                const text = value.trim()
570                if (text === '') return
571                act($, () => answer($, i, { kind: 'other', text }))()
572              }}
573            />
574          </Box>
575        )}
576      </Box>,
577    )
578  })
579  const answered = flow.answers.filter(a => a !== null).length
580  const total = flow.questions.length
581  const isAll = answered === total
582  out.push(<Text>{' '}</Text>)
583  out.push(heading(els, `${answered} of ${total} answered`, width))
584  out.push(
585    chipRows(
586      els,
587      [
588        {
589          key: 'upgrade',
590          label: isAll ? 'Upgrade now' : 'Skip the rest, upgrade now',
591          kind: 'gold',
592          onPress: act($, () => writeRewrite($, flow.id)),
593        },
594        { key: 'cancel', label: 'Cancel', kind: 'quiet', onPress: act($, () => cancel($)) },
595      ],
596      width,
597    ),
598  )
599  return out
600}
601
602function drawReview($: Engine, els: Els, flow: PromptUpgradeFlow, width: number, hasInput: boolean): RenderElement[] {
603  const { Box, Text, Input } = els
604  const out: RenderElement[] = [stats(els, flow, width)]
605  if (flow.isEditing && hasInput) {
606    out.push(heading(els, 'Edit the upgrade', width))
607    out.push(
608      <Input
609        key="edit"
610        value={flow.upgraded}
611        submitLabel="run it"
612        autoFocus
613        onSubmit={value => act($, () => runUpgrade($, value.trim()))()}
614      />,
615    )
616    out.push(
617      chipRows(
618        els,
619        [
620          {
621            key: 'stop-edit',
622            label: 'Stop editing',
623            kind: 'slate',
624            onPress: act($, async () => {
625              await update($, flowAtom, f => (f === null ? f : { ...f, isEditing: false }))
626            }),
627          },
628          { key: 'cancel', label: 'Cancel', kind: 'quiet', onPress: act($, () => cancel($)) },
629        ],
630        width,
631      ),
632    )
633  } else {
634    const buttons: Chip[] = [
635      { key: 'run', label: 'Run it', kind: 'gold', onPress: act($, () => runUpgrade($, flow.upgraded)) },
636    ]
637    if (hasInput) {
638      buttons.push({
639        key: 'edit',
640        label: 'Edit',
641        kind: 'slate',
642        onPress: act($, async () => {
643          await update($, flowAtom, f => (f === null ? f : { ...f, isEditing: true, isCopied: false }))
644        }),
645      })
646    }
647    buttons.push({
648      key: 'copy',
649      label: flow.isCopied ? 'Copied' : 'Copy',
650      kind: flow.isCopied ? 'green' : 'slate',
651      onPress: press =>
652        act($, async () => {
653          const copied = await $.ui.copy({ text: flow.upgraded, surface: press.surface })
654          if (!copied.isCopied) throw new Error('the clipboard did not take it')
655          await update($, flowAtom, f => (f === null ? f : { ...f, isCopied: true }))
656        })(),
657    })
658    buttons.push({ key: 'cancel', label: 'Cancel', kind: 'quiet', onPress: act($, () => cancel($)) })
659    out.push(chipRows(els, buttons, width))
660  }
661  if (flow.added.length > 0) {
662    out.push(heading(els, 'What I added', width))
663    for (const item of flow.added) {
664      const lines = wrapText(item, Math.max(1, width - 2))
665      lines.forEach((line, k) => {
666        out.push(
667          <Text wrap="truncate">
668            <Text color={GREEN}>{k === 0 ? '+ ' : '  '}</Text>
669            <Text color={TEXT}>{line}</Text>
670          </Text>,
671        )
672      })
673    }
674  }
675  out.push(heading(els, 'Side by side', width))
676  const isSide = width >= 60
677  const half = Math.floor((width - 1) / 2)
678  const left = card(els, { key: 'left', label: 'Your prompt', text: flow.original, width: isSide ? half : width })
679  const right = card(els, {
680    key: 'right',
681    label: 'The upgrade',
682    text: flow.upgraded,
683    width: isSide ? width - half - 1 : width,
684    color: GOLD,
685  })
686  out.push(
687    isSide ? (
688      <Box flexDirection="row" gap={1} alignItems="flex-start">
689        {left}
690        {right}
691      </Box>
692    ) : (
693      <Box flexDirection="column">
694        {left}
695        {right}
696      </Box>
697    ),
698  )
699  return out
700}
701
702function drawError($: Engine, els: Els, flow: PromptUpgradeFlow, width: number): RenderElement[] {
703  const { Box, Text } = els
704  const what = flow.failedDuring === 'rewrite' ? 'writing the upgrade' : 'choosing the questions'
705  const inner = Math.max(1, width - 4)
706  return [
707    <Box key="error" flexDirection="column" borderStyle="round" borderColor={AMBER} paddingX={1} width={width}>
708      <Text color={AMBER} bold wrap="truncate">
709        {clip('THAT DID NOT WORK', inner)}
710      </Text>
711      {wrapText(`The model could not finish ${what}. Your prompt is safe and nothing has been sent.`, inner).map(
712        line => (
713          <Text color={TEXT} wrap="truncate">
714            {line}
715          </Text>
716        ),
717      )}
718    </Box>,
719    chipRows(
720      els,
721      [
722        { key: 'retry', label: 'Try again', kind: 'gold', onPress: act($, () => retry($)) },
723        { key: 'send-original', label: 'Send my original', kind: 'slate', onPress: act($, () => sendOriginal($)) },
724        { key: 'cancel', label: 'Cancel', kind: 'quiet', onPress: act($, () => cancel($)) },
725      ],
726      width,
727    ),
728    card(els, { key: 'original', label: 'Your prompt', text: flow.original, width, maxLines: 10 }),
729  ]
730}
731
732function formatWhen(at: number): string {
733  const d = new Date(at)
734  const months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
735  const two = (n: number) => String(n).padStart(2, '0')
736  return `${months[d.getMonth()]} ${two(d.getDate())} ${two(d.getHours())}:${two(d.getMinutes())}`
737}
738
739function drawHistory($: Engine, els: Els, list: readonly PromptUpgradeEntry[], width: number): RenderElement[] {
740  const { Box, Button, Text } = els
741  const out: RenderElement[] = [heading(els, `Last ${HISTORY_SIZE} upgrades`, width)]
742  if (list.length === 0) out.push(note(els, 'No upgrades yet. Try /better followed by a rough prompt.', width))
743  for (const entry of list) {
744    const when = formatWhen(entry.at)
745    const room = Math.max(1, width - when.length - 2)
746    out.push(
747      <Button key={`h-${entry.id}`} plain onPress={act($, () => openEntry($, entry))}>
748        <Text color={MUTED}>{when}</Text>
749        {'  '}
750        <Text color={TEXT}>{clip(entry.original.replace(/\s+/g, ' ').trim(), room)}</Text>
751      </Button>,
752    )
753  }
754  out.push(<Box>{chipRows(els, [{ key: 'close', label: 'Close', kind: 'quiet', onPress: act($, () => cancel($)) }], width)}</Box>)
755  return out
756}
757
758// ------------------------------------------------------------------ register
759
760export const register: Register = on => {
761  on('session.start', async ($, e, next) => {
762    await $.command.register({
763      name: 'better',
764      description: 'Upgrade a rough prompt before it is sent (on, off, history)',
765      argumentHint: '[prompt | on | off | history]',
766    })
767    try {
768      const auto = (await $.store.get('auto')) === true
769      const saved = await $.store.get('history')
770      await update($, autoAtom, () => auto)
771      if (Array.isArray(saved)) await update($, historyAtom, () => (saved as PromptUpgradeEntry[]).slice(0, HISTORY_SIZE))
772    } catch {
773      // A store that cannot be read leaves automatic mode off and the history empty.
774    }
775    void resume($).catch(() => undefined)
776    return next(e)
777  })
778
779  on('command.run', { command: 'better' }, async ($, e) => {
780    try {
781      const args = e.args.trim()
782      const word = args.toLowerCase()
783      if (word === 'on' || word === 'off') {
784        const isOn = word === 'on'
785        await update($, autoAtom, () => isOn)
786        await $.store.set('auto', isOn)
787        return {
788          text: isOn
789            ? 'Prompt Upgrade is on: prompts you type open in the pane first. Start one with ">" to send it as is.'
790            : 'Prompt Upgrade is off: prompts go straight through. /better <prompt> still works.',
791        }
792      }
793      if (word === 'history') {
794        await update($, viewAtom, () => 'history' as const)
795        void showPane($).catch(() => undefined)
796        return { text: 'Prompt Upgrade: your last upgrades are in the pane.' }
797      }
798      if (args === '') {
799        await update($, flowAtom, () => newFlow('', 'compose'))
800        await update($, viewAtom, () => 'flow' as const)
801        void showPane($).catch(() => undefined)
802        return { text: 'Prompt Upgrade: type your prompt in the pane.' }
803      }
804      void startFlow($, args)
805        .then(async isPlaced => {
806          if (isPlaced) return
807          await $.prompt.fill({ text: args })
808          $.ui.toast(NARROW)
809        })
810        .catch(err => $.ui.toast(`Prompt Upgrade could not open (${reason(err)}). Your prompt is safe.`))
811      return { text: 'Prompt Upgrade: upgrading in the pane. Nothing is sent until you press Run it.' }
812    } catch (err) {
813      return { text: `Prompt Upgrade could not do that: ${reason(err)}.` }
814    }
815  }).catch(() => ({ text: 'Prompt Upgrade hit a problem and did nothing. Your prompt is safe.' }))
816
817  on('prompt.submit', async ($, e, next) => {
818    const decision = decide(e, await read($, autoAtom))
819    if (decision.action === 'pass') return next(e)
820    if (decision.action === 'strip') return next({ ...e, text: decision.text })
821    const text = e.text.trim()
822    const isPlaced = await startFlow($, text)
823    if (!isPlaced) {
824      await $.prompt.fill({ text })
825      return { drop: NARROW }
826    }
827    return { drop: 'Prompt Upgrade: upgrading in the pane. Nothing is sent until you press Run it.' }
828  }).catch(($, e, next) => next(e))
829
830  on('ui.close', { id: PANE }, async ($, e, next) => {
831    // The person closing the pane by hand is a Cancel: their prompt goes back in the box.
832    if (e.origin.kind === 'person') {
833      const flow = await read($, flowAtom)
834      await update($, flowAtom, () => null)
835      await update($, viewAtom, () => 'flow' as const)
836      if (flow !== null && flow.original !== '') await $.prompt.fill({ text: flow.original })
837    }
838    return next(e)
839  }).catch(($, e, next) => next(e))
840
841  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
842    const els = $.ui.resolve(e) as unknown as Els
843    const { Box, Text } = els
844    const width = Math.max(16, e.props.bodyColumns)
845    const hasInput = e.surface !== 'mobile'
846    const view = await read($, viewAtom)
847    const flow = await read($, flowAtom)
848    const isAuto = await read($, autoAtom)
849    let body: RenderElement[]
850    if (view === 'history') {
851      body = drawHistory($, els, await read($, historyAtom), width)
852    } else if (flow === null) {
853      body = [note(els, 'Nothing to upgrade. Type /better followed by a rough prompt.', width)]
854    } else {
855      body = [tracker(els, stepIndex(flow), width), <Text>{' '}</Text>]
856      if (flow.step === 'compose') body.push(...drawCompose($, els, width, hasInput))
857      else if (flow.step === 'reading' || flow.step === 'rewriting') body.push(...drawWaiting($, els, flow, width))
858      else if (flow.step === 'questions') body.push(...drawQuestions($, els, flow, width, hasInput))
859      else if (flow.step === 'review') body.push(...drawReview($, els, flow, width, hasInput))
860      else body.push(...drawError($, els, flow, width))
861    }
862    return (
863      <Box flexDirection="column" width={width}>
864        {header(els, width, isAuto)}
865        {body}
866      </Box>
867    )
868  })
869}
870

1// Pure helpers: deciding what to intercept, talking to the model, reading its
2// replies, and laying text out in cells. Nothing here touches `$`.
3
4import type { PromptUpgradeAnswer, PromptUpgradeQuestion } from '../types'
5
6export type Decision = { action: 'pass' } | { action: 'strip'; text: string } | { action: 'intercept' }
7
8export type Submission = {
9  text: string
10  attachments?: readonly unknown[]
11  turnId?: string
12  origin: { kind: string }
13}
14
15const ACKS =
16  /^(y|yes|yeah|yep|yup|no|nope|nah|ok|okay|k|sure|thanks|thank you|thx|ty|go|go ahead|go on|do it|keep going|continue|carry on|proceed|next|sounds good|looks good|lgtm|perfect|great|nice|cool|done|stop|wait|retry|try again|again|undo|revert)[\s.!?]*$/i
17
18/** Fewer words than this counts as a short reply and goes straight through. */
19export const SHORT_REPLY_WORDS = 4
20
21export function decide(e: Submission, isAuto: boolean): Decision {
22  if (!isAuto) return { action: 'pass' }
23  // Only what the person typed at the prompt box. The mod's own submissions
24  // arrive as origin `plugin`, so an upgrade can never come back in here.
25  if (e.origin.kind !== 'composer') return { action: 'pass' }
26  if (e.turnId !== undefined) return { action: 'pass' }
27  if ((e.attachments?.length ?? 0) > 0) return { action: 'pass' }
28  const text = e.text.trim()
29  if (text === '') return { action: 'pass' }
30  if (text.startsWith('>')) return { action: 'strip', text: text.slice(1).trimStart() }
31  if (text.startsWith('/')) return { action: 'pass' }
32  if (ACKS.test(text)) return { action: 'pass' }
33  if (wordCount(text) < SHORT_REPLY_WORDS) return { action: 'pass' }
34  return { action: 'intercept' }
35}
36
37export function wordCount(text: string): number {
38  const words = text.trim().split(/\s+/)
39  return words[0] === '' ? 0 : words.length
40}
41
42// ---------------------------------------------------------------- the model
43
44export const QUESTIONS_SYSTEM = `You help a person improve a prompt they are about to send to an AI coding and writing assistant.
45Pick the at most three multiple-choice questions whose answers would most improve the result.
46Ask only about things the prompt leaves genuinely open and that change the outcome: the goal, who it is for, the format, scope, or what "done" looks like.
47Never ask about something the prompt already says. If no question would help, return none.
48Each question is one short sentence. Each has two to four short options of at most six words. Do not add an "Other" option.
49Reply with JSON only, no prose and no code fence, in exactly this shape:
50{"questions":[{"question":"...","options":["...","..."]}]}`
51
52export function questionsPrompt(original: string): string {
53  return `The prompt:\n<prompt>\n${original}\n</prompt>`
54}
55
56export const REWRITE_SYSTEM = `You rewrite a rough prompt into a great one, written as one request the person could have typed themselves, in plain language.
57Rules:
58- Keep exactly what the person meant. Do not change the task.
59- Add a clear goal, who it is for, the format of the result, and what "done" looks like.
60- Use the answers the person gave. Leave out anything they skipped.
61- Never invent facts, names, numbers, prices, file names or URLs. Where something needed is unknown, write a placeholder in square brackets, like [the API base URL].
62- No headings, no bullet lists of instructions to the AI about being an AI; just a clear request, a few short paragraphs at most.
63Also list, in a few words each, what you added that was not in the original (at most six items).
64Reply with JSON only, no prose and no code fence, in exactly this shape:
65{"prompt":"...","added":["...","..."]}`
66
67export function rewritePrompt(
68  original: string,
69  questions: readonly PromptUpgradeQuestion[],
70  answers: readonly (PromptUpgradeAnswer | null)[],
71): string {
72  const lines: string[] = []
73  questions.forEach((q, i) => {
74    const a = answers[i]
75    if (a === null || a === undefined || a.kind === 'skip') return
76    lines.push(`Q: ${q.question}\nA: ${a.text}`)
77  })
78  const answered = lines.length === 0 ? 'The person answered no questions.' : `The person's answers:\n${lines.join('\n')}`
79  return `The original prompt:\n<prompt>\n${original}\n</prompt>\n\n${answered}`
80}
81
82/** Finds the JSON object in a reply, tolerating a code fence or a stray line around it. */
83function jsonObject(reply: string): unknown {
84  const start = reply.indexOf('{')
85  const end = reply.lastIndexOf('}')
86  if (start === -1 || end <= start) throw new Error('The reply held no JSON object.')
87  return JSON.parse(reply.slice(start, end + 1))
88}
89
90function cleanLine(value: unknown, max: number): string {
91  if (typeof value !== 'string') return ''
92  const one = value.replace(/\s+/g, ' ').trim()
93  return one.length > max ? `${one.slice(0, max - 1).trimEnd()}…` : one
94}
95
96export function parseQuestions(reply: string): PromptUpgradeQuestion[] {
97  const data = jsonObject(reply) as { questions?: unknown }
98  if (!Array.isArray(data.questions)) throw new Error('The reply had no questions list.')
99  const out: PromptUpgradeQuestion[] = []
100  for (const raw of data.questions) {
101    if (typeof raw !== 'object' || raw === null) continue
102    const item = raw as { question?: unknown; options?: unknown }
103    const question = cleanLine(item.question, 140)
104    if (question === '' || !Array.isArray(item.options)) continue
105    const seen = new Set<string>()
106    const options: string[] = []
107    for (const option of item.options) {
108      const text = cleanLine(option, 40)
109      if (text === '' || /^other\b/i.test(text) || seen.has(text.toLowerCase())) continue
110      seen.add(text.toLowerCase())
111      options.push(text)
112    }
113    if (options.length < 2) continue
114    out.push({ question, options: options.slice(0, 4) })
115    if (out.length === 3) break
116  }
117  return out
118}
119
120export function parseRewrite(reply: string): { prompt: string; added: string[] } {
121  const data = jsonObject(reply) as { prompt?: unknown; added?: unknown }
122  if (typeof data.prompt !== 'string' || data.prompt.trim() === '') throw new Error('The reply had no rewritten prompt.')
123  const added = Array.isArray(data.added)
124    ? data.added.map(item => cleanLine(item, 80)).filter(item => item !== '').slice(0, 6)
125    : []
126  return { prompt: data.prompt.trim(), added }
127}
128
129// ---------------------------------------------------------------- layout
130
131/** Wraps text into lines no wider than `width`, keeping the author's line breaks. */
132export function wrapText(text: string, width: number): string[] {
133  const room = Math.max(1, width)
134  const out: string[] = []
135  for (const paragraph of text.replace(/\r\n?/g, '\n').split('\n')) {
136    if (paragraph.trim() === '') {
137      out.push('')
138      continue
139    }
140    let line = ''
141    for (const word of paragraph.split(/\s+/).filter(w => w !== '')) {
142      let rest = word
143      while (rest.length > room) {
144        if (line !== '') {
145          out.push(line)
146          line = ''
147        }
148        out.push(rest.slice(0, room))
149        rest = rest.slice(room)
150      }
151      if (line === '') line = rest
152      else if (line.length + 1 + rest.length <= room) line = `${line} ${rest}`
153      else {
154        out.push(line)
155        line = rest
156      }
157    }
158    if (line !== '') out.push(line)
159  }
160  while (out.length > 0 && out[out.length - 1] === '') out.pop()
161  return out
162}
163
164export function clip(text: string, width: number): string {
165  if (width <= 0) return ''
166  return text.length > width ? `${text.slice(0, Math.max(0, width - 1))}…` : text
167}
168
169/** Splits chips into rows that fit `width`, given each chip's drawn width and a gap of one. */
170export function packRows<T>(items: readonly T[], widthOf: (item: T) => number, width: number): T[][] {
171  const rows: T[][] = []
172  let row: T[] = []
173  let used = 0
174  for (const item of items) {
175    const w = widthOf(item)
176    const need = row.length === 0 ? w : used + 1 + w
177    if (row.length > 0 && need > width) {
178      rows.push(row)
179      row = [item]
180      used = w
181    } else {
182      row.push(item)
183      used = need
184    }
185  }
186  if (row.length > 0) rows.push(row)
187  return rows
188}
189
190const DIGITS: Record<string, [string, string, string]> = {
191  '0': ['┏━┓', '┃ ┃', '┗━┛'],
192  '1': ['╺┓ ', ' ┃ ', '╺┻╸'],
193  '2': ['╺━┓', '┏━┛', '┗━╸'],
194  '3': ['╺━┓', ' ━┫', '╺━┛'],
195  '4': ['╻ ╻', '┗━┫', '  ╹'],
196  '5': ['┏━╸', '┗━┓', '╺━┛'],
197  '6': ['┏━╸', '┣━┓', '┗━┛'],
198  '7': ['╺━┓', '  ┃', '  ╹'],
199  '8': ['┏━┓', '┣━┫', '┗━┛'],
200  '9': ['┏━┓', '┗━┫', '╺━┛'],
201}
202
203/** A number drawn three rows tall in heavy box-drawing characters. */
204export function bigNumber(n: number): [string, string, string] {
205  const digits = String(Math.max(0, Math.round(n))).split('')
206  const rows: [string, string, string] = ['', '', '']
207  digits.forEach((d, i) => {
208    const glyph = DIGITS[d] ?? DIGITS['0']!
209    for (let r = 0; r < 3; r++) rows[r] += (i > 0 ? ' ' : '') + glyph[r]
210  })
211  return rows
212}
213

1export type PromptUpgradeStep = 'compose' | 'reading' | 'questions' | 'rewriting' | 'review' | 'error'
2
3export type PromptUpgradeQuestion = { question: string; options: string[] }
4
5export type PromptUpgradeAnswer = { kind: 'pick' | 'other' | 'skip'; text: string }
6
7export type PromptUpgradeFlow = {
8  id: string
9  original: string
10  step: PromptUpgradeStep
11  questions: PromptUpgradeQuestion[]
12  answers: (PromptUpgradeAnswer | null)[]
13  otherOpen: number | null
14  upgraded: string
15  added: string[]
16  isEditing: boolean
17  isCopied: boolean
18  failedDuring: 'questions' | 'rewrite' | null
19  callId: string
20}
21
22export type PromptUpgradeEntry = {
23  id: string
24  at: number
25  original: string
26  upgraded: string
27  added: string[]
28}
29
30declare module 'claude-code' {
31  interface PluginState {
32    'prompt-upgrade': {
33      flow: PromptUpgradeFlow | null
34      view: 'flow' | 'history'
35      auto: boolean
36      history: PromptUpgradeEntry[]
37    }
38  }
39}
40