fix(desktop): prompt input not clearing, attachments flaky

2 files changed, 61 insertions, 165 deletions

diff --git a/packages/desktop/src/components/prompt-input.tsx b/packages/desktop/src/components/prompt-input.tsx
index 5ae56f827..063822148 100644
--- a/packages/desktop/src/components/prompt-input.tsx
+++ b/packages/desktop/src/components/prompt-input.tsx
@@ -64,7 +64,6 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
   const handleFileSelect = (path: string | undefined) => {
     if (!path) return
     addPart({ type: "file", path, content: "@" + getFilename(path), start: 0, end: 0 })
-    setStore("popoverIsOpen", false)
   }
 
   const { flat, active, onInput, onKeyDown, refetch } = useFilteredList<string>({
@@ -114,16 +113,6 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     ),
   )
 
-  createEffect(
-    on(
-      () => session.prompt.cursor(),
-      (cursor) => {
-        if (cursor === undefined) return
-        queueMicrotask(() => setCursorPosition(editorRef, cursor))
-      },
-    ),
-  )
-
   const parseFromDOM = (): Prompt => {
     const newParts: Prompt = []
     let position = 0
@@ -173,118 +162,70 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
   }
 
   const addPart = (part: ContentPart) => {
+    const selection = window.getSelection()
+    if (!selection || selection.rangeCount === 0) return
+
     const cursorPosition = getCursorPosition(editorRef)
     const prompt = session.prompt.current()
     const rawText = prompt.map((p) => p.content).join("")
     const textBeforeCursor = rawText.substring(0, cursorPosition)
     const atMatch = textBeforeCursor.match(/@(\S*)$/)
 
-    const startIndex = atMatch ? atMatch.index! : cursorPosition
-    const endIndex = atMatch ? cursorPosition : cursorPosition
-
-    const pushText = (acc: { parts: ContentPart[]; runningIndex: number }, value: string) => {
-      if (!value) return
-      const last = acc.parts[acc.parts.length - 1]
-      if (last && last.type === "text") {
-        acc.parts[acc.parts.length - 1] = {
-          type: "text",
-          content: last.content + value,
-          start: last.start,
-          end: last.end + value.length,
+    if (part.type === "file") {
+      const pill = document.createElement("span")
+      pill.textContent = part.content
+      pill.setAttribute("data-type", "file")
+      pill.setAttribute("data-path", part.path)
+      pill.setAttribute("contenteditable", "false")
+      pill.style.userSelect = "text"
+      pill.style.cursor = "default"
+
+      const gap = document.createTextNode(" ")
+      const range = selection.getRangeAt(0)
+
+      if (atMatch) {
+        let node: Node | null = range.startContainer
+        let offset = range.startOffset
+        let runningLength = 0
+
+        const walker = document.createTreeWalker(editorRef, NodeFilter.SHOW_TEXT, null)
+        let currentNode = walker.nextNode()
+        while (currentNode) {
+          const textContent = currentNode.textContent || ""
+          if (runningLength + textContent.length >= atMatch.index!) {
+            const localStart = atMatch.index! - runningLength
+            const localEnd = cursorPosition - runningLength
+            if (currentNode === range.startContainer || runningLength + textContent.length >= cursorPosition) {
+              range.setStart(currentNode, localStart)
+              range.setEnd(currentNode, Math.min(localEnd, textContent.length))
+              break
+            }
+          }
+          runningLength += textContent.length
+          currentNode = walker.nextNode()
         }
-        return
       }
-      acc.parts.push({ type: "text", content: value, start: acc.runningIndex, end: acc.runningIndex + value.length })
-    }
-
-    const {
-      parts: nextParts,
-      inserted,
-      cursorPositionAfter,
-    } = prompt.reduce(
-      (acc, item) => {
-        if (acc.inserted) {
-          acc.parts.push({ ...item, start: acc.runningIndex, end: acc.runningIndex + item.content.length })
-          acc.runningIndex += item.content.length
-          return acc
-        }
-
-        const nextIndex = acc.runningIndex + item.content.length
-        if (nextIndex <= startIndex) {
-          acc.parts.push({ ...item, start: acc.runningIndex, end: acc.runningIndex + item.content.length })
-          acc.runningIndex = nextIndex
-          return acc
-        }
-
-        if (item.type !== "text") {
-          acc.parts.push({ ...item, start: acc.runningIndex, end: acc.runningIndex + item.content.length })
-          acc.runningIndex = nextIndex
-          return acc
-        }
 
-        const headLength = Math.max(0, startIndex - acc.runningIndex)
-        const tailLength = Math.max(0, endIndex - acc.runningIndex)
-        const head = item.content.slice(0, headLength)
-        const tail = item.content.slice(tailLength)
-
-        pushText(acc, head)
-        acc.runningIndex += head.length
-
-        if (part.type === "text") {
-          pushText(acc, part.content)
-          acc.runningIndex += part.content.length
-        }
-        if (part.type !== "text") {
-          acc.parts.push({ ...part, start: acc.runningIndex, end: acc.runningIndex + part.content.length })
-          acc.runningIndex += part.content.length
-        }
-
-        const needsGap = Boolean(atMatch)
-        const rest = needsGap ? (tail ? (/^\s/.test(tail) ? tail : ` ${tail}`) : " ") : tail
-        pushText(acc, rest)
-        acc.runningIndex += rest.length
-
-        const baseCursor = startIndex + part.content.length
-        const cursorAddition = needsGap && rest.length > 0 ? 1 : 0
-        acc.cursorPositionAfter = baseCursor + cursorAddition
-
-        acc.inserted = true
-        return acc
-      },
-      {
-        parts: [] as ContentPart[],
-        runningIndex: 0,
-        inserted: false,
-        cursorPositionAfter: cursorPosition + part.content.length,
-      },
-    )
-
-    if (!inserted) {
-      const baseParts = prompt.filter((item) => !(item.type === "text" && item.content === ""))
-      const runningIndex = baseParts.reduce((sum, p) => sum + p.content.length, 0)
-      const appendedAcc = { parts: [...baseParts] as ContentPart[], runningIndex }
-      if (part.type === "text") {
-        pushText(appendedAcc, part.content)
-      }
-      if (part.type !== "text") {
-        appendedAcc.parts.push({
-          ...part,
-          start: appendedAcc.runningIndex,
-          end: appendedAcc.runningIndex + part.content.length,
-        })
-      }
-      const next = appendedAcc.parts.length > 0 ? appendedAcc.parts : DEFAULT_PROMPT
-      const nextCursor = rawText.length + part.content.length
-      session.prompt.set(next, nextCursor)
-      setStore("popoverIsOpen", false)
-      queueMicrotask(() => setCursorPosition(editorRef, nextCursor))
-      return
+      range.deleteContents()
+      range.insertNode(gap)
+      range.insertNode(pill)
+      range.setStartAfter(gap)
+      range.collapse(true)
+      selection.removeAllRanges()
+      selection.addRange(range)
+    } else if (part.type === "text") {
+      const textNode = document.createTextNode(part.content)
+      const range = selection.getRangeAt(0)
+      range.deleteContents()
+      range.insertNode(textNode)
+      range.setStartAfter(textNode)
+      range.collapse(true)
+      selection.removeAllRanges()
+      selection.addRange(range)
     }
 
-    session.prompt.set(nextParts, cursorPositionAfter)
+    handleInput()
     setStore("popoverIsOpen", false)
-
-    queueMicrotask(() => setCursorPosition(editorRef, cursorPositionAfter))
   }
 
   const abort = () =>
diff --git a/packages/desktop/src/context/session.tsx b/packages/desktop/src/context/session.tsx
index 02c859a5b..0a517e0d3 100644
--- a/packages/desktop/src/context/session.tsx
+++ b/packages/desktop/src/context/session.tsx
@@ -22,65 +22,20 @@ export const { use: useSession, provider: SessionProvider } = createSimpleContex
           active?: string
           opened: string[]
         }
+        prompt: Prompt
+        cursor?: number
       }>({
         tabs: {
           opened: [],
         },
+        prompt: clonePrompt(DEFAULT_PROMPT),
+        cursor: undefined,
       }),
       {
         name: seed,
       },
     )
 
