fix(tui): restore pasted text when undo brings back its placeholder (#63265)

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

@@ -799,6 +799,7 @@ export function Prompt(props: PromptProps) {
           virtual: true,
           styleId,
           typeId: promptPartTypeId,
+          data: ref.type === "pasted" ? { pasted: prompt.pasted[ref.index].text, label: mention.text } : undefined,
         })
         setStore("extmarkToPart", (map: Map<number, PromptPartRef>) => {
           const newMap = new Map(map)
@@ -821,6 +822,19 @@ export function Prompt(props: PromptProps) {
         const pasted: PromptInfo["pasted"] = []
 
         for (const extmark of allExtmarks) {
+          // Undo restores the extmark after its entry in the live prompt has been removed.
+          if (typeof extmark.data?.pasted === "string" && typeof extmark.data?.label === "string") {
+            if (input.getTextRange(extmark.start, extmark.end) !== extmark.data.label) {
+              input.extmarks.delete(extmark.id)
+              continue
+            }
+            newMap.set(extmark.id, { type: "pasted", index: pasted.length })
+            pasted.push({
+              text: extmark.data.pasted,
+              source: { start: extmark.start, end: extmark.end, text: extmark.data.label },
+            })
+            continue
+          }
           const ref = draft.extmarkToPart.get(extmark.id)
           if (!ref) continue
           if (ref.type === "file") {
@@ -850,15 +864,7 @@ export function Prompt(props: PromptProps) {
             const index = skills.length
             skills.push(part)
             newMap.set(extmark.id, { type: "skill", index })
-            continue
           }
-          const part = draft.prompt.pasted[ref.index]
-          if (!part) continue
-          part.source.start = extmark.start
-          part.source.end = extmark.end
-          const index = pasted.length
-          pasted.push(part)
-          newMap.set(extmark.id, { type: "pasted", index })
         }
 
         const nextFiles = preserveMentionlessPromptAttachments(draft.prompt.files, files)
@@ -1405,6 +1411,7 @@ export function Prompt(props: PromptProps) {
       virtual: true,
       styleId: pasteStyleId,
       typeId: promptPartTypeId,
+      data: { pasted: text, label: virtualText },
     })
 
     setStore(
@@ -1426,7 +1433,6 @@ export function Prompt(props: PromptProps) {
     const part = store.prompt.pasted[ref.index]
     if (!part) return false
 
-    input.extmarks.delete(extmarkId)
     input.setSelection(extmark.start, extmark.end)
     input.insertText(part.text)
     return true

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

@@ -1,6 +1,8 @@
 import { expect, test } from "bun:test"
 import { takeDraft } from "../src/component/prompt/draft-stash"
 import { createAppFixture } from "./fixture/app"
+import { agent, model, session } from "./fixture/local"
+import { directory, json } from "./fixture/tui-client"
 
 test("the prompt keeps pasted control characters and drops escape sequences", async () => {
   // Other app tests can leave a home draft in the process-wide stash.
@@ -14,3 +16,38 @@ test("the prompt keeps pasted control characters and drops escape sequences", as
   await setup.waitForFrame((frame) => frame.includes("xa^Lbc^@d"))
   expect(setup.renderer.currentFocusedEditor?.plainText).toBe("xa\fbc\0d")
 })
+
+test.each(["whole", "fragment"])("undo submits the restored %s paste chip", async (kind) => {
+  const sessionID = `ses_paste_${crypto.randomUUID()}`
+  const prompted = Promise.withResolvers<string>()
+  await using setup = await createAppFixture({
+    args: { sessionID },
+    fetch: async (url, request) => {
+      const location = { directory }
+      if (url.pathname === "/api/agent") return json({ location, data: [agent("build")] })
+      if (url.pathname === "/api/model") return json({ location, data: [model("Paste Model")] })
+      if (url.pathname === `/api/session/${sessionID}`) return json({ data: session(sessionID) })
+      if (url.pathname.endsWith("/message") || url.pathname.endsWith("/inbox")) return json({ data: [], cursor: {} })
+      if (url.pathname.endsWith("/prompt")) {
+        prompted.resolve((await request.json()).text)
+        return json({ data: {} })
+      }
+    },
+  })
+  await setup.ready
+  await setup.waitForFrame((frame) => frame.includes("Paste Model"))
+  const pasted = "one\ntwo\nthree\nfour"
+  const text = () => setup.renderer.currentFocusedEditor?.plainText
+  await setup.mockInput.pasteBracketedText(pasted)
+  await setup.waitFor(() => text() === "[Pasted ~4 lines] ")
+  if (kind === "fragment") {
+    setup.mockInput.pressBackspace() // Remove the space after the chip.
+    setup.mockInput.pressKey("w", { ctrl: true }) // Leave a fragment of the chip.
+  }
+  setup.mockInput.pressKey("u", { ctrl: true })
+  await setup.waitFor(() => text() === "")
+  setup.mockInput.pressKey("z", { super: true })
+  await setup.waitFor(() => text() === (kind === "whole" ? "[Pasted ~4 lines] " : "[Pasted ~4 "))
+  setup.mockInput.pressEnter()
+  expect(await prompted.promise).toBe(kind === "whole" ? `${pasted} ` : "[Pasted ~4 ")
+})