test(tui): isolate prompt recovery drafts

packages/tui/src/component/prompt/draft-stash.ts

@@ -1,3 +1,5 @@
+import { createSignal } from "solid-js"
+import { unwrap } from "solid-js/store"
 import type { PromptInfo } from "../../prompt/history"
 
 // Holds one in-progress draft per tab across Prompt remounts. A draft is
@@ -6,6 +8,7 @@ import type { PromptInfo } from "../../prompt/history"
 export type DraftEntry = { prompt: PromptInfo; cursor: number }
 
 const byTab = new Map<string | undefined, DraftEntry>()
+const [failed, setFailed] = createSignal(new Map<string | undefined, DraftEntry[]>())
 
 export function takeDraft(sessionID: string | undefined) {
   const entry = byTab.get(sessionID)
@@ -16,3 +19,25 @@ export function takeDraft(sessionID: string | undefined) {
 export function saveDraft(sessionID: string | undefined, entry: DraftEntry) {
   byTab.set(sessionID, entry)
 }
+
+export function failedDrafts(sessionID: string | undefined) {
+  return failed().get(sessionID) ?? []
+}
+
+export function saveFailedDraft(sessionID: string | undefined, entry: DraftEntry) {
+  const snapshot = structuredClone(unwrap(entry))
+  setFailed((current) => new Map(current).set(sessionID, [...(current.get(sessionID) ?? []), snapshot]))
+}
+
+export function takeFailedDraft(sessionID: string | undefined) {
+  const entries = failedDrafts(sessionID)
+  const entry = entries[0]
+  if (!entry) return
+  setFailed((current) => {
+    const next = new Map(current)
+    if (entries.length > 1) return next.set(sessionID, entries.slice(1))
+    next.delete(sessionID)
+    return next
+  })
+  return entry
+}

packages/tui/src/component/prompt/index.tsx

@@ -7,7 +7,20 @@ import {
   decodePasteBytes,
   type KeyEvent,
 } from "@opentui/core"
-import { createEffect, createMemo, onMount, createSignal, onCleanup, on, Show, Switch, Match, For } from "solid-js"
+import {
+  batch,
+  createEffect,
+  createMemo,
+  onMount,
+  createSignal,
+  onCleanup,
+  on,
+  Show,
+  Switch,
+  Match,
+  For,
+  untrack,
+} from "solid-js"
 import path from "path"
 import { useLocal } from "../../context/local"
 import { useTheme, useThemes } from "../../context/theme"
@@ -31,7 +44,7 @@ import { parseSlashHead } from "../../prompt/parse"
 import { stringWidth } from "../../util/string-width"
 import { createStore, produce, unwrap } from "solid-js/store"
 import { emptyPrompt, usePromptHistory, type PromptInfo, type PromptPartRef } from "../../prompt/history"
-import { saveDraft, takeDraft } from "./draft-stash"
+import { failedDrafts, saveDraft, saveFailedDraft, takeDraft, takeFailedDraft } from "./draft-stash"
 import { Skill } from "@opencode-ai/schema/skill"
 import { computePromptTraits } from "../../prompt/traits"
 import { expandPastedTextPlaceholders, expandTrackedPastedText } from "../../prompt/part"
@@ -680,7 +693,7 @@ export function Prompt(props: PromptProps) {
       return input.focused
     },
     get current() {
-      return store.prompt
+      return { ...store.prompt, mode: store.mode }
     },
     focus() {
       input.focus()
@@ -689,10 +702,12 @@ export function Prompt(props: PromptProps) {
       input.blur()
     },
     set(prompt) {
-      input.setText(prompt.text)
-      setStore("prompt", prompt)
-      restoreExtmarksFromPrompt(prompt)
-      input.gotoBufferEnd()
+      batch(() => {
+        input.setText(prompt.text)
+        setStore("prompt", prompt)
+        restoreExtmarksFromPrompt(prompt)
+        input.gotoBufferEnd()
+      })
     },
     reset() {
       resetComposer()
@@ -703,32 +718,60 @@ export function Prompt(props: PromptProps) {
   }
 
   function resetComposer() {
-    input.extmarks.clear()
-    setStore("prompt", emptyPrompt())
-    setStore("extmarkToPart", new Map())
-    input.clear()
+    batch(() => {
+      input.extmarks.clear()
+      setStore("prompt", emptyPrompt())
+      setStore("extmarkToPart", new Map())
+      input.clear()
+    })
+  }
+
+  function restorePrompt(prompt: PromptInfo, cursor?: number) {
+    batch(() => {
+      input.setText(prompt.text)
+      setStore("prompt", prompt)
+      setStore("mode", prompt.mode ?? "normal")
+      restoreExtmarksFromPrompt(prompt)
+      input.gotoBufferEnd()
+      if (cursor !== undefined) input.cursorOffset = cursor
+    })
   }
 
   // Captured once: the session route is keyed by sessionID, so this Prompt
   // instance belongs to exactly one tab. Reading props.sessionID lazily would
   // observe the *next* route during onCleanup and stash under the wrong tab.
   const stashSessionID = props.sessionID
+  const failed = () => failedDrafts(stashSessionID)
+  const hasDraft = () =>
+    !!store.prompt.text ||
+    [store.prompt.files, store.prompt.agents, store.prompt.skills, store.prompt.pasted].some((parts) => parts?.length)
+
+  function restoreFailedDraft() {
+    batch(() => {
+      const entry = takeFailedDraft(stashSessionID)
+      if (!entry) return
+      if (hasDraft()) stash.push({ prompt: { ...store.prompt, mode: store.mode } })
+      restorePrompt(entry.prompt, entry.cursor)
+    })
+  }
 
   onMount(() => {
     const saved = takeDraft(stashSessionID)
     if (store.prompt.text) return
-    if (saved && saved.prompt.text) {
-      input.setText(saved.prompt.text)
-      setStore("prompt", saved.prompt)
-      restoreExtmarksFromPrompt(saved.prompt)
-      input.cursorOffset = saved.cursor
-    }
+    if (saved && saved.prompt.text) restorePrompt(saved.prompt, saved.cursor)
+  })
+
+  // The submitted Prompt may have unmounted before rejection. Recovery belongs
+  // to the tab, and waits rather than overwriting text typed after submission.
+  createEffect(() => {
+    if (!inputTarget() || disposed || hasDraft() || !failed().length) return
+    untrack(restoreFailedDraft)
   })
 
   onCleanup(() => {
     disposed = true
     if (store.prompt.text) {
-      saveDraft(stashSessionID, { prompt: unwrap(store.prompt), cursor: input.cursorOffset })
+      saveDraft(stashSessionID, { prompt: { ...unwrap(store.prompt), mode: store.mode }, cursor: input.cursorOffset })
     }
     setInputTarget(undefined)
     props.ref?.(undefined)
@@ -876,6 +919,16 @@ export function Prompt(props: PromptProps) {
 
   const stashCommands = createMemo(() =>
     [
+      {
+        title: "Restore failed prompt",
+        name: "prompt.restore_failed",
+        category: "Prompt",
+        enabled: failed().length > 0,
+        run: () => {
+          restoreFailedDraft()
+          dialog.clear()
+        },
+      },
       {
         title: "Stash prompt",
         name: "prompt.stash",
@@ -883,7 +936,7 @@ export function Prompt(props: PromptProps) {
         enabled: !!store.prompt.text,
         run: () => {
           if (!store.prompt.text) return
-          stash.push({ prompt: store.prompt })
+          stash.push({ prompt: { ...store.prompt, mode: store.mode } })
           resetComposer()
           dialog.clear()
         },
@@ -895,12 +948,7 @@ export function Prompt(props: PromptProps) {
         enabled: stash.list().length > 0,
         run: () => {
           const entry = stash.pop()
-          if (entry) {
-            input.setText(entry.prompt.text)
-            setStore("prompt", entry.prompt)
-            restoreExtmarksFromPrompt(entry.prompt)
-            input.gotoBufferEnd()
-          }
+          if (entry) restorePrompt(entry.prompt)
           dialog.clear()
         },
       },
@@ -910,16 +958,7 @@ export function Prompt(props: PromptProps) {
         category: "Prompt",
         enabled: stash.list().length > 0,
         run: () => {
-          dialog.replace(() => (
-            <DialogStash
-              onSelect={(entry) => {
-                input.setText(entry.prompt.text)
-                setStore("prompt", entry.prompt)
-                restoreExtmarksFromPrompt(entry.prompt)
-                input.gotoBufferEnd()
-              }}
-            />
-          ))
+          dialog.replace(() => <DialogStash onSelect={(entry) => restorePrompt(entry.prompt)} />)
         },
       },
     ].map(
@@ -1038,11 +1077,7 @@ export function Prompt(props: PromptProps) {
 
             const item = history.move(-1, input.plainText)
             if (!item) return false
-            input.setText(item.text)
-            setStore("prompt", item)
-            setStore("mode", item.mode ?? "normal")
-            restoreExtmarksFromPrompt(item)
-            input.cursorOffset = 0
+            restorePrompt(item, 0)
           },
         },
       ],
@@ -1077,11 +1112,7 @@ export function Prompt(props: PromptProps) {
 
             const item = history.move(1, input.plainText)
             if (!item) return false
-            input.setText(item.text)
-            setStore("prompt", item)
-            setStore("mode", item.mode ?? "normal")
-            restoreExtmarksFromPrompt(item)
-            input.cursorOffset = input.plainText.length
+            restorePrompt(item, item.text.length)
           },
         },
       ],
@@ -1190,19 +1221,14 @@ export function Prompt(props: PromptProps) {
     // Everything below reads the snapshot: text typed while a request is in
     // flight lands in the already-empty composer and survives, and prompt
     // history records exactly what was submitted instead of the live store
-    // (which may have absorbed mid-flight typing). Failure paths restore the
-    // snapshot unless the user has started typing something new.
+    // (which may have absorbed mid-flight typing). Failure recovery is tab-owned
+    // so navigation and newer input cannot discard the submitted snapshot.
     const currentMode = store.mode
     const entry = { ...store.prompt, mode: currentMode }
     resetComposer()
     props.onSubmit?.()
     const restoreEntry = () => {
-      if (disposed || input.isDestroyed || input.plainText !== "") return
-      input.setText(entry.text)
-      setStore("prompt", entry)
-      setStore("mode", entry.mode ?? "normal")
-      restoreExtmarksFromPrompt(entry)
-      input.cursorOffset = entry.text.length
+      saveFailedDraft(stashSessionID, { prompt: entry, cursor: entry.text.length })
     }
 
     const variant = selection.variant
@@ -1351,8 +1377,7 @@ export function Prompt(props: PromptProps) {
       }
       // The data layer admits optimistically: the prompt renders immediately
       // and rolls back if the server rejects it, so submission does not wait
-      // on the network. On rejection the row is already rolled back; restore
-      // the composer unless the user has started typing something new.
+      // on the network. On rejection retain the submitted snapshot for recovery.
       data.session
         .prompt({
           sessionID: target,
@@ -1686,6 +1711,27 @@ export function Prompt(props: PromptProps) {
             flexGrow={1}
             width="100%"
           >
+            <Show when={failed()[0]}>
+              {(entry) => (
+                <box flexDirection="row" gap={1} paddingBottom={1}>
+                  <text fg={theme.text.feedback.error.default} flexGrow={1} wrapMode="none" truncate>
+                    Failed to send: {entry().prompt.text.replace(/\s+/g, " ")}
+                    {failed().length > 1 ? ` (+${failed().length - 1} more)` : ""}
+                  </text>
+                  <text
+                    fg={theme.text.action.primary.default}
+                    flexShrink={0}
+                    onMouseUp={(event: MouseEvent) => {
+                      if (event.button !== 0) return
+                      event.stopPropagation()
+                      restoreFailedDraft()
+                    }}
+                  >
+                    restore
+                  </text>
+                </box>
+              )}
+            </Show>
             <Show when={config.prompt?.image_preview && visibleImageAttachments().length > 0}>
               <box
                 width="100%"

packages/tui/test/prompt-recovery.test.tsx

@@ -0,0 +1,220 @@
+import { expect, test } from "bun:test"
+import { createTestRenderer } from "@opentui/core/testing"
+import { InputRenderable, TextareaRenderable } from "@opentui/core"
+import { Effect, FileSystem } from "effect"
+import { Global } from "@opencode-ai/util/global"
+import { takeDraft } from "../src/component/prompt/draft-stash"
+import { createEventStream, createFetch, directory, json } from "./fixture/tui-client"
+import { tmpdir } from "./fixture/fixture"
+
+test.each([
+  "untouched",
+  "away",
+  "returned",
+  "typing",
+  "away-draft",
+  "cleared",
+  "multiple",
+  "shell-draft",
+  "history-paste",
+  "setup-mode",
+])("failed prompt recovery preserves input after %s", async (scenario) => {
+  await using state = await tmpdir()
+  const setup = await createTestRenderer({ width: 100, height: 30, useThread: false, kittyKeyboard: true })
+  setup.renderer.start()
+  const ready = Promise.withResolvers<void>()
+  const release = Promise.withResolvers<void>()
+  const requested = Promise.withResolvers<void>()
+  let sessionID = `ses_recovery_${scenario}`
+  const location = { directory, project: { id: "project", directory, canonical: directory } }
+  const events = createEventStream()
+  const calls = createFetch(async (url, request) => {
+    if (scenario === "setup-mode" && url.pathname === "/api/session" && request.method === "POST") {
+      sessionID = (await request.json()).id
+    }
+    if (
+      url.pathname === `/api/session/${sessionID}` ||
+      (scenario === "setup-mode" && url.pathname === "/api/session" && request.method === "POST")
+    )
+      return json({
+        data: {
+          id: sessionID,
+          projectID: "project",
+          title: "Prompt recovery",
+          agent: "build",
+          model: { providerID: "demo", id: "model" },
+          location: { directory },
+          cost: 0,
+          tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+          time: { created: 0, updated: 0 },
+        },
+      })
+    if (url.pathname === `/api/session/${sessionID}/message`) return json({ data: [], cursor: {} })
+    if (/^\/api\/session\/[^/]+\/(inbox|permission)$/.test(url.pathname)) return json({ data: [] })
+    if (url.pathname === "/api/agent")
+      return json({ location, data: [{ id: "build", mode: "primary", hidden: false, permissions: [] }] })
+    if (url.pathname === "/api/provider") return json({ location, data: [{ id: "demo", name: "Demo" }] })
+    if (url.pathname === "/api/model")
+      return json({ location, data: [{ id: "model", providerID: "demo", name: "Demo Model", variants: [] }] })
+    if (url.pathname === `/api/session/${sessionID}/model`) {
+      requested.resolve()
+      await release.promise
+      return json({ message: "Fixture admission rejected" }, { status: 400 })
+    }
+    return undefined
+  }, events)
+  const server = Bun.serve({ port: 0, fetch: (request) => calls.fetch(request) })
+  const { run } = await import("../src/app")
+  const task = Effect.runPromise(
+    run({
+      app: { name: "test", version: "test", channel: "test" },
+      server: { endpoint: { url: server.url.toString() } },
+      config: {
+        get: async () => ({
+          animations: false,
+          keybinds: {
+            "session.new": "f6",
+            "session.tab.select.1": "f7",
+            "prompt.stash.pop": "f8",
+            "prompt.stash": "f9",
+          },
+        }),
+        update: async () => ({}),
+      },
+      packages: { resolve: async () => undefined },
+      terminalHandoff: async () => ({ renderer: setup.renderer, mode: "dark", complete: ready.resolve }),
+      args: scenario === "setup-mode" ? {} : { sessionID },
+      log: () => {},
+    }).pipe(Effect.provide(Global.layerWith({ state: state.path })), Effect.provide(FileSystem.layerNoop({}))),
+  )
+  const input = () => {
+    const focused = setup.renderer.currentFocusedRenderable
+    if (!(focused instanceof TextareaRenderable)) throw new Error("composer is not focused")
+    return focused
+  }
+  const navigate = async (key: "F6" | "F7") => {
+    const previous = input()
+    setup.mockInput.pressKey(key)
+    await setup.waitForFrame(
+      () =>
+        setup.renderer.currentFocusedRenderable instanceof TextareaRenderable &&
+        setup.renderer.currentFocusedRenderable !== previous,
+    )
+  }
+  try {
+    await ready.promise
+    await setup.waitForFrame((frame) => frame.includes("Demo Model"))
+    await setup.mockInput.typeText("Recover this failed prompt")
+    if (scenario === "history-paste") {
+      await setup.mockInput.pasteBracketedText("first pasted line\nsecond pasted line\nthird pasted line")
+      await setup.waitForFrame(() => input().plainText.includes("[Pasted"))
+      expect(input().extmarks.getAll()).toHaveLength(1)
+    }
+    setup.mockInput.pressEnter()
+    await requested.promise
+    expect(input().plainText).toBe("")
+    expect(setup.captureCharFrame()).toContain("Recover this failed prompt")
+
+    if (scenario === "setup-mode") {
+      await setup.mockInput.typeText("!")
+      await setup.waitForFrame((frame) => frame.includes("Shell"))
+      await setup.mockInput.typeText("Keep my newer draft")
+      setup.mockInput.pressKey("F9")
+      await setup.waitForFrame(() => input().plainText === "")
+      setup.mockInput.pressKey("F8")
+      await setup.waitForFrame(() => input().plainText === "Keep my newer draft")
+      setup.mockInput.pressKey("ESCAPE")
+      await setup.waitForFrame((frame) => !frame.includes("Shell"))
+      release.resolve()
+      await setup.waitForFrame((frame) => frame.includes("Fixture admission rejected"))
+      await setup.waitForFrame(() => input().plainText === "Keep my newer draft")
+      await setup.renderOnce()
+      expect(setup.captureCharFrame()).not.toContain("Shell")
+      return
+    }
+
+    if (scenario === "shell-draft") {
+      await setup.mockInput.typeText("!")
+      await setup.waitForFrame((frame) => frame.includes("Shell"))
+    }
+    if (scenario === "typing" || scenario === "away-draft" || scenario === "cleared" || scenario === "shell-draft")
+      await setup.mockInput.typeText("Keep my newer draft")
+    if (scenario === "history-paste") {
+      setup.mockInput.pressArrow("up")
+      await setup.waitForFrame(() => input().plainText.includes("[Pasted"))
+    }
+    if (scenario === "multiple") {
+      await setup.mockInput.typeText("Another failed prompt")
+      setup.mockInput.pressEnter()
+      await setup.waitForFrame((frame) => frame.includes("Another failed prompt") && input().plainText === "")
+    }
+    if (scenario === "away" || scenario === "returned" || scenario === "away-draft") await navigate("F6")
+    if (scenario === "returned") await navigate("F7")
+    release.resolve()
+    await setup.waitForFrame((frame) => frame.includes("Fixture admission rejected"))
+    if (scenario === "away" || scenario === "away-draft") {
+      expect(input().plainText).toBe("")
+      await navigate("F7")
+    }
+
+    if (scenario === "history-paste") {
+      await setup.waitForFrame((frame) => frame.includes("Failed to send:"))
+      input().gotoBufferEnd()
+      setup.mockInput.pressArrow("down")
+      await setup.waitForFrame((frame) => input().plainText.includes("[Pasted") && !frame.includes("Failed to send:"))
+      expect(input().extmarks.getAll()).toHaveLength(1)
+      return
+    }
+
+    if (scenario === "multiple") {
+      await setup.waitForFrame((frame) => frame.includes("Failed to send: Another failed prompt"))
+      expect(input().plainText).toBe("Recover this failed prompt")
+      input().clear()
+      await setup.waitForFrame(
+        (frame) => input().plainText === "Another failed prompt" && !frame.includes("Failed to send:"),
+      )
+      return
+    }
+
+    if (scenario === "cleared") {
+      await setup.waitForFrame((frame) => frame.includes("Failed to send:"))
+      expect(input().plainText).toBe("Keep my newer draft")
+      input().clear()
+    }
+
+    if (scenario === "typing" || scenario === "away-draft" || scenario === "shell-draft") {
+      const failed = await setup.waitForFrame((frame) => frame.includes("Failed to send:") && frame.includes("restore"))
+      expect(failed).toContain("Recover this failed prompt")
+      expect(input().plainText).toBe("Keep my newer draft")
+      if (scenario === "away-draft") {
+        setup.mockInput.pressKey("p", { ctrl: true })
+        await setup.waitForFrame(() => setup.renderer.currentFocusedRenderable instanceof InputRenderable)
+        await setup.mockInput.typeText("Restore failed prompt")
+        await setup.renderOnce()
+        setup.mockInput.pressEnter()
+      } else {
+        const lines = failed.split("\n")
+        const row = lines.findIndex((line) => line.includes("restore"))
+        await setup.mockMouse.click(lines[row].indexOf("restore"), row)
+      }
+      await setup.waitForFrame(
+        (frame) => input().plainText === "Recover this failed prompt" && !frame.includes("Failed to send:"),
+      )
+
+      // Restoring explicitly stashes the newer draft rather than discarding it.
+      input().clear()
+      setup.mockInput.pressKey("F8")
+      await setup.waitForFrame(() => input().plainText === "Keep my newer draft")
+      if (scenario === "shell-draft") await setup.waitForFrame((frame) => frame.includes("Shell"))
+    } else {
+      await setup.waitForFrame(() => input().plainText === "Recover this failed prompt")
+    }
+  } finally {
+    release.resolve()
+    setup.renderer.destroy()
+    await task
+    await server.stop()
+    takeDraft(sessionID)
+    if (scenario === "setup-mode") takeDraft(undefined)
+  }
+})

packages/tui/test/prompt/draft-stash.test.ts

@@ -1,5 +1,11 @@
 import { describe, expect, test } from "bun:test"
-import { saveDraft, takeDraft } from "../../src/component/prompt/draft-stash"
+import {
+  failedDrafts,
+  saveDraft,
+  saveFailedDraft,
+  takeDraft,
+  takeFailedDraft,
+} from "../../src/component/prompt/draft-stash"
 import { emptyPrompt } from "../../src/prompt/history"
 
 // The Prompt component stashes an unsent draft in onCleanup and takes it back
@@ -39,4 +45,36 @@ describe("prompt draft stash", () => {
     saveDraft("ses_a", second)
     expect(takeDraft("ses_a")).toBe(second)
   })
+
+  test("failed drafts retain independent snapshots, metadata, and recovery order per tab", () => {
+    const first = {
+      prompt: {
+        ...emptyPrompt(),
+        text: "First",
+        mode: "shell" as const,
+        files: [{ uri: "file:///note.txt", name: "note.txt" }],
+      },
+      cursor: 2,
+    }
+    saveFailedDraft("ses_failed_one", first)
+    saveFailedDraft("ses_failed_one", draft("Second"))
+    saveFailedDraft("ses_failed_two", draft("Other tab"))
+    first.prompt.text = "Changed"
+    first.prompt.files[0].name = "changed.txt"
+
+    expect(failedDrafts("ses_failed_one")).toHaveLength(2)
+    expect(takeFailedDraft("ses_failed_one")).toEqual({
+      prompt: {
+        ...emptyPrompt(),
+        text: "First",
+        mode: "shell",
+        files: [{ uri: "file:///note.txt", name: "note.txt" }],
+      },
+      cursor: 2,
+    })
+    expect(takeFailedDraft("ses_failed_one")).toEqual(draft("Second"))
+    expect(failedDrafts("ses_failed_one")).toEqual([])
+    expect(takeFailedDraft("ses_failed_one")).toBeUndefined()
+    expect(takeFailedDraft("ses_failed_two")).toEqual(draft("Other tab"))
+  })
 })