prompt-upgrade

1import { atom, read, update } from 'claude-code'
2import type { Elements, EngineInterface, ModelCompleteRequest, Register, RenderChildren } from 'claude-code'
3
4import type { Answer, Entry, Flow, Question, Step } from '../types'
5import {
6  bigNumber,
7  cap,
8  clip,
9  describeFailure,
10  gate,
11  parseQuestions,
12  parseRewrite,
13  questionsRequest,
14  rewriteRequest,
15  when,
16  wordCount,
17  wrap,
18} from './logic'
19
20type D = EngineInterface
21
22const PANE = 'prompt-upgrade'
23const HISTORY_SIZE = 20
24const flow = atom({ plugin: 'prompt-upgrade', key: 'flow' } as const, null)
25
26const GOLD = '#F2C46D'
27const LINE = '#262B33'
28const MUTED = '#7D8590'
29const TEXT = '#E6EDF3'
30const GREEN = '#3DD68C'
31const AMBER = '#E3B341'
32const INK = '#0D1117'
33
34// The module's own: lost on reload, which the token in `pending` covers.
35let inflight: AbortController | undefined
36
37const newId = () => Math.random().toString(36).slice(2, 10)
38
39const blank = (original: string, step: Step): Flow => ({
40  id: newId(),
41  step,
42  original,
43  questions: [],
44  answers: [],
45  otherFor: null,
46  rewrite: '',
47  added: [],
48  editing: false,
49  pending: null,
50  error: null,
51  note: null,
52  fromHistory: false,
53  entries: [],
54})
55
56const setFlow = ($: D, value: Flow | null) => update($, flow, () => value)
57const patch = ($: D, change: (f: Flow) => Flow) => update($, flow, f => (f ? change(f) : f))
58
59const isAuto = async ($: D) => (await $.store.get('auto')) === true
60
61const loadHistory = async ($: D): Promise<Entry[]> => {
62  const saved = await $.store.get('history')
63  return Array.isArray(saved) ? (saved as Entry[]) : []
64}
65
66const remember = async ($: D, f: Flow) => {
67  if (f.fromHistory || !f.rewrite) return
68  const entry: Entry = { id: f.id, at: await $.clock.now(), original: f.original, rewrite: f.rewrite, added: f.added }
69  const kept = (await loadHistory($)).filter(one => one.id !== f.id)
70  await $.store.set('history', [...kept, entry].slice(-HISTORY_SIZE))
71}
72
73/** Opens the pane; when it cannot be placed, gives the words back and says so once. */
74const show = async ($: D, original: string): Promise<boolean> => {
75  const opened = await $.ui.open({ id: PANE, title: 'Prompt Upgrade', focus: true, closeOnEscape: true, rows: 32, columns: 104 })
76  if (opened.isPlaced) return true
77  inflight?.abort()
78  await setFlow($, null)
79  await $.ui.close({ id: PANE }).catch(() => undefined)
80  if (original) await $.prompt.fill({ text: original })
81  $.ui.log('Prompt Upgrade needs a wider window, so your prompt is back in the prompt box.')
82  return false
83}
84
85/** Runs background work; a failure gives the words back with one plain line. */
86const safely = (work: Promise<unknown>, $: D, original: string) =>
87  void work.catch(async () => {
88    await setFlow($, null).catch(() => undefined)
89    if (original) await $.prompt.fill({ text: original }).catch(() => undefined)
90    $.ui.log('Prompt Upgrade hit a problem. Your prompt is back in the prompt box.')
91  })
92
93const ask = async ($: D, request: ModelCompleteRequest): Promise<{ text: string } | { failure: string }> => {
94  inflight = new AbortController()
95  try {
96    const result = await $.model.complete(request, { signal: inflight.signal })
97    return result.isAnswered ? { text: result.text } : { failure: describeFailure(result) }
98  } catch {
99    return { failure: 'Claude Code would not send the request.' }
100  }
101}
102
103/** Moves to a waiting step under a fresh token; null when there is no upgrade. */
104const claim = async ($: D, step: 'reading' | 'rewriting', token: string) => {
105  const f = await patch($, cur => ({ ...cur, step, pending: token, error: null, note: null, editing: false, otherFor: null }))
106  return f?.pending === token ? f : null
107}
108
109const askQuestions = async ($: D) => {
110  const token = newId()
111  const f = await claim($, 'reading', token)
112  if (!f) return
113  const reply = await ask($, questionsRequest(f.original))
114  const now = await patch($, cur => {
115    if (cur.pending !== token) return cur
116    if ('failure' in reply) return { ...cur, pending: null, step: 'error', error: { during: 'questions', detail: reply.failure } }
117    // Questions are optional: an unreadable list goes straight to the rewrite.
118    const questions = parseQuestions(reply.text) ?? []
119    return { ...cur, pending: null, step: questions.length ? 'questions' : 'rewriting', questions, answers: questions.map(() => null) }
120  })
121  if (now?.step === 'rewriting' && now.pending === null) await rewrite($)
122}
123
124const rewrite = async ($: D) => {
125  const token = newId()
126  const f = await claim($, 'rewriting', token)
127  if (!f) return
128  const reply = await ask($, rewriteRequest(f.original, f.questions, f.answers))
129  const parsed = 'text' in reply ? parseRewrite(reply.text) : null
130  const now = await patch($, cur => {
131    if (cur.pending !== token) return cur
132    if ('failure' in reply) return { ...cur, pending: null, step: 'error', error: { during: 'rewrite', detail: reply.failure } }
133    if (!parsed) return { ...cur, pending: null, step: 'error', error: { during: 'rewrite', detail: 'The rewrite came back in a shape I could not read.' } }
134    return { ...cur, pending: null, step: 'review', rewrite: parsed.prompt, added: parsed.added }
135  })
136  if (now?.step === 'review' && parsed && now.rewrite === parsed.prompt) await remember($, now)
137}
138
139const begin = async ($: D, original: string) => {
140  inflight?.abort()
141  await setFlow($, blank(original, 'reading'))
142  if (await show($, original)) await askQuestions($)
143}
144
145const openBlank = async ($: D, step: 'input' | 'history') => {
146  inflight?.abort()
147  const entries = step === 'history' ? [...(await loadHistory($))].reverse() : []
148  await setFlow($, { ...blank('', step), entries })
149  await show($, '')
150}
151
152/** Closes the pane and forgets the upgrade; true when it gave the original back. */
153const finish = async ($: D, giveBack: boolean) => {
154  const f = await read($, flow)
155  inflight?.abort()
156  await setFlow($, null)
157  await $.ui.close({ id: PANE })
158  if (giveBack && f && !f.fromHistory && f.original) await $.prompt.fill({ text: f.original })
159}
160
161const run = async ($: D, text: string) => {
162  const f = await read($, flow)
163  const final = text.trim()
164  if (!final) return
165  if (f) await remember($, { ...f, rewrite: final })
166  await finish($, false)
167  // A plugin origin: the gate never takes it, so it cannot loop back.
168  await $.prompt.submit({ text: final, asUser: true })
169}
170
171const answer = ($: D, index: number, value: Answer) =>
172  patch($, cur => ({ ...cur, otherFor: null, answers: cur.answers.map((a, i) => (i === index ? value : a)) }))
173
174export const register: Register = on => {
175  on('session.start', async ($, e, next) => {
176    try {
177      await $.command.register({
178        name: 'better',
179        description: 'Upgrade a rough prompt before it is sent',
180        argumentHint: '<prompt> | on | off | history',
181      })
182      $.ui.status((await isAuto($)) ? 'better: auto' : undefined)
183      // A reload mid-step: pick the waiting call back up.
184      const f = await read($, flow)
185      if (f?.pending && f.step === 'reading') safely(askQuestions($), $, f.original)
186      if (f?.pending && f.step === 'rewriting') safely(rewrite($), $, f.original)
187    } catch {
188      // The session goes on without the mod rather than failing to start.
189    }
190    return next(e)
191  })
192
193  // Answers at once; the pane opens in the background.
194  on('command.run', { command: 'better' }, async ($, e) => {
195    const typed = e.args.trim()
196    const word = typed.toLowerCase()
197    if (word === 'on' || word === 'off') {
198      await $.store.set('auto', word === 'on')
199      $.ui.status(word === 'on' ? 'better: auto' : undefined)
200      return {
201        text:
202          word === 'on'
203            ? 'Automatic mode is on: prompts you type open in Prompt Upgrade first. Start one with > to send it as is.'
204            : 'Automatic mode is off: prompts go straight through.',
205      }
206    }
207    if (word === 'history') {
208      safely(openBlank($, 'history'), $, '')
209      return { text: 'Opening your last upgrades.' }
210    }
211    if (!typed) {
212      safely(openBlank($, 'input'), $, '')
213      return { text: 'Type the prompt to upgrade in the pane.' }
214    }
215    safely(begin($, typed), $, typed)
216    return { text: 'Upgrading your prompt. Nothing is sent until you press Run.' }
217  }).catch(() => ({ text: 'Prompt Upgrade could not do that just now. Nothing was sent.' }))
218
219  on('prompt.submit', async ($, e, next) => {
220    const decision = gate(e, await isAuto($))
221    if (decision.kind === 'pass') return next(e)
222    if (decision.kind === 'strip') return next({ ...e, text: decision.text })
223    const original = e.text.trim()
224    safely(begin($, original), $, original)
225    // Answered without `next`: nothing enters the session until Run.
226    return { text: e.text }
227  }).catch(($, e, next) => next(e))
228
229  on('ui.close', { id: PANE }, async ($, e, next) => {
230    const closed = await next(e)
231    // The person's close mark or Escape is a Cancel: the words go back.
232    const f = e.origin.kind === 'person' ? await read($, flow) : null
233    if (f) {
234      inflight?.abort()
235      await setFlow($, null)
236      if (!f.fromHistory && f.original) await $.prompt.fill({ text: f.original })
237    }
238    return closed
239  }).catch(($, e, next) => next(e))
240
241  on('ui.render', { component: 'Pane', requestId: PANE }, async ($, e) => {
242    const table = $.ui.resolve(e)
243    const kit: Kit = { Box: table.Box, Text: table.Text, Button: table.Button, Input: 'Input' in table ? table.Input : undefined }
244    const width = Math.max(24, e.props.bodyColumns)
245    return draw($, kit, width, await read($, flow))
246  })
247}
248
249// Drawing. Reads only: every change happens in a press or an input handler.
250
251type Table = Elements['terminal']
252type Kit = { Box: Table['Box']; Text: Table['Text']; Button: Table['Button']; Input: Table['Input'] | undefined }
253
254const STEPS = [
255  { long: 'Your prompt', short: 'Prompt' },
256  { long: 'Questions', short: 'Ask' },
257  { long: 'Rewrite', short: 'Write' },
258  { long: 'Review', short: 'Review' },
259]
260
261const stepIndex = (f: Flow): number => {
262  if (f.step === 'error') return f.error?.during === 'rewrite' ? 2 : 1
263  return { input: 0, reading: 0, questions: 1, rewriting: 2, review: 3, history: 0 }[f.step]
264}
265
266const draw = ($: D, kit: Kit, width: number, f: Flow | null) => {
267  const { Box, Text, Button, Input } = kit
268
269  const title = 'PROMPT UPGRADE'
270  const header = (
271    <Box flexDirection="column" width={width}>
272      <Text color={TEXT} weight="bold" wrap="truncate-end">
273        {clip(title, width)}
274      </Text>
275      <Text wrap="truncate-end">
276        <Text color={GOLD}>{'━'.repeat(Math.min(title.length, width))}</Text>
277        <Text color={LINE}>{'─'.repeat(Math.max(0, width - title.length))}</Text>
278      </Text>
279    </Box>
280  )
281
282  const heading = (label: string, tone = MUTED) => (
283    <Text wrap="truncate-end">
284      <Text color={tone}>{clip(label, width)}</Text>
285      <Text color={LINE}>{` ${'─'.repeat(Math.max(0, width - label.length - 1))}`}</Text>
286    </Text>
287  )
288
289  const tracker = (current: number) => {
290    const mark = (i: number) => (i < current ? '✓' : i === current ? '●' : '○')
291    const tone = (i: number) => (i < current ? GREEN : i === current ? GOLD : MUTED)
292    const fits = (pick: 'long' | 'short', joint: string) =>
293      STEPS.reduce((n, s) => n + s[pick].length + 2, 0) + joint.length * 3 <= width
294    const layout = fits('long', ' ── ') ? (['long', ' ── '] as const) : fits('short', ' ─ ') ? (['short', ' ─ '] as const) : null
295    if (!layout) {
296      return (
297        <Text color={GOLD} wrap="truncate-end">
298          {clip(`Step ${current + 1} of 4  ${STEPS[current]!.long}`, width)}
299        </Text>
300      )
301    }
302    const [pick, joint] = layout
303    return (
304      <Text wrap="truncate-end">
305        {STEPS.map((s, i) => (
306          <Text>
307            <Text color={tone(i)} weight={i === current ? 'bold' : 'regular'}>{`${mark(i)} ${s[pick]}`}</Text>
308            {i < STEPS.length - 1 ? <Text color={LINE}>{joint}</Text> : ''}
309          </Text>
310        ))}
311      </Text>
312    )
313  }
314
315  const card = (text: string, border: string, w: number, maxLines: number, tone = TEXT) => (
316    <Box borderStyle="round" borderColor={border} paddingX={1} width={w} flexDirection="column">
317      {cap(wrap(text, w - 4), maxLines, w - 4).map(line => (
318        <Text color={tone} wrap="truncate-end">
319          {line || ' '}
320        </Text>
321      ))}
322    </Box>
323  )
324
325  const chip = (key: string, label: string, tone: 'gold' | 'slate' | 'green', onPress: () => unknown, autoFocus?: true) => (
326    <Button key={key} plain autoFocus={autoFocus} onPress={() => void onPress()}>
327      <Text
328        backgroundColor={tone === 'gold' ? GOLD : tone === 'green' ? GREEN : LINE}
329        color={tone === 'slate' ? TEXT : INK}
330        weight={tone === 'slate' ? 'regular' : 'bold'}
331      >
332        {` ${clip(label, width - 2)} `}
333      </Text>
334    </Button>
335  )
336
337  const quiet = (key: string, label: string, onPress: () => unknown) => (
338    <Button key={key} plain onPress={() => void onPress()}>
339      <Text color={MUTED}>{label}</Text>
340    </Button>
341  )
342
343  const row = (...children: RenderChildren[]) => (
344    <Box flexDirection="row" flexWrap="wrap" columnGap={1} width={width}>
345      {children}
346    </Box>
347  )
348
349  const line = (text: string, tone = MUTED) => (
350    <Text color={tone} wrap="truncate-end">
351      {clip(text, width)}
352    </Text>
353  )
354
355  const frame = (current: number | null, ...body: RenderChildren[]) => (
356    <Box flexDirection="column" width={width} rowGap={1}>
357      <Box flexDirection="column">
358        {header}
359        {current === null ? '' : tracker(current)}
360      </Box>
361      {body}
362    </Box>
363  )
364
365  if (!f) return frame(null, line('Nothing to upgrade. Type /better followed by a rough prompt.'))
366
367  const cancel = quiet('cancel', 'Cancel', () => finish($, true))
368  const yours = (
369    <Box flexDirection="column">
370      {heading('YOUR PROMPT')}
371      {card(f.original, LINE, width, 8)}
372    </Box>
373  )
374
375  if (f.step === 'input') {
376    return frame(
377      0,
378      <Box flexDirection="column">
379        {heading('WHAT DO YOU WANT TO ASK?')}
380        {Input ? (
381          <Input
382            key="draft"
383            rows={3}
384            autoFocus
385            placeholder="Type a rough prompt. Enter upgrades it."
386            submitLabel="upgrade"
387            onSubmit={value => void (value.trim() && begin($, value.trim()))}
388          />
389        ) : (
390          line('Type /better followed by your prompt.', TEXT)
391        )}
392      </Box>,
393      row(cancel),
394    )
395  }
396
397  if (f.step === 'reading' || f.step === 'rewriting') {
398    const doing =
399      f.step === 'reading'
400        ? 'Reading your prompt and choosing the questions that matter most.'
401        : 'Rewriting your prompt with your answers.'
402    return frame(stepIndex(f), yours, <Box flexDirection="column">{wrap(doing, width).map(text => line(text))}</Box>, row(cancel))
403  }
404
405  if (f.step === 'questions') {
406    const answered = f.answers.filter(a => a && 'text' in a).length
407    const skipped = f.answers.filter(a => a && 'skipped' in a).length
408    const isDone = answered + skipped === f.questions.length
409    return frame(
410      1,
411      yours,
412      <Box flexDirection="column" rowGap={1}>
413        {heading('QUESTIONS')}
414        {f.questions.map((question, i) => questionBlock(question, i, f.answers[i] ?? null, f.otherFor === i))}
415      </Box>,
416      <Box flexDirection="column">
417        {line(`${answered} of ${f.questions.length} answered${skipped ? `, ${skipped} skipped` : ''}`, answered ? GREEN : MUTED)}
418        {row(chip('upgrade', isDone ? 'Upgrade now' : 'Skip the rest, upgrade now', 'gold', () => rewrite($)), cancel)}
419      </Box>,
420    )
421  }
422
423  function questionBlock(question: Question, i: number, given: Answer, isTyping: boolean) {
424    const isAnswered = given !== null && 'text' in given
425    const isSkipped = given !== null && 'skipped' in given
426    const badge = isAnswered ? ' ✓ ' : isSkipped ? ' – ' : ` ${i + 1} `
427    const badgeTone = isAnswered ? GREEN : isSkipped ? LINE : GOLD
428    const asked = wrap(question.q, width - 4)
429    const custom = isAnswered && !question.options.includes(given.text) ? given.text : null
430    return (
431      <Box flexDirection="column" width={width}>
432        {asked.map((text, n) => (
433          <Text wrap="truncate-end">
434            {n === 0 ? (
435              <Text backgroundColor={badgeTone} color={isSkipped ? MUTED : INK} weight="bold">
436                {badge}
437              </Text>
438            ) : (
439              '   '
440            )}
441            <Text color={isSkipped ? MUTED : TEXT}>{` ${text}`}</Text>
442          </Text>
443        ))}
444        <Box flexDirection="row" flexWrap="wrap" columnGap={1} paddingLeft={4} width={width}>
445          {question.options.map((option, n) =>
446            chip(`q${i}-o${n}`, clip(option, width - 8), isAnswered && given.text === option ? 'gold' : 'slate', () =>
447              answer($, i, { text: option }),
448            ),
449          )}
450          {Input
451            ? chip(`q${i}-other`, custom ? clip(`Other: ${custom}`, width - 8) : 'Other (type it)', custom ? 'gold' : 'slate', () =>
452                patch($, cur => ({ ...cur, otherFor: i })),
453              )
454            : ''}
455          {quiet(`q${i}-skip`, isSkipped ? 'Skipped' : 'Skip', () => answer($, i, { skipped: true }))}
456        </Box>
457        {isTyping && Input ? (
458          <Box paddingLeft={4} width={width}>
459            <Input key={`q${i}-text`} autoFocus placeholder="Type your answer" submitLabel="save" onSubmit={value => void (value.trim() && answer($, i, { text: value.trim() }))} />
460          </Box>
461        ) : (
462          ''
463        )}
464      </Box>
465    )
466  }
467
468  if (f.step === 'error') {
469    const during = f.error?.during === 'rewrite' ? 'rewriting your prompt' : 'choosing questions'
470    return frame(
471      stepIndex(f),
472      <Box borderStyle="round" borderColor={AMBER} paddingX={1} width={width} flexDirection="column">
473        {wrap(`Something went wrong while ${during}.`, width - 4).map(text => line(text, AMBER))}
474        {line('Your prompt is safe.', TEXT)}
475        {wrap(f.error?.detail ?? '', width - 4).map(text => line(text))}
476      </Box>,
477      row(
478        chip('retry', 'Try again', 'gold', () => (f.error?.during === 'rewrite' ? rewrite($) : askQuestions($)), true),
479        chip('original', 'Send my original', 'slate', () => run($, f.original)),
480        cancel,
481      ),
482    )
483  }
484
485  if (f.step === 'history') {
486    return frame(
487      null,
488      <Box flexDirection="column">
489        {heading(`LAST UPGRADES  ${f.entries.length} of ${HISTORY_SIZE}`)}
490        {f.entries.length === 0
491          ? line('No upgrades yet. Try /better followed by a rough prompt.')
492          : f.entries.map(entry => (
493              <Button
494                key={`h-${entry.id}`}
495                plain
496                onPress={() =>
497                  void setFlow($, {
498                    ...blank(entry.original, 'review'),
499                    id: entry.id,
500                    rewrite: entry.rewrite,
501                    added: entry.added,
502                    fromHistory: true,
503                  })
504                }
505              >
506                <Text color={MUTED}>{`${when(entry.at)}  `}</Text>
507                <Text color={TEXT}>{clip(entry.original.replace(/\s+/g, ' '), width - 14)}</Text>
508              </Button>
509            ))}
510      </Box>,
511      row(quiet('close', 'Close', () => finish($, false))),
512    )
513  }
514
515  // Review.
516  const before = wordCount(f.original)
517  const after = wordCount(f.rewrite)
518  const tiles = [
519    { label: 'WORDS BEFORE', value: before, tone: TEXT },
520    { label: 'WORDS AFTER', value: after, tone: GOLD },
521    { label: 'DETAILS ADDED', value: f.added.length, tone: GREEN },
522  ]
523  const tileWidth = (t: (typeof tiles)[number]) => Math.max(t.label.length, bigNumber(t.value)[0].length) + 4
524  const tilesFit = tiles.reduce((n, t) => n + tileWidth(t) + 1, -1) <= width
525  const numbers = tilesFit ? (
526    <Box flexDirection="row" columnGap={1}>
527      {tiles.map(t => (
528        <Box borderStyle="round" borderColor={LINE} paddingX={1} width={tileWidth(t)} flexDirection="column">
529          {bigNumber(t.value).map(text => (
530            <Text color={t.tone}>{text}</Text>
531          ))}
532          <Text color={MUTED}>{t.label}</Text>
533        </Box>
534      ))}
535    </Box>
536  ) : (
537    <Text wrap="truncate-end">
538      <Text color={TEXT}>{`${before}`}</Text>
539      <Text color={MUTED}>{' → '}</Text>
540      <Text color={GOLD}>{`${after}`}</Text>
541      <Text color={MUTED}>{' words, '}</Text>
542      <Text color={GREEN}>{`${f.added.length}`}</Text>
543      <Text color={MUTED}>{' added'}</Text>
544    </Text>
545  )
546
547  const buttons = row(
548    chip('run', 'Run it', 'gold', () => run($, f.rewrite), true),
549    Input ? chip('edit', f.editing ? 'Stop editing' : 'Edit', 'slate', () => patch($, cur => ({ ...cur, editing: !cur.editing, note: null }))) : '',
550    chip('copy', 'Copy', 'slate', async () => {
551      const copied = await $.ui.copy({ text: f.rewrite })
552      await patch($, cur => ({
553        ...cur,
554        note: copied.isCopied ? { text: 'Copied to the clipboard.', tone: 'good' } : { text: 'Could not reach the clipboard here.', tone: 'warn' },
555      }))
556    }),
557    cancel,
558  )
559
560  const added = (
561    <Box flexDirection="column">
562      {heading('WHAT I ADDED')}
563      {f.added.length === 0
564        ? line('Nothing new: your prompt was already clear.')
565        : f.added.map(item => (
566            <Text wrap="truncate-end">
567              <Text color={GREEN}>{'+ '}</Text>
568              <Text color={TEXT}>{clip(item, width - 2)}</Text>
569            </Text>
570          ))}
571    </Box>
572  )
573
574  const editor =
575    f.editing && Input ? (
576      <Box flexDirection="column">
577        {heading('EDIT THE UPGRADE', GOLD)}
578        <Input
579          key="edit-text"
580          autoFocus
581          rows={Math.min(12, Math.max(3, wrap(f.rewrite, width - 2).length))}
582          value={f.rewrite}
583          submitLabel="run"
584          onSubmit={value => void run($, value)}
585        />
586      </Box>
587    ) : (
588      ''
589    )
590
591  const isSideBySide = width >= 60
592  const half = isSideBySide ? Math.floor((width - 1) / 2) : width
593  const side = (label: string, tone: string, text: string, border: string) => (
594    <Box flexDirection="column" width={half}>
595      <Text color={tone} wrap="truncate-end">
596        {clip(label, half)}
597      </Text>
598      {card(text, border, half, 40)}
599    </Box>
600  )
601  const compare = (
602    <Box flexDirection={isSideBySide ? 'row' : 'column'} columnGap={1} rowGap={1} width={width}>
603      {side('YOUR PROMPT', MUTED, f.original, LINE)}
604      {side('UPGRADED', GOLD, f.rewrite, GOLD)}
605    </Box>
606  )
607
608  return frame(
609    3,
610    numbers,
611    <Box flexDirection="column">
612      {buttons}
613      {f.note ? line(f.note.text, f.note.tone === 'good' ? GREEN : AMBER) : ''}
614    </Box>,
615    added,
616    editor,
617    compare,
618  )
619}
620

