Browse Source

fix(tui): deduplicate repeated image attachments (#41651)

Kit Langton 5 days ago
parent
commit
a5f7f8d3b5

+ 20 - 1
packages/core/src/session/runner/to-llm-message.ts

@@ -59,6 +59,25 @@ const attachmentContent = (file: FileAttachment): ContentPart[] => {
   return []
 }
 
+const userAttachmentContent = (files: readonly FileAttachment[]) => {
+  const eligible = files.filter(
+    (file) => imageMimes.has(file.mime) && file.source.type === "inline" && file.mention?.text,
+  )
+  if (eligible.length < 2) return files.flatMap(attachmentContent)
+
+  const seen = new Map<string, string[]>()
+  return files.flatMap((file) => {
+    if (!imageMimes.has(file.mime) || file.source.type !== "inline" || !file.mention?.text)
+      return attachmentContent(file)
+    const metadata = JSON.stringify([file.mime, file.name ?? null, file.description ?? null, file.mention.text])
+    const matches = seen.get(metadata)
+    if (matches?.includes(file.data)) return []
+    if (matches) matches.push(file.data)
+    if (!matches) seen.set(metadata, [file.data])
+    return attachmentContent(file)
+  })
+}
+
 const decodeToolInput = Schema.decodeUnknownOption(Schema.UnknownFromJsonString)
 
 const providerMetadata = (
@@ -186,7 +205,7 @@ function toLLMMessage(message: SessionMessage.Info, model: Model.Ref, providerMe
       const content = [
         ...(message.skills ?? []).map((skill) => Message.text(skill.text)),
         ...(message.text === "" ? [] : [Message.text(message.text)]),
-        ...(message.files ?? []).flatMap(attachmentContent),
+        ...userAttachmentContent(message.files ?? []),
       ]
       if (content.length === 0) return []
       return [

+ 97 - 0
packages/core/test/session-runner-message.test.ts

@@ -373,6 +373,103 @@ Recent work
     ])
   })
 
