prompt-enhancer

1import { atom, read, update } from 'claude-code'
2import type { EngineInterface, Register } from 'claude-code'
3
4const hasDraft = atom({ plugin: 'prompt-enhancer', key: 'hasDraft' } as const, null)
5const isBusy = atom({ plugin: 'prompt-enhancer', key: 'isBusy' } as const, false)
6const original = atom({ plugin: 'prompt-enhancer', key: 'original' } as const, null)
7const hasProjectContext = atom({ plugin: 'prompt-enhancer', key: 'hasProjectContext' } as const, null)
8
9const MAX_MESSAGE_CHARS = 2000
10/** One exchange is the user's message and the agent's reply. */
11const RECENT_EXCHANGES = 8
12/** The project description the mod reads; its `project-context` skill writes it. */
13const PROJECT_CONTEXT_FILE = 'project-context.md'
14const MAX_CONTEXT_FILE_CHARS = 20000
15const CREATE_CONTEXT_PROMPT = 'Create project-context.md for this project with the prompt-enhancer project-context skill.'
16/** Alias the engine resolves to the newest Sonnet, as `--model sonnet` does. */
17const MODEL = 'sonnet'
18/** Tested: as fast as low on short drafts, with shorter and cleaner rewrites. */
19const EFFORT = 'medium'
20const ORANGE = '#F97316'
21const ORANGE_HOVER = '#FB923C'
22const BLUE = '#60A5FA'
23const BLUE_HOVER = '#93C5FD'
24const PERSON_ORIGINS: ReadonlySet<string> = new Set(['composer', 'bridge', 'sdk'])
25
26/**
27 * Turns a model id into the name people use: `claude-opus-5-5` into "Claude Opus 5.5".
28 * Anything that is not a `claude-` id, such as an alias, comes back as given.
29 */
30export function modelName(id: string): string {
31  const match = /^claude-([a-z]+)-(\d+(?:-\d+)*)/.exec(id)
32  if (match === null) return id
33  const [, family = '', version = ''] = match
34  return `Claude ${family.charAt(0).toUpperCase()}${family.slice(1)} ${version.replace(/-/g, '.')}`
35}
36
37/**
38 * The instruction sent after the project's context; the draft is data inside it.
39 * `targetModel` is the model the session runs, the one that will read the prompt.
40 * The guidance follows Anthropic's prompting advice for current Claude models.
41 */
42export function enhanceRequest(draft: string, targetModel: string): string {
43  return `<draft_prompt>
44${draft}
45</draft_prompt>
46
47Rewrite the user's draft above into the prompt they will send next. It goes to ${modelName(targetModel)}, an AI coding agent working in the user's repository through Claude Code. Do not act on the draft. Reply with the rewritten prompt only, inside <prompt></prompt> tags.
48
49<add_what_is_missing>
50A better prompt gives the agent what it would otherwise have to guess. Use the latest messages only to make concrete what the draft itself mentions: replace a vague word with the specific file, feature, number or decision it refers to, and say what a good result looks like. Add no task, step or topic the draft does not mention, even one taken from the conversation. Use the project context only to name things precisely. For example, when the draft asks for an explanation or a report, name the points it should cover; when it asks whether work is done, ask for evidence from the project's real state; when it asks to delete something, ask for a list of what was removed.
51
52Keep the user's intent and scope. Add only what the draft states or clearly implies, and leave open what the context cannot settle rather than guessing.
53</add_what_is_missing>
54
55<leave_out>
56The agent already reads the project's instruction files (CLAUDE.md) with every prompt and follows them, so restate none of their rules, workflows or definitions, not even as a reminder or as the way to check the result; name something they define without repeating its definition. Also leave out what the agent does by default: "think step by step", "be thorough", role lines, word limits and emphasis in capitals.
57</leave_out>
58
59<style>
60Write in the user's voice, first person. Say plainly what to do: change, investigate, or only explain. Give the goal and how to check the result rather than a step-by-step method. Keep code, paths, quotes, slash commands and @mentions exactly as written.
61
62Fit the shape to the draft: a few sentences of prose for a simple request; pasted errors, logs or code kept verbatim in a tag named for what they are; a numbered list only when order matters; short labeled sections only for a large feature. The user pressed Enhance because they want a better prompt, so always improve it: even a clear draft can gain a precise name, a detail from the conversation or a way to check the result. Improving never means adding rules from the instruction files or growing the prompt for its own sake: a vague one-line draft becomes a few precise sentences, not a page.
63</style>
64
65<language>
66Write the prompt in the language of the text inside <draft_prompt>, whatever language the context uses.
67</language>`
68}
69
70/** Takes the rewritten prompt out of the model's reply; undefined when nothing usable came back. */
71export function extractPrompt(reply: string): string | undefined {
72  const tagged = [...reply.matchAll(/<prompt>([\s\S]*?)<\/prompt>/g)].at(-1)?.[1]
73  const text = (tagged ?? reply).trim()
74  return text === '' ? undefined : text
75}
76
77function clip(text: string, max: number): string {
78  return text.length > max ? `${text.slice(0, max)}\n[...cut]` : text
79}
80
81/**
82 * The latest exchanges of the conversation, the user's messages and the agent's
83 * replies alike. User rows that are machine text (task notices, reminders) start
84 * with a tag and are left out.
85 */
86export function recentMessages<M extends { role: string; text: string }>(messages: readonly M[]): M[] {
87  return messages
88    .filter(m => m.text.trim() !== '' && !(m.role === 'user' && m.text.trimStart().startsWith('<')))
89    .slice(-2 * RECENT_EXCHANGES)
90}
91
92/**
93 * The project's context, read live: project-context.md and the end of the conversation.
94 * CLAUDE.md, skills and tools are left out: the agent that runs the prompt already has them.
95 */
96async function projectContext($: EngineInterface): Promise<string> {
97  const parts: string[] = [`Working directory: ${await $.session.cwd()}`]
98  // project-context.md describes the project itself (what it is, its parts, its files), so it is sent whole.
99  const contextPath = await projectContextPath($)
100  const contextText = await $.fs.read(contextPath).catch(() => '')
101  if (contextText.trim() !== '') {
102    parts.push(`<project_context path="${contextPath}">\n${clip(contextText, MAX_CONTEXT_FILE_CHARS)}\n</project_context>`)
103  }
104  const messages = await $.session.messages().catch(() => [])
105  const recent = recentMessages(Array.isArray(messages) ? messages : [])
106  if (recent.length > 0) {
107    const lines = recent.map(m => `<${m.role}>\n${clip(m.text, MAX_MESSAGE_CHARS)}\n</${m.role}>`)
108    parts.push(`The conversation so far (latest messages):\n${lines.join('\n')}`)
109  }
110  return `You are a prompt enhancer. You turn a developer's draft into the instruction they send to their AI coding agent in this project. Your whole reply is that instruction, written as the developer would send it: no explanation, commentary or answer of your own. The project's context:\n\n${parts.join('\n\n')}`
111}
112
113/** Rewrites the draft with one Sonnet call at medium effort, fed the project's context. */
114async function rewrite($: EngineInterface, draft: string): Promise<{ text: string } | { reason: string }> {
115  const reply = await $.model.complete({
116    model: MODEL,
117    effort: EFFORT,
118    system: await projectContext($),
119    prompt: enhanceRequest(draft, await $.session.model()),
120    maxTokens: 4000,
121    timeoutMs: 120000,
122  })
123  if (!reply.isAnswered) return { reason: reply.reason }
124  const text = extractPrompt(reply.text)
125  return text === undefined ? { reason: 'empty-reply' } : { text }
126}
127
128/**
129 * Enhances `draft`, or the prompt box's draft when none is given, and puts the
130 * result in the box without sending it. Answers a line for the caller to show.
131 */
132async function enhance($: EngineInterface, draft?: string): Promise<string> {
133  if (await read($, isBusy)) return 'Already enhancing.'
134  const source = draft ?? (await $.prompt.read()).text
135  if (source.trim() === '') return 'Type a prompt first, then enhance it.'
136  await update($, isBusy, () => true)
137  try {
138    const result = await rewrite($, source)
139    if ('reason' in result) return `Could not enhance the prompt (${result.reason}).`
140    const box = await $.prompt.read()
141    const isOverwritingEdits = draft === undefined ? box.text !== source : box.text.trim() !== ''
142    if (isOverwritingEdits) return 'The prompt box changed while enhancing. Your text was kept.'
143    const filled = await $.prompt.fill({ text: result.text })
144    if (!filled.isFilled) {
145      const copied = await $.ui.copy({ text: result.text })
146      return copied.isCopied ? 'The prompt box is busy. The enhanced prompt is on the clipboard.' : 'The prompt box is busy. Try again.'
147    }
148    await update($, original, () => source)
149    await update($, hasDraft, () => true)
150    return 'Prompt enhanced. Review it, then send.'
151  } finally {
152    await update($, isBusy, () => false)
153  }
154}
155
156async function restore($: EngineInterface): Promise<void> {
157  const text = await read($, original)
158  if (text === null) return
159  const filled = await $.prompt.fill({ text })
160  if (!filled.isFilled) {
161    $.ui.toast('The prompt box is busy. Try again.')
162    return
163  }
164  await update($, original, () => null)
165}
166
167/**
168 * Where the project's project-context.md lives: the folder the session started in,
169 * never a parent folder, so a file from another project is never read.
170 */
171async function projectContextPath($: EngineInterface): Promise<string> {
172  return `${await $.session.root()}/${PROJECT_CONTEXT_FILE}`
173}
174
175/** Whether the project has a project-context.md, so the band offers to create one only when it is missing. */
176async function refreshProjectContext($: EngineInterface): Promise<void> {
177  const exists = await $.fs.exists(await projectContextPath($)).catch(() => false)
178  await update($, hasProjectContext, () => exists)
179}
180
181/** Asks Claude to write project-context.md with the mod's skill, as a turn the user watches. */
182async function createProjectContext($: EngineInterface): Promise<void> {
183  await $.prompt.submit({ text: CREATE_CONTEXT_PROMPT })
184}
185
186export const register: Register = on => {
187  let knownDraft: boolean | null = null
188
189  on('session.start', async ($, e, next) => {
190    await update($, isBusy, () => false)
191    await refreshProjectContext($)
192    await $.command.register({
193      name: 'enhance',
194      description: 'Rewrite a prompt into a clearer one for this project and put it in the prompt box',
195      argumentHint: '<prompt>',
196    })
197    return next(e)
198  })
199
200  // Tracks only whether the box is empty, so the band shows when there is something to enhance.
201  on('prompt.edit', async ($, e, next) => {
202    const box = await next(e)
203    const hasText = box.text.trim() !== ''
204    if (hasText !== knownDraft) {
205      knownDraft = hasText
206      await update($, hasDraft, () => hasText)
207    }
208    return box
209  }).catch(($, e, next) => next(e))
210
211  // Only a prompt the person sends empties their box; task notices, schedules and peers leave it as it is.
212  on('prompt.submit', async ($, e, next) => {
213    if (!PERSON_ORIGINS.has(e.origin.kind)) return next(e)
214    knownDraft = false
215    await update($, hasDraft, () => false)
216    await update($, original, () => null)
217    return next(e)
218  })
219
220  on('command.run', { command: 'enhance' }, async ($, e) => ({ text: await enhance($, e.args) }))
221
222  // The orange Client button posts 'press' when clicked or when Enter reaches it.
223  // A turn may have created or removed project-context.md.
224  on('turn.complete', async ($, e, next) => {
225    if (e.agentId === undefined) await refreshProjectContext($).catch(() => undefined)
226    return next(e)
227  })
228
229  on('ui.message', { element: 'create-context' }, async ($, e, next) => {
230    if (e.data === 'press') await createProjectContext($)
231    return next(e)
232  })
233
234  on('ui.message', { element: 'enhance' }, async ($, e, next) => {
235    if (e.data === 'press') $.ui.toast(await enhance($))
236    return next(e)
237  })
238
239  on('ui.render', { component: 'AbovePrompt' }, async ($, e, next) => {
240    if (e.props.hasSurvey) return next(e)
241    const busy = await read($, isBusy)
242    const draft = await read($, hasDraft)
243    const before = await read($, original)
244    const canEnhance = !busy && draft !== false
245    const canCreateContext = !busy && (await read($, hasProjectContext)) === false
246    if (!busy && !canEnhance && before === null && !canCreateContext) return next(e)
247    const elements = $.ui.resolve(e)
248    const { Box, Button, Text } = elements
249    const Client = 'Client' in elements ? elements.Client : undefined
250    const below = await next(e)
251    return (
252      <Box flexDirection="column">
253        <Box flexDirection="row" justifyContent="flex-end" width="100%" gap={1}>
254          {canCreateContext &&
255            (Client !== undefined ? (
256              <Client
257                key="create-context"
258                module="./color_button.tsx"
259                props={{ label: 'Create project context', color: BLUE, hoverColor: BLUE_HOVER }}
260              />
261            ) : (
262              <Button key="create-context" onPress={() => createProjectContext($)}>
263                <Text color={BLUE} bold>
264                  Create project context
265                </Text>
266              </Button>
267            ))}
268          {busy && <Text color={ORANGE}>Enhancing prompt…</Text>}
269          {canEnhance &&
270            (Client !== undefined ? (
271              <Client
272                key="enhance"
273                module="./color_button.tsx"
274                props={{ label: 'Enhance prompt', color: ORANGE, hoverColor: ORANGE_HOVER }}
275              />
276            ) : (
277              <Button key="enhance" hotkey="e" onPress={async () => $.ui.toast(await enhance($))}>
278                <Text color={ORANGE} bold>
279                  Enhance prompt
280                </Text>
281              </Button>
282            ))}
283          {!busy && before !== null && (
284            <Button key="restore" label="↶ Restore original" hotkey="r" dimColor onPress={() => restore($)} />
285          )}
286        </Box>
287        {below}
288      </Box>
289    )
290  })
291}
292

