fix: ensure that tool attachments arent sent as user messages (#8944)

3 files changed, 89 insertions, 43 deletions

diff --git a/packages/opencode/src/session/message-v2.ts b/packages/opencode/src/session/message-v2.ts
index 9f2e0ba06..da714c437 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 ToolSet,
+  type UIMessage,
+} 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
+                  ? { output: "[Old tool result content cleared]" }
+                  : { output: part.state.output, attachments: part.state.attachments },
                 callProviderMetadata: part.metadata,
               })
             }
@@ -565,7 +556,10 @@ 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..151b2d62f 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
             ? [
                 {
@@ -718,8 +718,22 @@ export namespace SessionPrompt {
         },
         toModelOutput(result) {
           return {
-            type: "text",
-            value: result.output,
+            type: "content",
+            value: [
+              {
+                type: "text",
+                text: result.output,
+              },
+              ...(result.attachments?.map((attachment: MessageV2.FilePart) => {
+                const base64 = attachment.url.startsWith("data:") ? attachment.url.split(",", 2)[1] : attachment.url
+
+                return {
+                  type: "media",
+                  data: base64,
+                  mediaType: attachment.mime,
+                }
+              }) ?? []),
+            ],
           }
         },
       })
@@ -808,8 +822,22 @@ export namespace SessionPrompt {
       }
       item.toModelOutput = (result) => {
         return {
-          type: "text",
-          value: result.output,
+          type: "content",
+          value: [
+            {
+              type: "text",
+              text: result.output,
+            },
+            ...(result.attachments?.map((attachment: MessageV2.FilePart) => {
+              const base64 = attachment.url.startsWith("data:") ? attachment.url.split(",", 2)[1] : attachment.url
+
+              return {
+                type: "media",
+                data: base64,
+                mediaType: attachment.mime,
+              }
+            }) ?? []),
+          ],
         }
       }
       tools[key] = item
diff --git a/packages/opencode/test/session/message-v2.test.ts b/packages/opencode/test/session/message-v2.test.ts
index f069f6ba6..376c189ba 100644
--- a/packages/opencode/test/session/message-v2.test.ts
+++ b/packages/opencode/test/session/message-v2.test.ts
@@ -1,8 +1,35 @@
 import { describe, expect, test } from "bun:test"
 import { MessageV2 } from "../../src/session/message-v2"
+import type { ToolSet } from "ai"
 
 const sessionID = "session"
 
+// Mock tool that transforms output to content format with media support
+function createMockTools(): ToolSet {
+  return {
+    bash: {
+      description: "mock bash tool",
+      inputSchema: { type: "object", properties: {} } as any,
+      toModelOutput(result: { output: string; attachments?: MessageV2.FilePart[] }) {
+        return {
+          type: "content" as const,
+          value: [
+            { type: "text" as const, text: result.output },
+            ...(result.attachments?.map((attachment) => {
+              const base64 = attachment.url.startsWith("data:") ? attachment.url.split(",", 2)[1] : attachment.url
+              return {
+                type: "media" as const,
+                data: base64,
+                mediaType: attachment.mime,
+              }
+            }) ?? []),
+          ],
+        }
+      },
+    },
+  } as ToolSet
+}
+
 function userInfo(id: string): MessageV2.User {
   return {
     id,
@@ -259,24 +286,12 @@ describe("session.message-v2.toModelMessage", () => {
       },
     ]
 
-    expect(MessageV2.toModelMessage(input)).toStrictEqual([
+    expect(MessageV2.toModelMessage(input, { tools: createMockTools() })).toStrictEqual([
       {
         role: "user",
         content: [{ type: "text", text: "run tool" }],
       },
       {
-        role: "user",
-        content: [
-          { type: "text", text: "Tool bash returned an attachment:" },
-          {
-            type: "file",
-            mediaType: "image/png",
-            filename: "attachment.png",
-            data: "https://example.com/attachment.png",
-          },
-        ],
-      },
-      {
         role: "assistant",
         content: [
           { type: "text", text: "done", providerOptions: { openai: { assistant: "meta" } } },
@@ -297,7 +312,13 @@ describe("session.message-v2.toModelMessage", () => {
             type: "tool-result",
             toolCallId: "call-1",
             toolName: "bash",
-            output: { type: "text", value: "ok" },
+            output: {
+              type: "content",
+              value: [
+                { type: "text", text: "ok" },
+                { type: "media", data: "https://example.com/attachment.png", mediaType: "image/png" },
+              ],
+            },
             providerOptions: { openai: { tool: "meta" } },
           },
         ],
@@ -341,7 +362,7 @@ describe("session.message-v2.toModelMessage", () => {
       },
     ]
 
-    expect(MessageV2.toModelMessage(input)).toStrictEqual([
+    expect(MessageV2.toModelMessage(input, { tools: createMockTools() })).toStrictEqual([
       {
         role: "user",
         content: [{ type: "text", text: "run tool" }],
@@ -365,7 +386,10 @@ describe("session.message-v2.toModelMessage", () => {
             type: "tool-result",
             toolCallId: "call-1",
             toolName: "bash",
-            output: { type: "text", value: "[Old tool result content cleared]" },
+            output: {
+              type: "content",
+              value: [{ type: "text", text: "[Old tool result content cleared]" }],
+            },
           },
         ],
       },