+  test("deduplicates provider media while preserving durable attachment references", () => {
+    const data = Base64.make("AAECAw==")
+    const messages = toLLMMessages(
+      [
+        SessionMessage.User.make({
+          id: id("user-duplicate-image"),
+          type: "user",
+          text: "[Image 1] [Image 1] [Image 2]",
+          files: [
+            FileAttachment.make({
+              data,
+              mime: "image/png",
+              source: { type: "inline" },
+              name: "image.png",
+              mention: { start: 0, end: 9, text: "[Image 1]" },
+            }),
+            FileAttachment.make({
+              data,
+              mime: "image/png",
+              source: { type: "inline" },
+              name: "image.png",
+              mention: { start: 10, end: 19, text: "[Image 1]" },
+            }),
+            FileAttachment.make({
+              data,
+              mime: "image/png",
+              source: { type: "inline" },
+              name: "image.png",
+              description: "alternate use",
+              mention: { start: 20, end: 29, text: "[Image 2]" },
+            }),
+          ],
+          time: { created },
+        }),
+      ],
+      model,
+    )
+
+    expect(messages[0]?.content).toEqual([
+      { type: "text", text: "[Image 1] [Image 1] [Image 2]" },
+      { type: "media", mediaType: "image/png", data, filename: "image.png" },
+      {
+        type: "media",
+        mediaType: "image/png",
+        data,
+        filename: "image.png",
+        metadata: { description: "alternate use" },
+      },
+    ])
+  })
+
+  test("preserves provider media with distinct labels or URI sources", () => {
+    const data = Base64.make("AAECAw==")
+    const messages = toLLMMessages(
+      [
+        SessionMessage.User.make({
+          id: id("user-distinct-images"),
+          type: "user",
+          text: "[Image 1] [Image 2]",
+          files: [
+            FileAttachment.make({
+              data,
+              mime: "image/png",
+              source: { type: "inline" },
+              name: "image.png",
+              mention: { start: 0, end: 9, text: "[Image 1]" },
+            }),
+            FileAttachment.make({
+              data,
+              mime: "image/png",
+              source: { type: "inline" },
+              name: "image.png",
+              mention: { start: 10, end: 19, text: "[Image 2]" },
+            }),
+            FileAttachment.make({
+              data,
+              mime: "image/png",
+              source: { type: "uri", uri: "file:///project/image.png" },
+              name: "image.png",
+              mention: { start: 0, end: 9, text: "[Image 1]" },
+            }),
+            FileAttachment.make({
+              data,
+              mime: "image/png",
+              source: { type: "inline" },
+              name: "image.png",
+            }),
+          ],
+          time: { created },
+        }),
+      ],
+      model,
+    )
+
+    expect(messages[0]?.content.filter((part) => part.type === "media")).toHaveLength(4)
+  })
+
   test("replays durable tool media into canonical tool messages without structured base64", () => {
     const messages = toLLMMessages(
       [

+ 21 - 12
packages/tui/src/component/prompt/index.tsx

@@ -60,6 +60,11 @@ import { Keymap, type KeymapCommand } from "../../context/keymap"
 import { abbreviateHome } from "../../runtime"
 import { PluginSlot } from "../../plugin/render"
 import type { SessionPending } from "@opencode-ai/schema/session-pending"
+import {
+  deduplicatePromptImages,
+  preserveMentionlessPromptAttachments,
+  promptAttachmentLabel,
+} from "../../prompt/attachment"
 import { DialogImagePreview } from "../dialog-image-preview"
 
 export type PromptProps = {
@@ -331,7 +336,7 @@ export function Prompt(props: PromptProps) {
   }
 
   const imageAttachments = createMemo(() =>
-    (store.prompt.files ?? []).filter((file) => typeof file.uri === "string" && file.uri.startsWith("data:image/")),
+    (deduplicatePromptImages(store.prompt.files) ?? []).filter((file) => file.uri.startsWith("data:image/")),
   )
   const imagePreviewHeight = createMemo(() => Math.max(4, Math.min(8, Math.floor(dimensions().height / 4))))
   const imagePreviewWidth = createMemo(() => imagePreviewHeight() * 2)
@@ -736,6 +741,7 @@ export function Prompt(props: PromptProps) {
     setStore(
       produce((draft) => {
         const newMap = new Map<number, PromptPartRef>()
+        const fileExtmarks = new Map<number, NonNullable<PromptInfo["files"]>[number]>()
         const files: NonNullable<PromptInfo["files"]> = []
         const agents: NonNullable<PromptInfo["agents"]> = []
         const skills: NonNullable<PromptInfo["skills"]> = []
@@ -749,9 +755,8 @@ export function Prompt(props: PromptProps) {
             if (!part?.mention) continue
             part.mention.start = extmark.start
             part.mention.end = extmark.end
-            const index = files.length
             files.push(part)
-            newMap.set(extmark.id, { type: "file", index })
+            fileExtmarks.set(extmark.id, part)
             continue
           }
           if (ref.type === "agent") {
@@ -783,8 +788,19 @@ export function Prompt(props: PromptProps) {
           newMap.set(extmark.id, { type: "pasted", index })
         }
 
+        const nextFiles = preserveMentionlessPromptAttachments(draft.prompt.files, files)
+        const fileIndices = new Map(nextFiles.map((file, index) => [file, index]))
+        for (const [extmark, file] of fileExtmarks) {
+          const index = fileIndices.get(file)
+          if (index !== undefined) newMap.set(extmark, { type: "file", index })
+        }
+
         draft.extmarkToPart = newMap
-        draft.prompt.files = files
+        if (
+          nextFiles.length !== draft.prompt.files?.length ||
+          nextFiles.some((file, index) => file !== draft.prompt.files?.[index])
+        )
+          draft.prompt.files = nextFiles
         draft.prompt.agents = agents
         draft.prompt.skills = skills
         draft.prompt.pasted = pasted
@@ -1138,7 +1154,6 @@ export function Prompt(props: PromptProps) {
 
     // Capture mode before it gets reset
     const currentMode = store.mode
-
     if (store.mode === "shell") {
       move.startSubmit()
       void client.api.session.shell({
@@ -1376,13 +1391,7 @@ export function Prompt(props: PromptProps) {
   function pasteAttachment(file: { filename?: string; uri: string }) {
     const currentOffset = input.cursorOffset
     const extmarkStart = currentOffset
-    const pdf = file.uri.startsWith("data:application/pdf;")
-    const count = pdf
-      ? (store.prompt.files?.filter(
-          (attachment) => typeof attachment.uri === "string" && attachment.uri.startsWith("data:application/pdf;"),
-        ).length ?? 0)
-      : imageAttachments().length
-    const virtualText = pdf ? `[PDF ${count + 1}]` : `[Image ${count + 1}]`
+    const virtualText = promptAttachmentLabel(store.prompt.files, { uri: file.uri, name: file.filename })
     const extmarkEnd = extmarkStart + virtualText.length
     const textToInsert = virtualText + " "
 

+ 96 - 0
packages/tui/src/prompt/attachment.ts

@@ -0,0 +1,96 @@
+import type { PromptInput } from "@opencode-ai/schema"
+
+type PromptFile = PromptInput.FileAttachment
+type PromptFileIdentity = Pick<PromptFile, "uri" | "name" | "description">
+type ProjectedFile = Readonly<{
+  data: string
+  mime: string
+  source: { type: string }
+  name?: string
+  description?: string
+  mention?: { text: string }
+}>
+
+function attachmentKind(uri: string) {
+  if (uri.startsWith("data:image/")) return "Image"
+  if (uri.startsWith("data:application/pdf;")) return "PDF"
+  return undefined
+}
+
+function attachmentMetadata(file: PromptFileIdentity) {
+  return JSON.stringify([file.name ?? null, file.description ?? null])
+}
+
+function deduplicateByIdentity<T>(
+  items: readonly T[],
+  identity: (item: T) => { metadata: string; payload: string } | undefined,
+) {
+  const seen = new Map<string, string[]>()
+  return items.filter((item) => {
+    const key = identity(item)
+    if (!key) return true
+    const matches = seen.get(key.metadata)
+    if (matches?.includes(key.payload)) return false
+    if (matches) matches.push(key.payload)
+    if (!matches) seen.set(key.metadata, [key.payload])
+    return true
+  })
+}
+
+export function deduplicatePromptImages(files: readonly PromptFile[] | undefined) {
+  if (!files || files.length < 2) return files
+  return deduplicateByIdentity(files, (file) =>
+    file.uri.startsWith("data:image/") && file.mention?.text
+      ? {
+          metadata: JSON.stringify([attachmentMetadata(file), file.mention.text]),
+          payload: file.uri,
+        }
+      : undefined,
+  )
+}
+
+export function preserveMentionlessPromptAttachments(
+  files: readonly PromptFile[] | undefined,
+  mentioned: PromptFile[],
+) {
+  if (!files) return mentioned
+  const tracked = mentioned.values()
+  return files.flatMap((file) => {
+    if (!file.mention?.text) return [file]
+    const next = tracked.next()
+    return next.done ? [] : [next.value]
+  })
+}
+
+export function deduplicateVisibleImages<T extends ProjectedFile>(files: readonly T[]) {
+  return deduplicateByIdentity(files, (file) =>
+    file.mime.startsWith("image/") && file.source.type === "inline" && file.mention?.text
+      ? {
+          metadata: JSON.stringify([file.mime, file.name ?? null, file.description ?? null, file.mention.text]),
+          payload: file.data,
+        }
+      : undefined,
+  )
+}
+
+export function promptAttachmentLabel(files: readonly PromptFile[] | undefined, file: PromptFileIdentity) {
+  const kind = attachmentKind(file.uri)
+  if (!kind) throw new Error(`Unsupported inline attachment: ${file.uri}`)
+  const metadata = attachmentMetadata(file)
+  const existing =
+    kind === "Image"
+      ? files?.find(
+          (candidate) =>
+            candidate.uri === file.uri && attachmentMetadata(candidate) === metadata && candidate.mention?.text,
+        )?.mention?.text
+      : undefined
+  if (existing) return existing
+
+  const pattern = new RegExp(`^\\[${kind} (\\d+)\\]$`)
+  const count =
+    files?.reduce((highest, candidate) => {
+      const match = candidate.mention?.text.match(pattern)
+      return match ? Math.max(highest, Number(match[1])) : highest
+    }, 0) ?? 0
+  return `[${kind} ${count + 1}]`
+}

+ 2 - 1
packages/tui/src/routes/session/index.tsx

@@ -70,6 +70,7 @@ import stripAnsi from "strip-ansi"
 import { usePromptRef } from "../../context/prompt"
 import { sessionTabsFitVertically, SESSION_SIDEBAR_WIDTH } from "../../ui/layout"
 import { projectedPromptInput } from "../../prompt/codec"
+import { deduplicateVisibleImages } from "../../prompt/attachment"
 import { useEpilogue } from "../../context/epilogue"
 import { normalizePath } from "../../util/path"
 import { PermissionPrompt } from "./permission"
@@ -1899,7 +1900,7 @@ function UserMessage(props: { message: SessionMessageUser }) {
   const ctx = use()
   const data = useData()
   const local = useLocal()
-  const files = createMemo(() => props.message.files ?? [])
+  const files = createMemo(() => deduplicateVisibleImages(props.message.files ?? []))
   const skills = createMemo(() => props.message.skills ?? [])
   const images = createMemo(() =>
     files().flatMap((file) =>

+ 125 - 0
packages/tui/test/prompt/attachment.test.ts

@@ -0,0 +1,125 @@
+import { describe, expect, test } from "bun:test"
+import {
+  deduplicatePromptImages,
+  deduplicateVisibleImages,
+  preserveMentionlessPromptAttachments,
+  promptAttachmentLabel,
+} from "../../src/prompt/attachment"
+
+describe("prompt attachments", () => {
+  test("deduplicates identical inline images while preserving other attachments", () => {
+    const files = [
+      {
+        uri: "data:image/png;base64,AAA",
+        name: "first.png",
+        mention: { start: 0, end: 9, text: "[Image 1]" },
+      },
+      { uri: "file:///same", name: "first.txt" },
+      { uri: "data:application/pdf;base64,CCC", name: "first.pdf" },
+      {
+        uri: "data:image/png;base64,BBB",
+        name: "second.png",
+        mention: { start: 10, end: 19, text: "[Image 2]" },
+      },
+      {
+        uri: "data:image/png;base64,AAA",
+        name: "first.png",
+        mention: { start: 20, end: 29, text: "[Image 1]" },
+      },
+      {
+        uri: "data:image/png;base64,AAA",
+        name: "first.png",
+        description: "alternate use",
+        mention: { start: 30, end: 39, text: "[Image 1]" },
+      },
+      { uri: "file:///same", name: "second.txt" },
+      { uri: "data:application/pdf;base64,CCC", name: "first.pdf" },
+    ]
+
+    expect(deduplicatePromptImages(files)).toEqual([
+      files[0],
+      files[1],
+      files[2],
+      files[3],
+      files[5],
+      files[6],
+      files[7],
+    ])
+    expect(files).toHaveLength(8)
+  })
+
+  test("reuses labels for identical image data", () => {
+    const first = "data:image/png;base64,AAA"
+    const second = "data:image/png;base64,BBB"
+    const files = [{ uri: first, mention: { start: 0, end: 9, text: "[Image 1]" } }]
+
+    expect(promptAttachmentLabel(files, { uri: first })).toBe("[Image 1]")
+    expect(promptAttachmentLabel([...files, { ...files[0], mention: undefined }], { uri: second })).toBe("[Image 2]")
+    expect(promptAttachmentLabel([{ uri: first }], { uri: first })).toBe("[Image 1]")
+  })
+
+  test("numbers PDFs independently from images", () => {
+    const files = [{ uri: "data:image/png;base64,AAA" }]
+
+    expect(promptAttachmentLabel(files, { uri: "data:application/pdf;base64,BBB" })).toBe("[PDF 1]")
+  })
+
+  test("does not reuse a label when attachment metadata differs", () => {
+    const uri = "data:image/png;base64,AAA"
+    const files = [{ uri, name: "one.png", mention: { start: 0, end: 9, text: "[Image 1]" } }]
+
+    expect(promptAttachmentLabel(files, { uri, name: "two.png" })).toBe("[Image 2]")
+  })
+
+  test("does not reuse numbers after an earlier attachment is removed", () => {
+    const files = [{ uri: "data:image/png;base64,BBB", mention: { start: 0, end: 9, text: "[Image 2]" } }]
+
+    expect(promptAttachmentLabel(files, { uri: "data:image/png;base64,CCC" })).toBe("[Image 3]")
+  })
+
+  test("preserves mentionless attachments when tracked mentions are synchronized", () => {
+    const mentionless = { uri: "data:image/png;base64,AAA" }
+    const emptyMention = {
+      uri: "data:image/png;base64,CCC",
+      mention: { start: 0, end: 0, text: "" },
+    }
+    const mentioned = {
+      uri: "data:image/png;base64,BBB",
+      mention: { start: 0, end: 9, text: "[Image 1]" },
+    }
+
+    const restored = preserveMentionlessPromptAttachments([mentionless, emptyMention, mentioned], [mentioned])
+    expect(restored).toEqual([mentionless, emptyMention, mentioned])
+    expect(restored.indexOf(mentioned)).toBe(2)
+
+    const another = {
+      uri: "data:image/png;base64,DDD",
+      mention: { start: 10, end: 19, text: "[Image 2]" },
+    }
+    expect(preserveMentionlessPromptAttachments([mentioned, mentionless, another], [another, mentioned])).toEqual([
+      another,
+      mentionless,
+      mentioned,
+    ])
+  })
+
+  test("deduplicates visible inline image cards without dropping durable references", () => {
+    const file = {
+      data: "AAA",
+      mime: "image/png",
+      source: { type: "inline" },
+      name: "clipboard",
+      mention: { text: "[Image 1]" },
+    }
+    const files = [file, { ...file, mention: { text: "[Image 1]" } }]
+
+    expect(deduplicateVisibleImages(files)).toEqual([file])
+    expect(files).toHaveLength(2)
+
+    const distinct = [
+      { ...file, mention: { text: "[Image 2]" } },
+      { ...file, mention: undefined },
+    ]
+    expect(deduplicateVisibleImages([file, ...distinct])).toEqual([file, ...distinct])
+  })
+})

+ 17 - 0
packages/tui/test/prompt/history.test.ts

@@ -41,4 +41,21 @@ describe("prompt history", () => {
     const b = entry("describe this", [{ name: "b.png", uri: "data:image/png;base64,BBB" }])
     expect(isDuplicateEntry(a, b)).toBe(false)
   })
+
+  test("preserves duplicate attachment mentions for prompt restoration", () => {
+    const value = entry("[Image 1] [Image 1]", [
+      {
+        name: "clipboard",
+        uri: "data:image/png;base64,AAA",
+        mention: { start: 0, end: 9, text: "[Image 1]" },
+      },
+      {
+        name: "clipboard",
+        uri: "data:image/png;base64,AAA",
+        mention: { start: 10, end: 19, text: "[Image 1]" },
+      },
+    ])
+
+    expect(parsePromptHistory(JSON.stringify(value))).toEqual([value])
+  })
 })