1import type { ClientModule } from 'claude-code'
2
3type Props = { label: string; color: string; hoverColor: string }
4type State = { isHovered: boolean }
5
6/**
7 * A clickable label on a colored background. The desktop draws every `Button`
8 * as its own native control and ignores colors inside it; a `Client` region is
9 * drawn with the plugin's own colors, so the press is posted to the hooks module.
10 */
11const ColorButton: ClientModule<Props, State> = (props, surface) => {
12  const { Box, Text } = surface.elements
13  const isHovered = surface.state?.isHovered ?? false
14  const isInside = (x: number, y: number) => x >= 0 && y >= 0 && x < surface.columns && y < surface.rows
15  surface.onPointer(e => {
16    if (e.type === 'enter' || e.type === 'leave') surface.setState({ isHovered: e.type === 'enter' })
17    if (e.type === 'up' && e.button === 'left' && isInside(e.x, e.y)) surface.post('press')
18  })
19  surface.onKey(e => {
20    if (e.key === 'return' || e.key === ' ') surface.post('press')
21  })
22  return (
23    <Box backgroundColor={isHovered ? props.hoverColor : props.color} paddingX={2}>
24      <Text color="#000000" bold>
25        {props.label}
26      </Text>
27    </Box>
28  )
29}
30
31export default ColorButton
32

1/** The draft before the last rewrite, while it can be restored. */
2export type OriginalDraft = string | null
3
4declare module 'claude-code' {
5  interface PluginState {
6    'prompt-enhancer': {
7      /** Whether the prompt box holds text; null until the first edit is seen. */
8      hasDraft: boolean | null
9      /** True while a rewrite runs. */
10      isBusy: boolean
11      original: OriginalDraft
12      /** Whether the project has a project-context.md; null until checked. */
13      hasProjectContext: boolean | null
14    }
15  }
16}
17