prompt-deepener

1import type { Register } from 'claude-code'
2
3import { DEFAULT_IMPROVER_TEXT, buildReplyAndStopRequest, unwrapWake } from './deepen'
4import { cancel, detach, improve, probeFill, putInDraft, readBand, reset, restore, startDeepening } from './flow'
5
6/** Borrowed so the chord the person binds to it (ctrl+x d, see the README) presses Improve at the prompt. */
7const CHORD_ACTION = 'app:diffFileListDown'
8
9export const register: Register = (on, options) => {
10  const improverText = typeof options.improverText === 'string' && options.improverText.trim() !== '' ? options.improverText : DEFAULT_IMPROVER_TEXT
11
12  on('prompt.submit', async ($, e, next) => {
13    const isOwn = e.origin.kind === 'composer' || e.origin.kind === 'bridge'
14    if (isOwn && (await readBand($)).kind === 'deepening') return { drop: 'Deepening in progress, Cancel first.' }
15
16    // `++ prompt` typed by the person: deepen it as if Improve were pressed, and send nothing.
17    // In a project thread the message arrives wrapped in a `<wake>` envelope, so the marker is read from inside it.
18    const wrapped = unwrapWake(e.text)
19    const typed = wrapped?.text ?? e.text
20    const marked = isOwn && e.command === undefined && e.isMeta === undefined ? /^\+\+\s*([\s\S]*\S)/.exec(typed) : null
21    if (marked !== null) {
22      const text = marked[1]!
23      // Where the box takes no text (a thread) the main model gets the request as this turn and stops after answering.
24      if ((await probeFill($)) === 'no_composer') {
25        const request = buildReplyAndStopRequest(improverText, text)
26        return next({ ...e, text: wrapped === null ? request : wrapped.rewrap(request) })
27      }
28      // Started unawaited: a hook that awaits a model call inside the turn's own chain can hang the session.
29      detach($, startDeepening($, improverText, text, e.text, true))
30      return { drop: 'Deepening your prompt…' }
31    }
32
33    await reset($)
34    return next(e)
35  })
36
37  // A survey keeps the band to itself.
38  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
39    if (e.props.hasSurvey) return next(e)
40    const state = await readBand($)
41    const { Box, Text, Button } = $.ui.resolve(e)
42
43    const improveButton = (label: string) => (
44      <Button key="improve" label={label} hotkey="d" action={CHORD_ACTION} onPress={() => detach($, improve($, improverText))} />
45    )
46    const restoreButton = <Button key="restore" label="Restore original" hotkey="r" onPress={() => detach($, restore($))} />
47
48    // Every press returns at once and lets the work run unawaited: a hook that outlives its ten-second budget is skipped,
49    // and a deepening takes longer than that.
50    if (state.kind === 'deepening') {
51      return (
52        <Box flexDirection="row">
53          <Text key="status">deepening… </Text>
54          <Button key="cancel" label="Cancel" hotkey="c" onPress={() => detach($, cancel($))} />
55        </Box>
56      )
57    }
58    if (state.kind === 'deepened') {
59      return (
60        <Box flexDirection="column">
61          <Box flexDirection="row">
62            <Text key="status">{`deepened ×${state.count} `}</Text>
63            {improveButton('Deepen again')}
64            {state.isInDraft ? null : <Button key="put" label="Put in draft" onPress={() => detach($, putInDraft($))} />}
65            {restoreButton}
66          </Box>
67          {state.answers.map((answer, i) => (
68            <Text key={`answer-${i + 1}`} dimColor>{`${i + 1}. ${answer}`}</Text>
69          ))}
70        </Box>
71      )
72    }
73    // A failed deepening leaves the draft as it was, so a deepening before it can still be restored.
74    const isRestorable = state.kind === 'failed' && state.previous.kind === 'deepened'
75    return (
76      <Box flexDirection="row">
77        {improveButton(isRestorable ? 'Deepen again' : 'Improve')}
78        {isRestorable ? restoreButton : null}
79        {state.kind === 'failed' ? <Text key="failed" dimColor>{` failed: ${state.reason}`}</Text> : null}
80      </Box>
81    )
82  })
83}
84

