fix(app): preserve mentions during paste (#37940)

Co-authored-by: Brendan Allan <14191578+Brendonovich@users.noreply.github.com>

packages/session-ui/src/v2/components/prompt-input/interaction.ts

@@ -82,9 +82,7 @@ export function createPromptInputV2Controller(input: {
       draft.addMention(part)
       return true
     }
-    const text = draft.state.prompt.map((item) => ("content" in item ? item.content : "")).join("")
-    const cursor = draft.state.cursor ?? text.length
-    draft.setText(text.slice(0, cursor) + part.content + text.slice(cursor))
+    draft.addText(part.content)
     return true
   }
   const attachments = input.attachments

packages/session-ui/src/v2/components/prompt-input/store.test.ts

@@ -56,6 +56,28 @@ describe("prompt input v2 store", () => {
     expect(prompt.state.cursor).toBe(7)
   })
 
+  test("inserts text without flattening structured mentions", () => {
+    const [state, setState] = createStore<PromptInputV2PersistedState>({
+      prompt: [
+        { type: "text", content: "A ", start: 0, end: 2 },
+        { type: "file", path: "one", content: "@one", start: 2, end: 6 },
+        { type: "text", content: " B", start: 6, end: 8 },
+      ],
+      cursor: 2,
+      context: { items: [] },
+    })
+    const prompt = createPromptInputV2Store([state, setState])
+
+    prompt.addText("X\nY")
+
+    expect(prompt.state.prompt).toEqual([
+      { type: "text", content: "A X\nY", start: 0, end: 5 },
+      { type: "file", path: "one", content: "@one", start: 5, end: 9 },
+      { type: "text", content: " B", start: 9, end: 11 },
+    ])
+    expect(prompt.state.cursor).toBe(5)
+  })
+
   test("mutates context, attachments, and model through shared actions", () => {
     const prompt = createPromptStore()
     const context = { key: "file:src/index.ts", type: "file" as const, path: "src/index.ts" }

packages/session-ui/src/v2/components/prompt-input/store.ts

@@ -47,6 +47,13 @@ export function createPromptInputV2Store(input: PromptInputV2StoreInput) {
         setStore()("cursor", content.length)
       })
     },
+    addText(content: string) {
+      const cursor = store().cursor ?? promptLength(store().prompt)
+      batch(() => {
+        setStore()("prompt", (prompt) => insertText(prompt, cursor, content))
+        setStore()("cursor", cursor + content.length)
+      })
+    },
     reset() {
       batch(() => {
         setStore()("prompt", [{ type: "text", content: "", start: 0, end: 0 }])
@@ -86,6 +93,27 @@ export function createPromptInputV2Store(input: PromptInputV2StoreInput) {
 
 export type PromptInputV2Store = ReturnType<typeof createPromptInputV2Store>
 
+function insertText(prompt: PromptInputV2Prompt, cursor: number, content: string): PromptInputV2Prompt {
+  let position = 0
+  let inserted = false
+  const parts = prompt.flatMap<PromptInputV2Prompt[number]>((part) => {
+    if (part.type === "image") return [part]
+    const start = position
+    position += part.content.length
+    if (inserted) return [part]
+    if (part.type === "text" && cursor >= start && cursor <= position) {
+      inserted = true
+      const offset = cursor - start
+      return [{ ...part, content: part.content.slice(0, offset) + content + part.content.slice(offset) }]
+    }
+    if (cursor > start) return [part]
+    inserted = true
+    return [{ type: "text", content, start: 0, end: 0 }, part]
+  })
+  if (!inserted) parts.push({ type: "text", content, start: 0, end: 0 })
+  return withOffsets(parts)
+}
+
 function insertMention(
   prompt: PromptInputV2Prompt,
   start: number,
@@ -106,11 +134,19 @@ function insertMention(
       { type: "text" as const, content: ` ${after}`, start: 0, end: 0 },
     ]
   })
+  return withOffsets(parts)
+}
+
+function withOffsets(prompt: PromptInputV2Prompt): PromptInputV2Prompt {
   let offset = 0
-  return parts.map((part) => {
+  return prompt.map((part) => {
     if (part.type === "image") return part
     const next = { ...part, start: offset, end: offset + part.content.length }
     offset = next.end
     return next
   })
 }
+
+function promptLength(prompt: PromptInputV2Prompt) {
+  return prompt.reduce((length, part) => length + ("content" in part ? part.content.length : 0), 0)
+}