fix(app): store image attachments

3 files changed, 192 insertions, 60 deletions

diff --git a/packages/app/src/components/prompt-input.tsx b/packages/app/src/components/prompt-input.tsx
index d269976f3..fc4a3d1e6 100644
--- a/packages/app/src/components/prompt-input.tsx
+++ b/packages/app/src/components/prompt-input.tsx
@@ -165,6 +165,9 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       },
   )
   const working = createMemo(() => status()?.type !== "idle")
+  const imageAttachments = createMemo(
+    () => prompt.current().filter((part) => part.type === "image") as ImageAttachmentPart[],
+  )
 
   const [store, setStore] = createStore<{
     popover: "at" | "slash" | null
@@ -172,7 +175,6 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     savedPrompt: Prompt | null
     placeholder: number
     dragging: boolean
-    imageAttachments: ImageAttachmentPart[]
     mode: "normal" | "shell"
     applyingHistory: boolean
   }>({
@@ -181,7 +183,6 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     savedPrompt: null,
     placeholder: Math.floor(Math.random() * PLACEHOLDERS.length),
     dragging: false,
-    imageAttachments: [],
     mode: "normal",
     applyingHistory: false,
   })
@@ -274,21 +275,16 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
         mime: file.type,
         dataUrl,
       }
-      setStore(
-        produce((draft) => {
-          draft.imageAttachments.push(attachment)
-        }),
-      )
+      const cursorPosition = prompt.cursor() ?? getCursorPosition(editorRef)
+      prompt.set([...prompt.current(), attachment], cursorPosition)
     }
     reader.readAsDataURL(file)
   }
 
   const removeImageAttachment = (id: string) => {
-    setStore(
-      produce((draft) => {
-        draft.imageAttachments = draft.imageAttachments.filter((a) => a.id !== id)
-      }),
-    )
+    const current = prompt.current()
+    const next = current.filter((part) => part.type !== "image" || part.id !== id)
+    prompt.set(next, prompt.cursor())
   }
 
   const handlePaste = async (event: ClipboardEvent) => {
@@ -538,8 +534,9 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
     on(
       () => prompt.current(),
       (currentParts) => {
+        const inputParts = currentParts.filter((part) => part.type !== "image") as Prompt
         const domParts = parseFromDOM()
-        if (isNormalizedEditor() && isPromptEqual(currentParts, domParts)) return
+        if (isNormalizedEditor() && isPromptEqual(inputParts, domParts)) return
 
         const selection = window.getSelection()
         let cursorPosition: number | null = null
@@ -547,7 +544,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
           cursorPosition = getCursorPosition(editorRef)
         }
 
-        renderEditor(currentParts)
+        renderEditor(inputParts)
 
         if (cursorPosition !== null) {
           setCursorPosition(editorRef, cursorPosition)
@@ -638,11 +635,12 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
 
   const handleInput = () => {
     const rawParts = parseFromDOM()
+    const images = imageAttachments()
     const cursorPosition = getCursorPosition(editorRef)
     const rawText = rawParts.map((p) => ("content" in p ? p.content : "")).join("")
     const trimmed = rawText.replace(/\u200B/g, "").trim()
     const hasNonText = rawParts.some((part) => part.type !== "text")
-    const shouldReset = trimmed.length === 0 && !hasNonText
+    const shouldReset = trimmed.length === 0 && !hasNonText && images.length === 0
 
     if (shouldReset) {
       setStore("popover", null)
@@ -681,7 +679,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       setStore("savedPrompt", null)
     }
 
-    prompt.set(rawParts, cursorPosition)
+    prompt.set([...rawParts, ...images], cursorPosition)
     queueScroll()
   }
 
@@ -784,16 +782,14 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
       .map((p) => ("content" in p ? p.content : ""))
       .join("")
       .trim()
-    if (!text) return
+    const hasImages = prompt.some((part) => part.type === "image")
+    if (!text && !hasImages) return
 
     const entry = clonePromptParts(prompt)
     const currentHistory = mode === "shell" ? shellHistory : history
     const setCurrentHistory = mode === "shell" ? setShellHistory : setHistory
     const lastEntry = currentHistory.entries[0]
-    if (lastEntry) {
-      const lastText = lastEntry.map((p) => ("content" in p ? p.content : "")).join("")
-      if (lastText === text) return
-    }
+    if (lastEntry && isPromptEqual(lastEntry, entry)) return
 
     setCurrentHistory("entries", (entries) => [entry, ...entries].slice(0, MAX_HISTORY))
   }
@@ -967,7 +963,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
 
     const currentPrompt = prompt.current()
     const text = currentPrompt.map((part) => ("content" in part ? part.content : "")).join("")
-    const images = store.imageAttachments.slice()
+    const images = imageAttachments().slice()
     const mode = store.mode
 
     if (text.trim().length === 0 && images.length === 0) {
@@ -1061,14 +1057,12 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
 
     const clearInput = () => {
       prompt.reset()
-      setStore("imageAttachments", [])
       setStore("mode", "normal")
       setStore("popover", null)
     }
 
     const restoreInput = () => {
       prompt.set(currentPrompt, promptLength(currentPrompt))
-      setStore("imageAttachments", images)
       setStore("mode", mode)
       setStore("popover", null)
       requestAnimationFrame(() => {
@@ -1471,9 +1465,9 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
             </For>
           </div>
         </Show>
-        <Show when={store.imageAttachments.length > 0}>
+        <Show when={imageAttachments().length > 0}>
           <div class="flex flex-wrap gap-2 px-3 pt-3">
-            <For each={store.imageAttachments}>
+            <For each={imageAttachments()}>
               {(attachment) => (
                 <div class="relative group">
                   <Show
@@ -1525,7 +1519,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
               "font-mono!": store.mode === "shell",
             }}
           />
-          <Show when={!prompt.dirty() && store.imageAttachments.length === 0}>
+          <Show when={!prompt.dirty()}>
             <div class="absolute top-0 inset-x-0 px-5 py-3 pr-12 text-14-regular text-text-weak pointer-events-none whitespace-nowrap truncate">
               {store.mode === "shell"
                 ? "Enter shell command..."
@@ -1658,7 +1652,7 @@ export const PromptInput: Component<PromptInputProps> = (props) => {
             >
               <IconButton
                 type="submit"
-                disabled={!prompt.dirty() && store.imageAttachments.length === 0 && !working()}
+                disabled={!prompt.dirty() && !working()}
                 icon={working() ? "stop" : "arrow-up"}
                 variant="primary"
                 class="h-6 w-4.5"
diff --git a/packages/app/src/pages/layout.tsx b/packages/app/src/pages/layout.tsx
index aabaa1237..932026794 100644
--- a/packages/app/src/pages/layout.tsx
+++ b/packages/app/src/pages/layout.tsx
@@ -644,13 +644,13 @@ export default function Layout(props: ParentProps) {
     const updated = createMemo(() => DateTime.fromMillis(props.session.time.updated))
     const notifications = createMemo(() => notification.session.unseen(props.session.id))
     const hasError = createMemo(() => notifications().some((n) => n.type === "error"))
+    const [sessionStore] = globalSync.child(props.session.directory)
     const hasPermissions = createMemo(() => {
-      const store = globalSync.child(props.project.worktree)[0]
-      const permissions = store.permission?.[props.session.id] ?? []
+      const permissions = sessionStore.permission?.[props.session.id] ?? []
       if (permissions.length > 0) return true
-      const childSessions = store.session.filter((s) => s.parentID === props.session.id)
+      const childSessions = sessionStore.session.filter((s) => s.parentID === props.session.id)
       for (const child of childSessions) {
-        const childPermissions = store.permission?.[child.id] ?? []
+        const childPermissions = sessionStore.permission?.[child.id] ?? []
         if (childPermissions.length > 0) return true
       }
       return false
@@ -658,7 +658,7 @@ export default function Layout(props: ParentProps) {
     const isWorking = createMemo(() => {
       if (props.session.id === params.id) return false
       if (hasPermissions()) return false
-      const status = globalSync.child(props.project.worktree)[0].session_status[props.session.id]
+      const status = sessionStore.session_status[props.session.id]
       return status?.type === "busy" || status?.type === "retry"
     })
     return (
diff --git a/packages/app/src/utils/prompt.ts b/packages/app/src/utils/prompt.ts
index 45c5ce1f3..29a774c2a 100644
--- a/packages/app/src/utils/prompt.ts
+++ b/packages/app/src/utils/prompt.ts
@@ -1,47 +1,185 @@
-import type { Part, TextPart, FilePart } from "@opencode-ai/sdk/v2"
-import type { Prompt, FileAttachmentPart } from "@/context/prompt"
+import type { AgentPart as MessageAgentPart, FilePart, Part, TextPart } from "@opencode-ai/sdk/v2"
+import type { AgentPart, FileAttachmentPart, ImageAttachmentPart, Prompt } from "@/context/prompt"
+
+type Inline =
+  | {
+      type: "file"
+      start: number
+      end: number
+      value: string
+      path: string
+      selection?: {
+        startLine: number
+        endLine: number
+        startChar: number
+        endChar: number
+      }
+    }
+  | {
+      type: "agent"
+      start: number
+      end: number
+      value: string
+      name: string
+    }
+
+function selectionFromFileUrl(url: string): Extract<Inline, { type: "file" }>["selection"] {
+  const queryIndex = url.indexOf("?")
+  if (queryIndex === -1) return undefined
+  const params = new URLSearchParams(url.slice(queryIndex + 1))
+  const startLine = Number(params.get("start"))
+  const endLine = Number(params.get("end"))
+  if (!Number.isFinite(startLine) || !Number.isFinite(endLine)) return undefined
+  return {
+    startLine,
+    endLine,
+    startChar: 0,
+    endChar: 0,
+  }
+}
+
+function textPartValue(parts: Part[]) {
+  const candidates = parts
+    .filter((part): part is TextPart => part.type === "text")
+    .filter((part) => !part.synthetic && !part.ignored)
+  return candidates.reduce((best: TextPart | undefined, part) => {
+    if (!best) return part
+    if (part.text.length > best.text.length) return part
+    return best
+  }, undefined)
+}
 
 /**
  * Extract prompt content from message parts for restoring into the prompt input.
  * This is used by undo to restore the original user prompt.
  */
 export function extractPromptFromParts(parts: Part[]): Prompt {
-  const result: Prompt = []
-  let position = 0
+  const textPart = textPartValue(parts)
+  const text = textPart?.text ?? ""
+
+  const inline: Inline[] = []
+  const images: ImageAttachmentPart[] = []
 
   for (const part of parts) {
-    if (part.type === "text") {
-      const textPart = part as TextPart
-      if (!textPart.synthetic && textPart.text) {
-        result.push({
-          type: "text",
-          content: textPart.text,
-          start: position,
-          end: position + textPart.text.length,
-        })
-        position += textPart.text.length
-      }
-    } else if (part.type === "file") {
+    if (part.type === "file") {
       const filePart = part as FilePart
-      if (filePart.source?.type === "file") {
-        const path = filePart.source.path
-        const content = "@" + path
-        const attachment: FileAttachmentPart = {
+      const sourceText = filePart.source?.text
+      if (sourceText) {
+        const value = sourceText.value
+        const start = sourceText.start
+        const end = sourceText.end
+        let path = value
+        if (value.startsWith("@")) path = value.slice(1)
+        if (!value.startsWith("@") && filePart.source && "path" in filePart.source) {
+          path = filePart.source.path
+        }
+        inline.push({
           type: "file",
+          start,
+          end,
+          value,
           path,
-          content,
-          start: position,
-          end: position + content.length,
-        }
-        result.push(attachment)
-        position += content.length
+          selection: selectionFromFileUrl(filePart.url),
+        })
+        continue
       }
+
+      if (filePart.url.startsWith("data:")) {
+        images.push({
+          type: "image",
+          id: filePart.id,
+          filename: filePart.filename ?? "attachment",
+          mime: filePart.mime,
+          dataUrl: filePart.url,
+        })
+      }
+    }
+
+    if (part.type === "agent") {
+      const agentPart = part as MessageAgentPart
+      const source = agentPart.source
+      if (!source) continue
+      inline.push({
+        type: "agent",
+        start: source.start,
+        end: source.end,
+        value: source.value,
+        name: agentPart.name,
+      })
     }
   }
 
+  inline.sort((a, b) => {
+    if (a.start !== b.start) return a.start - b.start
+    return a.end - b.end
+  })
+
+  const result: Prompt = []
+  let position = 0
+  let cursor = 0
+
+  const pushText = (content: string) => {
+    if (!content) return
+    result.push({
+      type: "text",
+      content,
+      start: position,
+      end: position + content.length,
+    })
+    position += content.length
+  }
+
+  const pushFile = (item: Extract<Inline, { type: "file" }>) => {
+    const content = item.value
+    const attachment: FileAttachmentPart = {
+      type: "file",
+      path: item.path,
+      content,
+      start: position,
+      end: position + content.length,
+      selection: item.selection,
+    }
+    result.push(attachment)
+    position += content.length
+  }
+
+  const pushAgent = (item: Extract<Inline, { type: "agent" }>) => {
+    const content = item.value
+    const mention: AgentPart = {
+      type: "agent",
+      name: item.name,
+      content,
+      start: position,
+      end: position + content.length,
+    }
+    result.push(mention)
+    position += content.length
+  }
+
+  for (const item of inline) {
+    if (item.start < 0 || item.end < item.start) continue
+    if (item.end > text.length) continue
+    if (item.start < cursor) continue
+
+    pushText(text.slice(cursor, item.start))
+
+    if (item.type === "file") {
+      pushFile(item)
+    }
+
+    if (item.type === "agent") {
+      pushAgent(item)
+    }
+
+    cursor = item.end
+  }
+
+  pushText(text.slice(cursor))
+
   if (result.length === 0) {
     result.push({ type: "text", content: "", start: 0, end: 0 })
   }
 
-  return result
+  if (images.length === 0) return result
+  return [...result, ...images]
 }