1/*
2 * The pure part of deepening: the request sent to the model and the parse of its reply.
3 * Imports nothing, so the tests and the flow can both use it.
4 */
5
6export const DEFAULT_IMPROVER_TEXT =
7  "I want a deeper version of this task, not a bigger one. Before rewriting it: (1) say what a top professional in this field would do here that I'm not asking for; (2) name the rawest material they would start from instead of what has already been digested for me; (3) name the specific methods and tools they would use; (4) say what, other than you or me, would show the result is right. Then rewrite my prompt to ask for exactly that."
8
9/** Fixed by the mod and appended after the improver text, so a custom improver text cannot break the parse. */
10export const OUTPUT_FORMAT = [
11  'Reply in exactly this format and write nothing outside it. Write everything in the language of the original prompt.',
12  'Each of the four answers is at most 3 lines. The deepened prompt may be as long as it needs.',
13  '<answer-1>your answer to (1)</answer-1>',
14  '<answer-2>your answer to (2)</answer-2>',
15  '<answer-3>your answer to (3)</answer-3>',
16  '<answer-4>your answer to (4)</answer-4>',
17  '<deepened-prompt>the rewritten prompt, ready to send</deepened-prompt>',
18  'Do not carry out the task itself.',
19].join('\n')
20
21/** For the turn that replaces a fill the draft box cannot take: the main model answers in the chat, readable, and stops. */
22export const REPLY_FORMAT = [
23  'Reply with exactly two parts and write nothing else. Write everything in the language of the original prompt.',
24  'Part 1, the deepening: four numbered answers to (1) to (4), each at most 3 lines.',
25  'Part 2, the deepened prompt: the rewritten prompt, ready to send, in one fenced block.',
26  'Then stop. Do not carry out the task itself, do not use any tools, and wait until I send the deepened prompt or say "go".',
27].join('\n')
28
29export const MAX_ANSWER_LINES = 3
30
31export type Answers = [string, string, string, string]
32
33export type Parsed =
34  | { isOk: true; answers: Answers; deepenedPrompt: string }
35  | { isOk: false; reason: 'no-deepened-prompt' | 'wrong-shape' }
36
37/** The improver text, then the output format, then the original prompt. */
38export function buildRequest(improverText: string, originalPrompt: string): string {
39  const improver = improverText.trim() === '' ? DEFAULT_IMPROVER_TEXT : improverText.trim()
40  return `${improver}\n\n${OUTPUT_FORMAT}\n\nOriginal prompt:\n${originalPrompt}`
41}
42
43/** The turn submitted where the draft cannot be filled: improver text, reply format, original prompt, and a stop. */
44export function buildReplyAndStopRequest(improverText: string, originalPrompt: string): string {
45  const improver = improverText.trim() === '' ? DEFAULT_IMPROVER_TEXT : improverText.trim()
46  return `${improver}\n\n${REPLY_FORMAT}\n\nOriginal prompt:\n${originalPrompt}`
47}
48
49function tag(text: string, name: string): string | null {
50  const match = new RegExp(`<${name}>([\\s\\S]*?)</${name}>`).exec(text)
51  return match === null ? null : match[1]!.trim()
52}
53
54/** Keeps the first MAX_ANSWER_LINES non-empty lines. */
55export function cutLines(answer: string): string {
56  return answer
57    .split('\n')
58    .map(line => line.trimEnd())
59    .filter(line => line.trim() !== '')
60    .slice(0, MAX_ANSWER_LINES)
61    .join('\n')
62}
63
64export function parseReply(reply: string): Parsed {
65  const deepenedPrompt = tag(reply, 'deepened-prompt')
66  if (deepenedPrompt === null || deepenedPrompt === '') return { isOk: false, reason: 'no-deepened-prompt' }
67  const found = [1, 2, 3, 4].map(n => tag(reply, `answer-${n}`))
68  if (found.some(answer => answer === null || cutLines(answer) === '')) return { isOk: false, reason: 'wrong-shape' }
69  const [a, b, c, d] = found.map(answer => cutLines(answer!))
70  return { isOk: true, answers: [a!, b!, c!, d!], deepenedPrompt }
71}
72
73export function describeFailure(reason: 'no-deepened-prompt' | 'wrong-shape'): string {
74  return reason === 'no-deepened-prompt' ? 'the reply had no deepened prompt' : 'the reply was not four answers and a deepened prompt'
75}
76
77const ENTITIES: Record<string, string> = { '&amp;': '&', '&lt;': '<', '&gt;': '>', '&quot;': '"', '&#39;': "'" }
78const unescapeText = (s: string) => s.replace(/&(?:amp|lt|gt|quot|#39);/g, m => ENTITIES[m]!)
79const escapeText = (s: string) => s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/>/g, '&gt;')
80
81export type Wrapped = { text: string; rewrap: (text: string) => string }
82
83/**
84 * A project thread hands the model the person's message inside a `<wake>` envelope, not as the bare text.
85 * Finds the triggering message in it so the `++` marker can be read, and rewraps a replacement into the same envelope.
86 */
87export function unwrapWake(raw: string): Wrapped | null {
88  if (!raw.trimStart().startsWith('<wake')) return null
89  const found = [...raw.matchAll(/(<message\b[^>]*\btrigger="true"[^>]*>)([\s\S]*?)(<\/message>)/g)].at(-1)
90  if (found === undefined) return null
91  const [whole, open, body, close] = found as unknown as [string, string, string, string]
92  const start = found.index!
93  return {
94    text: unescapeText(body),
95    rewrap: text => raw.slice(0, start) + open + escapeText(text) + close + raw.slice(start + whole.length),
96  }
97}
98

1/*
2 * Deepening as the engine runs it: the side call, the draft fill and the band's state.
3 * `deepen.ts` holds the pure parts.
4 */
5import { atom, read, update } from 'claude-code'
6import type { EngineInterface } from 'claude-code'
7
8import type { Band, Settled } from '../types'
9import { buildReplyAndStopRequest, buildRequest, describeFailure, parseReply } from './deepen'
10import type { Answers } from './deepen'
11
12export const PLUGIN = 'prompt-deepener'
13
14export const band = atom({ plugin: 'prompt-deepener', key: 'band' } as const, { kind: 'idle' } as Band)
15
16export type Deepening = { isOk: true; answers: Answers; deepenedPrompt: string } | { isOk: false; reason: string }
17
18/** The band's state, for the drawing. */
19export async function readBand($: EngineInterface): Promise<Band> {
20  return read($, band)
21}
22
23function failure(reason: string): Deepening {
24  return { isOk: false, reason }
25}
26
27/** Why a model call gave no text, in a sentence. */
28function describeNoReply(r: { reason: string; status?: number | null; error?: string }): string {
29  if (r.reason === 'api-error') return `the model call failed (${r.error ?? 'error'}${r.status == null ? '' : ` ${r.status}`})`
30  if (r.reason === 'empty-reply') return 'the model sent no text'
31  if (r.reason === 'aborted') return 'the model call was cut short'
32  return `no reply (${r.reason})`
33}
34
35/**
36 * Runs the original prompt through the improver text as a side call over the session's transcript;
37 * with no transcript yet (a fresh session, after /clear) as a plain completion on the prompt alone.
38 */
39export async function deepen($: EngineInterface, improverText: string, original: string, signal?: AbortSignal): Promise<Deepening> {
40  const prompt = buildRequest(improverText, original)
41  const forked = await $.model.fork({ prompt })
42  let text: string
43  if (forked.isAnswered) {
44    text = forked.text
45  } else if (forked.reason === 'nothing-to-fork') {
46    const completed = await $.model.complete({ model: await $.session.model(), prompt, maxTokens: 8192 }, { signal })
47    if (!completed.isAnswered) return failure(describeNoReply(completed))
48    text = completed.text
49  } else {
50    return failure(describeNoReply(forked))
51  }
52  const parsed = parseReply(text)
53  return parsed.isOk ? parsed : failure(describeFailure(parsed.reason))
54}
55
56function settled(current: Band): Settled {
57  return current.kind === 'idle' || current.kind === 'deepened' ? current : current.previous
58}
59
60const DEEPEN_TIMEOUT_MS = 90_000
61
62/** The deepening in flight: which one counts, and the signal that stops its timer and its completion. */
63const run = { id: 0, active: null as number | null, controller: null as AbortController | null }
64
65/** Cancel: the band goes back at once and the draft is as it was; a result arriving later changes nothing. */
66export async function cancel($: EngineInterface): Promise<void> {
67  const current = await read($, band)
68  if (current.kind !== 'deepening') return
69  run.id += 1
70  run.active = null
71  run.controller?.abort()
72  if (current.typed !== null) await $.prompt.fill({ text: current.typed, mode: 'replace' })
73  await update($, band, (): Band => current.previous)
74}
75
76/** What the draft box says to a write that changes nothing: `ok`, or why it refuses (`no_composer`: threads, maybe desktop). */
77export async function probeFill($: EngineInterface): Promise<'ok' | 'no_composer' | 'dialog'> {
78  const box = await $.prompt.read()
79  const filled = await $.prompt.fill({ text: box.text, mode: 'replace', cursor: box.cursor })
80  if (filled.isFilled) return 'ok'
81  return filled.refusal === 'no_composer' ? 'no_composer' : 'dialog'
82}
83
84/** Improve: deepens the draft and puts the deepened prompt in its place. Nothing is sent. */
85export async function improve($: EngineInterface, improverText: string): Promise<void> {
86  const draft = (await $.prompt.read()).text
87  if (draft.trim() === '') return
88  await startDeepening($, improverText, draft, null)
89}
90
91/**
92 * Deepens `text`. `typed` is what the person typed when `++` took it out of the box: put back if the
93 * deepening fails or is cancelled, so the retry is one Enter away; null when the text is still in the draft.
94 */
95export async function startDeepening($: EngineInterface, improverText: string, text: string, typed: string | null, isProbed = false): Promise<void> {
96  // Claimed before the first await: presses return at once, so a second press must find this one running.
97  if (run.active !== null) return
98  const mine = ++run.id
99  run.active = mine
100  try {
101    await deepenOnce($, improverText, text, typed, isProbed, mine)
102  } catch (error) {
103    // Whatever threw, the band must not stay at `deepening` and a taken-out `++` text must come back.
104    const current = await read($, band)
105    const previous = settled(current)
106    await fail($, `unexpected error: ${String(error)}`, previous, typed)
107  } finally {
108    if (run.active === mine) run.active = null
109  }
110}
111
112async function deepenOnce($: EngineInterface, improverText: string, text: string, typed: string | null, isProbed: boolean, mine: number): Promise<void> {
113  const current = await read($, band)
114  if (current.kind === 'deepening') return
115  const previous = settled(current)
116
117  // A box that takes no text (a claude.ai thread): the main model answers in the chat instead, and stops.
118  if (!isProbed && (await probeFill($)) === 'no_composer') {
119    await $.prompt.submit({ text: buildReplyAndStopRequest(improverText, text) })
120    return
121  }
122
123  await update($, band, (): Band => ({ kind: 'deepening', text, previous, typed }))
124
125  // `$.model.fork` takes no abort signal, so a cancelled or timed-out deepening is left to finish and its result discarded;
126  // only the plain-completion fallback is really aborted.
127  const controller = new AbortController()
128  run.controller = controller
129  const timeout = $.clock.sleep(DEEPEN_TIMEOUT_MS, { signal: controller.signal }).then(
130    (): Deepening => failure(`timed out after ${DEEPEN_TIMEOUT_MS / 1000} seconds`),
131    () => new Promise<never>(() => {}),
132  )
133  let result: Deepening
134  try {
135    result = await Promise.race([deepen($, improverText, text, controller.signal), timeout])
136  } finally {
137    controller.abort()
138  }
139  if (mine !== run.id) return
140  if (!result.isOk) return fail($, result.reason, previous, typed)
141
142  const original = previous.kind === 'deepened' ? previous.original : text
143  const filled = await $.prompt.fill({ text: result.deepenedPrompt, mode: 'replace' })
144  // A refusal that names no cause counts as a dialog: the result is kept, never lost.
145  if (!filled.isFilled && filled.refusal === 'no_composer') return fail($, 'the draft box could not take the deepened prompt', previous, typed)
146  if (!filled.isFilled) $.ui.toast('The draft is blocked, probably by a dialog: close it, then press Put in draft.')
147  await update($, band, (): Band => ({
148    kind: 'deepened',
149    original,
150    answers: result.answers,
151    count: previous.kind === 'deepened' ? previous.count + 1 : 1,
152    deepenedPrompt: result.deepenedPrompt,
153    isInDraft: filled.isFilled,
154  }))
155}
156
157/** Put in draft: the deepened prompt, once the dialog that held the keys is gone. */
158export async function putInDraft($: EngineInterface): Promise<void> {
159  const current = await read($, band)
160  if (current.kind !== 'deepened' || current.isInDraft) return
161  const filled = await $.prompt.fill({ text: current.deepenedPrompt, mode: 'replace' })
162  if (!filled.isFilled) {
163    $.ui.toast('The draft box still cannot take it: close any dialog and press Put in draft again.')
164    return
165  }
166  await update($, band, (): Band => ({ ...current, isInDraft: true }))
167}
168
169async function fail($: EngineInterface, reason: string, previous: Settled, typed: string | null): Promise<void> {
170  if (typed !== null) await $.prompt.fill({ text: typed, mode: 'replace' })
171  await update($, band, (): Band => ({ kind: 'failed', reason, previous }))
172  $.ui.toast(`Deepening failed: ${reason}. Your draft is unchanged.`)
173}
174
175/** Any prompt sent: the band forgets the deepening. */
176export async function reset($: EngineInterface): Promise<void> {
177  await update($, band, (): Band => ({ kind: 'idle' }))
178}
179
180/** Restore original: puts the prompt as first typed back in the draft and forgets the deepening. */
181export async function restore($: EngineInterface): Promise<void> {
182  const current = settled(await read($, band))
183  if (current.kind !== 'deepened') return
184  const filled = await $.prompt.fill({ text: current.original, mode: 'replace' })
185  if (!filled.isFilled) {
186    $.ui.toast('Could not restore the original: the draft box did not take it.')
187    return
188  }
189  await update($, band, (): Band => ({ kind: 'idle' }))
190}
191
192/** Runs a press's work unawaited, so the press returns at once; a failure becomes a toast and never an unhandled rejection. */
193export function detach($: EngineInterface, work: Promise<void>): void {
194  work.catch((error: unknown) => {
195    $.ui.toast(`Prompt deepener: ${String(error)}`)
196  })
197}
198

1export type Answers = [string, string, string, string]
2
3/** A band that is not waiting for the model. `original` is the prompt as first typed (what Restore goes back to). */
4export type Settled =
5  | { kind: 'idle' }
6  | { kind: 'deepened'; original: string; answers: Answers; count: number; deepenedPrompt: string; isInDraft: boolean }
7
8export type Band =
9  | Settled
10  | { kind: 'deepening'; text: string; previous: Settled; typed: string | null }
11  | { kind: 'failed'; reason: string; previous: Settled }
12
13declare module 'claude-code' {
14  interface PluginState {
15    'prompt-deepener': { band: Band }
16  }
17}
18