-    const [promptStore, setPromptStore] = createStore<{
-      prompt: Prompt
-      cursor?: number
-    }>({
-      prompt: clonePrompt(DEFAULT_PROMPT),
-    })
-
-    const key = createMemo(() => props.sessionId ?? "new-session")
-    const [ready, setReady] = createSignal(false)
-    const prefix = "session-prompt:"
-
-    createEffect(
-      on(
-        key,
-        (value) => {
-          setReady(false)
-          const record = localStorage.getItem(prefix + value)
-          if (!record) {
-            setPromptStore("prompt", clonePrompt(DEFAULT_PROMPT))
-            setPromptStore("cursor", undefined)
-            setReady(true)
-            return
-          }
-          const payload = JSON.parse(record) as { prompt?: Prompt; cursor?: number }
-          const parts = payload.prompt ?? DEFAULT_PROMPT
-          const cursor = typeof payload.cursor === "number" ? payload.cursor : undefined
-          setPromptStore("prompt", clonePrompt(parts))
-          setPromptStore("cursor", cursor)
-          setReady(true)
-        },
-        { defer: true },
-      ),
-    )
-
-    createEffect(() => {
-      if (!ready()) return
-      const value = key()
-      const isDefault = isPromptEqual(promptStore.prompt, DEFAULT_PROMPT)
-      if (isDefault && (promptStore.cursor === undefined || promptStore.cursor <= 0)) {
-        localStorage.removeItem(prefix + value)
-        return
-      }
-      const next = JSON.stringify({
-        prompt: clonePrompt(promptStore.prompt),
-        cursor: promptStore.cursor,
-      })
-      localStorage.setItem(prefix + value, next)
-    })
-
     createEffect(() => {
       if (!props.sessionId) return
       sync.session.sync(props.sessionId)
@@ -149,14 +104,14 @@ export const { use: useSession, provider: SessionProvider } = createSimpleContex
       working,
       diffs,
       prompt: {
-        current: createMemo(() => promptStore.prompt),
-        cursor: createMemo(() => promptStore.cursor),
-        dirty: createMemo(() => !isPromptEqual(promptStore.prompt, DEFAULT_PROMPT)),
+        current: createMemo(() => persist.prompt),
+        cursor: createMemo(() => persist.cursor),
+        dirty: createMemo(() => !isPromptEqual(persist.prompt, DEFAULT_PROMPT)),
         set(prompt: Prompt, cursorPosition?: number) {
           const next = clonePrompt(prompt)
           batch(() => {
-            setPromptStore("prompt", next)
-            if (cursorPosition !== undefined) setPromptStore("cursor", cursorPosition)
+            setPersist("prompt", next)
+            if (cursorPosition !== undefined) setPersist("cursor", cursorPosition)
           })
         },
       },