1import type { ModelCompleteRequest, ModelCompleteResult } from 'claude-code'
2
3import type { Answer, Question } from '../types'
4
5// What decides whether a typed prompt goes to the pane.
6
7export type Gate = { kind: 'pass' } | { kind: 'strip'; text: string } | { kind: 'take' }
8
9type Submitted = {
10  text: string
11  command?: unknown
12  isMeta?: true
13  attachments?: readonly unknown[]
14  turnId?: string
15  origin: { kind: string }
16}
17
18const ACK =
19  /^(y|yes|yep|yeah|yup|n|no|nope|ok|okay|k|sure|thanks|thank you|thx|ty|go|go on|go ahead|continue|keep going|carry on|do it|proceed|sounds good|looks good|lgtm|perfect|great|nice|good|stop|done|next|again|retry|try again)[\s.!?]*$/i
20
21/** Short replies that should never be upgraded. */
22export const isShortReply = (text: string): boolean => {
23  const trimmed = text.trim()
24  return trimmed.split(/\s+/).length <= 2 || ACK.test(trimmed)
25}
26
27export const gate = (e: Submitted, isAuto: boolean): Gate => {
28  // Only the person's own Enter is ever taken; this plugin's own Run submits
29  // with a plugin origin, so it can never come back here.
30  if (!isAuto || e.origin.kind !== 'composer' || e.command !== undefined || e.isMeta) return { kind: 'pass' }
31  const trimmed = e.text.trimStart()
32  if (trimmed.startsWith('>')) {
33    const rest = trimmed.slice(1).trimStart()
34    return rest ? { kind: 'strip', text: rest } : { kind: 'pass' }
35  }
36  if ((e.attachments?.length ?? 0) > 0 || e.turnId !== undefined) return { kind: 'pass' }
37  if (trimmed.startsWith('/') || trimmed.startsWith('!') || isShortReply(trimmed)) return { kind: 'pass' }
38  return { kind: 'take' }
39}
40
41// The two model calls.
42
43export const QUESTION_MODEL = 'haiku'
44export const REWRITE_MODEL = 'sonnet'
45
46const QUESTIONS_SYSTEM = `You help a person sharpen a rough prompt before they send it to Claude Code, an AI coding assistant.
47Choose at most three multiple-choice questions whose answers would most improve the result.
48Ask only about what the prompt leaves open and what would change the outcome: scope, audience, format, constraints, what done looks like.
49Never ask about something the prompt already answers. If the prompt is already clear enough, return no questions.
50Each question is under 12 words. Each has 2 to 4 options of 1 to 5 words.
51Reply with JSON only, no other text: {"questions":[{"q":"...","options":["...","..."]}]}`
52
53const REWRITE_SYSTEM = `You rewrite a person's rough prompt into a clear, complete request for Claude Code, an AI coding assistant.
54Rules:
55- Keep exactly what they meant. Do not change or widen the task.
56- Add a clear goal, who it is for, the format of the result, and what done looks like.
57- Use the answers given. Leave out anything they skipped.
58- Never invent facts, names, numbers, prices, files or tools. Where something needed is unknown, write a placeholder in square brackets, like [file name].
59- Write it as one request the person could have typed themselves: plain language, first person, no headings, no markdown, no preamble.
60Reply with JSON only, no other text: {"prompt":"the upgraded prompt","added":["short phrase for each thing you added"]}
61Give at most 6 items in "added", each under 8 words.`
62
63export const questionsRequest = (original: string): ModelCompleteRequest => ({
64  model: QUESTION_MODEL,
65  system: QUESTIONS_SYSTEM,
66  prompt: `<prompt>\n${original}\n</prompt>`,
67  maxTokens: 600,
68  effort: 'low',
69  timeoutMs: 45_000,
70})
71
72export const rewriteRequest = (original: string, questions: readonly Question[], answers: readonly Answer[]): ModelCompleteRequest => {
73  const answered = questions.flatMap((question, i) => {
74    const answer = answers[i]
75    return answer && 'text' in answer ? [`- ${question.q} ${answer.text}`] : []
76  })
77  const facts = answered.length ? `\n\n<answers>\n${answered.join('\n')}\n</answers>` : ''
78  return {
79    model: REWRITE_MODEL,
80    system: REWRITE_SYSTEM,
81    prompt: `<prompt>\n${original}\n</prompt>${facts}`,
82    maxTokens: 2000,
83    timeoutMs: 90_000,
84  }
85}
86
87const jsonIn = (text: string): unknown => {
88  const start = text.indexOf('{')
89  const end = text.lastIndexOf('}')
90  if (start < 0 || end <= start) return undefined
91  try {
92    return JSON.parse(text.slice(start, end + 1))
93  } catch {
94    return undefined
95  }
96}
97
98const isText = (value: unknown): value is string => typeof value === 'string' && value.trim() !== ''
99
100/** The questions in a reply: at most 3, each with 2 to 4 options; null when unreadable. */
101export const parseQuestions = (text: string): Question[] | null => {
102  const data = jsonIn(text)
103  if (typeof data !== 'object' || data === null || !Array.isArray((data as { questions?: unknown }).questions)) return null
104  const list = (data as { questions: unknown[] }).questions
105  return list
106    .flatMap(item => {
107      const { q, options } = (item ?? {}) as { q?: unknown; options?: unknown }
108      if (!isText(q) || !Array.isArray(options)) return []
109      const kept = [...new Set(options.filter(isText).map(o => o.trim()))].slice(0, 4)
110      return kept.length >= 2 ? [{ q: q.trim(), options: kept }] : []
111    })
112    .slice(0, 3)
113}
114
115/** The upgraded prompt and what it added; null when unreadable. */
116export const parseRewrite = (text: string): { prompt: string; added: string[] } | null => {
117  const data = jsonIn(text) as { prompt?: unknown; added?: unknown } | undefined
118  if (!data || !isText(data.prompt)) return null
119  const added = Array.isArray(data.added) ? data.added.filter(isText).map(a => a.trim()).slice(0, 6) : []
120  return { prompt: data.prompt.trim(), added }
121}
122
123/** A calm sentence for a model call that came back without text. */
124export const describeFailure = (result: Exclude<ModelCompleteResult, { isAnswered: true }>): string => {
125  if (result.reason === 'aborted') return 'The model took too long to answer.'
126  if (result.reason === 'empty-reply') return 'The model answered with nothing.'
127  if (result.error === 'rate_limit') return 'You have hit a rate limit for now.'
128  if (result.error === 'overloaded') return 'The model service is busy right now.'
129  if (result.error === 'authentication_failed') return 'Claude Code could not sign in to the model.'
130  return `The model call failed${result.status === null ? '' : ` (status ${result.status})`}.`
131}
132
133// Text layout.
134
135export const wordCount = (text: string): number => text.split(/\s+/).filter(Boolean).length
136
137export const clip = (text: string, width: number): string =>
138  width <= 0 ? '' : text.length <= width ? text : `${text.slice(0, Math.max(0, width - 1))}…`
139
140/** Word-wraps to `width`, keeping line breaks and splitting words longer than a line. */
141export const wrap = (text: string, width: number): string[] => {
142  const room = Math.max(1, width)
143  return text.split('\n').flatMap(paragraph => {
144    const lines: string[] = []
145    let line = ''
146    for (const raw of paragraph.split(/\s+/).filter(Boolean)) {
147      let word = raw
148      while (word.length > room) {
149        if (line) lines.push(line)
150        lines.push(word.slice(0, room))
151        word = word.slice(room)
152        line = ''
153      }
154      if (!line) line = word
155      else if (line.length + 1 + word.length <= room) line += ` ${word}`
156      else {
157        lines.push(line)
158        line = word
159      }
160    }
161    lines.push(line)
162    return lines
163  })
164}
165
166/** Caps a list of lines, the last kept line marking the cut. */
167export const cap = (lines: string[], max: number, width: number): string[] =>
168  lines.length <= max ? lines : [...lines.slice(0, max - 1), clip(`${lines[max - 1]} …`, width)]
169
170const DIGITS: Record<string, [string, string, string]> = {
171  '0': ['┏━┓', '┃ ┃', '┗━┛'],
172  '1': ['╺┓ ', ' ┃ ', '╺┻╸'],
173  '2': ['╺━┓', '┏━┛', '┗━╸'],
174  '3': ['╺━┓', ' ━┫', '╺━┛'],
175  '4': ['╻ ╻', '┗━┫', '  ╹'],
176  '5': ['┏━╸', '┗━┓', '╺━┛'],
177  '6': ['┏━╸', '┣━┓', '┗━┛'],
178  '7': ['╺━┓', '  ┃', '  ╹'],
179  '8': ['┏━┓', '┣━┫', '┗━┛'],
180  '9': ['┏━┓', '┗━┫', '╺━┛'],
181}
182
183/** A number drawn three rows tall in heavy box-drawing characters. */
184export const bigNumber = (n: number): [string, string, string] => {
185  const glyphs = String(Math.max(0, Math.round(n))).split('').map(d => DIGITS[d] ?? DIGITS['0']!)
186  return [0, 1, 2].map(row => glyphs.map(g => g[row]).join(' ')) as [string, string, string]
187}
188
189const MONTHS = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
190
191export const when = (at: number): string => {
192  const d = new Date(at)
193  const pad = (n: number) => String(n).padStart(2, '0')
194  return `${MONTHS[d.getMonth()]} ${pad(d.getDate())} ${pad(d.getHours())}:${pad(d.getMinutes())}`
195}
196

1export type Question = { q: string; options: string[] }
2
3/** A chosen or typed answer, a skip, or nothing yet. */
4export type Answer = { text: string } | { skipped: true } | null
5
6export type Step = 'input' | 'reading' | 'questions' | 'rewriting' | 'review' | 'error' | 'history'
7
8export type Entry = { id: string; at: number; original: string; rewrite: string; added: string[] }
9
10export type Flow = {
11  id: string
12  step: Step
13  original: string
14  questions: Question[]
15  answers: Answer[]
16  /** Which question's "Other" field is open. */
17  otherFor: number | null
18  rewrite: string
19  added: string[]
20  editing: boolean
21  /** Token of the model call in flight; a reply under another token is dropped. */
22  pending: string | null
23  error: { during: 'questions' | 'rewrite'; detail: string } | null
24  note: { text: string; tone: 'good' | 'warn' } | null
25  /** Reopened from history: Cancel then just closes. */
26  fromHistory: boolean
27  entries: Entry[]
28}
29
30declare module 'claude-code' {
31  interface PluginState {
32    'prompt-upgrade': { flow: Flow | null }
33  }
34}
35