fix(app): preserve prompt input editing

packages/app/e2e/regression/prompt-input-v2-editing.spec.ts

@@ -0,0 +1,50 @@
+import { expect, test } from "@playwright/test"
+import { base64Encode } from "@opencode-ai/core/util/encode"
+import { mockOpenCodeServer } from "../utils/mock-server"
+import { expectAppVisible } from "../utils/waits"
+
+const directory = "C:/OpenCode/PromptInputV2Editing"
+const projectID = "proj_prompt_input_v2_editing"
+const sessionID = "ses_prompt_input_v2_editing"
+
+test("preserves the draft when a populated command menu triggers a built-in", async ({ page }) => {
+  await mockOpenCodeServer(page, {
+    directory,
+    project: {
+      id: projectID,
+      worktree: directory,
+      vcs: "git",
+      name: "prompt-input-v2-editing",
+      time: { created: 1700000000000, updated: 1700000000000 },
+      sandboxes: [],
+    },
+    provider: { all: [], connected: [], default: {} },
+    sessions: [
+      {
+        id: sessionID,
+        slug: "prompt-input-v2-editing",
+        projectID,
+        directory,
+        title: "Prompt input V2 editing",
+        version: "dev",
+        time: { created: 1700000000000, updated: 1700000000000 },
+      },
+    ],
+    pageMessages: () => ({ items: [] }),
+  })
+  await page.addInitScript(() => {
+    localStorage.setItem("settings.v3", JSON.stringify({ general: { newLayoutDesigns: true } }))
+  })
+
+  await page.goto(`/${base64Encode(directory)}/session/${sessionID}`)
+  const composer = page.locator('[data-component="prompt-input-v2"]')
+  const input = composer.locator('[data-component="prompt-input"]')
+  await expectAppVisible(composer)
+
+  await input.fill("keep me")
+  await composer.getByRole("button", { name: "Add images and files" }).click()
+  await page.getByRole("menuitem", { name: "Commands" }).click()
+  await page.locator('[data-suggestion-id="model.choose"]').click()
+
+  await expect(input).toHaveText("keep me")
+})

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
@@ -163,6 +161,10 @@ export function createPromptInputV2Controller(input: {
 
   function dispatch(event: PromptInputV2InteractionEvent) {
     const mode = state.mode
+    const menuDraft =
+      event.type === "popover.select" && state.popover.type === "command-menu"
+        ? { prompt: clonePrompt(draft.state.prompt), cursor: draft.state.cursor }
+        : undefined
     const result = transitionPromptInputV2(state, event, draft.state)
     setState(reconcile(result.state))
     result.commands.forEach(execute)
@@ -174,10 +176,13 @@ export function createPromptInputV2Controller(input: {
       const action = input.onSuggestionSelect?.(event.item)
       if (!action) return result.handled
       if (event.item.kind === "command") {
-        draft.setPrompt(
-          draft.state.prompt.filter((part): part is PromptInputV2Attachment => part.type === "image"),
-          0,
-        )
+        if (menuDraft) draft.setPrompt(menuDraft.prompt, menuDraft.cursor)
+        if (!menuDraft) {
+          draft.setPrompt(
+            draft.state.prompt.filter((part): part is PromptInputV2Attachment => part.type === "image"),
+            0,
+          )
+        }
       }
       action()
     }

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

@@ -26,6 +26,20 @@ describe("prompt input v2 interaction machine", () => {
     expect(closed.state.popover).toEqual({ type: "closed" })
   })
 
+  test("opens context completion at the cursor", () => {
+    const value = "alpha @sr omega"
+    const input = persisted(value)
+    input.cursor = 9
+
+    const result = transitionPromptInputV2(
+      createPromptInputV2InteractionState(),
+      { type: "input.changed", value, persist: false },
+      input,
+    )
+
+    expect(result.state.popover).toEqual({ type: "context", query: "sr" })
+  })
+
   test("enters shell mode from an initial exclamation mark", () => {
     const result = transitionPromptInputV2(
       createPromptInputV2InteractionState(),

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

@@ -61,7 +61,7 @@ export function transitionPromptInputV2(
   event: PromptInputV2InteractionEvent,
   persisted: PromptInputV2PersistedState,
 ): PromptInputV2Transition {
-  if (event.type === "input.changed") return inputChanged(state, event.value, event.persist !== false)
+  if (event.type === "input.changed") return inputChanged(state, event.value, event.persist !== false, persisted.cursor)
   if (event.type === "commands.open") return openCommands(state, persisted)
   if (event.type === "context.open") return openContext(state, persisted)
   if (event.type === "popover.query") return queryChanged(state, event.value)
@@ -81,14 +81,19 @@ export function transitionPromptInputV2(
   return changed({ ...state, focus: "external" })
 }
 
-function inputChanged(state: PromptInputV2InteractionState, value: string, persist: boolean): PromptInputV2Transition {
+function inputChanged(
+  state: PromptInputV2InteractionState,
+  value: string,
+  persist: boolean,
+  cursor: number | undefined,
+): PromptInputV2Transition {
   const setText: PromptInputV2InteractionCommand[] = persist ? [{ type: "draft.setText", value }] : []
   if (state.mode === "normal" && value === "!") {
     return changed({ ...state, mode: "shell", popover: { type: "closed" }, focus: "editor" }, [
       { type: "draft.setText", value: "" },
     ])
   }
-  const context = value.match(/(?:^|\s)@([^\s@]*)$/)
+  const context = value.slice(0, cursor ?? value.length).match(/(?:^|\s)@([^\s@]*)$/)
   if (context) {
     const query = context[1] ?? ""
     return changed({ ...state, popover: { type: "context", query }, focus: "editor" }, [

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)
+}