fix: rm user message when dealing w/ image attachments, use proper tool attachment instead

2 files changed, 41 insertions, 29 deletions

diff --git a/packages/opencode/src/session/message-v2.ts b/packages/opencode/src/session/message-v2.ts
index 9f2e0ba06..4b081b5b4 100644
--- a/packages/opencode/src/session/message-v2.ts
+++ b/packages/opencode/src/session/message-v2.ts
@@ -1,7 +1,14 @@
 import { BusEvent } from "@/bus/bus-event"
 import z from "zod"
 import { NamedError } from "@opencode-ai/util/error"
-import { APICallError, convertToModelMessages, LoadAPIKeyError, type ModelMessage, type UIMessage } from "ai"
+import {
+  APICallError,
+  convertToModelMessages,
+  LoadAPIKeyError,
+  type ModelMessage,
+  type UIMessage,
+  type ToolSet,
+} from "ai"
 import { Identifier } from "../id/id"
 import { LSP } from "../lsp"
 import { Snapshot } from "@/snapshot"
@@ -432,7 +439,7 @@ export namespace MessageV2 {
   })
   export type WithParts = z.infer<typeof WithParts>
 
-  export function toModelMessage(input: WithParts[]): ModelMessage[] {
+  export function toModelMessage(input: WithParts[], options?: { tools?: ToolSet }): ModelMessage[] {
     const result: UIMessage[] = []
 
     for (const msg of input) {
@@ -503,30 +510,14 @@ export namespace MessageV2 {
             })
           if (part.type === "tool") {
             if (part.state.status === "completed") {
-              if (part.state.attachments?.length) {
-                result.push({
-                  id: Identifier.ascending("message"),
-                  role: "user",
-                  parts: [
-                    {
-                      type: "text",
-                      text: `Tool ${part.tool} returned an attachment:`,
-                    },
-                    ...part.state.attachments.map((attachment) => ({
-                      type: "file" as const,
-                      url: attachment.url,
-                      mediaType: attachment.mime,
-                      filename: attachment.filename,
-                    })),
-                  ],
-                })
-              }
               assistantMessage.parts.push({
                 type: ("tool-" + part.tool) as `tool-${string}`,
                 state: "output-available",
                 toolCallId: part.callID,
                 input: part.state.input,
-                output: part.state.time.compacted ? "[Old tool result content cleared]" : part.state.output,
+                output: part.state.time.compacted
+                  ? "[Old tool result content cleared]"
+                  : { output: part.state.output, attachments: part.state.attachments },
                 callProviderMetadata: part.metadata,
               })
             }
@@ -565,7 +556,12 @@ export namespace MessageV2 {
       }
     }
 
-    return convertToModelMessages(result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")))
+    return convertToModelMessages(
+      result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")),
+      {
+        tools: options?.tools,
+      },
+    )
   }
 
   export const stream = fn(Identifier.schema("session"), async function* (sessionID) {
diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts
index 8327698fd..663f5660f 100644
--- a/packages/opencode/src/session/prompt.ts
+++ b/packages/opencode/src/session/prompt.ts
@@ -597,7 +597,7 @@ export namespace SessionPrompt {
         sessionID,
         system: [...(await SystemPrompt.environment()), ...(await SystemPrompt.custom())],
         messages: [
-          ...MessageV2.toModelMessage(sessionMessages),
+          ...MessageV2.toModelMessage(sessionMessages, { tools }),
           ...(isLastStep
             ? [
                 {
@@ -716,10 +716,18 @@ export namespace SessionPrompt {
           )
           return result
         },
-        toModelOutput(result) {
+        toModelOutput(result: { output: string; attachments?: MessageV2.FilePart[] }) {
+          if (!result.attachments?.length) return { type: "text", value: result.output }
           return {
-            type: "text",
-            value: result.output,
+            type: "content",
+            value: [
+              { type: "text", text: result.output },
+              ...result.attachments.map((a) => ({
+                type: "media" as const,
+                data: a.url.slice(a.url.indexOf(",") + 1),
+                mediaType: a.mime,
+              })),
+            ],
           }
         },
       })
@@ -806,10 +814,18 @@ export namespace SessionPrompt {
           content: result.content, // directly return content to preserve ordering when outputting to model
         }
       }
-      item.toModelOutput = (result) => {
+      item.toModelOutput = (result: { output: string; attachments?: MessageV2.FilePart[] }) => {
+        if (!result.attachments?.length) return { type: "text", value: result.output }
         return {
-          type: "text",
-          value: result.output,
+          type: "content",
+          value: [
+            { type: "text", text: result.output },
+            ...result.attachments.map((a) => ({
+              type: "media" as const,
+              data: a.url.slice(a.url.indexOf(",") + 1),
+              mediaType: a.mime,
+            })),
+          ],
         }
       }
       tools[key] = item