fix(tui): restore the live draft after browsing prompt history (#63262)

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

@@ -211,7 +211,6 @@ export function Prompt(props: PromptProps) {
   const dialog = useDialog()
   const toast = useToast()
   const status = createMemo(() => data.session.status(props.sessionID ?? ""))
-  const history = usePromptHistory()
   const stash = usePromptStash()
   const keymap = Keymap.use()
   const renderer = useRenderer()
@@ -369,6 +368,7 @@ export function Prompt(props: PromptProps) {
     extmarkToPart: new Map(),
     interrupt: 0,
   })
+  const history = usePromptHistory(() => ({ ...unwrap(store.prompt), mode: store.mode }))
   let disposed = false
   let pasteQueue = Promise.resolve()
 
@@ -691,6 +691,7 @@ export function Prompt(props: PromptProps) {
       setStore("mode", mode)
     },
     set(prompt) {
+      history.reset()
       input.setText(prompt.text)
       setStore("prompt", prompt)
       restoreExtmarksFromPrompt(prompt)
@@ -705,6 +706,7 @@ export function Prompt(props: PromptProps) {
   }
 
   function resetComposer() {
+    history.reset()
     input.extmarks.clear()
     setStore("prompt", emptyPrompt())
     setStore("extmarkToPart", new Map())
@@ -900,12 +902,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) ref.set(entry.prompt)
           dialog.clear()
         },
       },
