feat: add AI-powered prompt enhancement (#144)

anandgupta42 · claude · web-flow · commit c189028a64f2 · 2026-03-15T11:39:01.000-07:00
* feat: add prompt enhancement feature

Add AI-powered prompt enhancement that rewrites rough user prompts into
clearer, more specific versions before sending to the main model.

- Add `enhancePrompt()` utility using a small/cheap model to polish prompts
- Register `prompt.enhance` TUI command with `&lt;leader&gt;i` keybind
- Show "enhance" hint in the bottom bar alongside agents/commands
- Add `prompt_enhance` keybind to config schema
- Add unit tests for the `clean()` text sanitization function

Inspired by KiloCode's prompt enhancement feature.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* feat: improve enhancement prompt with research-backed approach and add auto-enhance config

- Rewrite system prompt based on AutoPrompter research (5 missing info
  categories: specifics, action plan, scope, verification, intent)
- Add few-shot examples for data engineering tasks (dbt, SQL, migrations)
- Add `experimental.auto_enhance_prompt` config flag (default: false)
- Auto-enhance normal prompts on submit when enabled (skips shell/slash)
- Export `isAutoEnhanceEnabled()` for config-driven behavior

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* fix: address code review findings for prompt enhancement

- Add 15s timeout via `AbortController` to prevent indefinite hangs
- Extract `ENHANCE_ID` constant and document synthetic `as any` casts
- Fix `clean()` regex to match full-string code fences only (avoids
  stripping inner code blocks)
- Export `stripThinkTags()` as separate utility for testability
- Move auto-enhance before extmark expansion (prevents sending
  expanded paste content to the small model)
- Add toast feedback and error logging for auto-enhance path
- Update `store.prompt.input` after enhancement so history is accurate
- Add outer try/catch with logging to `enhancePrompt()`
- Expand tests from 9 to 30: `stripThinkTags()`, `clean()` edge
  cases, combined pipeline tests

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* fix: handle unclosed `&lt;think&gt;` tags from truncated model output

When the small model hits its token limit mid-generation, `&lt;think&gt;` tags
may not have a closing `&lt;/think&gt;`. The previous regex required a closing
tag, which would leak the entire reasoning block into the enhanced prompt.

Now `stripThinkTags()` matches both closed and unclosed think blocks.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* fix: address remaining review findings — history, debounce, tests

- Fix history storing original text instead of enhanced text by passing
  `inputText` explicitly to `history.append()` instead of spreading
  `store.prompt` which may contain stale state
- Add concurrency guard (`enhancingInProgress` flag) to prevent multiple
  concurrent auto-enhance LLM calls from rapid submissions
- Consolidate magic string into `ENHANCE_NAME` constant used across
  agent name, user agent, log service, and ID derivation
- Add justifying comment for `as any` cast on synthetic IDs explaining
  why branded types are safely bypassed
- Add `isAutoEnhanceEnabled()` tests (5 cases): config absent, present
  but missing flag, false, true, undefined
- Add `enhancePrompt()` tests (10 cases): empty input, whitespace,
  successful enhancement, think tag stripping, code fence stripping,
  stream.text failure, stream init failure, empty LLM response, think
  tags with no content, combined pipeline

Test count: 32 -&gt; 48

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

* fix: address Sentry findings — stream consumption and race condition

- Explicitly consume `stream.fullStream` before awaiting `stream.text`
  to prevent potential hangs from Vercel AI SDK stream not being drained
- Add race condition guard to manual enhance command: if user edits the
  prompt while enhancement is in-flight, discard the stale result
- Add same guard to auto-enhance path in `submit()` for consistency
- Update LLM mock to include `fullStream` async iterable

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;

---------

Co-authored-by: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/packages/opencode/src/altimate/enhance-prompt.ts b/packages/opencode/src/altimate/enhance-prompt.ts
@@ -0,0 +1,157 @@
+// altimate_change - new file
+import { Provider } from "@/provider/provider"
+import { LLM } from "@/session/llm"
+import { Agent } from "@/agent/agent"
+import { Config } from "@/config/config"
+import { Log } from "@/util/log"
+import { MessageV2 } from "@/session/message-v2"
+
+const ENHANCE_NAME = "enhance-prompt"
+const ENHANCE_TIMEOUT_MS = 15_000
+// MessageV2.User requires branded MessageID/SessionID types, but this is a
+// synthetic message that never enters the session store — cast is safe here.
+const ENHANCE_ID = ENHANCE_NAME as any
+
+const log = Log.create({ service: ENHANCE_NAME })
+
+// Research-backed enhancement prompt based on:
+// - AutoPrompter (arxiv 2504.20196): 5 missing info categories that cause 27% lower edit correctness
+// - Meta-prompting best practices: clear role, structural scaffolding, few-shot examples
+// - KiloCode's enhance-prompt implementation: lightweight model, preserve intent, no wrapping
+const ENHANCE_SYSTEM_PROMPT = `You are a prompt rewriter for a data engineering coding agent. The agent can read/write files, run SQL, manage dbt models, inspect schemas, and execute shell commands.
+
+Your task: rewrite the user's rough prompt into a clearer version that will produce better results. Reply with ONLY the enhanced prompt — no explanations, no wrapping in quotes or code fences.
+
+## What to improve
+
+Research shows developer prompts commonly lack these five categories of information. Add them when missing:
+
+1. **Specifics** — Add concrete details the agent needs: table names, column names, file paths, SQL dialects, error messages. If the user references "the model" or "the table", keep the reference but clarify what the agent should look for.
+2. **Action plan** — When the prompt is vague ("fix this"), add explicit steps: investigate first, then modify, then verify. Structure as a logical sequence.
+3. **Scope** — Clarify what files, models, or queries are in scope. If ambiguous, instruct the agent to identify the scope first.
+4. **Verification** — Add a verification step when the user implies correctness matters (fixes, migrations, refactors). E.g. "run the query to confirm results" or "run dbt test after changes".
+5. **Intent clarification** — When the request could be interpreted multiple ways, pick the most likely interpretation and make it explicit.
+
+## Rules
+
+- Preserve the user's intent exactly — never add requirements they didn't ask for
+- Keep it concise — a good enhancement adds 1-3 sentences, not paragraphs
+- If the prompt is already clear and specific, return it unchanged
+- Write in the same tone/style as the user (casual stays casual, technical stays technical)
+- Never add generic filler like "please ensure best practices" or "follow coding standards"
+- Do not mention yourself or the enhancement process
+
+## Examples
+
+User: "fix the failing test"
+Enhanced: "Investigate the failing test — run the test suite first to identify which test is failing and why, then examine the relevant source code, apply a fix, and re-run the test to confirm it passes."
+
+User: "add a created_at column to the users model"
+Enhanced: "Add a created_at timestamp column to the users dbt model. Update the SQL definition and the schema.yml entry. Use the appropriate timestamp type for the target warehouse."
+
+User: "why is this query slow"
+Enhanced: "Analyze why the query is slow. Run EXPLAIN/query profile to identify bottlenecks (full table scans, missing indexes, expensive joins). Suggest specific optimizations based on the findings."
+
+User: "migrate this from snowflake to bigquery"
+Enhanced: "Migrate the SQL from Snowflake dialect to BigQuery dialect. Convert Snowflake-specific functions (e.g. DATEADD, IFF, QUALIFY) to BigQuery equivalents. Preserve the query logic and verify the translated query is syntactically valid."`
+
+export function stripThinkTags(text: string) {
+  // Match closed <think>...</think> blocks, and also unclosed <think>... to end of string
+  // (unclosed tags happen when the model hits token limit mid-generation)
+  return text.replace(/<think>[\s\S]*?(?:<\/think>\s*|$)/g, "")
+}
+
+export function clean(text: string) {
+  return text
+    .trim()
+    .replace(/^```\w*\n([\s\S]*?)\n```$/, "$1")
+    .trim()
+    .replace(/^(['"])([\s\S]*)\1$/, "$2")
+    .trim()
+}
+
+/**
+ * Check if auto-enhance is enabled in config.
+ * Defaults to false — user must explicitly opt in.
+ */
+export async function isAutoEnhanceEnabled(): Promise<boolean> {
+  const cfg = await Config.get()
+  return cfg.experimental?.auto_enhance_prompt === true
+}
+
+export async function enhancePrompt(text: string): Promise<string> {
+  const trimmed = text.trim()
+  if (!trimmed) return text
+
+  log.info("enhancing", { length: trimmed.length })
+
+  const controller = new AbortController()
+  const timeout = setTimeout(() => controller.abort(), ENHANCE_TIMEOUT_MS)
+
+  try {
+    const defaultModel = await Provider.defaultModel()
+    const model =
+      (await Provider.getSmallModel(defaultModel.providerID)) ??
+      (await Provider.getModel(defaultModel.providerID, defaultModel.modelID))
+
+    const agent: Agent.Info = {
+      name: ENHANCE_NAME,
+      mode: "primary",
+      hidden: true,
+      options: {},
+      permission: [],
+      prompt: ENHANCE_SYSTEM_PROMPT,
+      temperature: 0.7,
+    }
+
+    const user: MessageV2.User = {
+      id: ENHANCE_ID,
+      sessionID: ENHANCE_ID,
+      role: "user",
+      time: { created: Date.now() },
+      agent: ENHANCE_NAME,
+      model: {
+        providerID: model.providerID,
+        modelID: model.id,
+      },
+    }
+
+    const stream = await LLM.stream({
+      agent,
+      user,
+      system: [],
+      small: true,
+      tools: {},
+      model,
+      abort: controller.signal,
+      sessionID: ENHANCE_ID,
+      retries: 2,
+      messages: [
+        {
+          role: "user",
+          content: trimmed,
+        },
+      ],
+    })
+
+    // Consume the stream explicitly to avoid potential SDK hangs where
+    // .text never resolves if the stream isn't drained (Vercel AI SDK caveat)
+    for await (const _ of stream.fullStream) {
+      // drain
+    }
+    const result = await stream.text.catch((err) => {
+      log.error("failed to enhance prompt", { error: err })
+      return undefined
+    })
+
+    if (!result) return text
+
+    const cleaned = clean(stripThinkTags(result).trim())
+    return cleaned || text
+  } catch (err) {
+    log.error("enhance prompt failed", { error: err })
+    return text
+  } finally {
+    clearTimeout(timeout)
+  }
+}
diff --git a/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx b/packages/opencode/src/cli/cmd/tui/component/prompt/index.tsx
@@ -34,6 +34,10 @@ import { useToast } from "../../ui/toast"
 import { useKV } from "../../context/kv"
 import { useTextareaKeybindings } from "../textarea-keybindings"
 import { DialogSkill } from "../dialog-skill"
+// altimate_change start - import prompt enhancement
+import { enhancePrompt, isAutoEnhanceEnabled } from "@/altimate/enhance-prompt"
+let enhancingInProgress = false
+// altimate_change end
 
 export type PromptProps = {
   sessionID?: string
@@ -194,6 +198,53 @@ export function Prompt(props: PromptProps) {
           dialog.clear()
         },
       },
+      // altimate_change start - add prompt enhance command
+      {
+        title: "Enhance prompt",
+        value: "prompt.enhance",
+        keybind: "prompt_enhance",
+        category: "Prompt",
+        enabled: !!store.prompt.input,
+        onSelect: async (dialog) => {
+          if (!store.prompt.input.trim()) return
+          dialog.clear()
+          const original = store.prompt.input
+          toast.show({
+            message: "Enhancing prompt...",
+            variant: "info",
+            duration: 2000,
+          })
+          try {
+            const enhanced = await enhancePrompt(original)
+            // Guard against race condition: if user edited the prompt while
+            // enhancement was in-flight, discard the stale enhanced result
+            if (store.prompt.input !== original) return
+            if (enhanced !== original) {
+              input.setText(enhanced)
+              setStore("prompt", "input", enhanced)
+              input.gotoBufferEnd()
+              toast.show({
+                message: "Prompt enhanced",
+                variant: "success",
+                duration: 2000,
+              })
+            } else {
+              toast.show({
+                message: "Prompt already looks good",
+                variant: "info",
+                duration: 2000,
+              })
+            }
+          } catch {
+            toast.show({
+              message: "Failed to enhance prompt",
+              variant: "error",
+              duration: 3000,
+            })
+          }
+        },
+      },
+      // altimate_change end
       {
         title: "Paste",
         value: "prompt.paste",
@@ -564,6 +615,32 @@ export function Prompt(props: PromptProps) {
     const messageID = MessageID.ascending()
     let inputText = store.prompt.input
 
+    // altimate_change start - auto-enhance prompt before expanding paste text
+    // Only enhance the raw user text, not shell commands or slash commands
+    // Guard prevents concurrent enhancement calls from rapid submissions
+    if (store.mode === "normal" && !inputText.startsWith("/") && !enhancingInProgress) {
+      try {
+        const autoEnhance = await isAutoEnhanceEnabled()
+        if (autoEnhance) {
+          enhancingInProgress = true
+          toast.show({ message: "Enhancing prompt...", variant: "info", duration: 2000 })
+          const enhanced = await enhancePrompt(inputText)
+          // Discard if user changed the prompt during enhancement
+          if (store.prompt.input !== inputText) return
+          if (enhanced !== inputText) {
+            inputText = enhanced
+            setStore("prompt", "input", enhanced)
+          }
+        }
+      } catch (err) {
+        // Enhancement failure should never block prompt submission
+        console.error("auto-enhance failed, using original prompt", err)
+      } finally {
+        enhancingInProgress = false
+      }
+    }
+    // altimate_change end
+
     // Expand pasted text inline before submitting
     const allExtmarks = input.extmarks.getAllForTypeId(promptPartTypeId)
     const sortedExtmarks = allExtmarks.sort((a: { start: number }, b: { start: number }) => b.start - a.start)
@@ -653,6 +730,7 @@ export function Prompt(props: PromptProps) {
     }
     history.append({
       ...store.prompt,
+      input: inputText,
       mode: currentMode,
     })
     input.extmarks.clear()
@@ -1155,6 +1233,11 @@ export function Prompt(props: PromptProps) {
                   <text fg={theme.text}>
                     {keybind.print("command_list")} <span style={{ fg: theme.textMuted }}>commands</span>
                   </text>
+                  {/* altimate_change start - show enhance hint */}
+                  <text fg={theme.text}>
+                    {keybind.print("prompt_enhance")} <span style={{ fg: theme.textMuted }}>enhance</span>
+                  </text>
+                  {/* altimate_change end */}
                 </Match>
                 <Match when={store.mode === "shell"}>
                   <text fg={theme.text}>
diff --git a/packages/opencode/src/config/config.ts b/packages/opencode/src/config/config.ts
@@ -866,6 +866,9 @@ export namespace Config {
       agent_cycle: z.string().optional().default("tab").describe("Next agent"),
       agent_cycle_reverse: z.string().optional().default("shift+tab").describe("Previous agent"),
       variant_cycle: z.string().optional().default("ctrl+t").describe("Cycle model variants"),
+      // altimate_change start - add prompt enhance keybind
+      prompt_enhance: z.string().optional().default("<leader>i").describe("Enhance prompt with AI before sending"),
+      // altimate_change end
       input_clear: z.string().optional().default("ctrl+c").describe("Clear input field"),
       input_paste: z.string().optional().default("ctrl+v").describe("Paste from clipboard"),
       input_submit: z.string().optional().default("return").describe("Submit input"),
@@ -1226,6 +1229,14 @@ export namespace Config {
             .positive()
             .optional()
             .describe("Timeout in milliseconds for model context protocol (MCP) requests"),
+          // altimate_change start - auto-enhance prompt config
+          auto_enhance_prompt: z
+            .boolean()
+            .optional()
+            .describe(
+              "Automatically enhance prompts with AI before sending (default: false). Uses a small model to rewrite rough prompts into clearer versions.",
+            ),
+          // altimate_change end
         })
         .optional(),
     })
diff --git a/packages/opencode/test/altimate/enhance-prompt.test.ts b/packages/opencode/test/altimate/enhance-prompt.test.ts
