fix(app): store image attachments

1 files changed, 165 insertions, 27 deletions

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]
 }