@@ -915,16 +912,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) => ref.set(entry.prompt)} />)
         },
       },
     ].map(

packages/tui/src/prompt/history.tsx

@@ -90,7 +90,7 @@ export function parsePromptInfo(value: unknown): PromptInfo | undefined {
   return input as PromptInfo
 }
 
-export const { use: usePromptHistory, provider: PromptHistoryProvider } = createSimpleContext({
+const { use: useHistoryEntries, provider: PromptHistoryProvider } = createSimpleContext({
   name: "PromptHistory",
   init: () => {
     const paths = useTuiPaths()
@@ -104,26 +104,15 @@ export const { use: usePromptHistory, provider: PromptHistoryProvider } = create
         writeText(historyPath, lines.map((line) => JSON.stringify(line)).join("\n") + "\n").catch(() => {})
     })
 
-    const [store, setStore] = createStore({ index: 0, history: [] as PromptInfo[] })
+    const [store, setStore] = createStore({ history: [] as PromptInfo[] })
 
     return {
-      move(direction: 1 | -1, input: string) {
-        if (!store.history.length) return undefined
-        const current = store.history.at(store.index)
-        if (!current) return undefined
-        if (current.text !== input && input.length) return
-        const next = store.index + direction
-        if (Math.abs(next) > store.history.length || next > 0) return
-        setStore("index", next)
-        if (next === 0) return emptyPrompt()
-        return store.history.at(next)
+      get entries() {
+        return store.history
       },
       append(item: PromptInfo) {
         const entry = structuredClone(unwrap(item))
-        if (isDuplicateEntry(store.history.at(-1), entry)) {
-          setStore("index", 0)
-          return
-        }
+        if (isDuplicateEntry(store.history.at(-1), entry)) return
         let trimmed = false
         setStore(
           produce((draft) => {
@@ -132,7 +121,6 @@ export const { use: usePromptHistory, provider: PromptHistoryProvider } = create
               draft.history = draft.history.slice(-MAX_HISTORY_ENTRIES)
               trimmed = true
             }
-            draft.index = 0
           }),
         )
 
@@ -145,3 +133,36 @@ export const { use: usePromptHistory, provider: PromptHistoryProvider } = create
     }
   },
 })
+
+export { PromptHistoryProvider }
+
+export function usePromptHistory(currentDraft = emptyPrompt) {
+  const history = useHistoryEntries()
+  let index = 0
+  let draft: PromptInfo | undefined
+  function reset() {
+    index = 0
+    draft = undefined
+  }
+  return {
+    move(direction: 1 | -1, input: string) {
+      const next = index + direction
+      if (next < -history.entries.length || next > 0) return
+      if (index < 0 && input && history.entries.at(index)?.text !== input) return
+      if (index === 0) draft = structuredClone({ ...currentDraft(), text: input })
+      index = next
+      if (next !== 0) {
+        const entry = history.entries.at(next)
+        return entry ? { ...emptyPrompt(), ...structuredClone(unwrap(entry)) } : undefined
+      }
+      const restored = draft
+      draft = undefined
+      return restored
+    },
+    reset,
+    append(item: PromptInfo) {
+      reset()
+      history.append(item)
+    },
+  }
+}

packages/tui/test/prompt/history-composer.test.tsx

@@ -0,0 +1,52 @@
+import { expect, test } from "bun:test"
+import path from "node:path"
+import { takeDraft } from "../../src/component/prompt/draft-stash"
+import { createAppFixture } from "../fixture/app"
+import { tmpdir } from "../fixture/fixture"
+
+async function composerWithHistory(state: string) {
+  takeDraft(undefined)
+  const file = path.join(state, "prompt-history.jsonl")
+  // Matching the oldest entry lets Up recall "previous" even before the draft restoration fix.
+  const entries = ["echo pen", "previous"].map((text) => JSON.stringify({ text, mode: "normal", pasted: [] }))
+  await Bun.write(file, entries.join("\n"))
+  const setup = await createAppFixture({ state })
+  await setup.ready
+  await setup.waitForFrame((frame) => frame.includes("commands"))
+  // Loading history rewrites the file with a trailing newline.
+  await setup.waitFor(async () => (await Bun.file(file).text()).endsWith("\n"))
+  return setup
+}
+
+test("down past the newest entry restores the draft text and shell mode", async () => {
+  await using state = await tmpdir()
+  await using setup = await composerWithHistory(state.path)
+  const editor = setup.renderer.currentFocusedEditor!
+
+  await setup.mockInput.typeText("!echo pen")
+  setup.mockInput.pressKey("HOME")
+  setup.mockInput.pressArrow("up")
+  expect(editor.plainText).toBe("previous")
+  await setup.waitForFrame((frame) => frame.includes("previous") && !frame.includes("Shell"))
+
+  setup.mockInput.pressKey("END")
+  setup.mockInput.pressArrow("down")
+  expect(editor.plainText).toBe("echo pen")
+  await setup.waitForFrame((frame) => frame.includes("Shell"))
+})
+
+test("clearing a recalled entry drops the draft so down stays empty", async () => {
+  await using state = await tmpdir()
+  await using setup = await composerWithHistory(state.path)
+  const editor = setup.renderer.currentFocusedEditor!
+
+  await setup.mockInput.typeText("unsent")
+  setup.mockInput.pressKey("HOME")
+  setup.mockInput.pressArrow("up")
+  expect(editor.plainText).toBe("previous")
+
+  setup.mockInput.pressKey("c", { ctrl: true })
+  expect(editor.plainText).toBe("")
+  setup.mockInput.pressArrow("down")
+  expect(editor.plainText).toBe("")
+})

packages/tui/test/prompt/history-provider.test.tsx

@@ -50,6 +50,33 @@ test("keeps legacy unscoped history on the home composer", async () => {
   }
 })
 
+test("reset clears browsing position and live draft, and move returns a deep copy", async () => {
+  await using tmp = await tmpdir()
+  const setup = await renderHistory(tmp.path)
+  try {
+    setup.history.append({
+      text: "[Pasted ~3 lines] ",
+      files: [],
+      agents: [],
+      pasted: [{ text: "a\nb\nc", source: { start: 0, end: 17, text: "[Pasted ~3 lines]" } }],
+    })
+
+    const recalled = setup.history.move(-1, "unsent draft")
+    expect(recalled?.text).toBe("[Pasted ~3 lines] ")
+    recalled!.pasted[0]!.source.start = 99
+    recalled!.pasted.push({ text: "extra", source: { start: 0, end: 5, text: "extra" } })
+
+    setup.history.reset()
+    expect(setup.history.move(1, "")).toBeUndefined()
+
+    const again = setup.history.move(-1, "fresh draft")
+    expect(again?.pasted).toEqual([{ text: "a\nb\nc", source: { start: 0, end: 17, text: "[Pasted ~3 lines]" } }])
+    expect(setup.history.move(1, "[Pasted ~3 lines] ")?.text).toBe("fresh draft")
+  } finally {
+    setup.app.renderer.destroy()
+  }
+})
+
 async function renderHistory(root: string, persisted?: string) {
   const state = path.join(root, "state")
   await mkdir(state, { recursive: true })