fix: ensure images are properly returned as tool results

pull/9961/head
Aiden Cline 2026-01-21 23:54:39 -06:00
parent fc0210c2fd
commit c2844697f3
3 changed files with 64 additions and 47 deletions

View File

@ -435,6 +435,40 @@ export namespace MessageV2 {
export function toModelMessages(input: WithParts[], model: Provider.Model): ModelMessage[] { export function toModelMessages(input: WithParts[], model: Provider.Model): ModelMessage[] {
const result: UIMessage[] = [] const result: UIMessage[] = []
const toolNames = new Set<string>()
const toModelOutput = (output: unknown) => {
if (typeof output === "string") {
return { type: "text", value: output }
}
if (typeof output === "object") {
const outputObject = output as {
text: string
attachments?: Array<{ mime: string; url: string }>
}
const attachments = (outputObject.attachments ?? []).filter((attachment) => {
return attachment.url.startsWith("data:") && attachment.url.includes(",")
})
return {
type: "content",
value: [
{ type: "text", text: outputObject.text },
...attachments.map((attachment) => ({
type: "media",
mediaType: attachment.mime,
data: iife(() => {
const commaIndex = attachment.url.indexOf(",")
return commaIndex === -1 ? attachment.url : attachment.url.slice(commaIndex + 1)
}),
})),
],
}
}
return { type: "json", value: output as never }
}
for (const msg of input) { for (const msg of input) {
if (msg.parts.length === 0) continue if (msg.parts.length === 0) continue
@ -505,31 +539,24 @@ export namespace MessageV2 {
type: "step-start", type: "step-start",
}) })
if (part.type === "tool") { if (part.type === "tool") {
toolNames.add(part.tool)
if (part.state.status === "completed") { if (part.state.status === "completed") {
if (part.state.attachments?.length) { const outputText = part.state.time.compacted ? "[Old tool result content cleared]" : part.state.output
result.push({ const attachments = part.state.time.compacted ? [] : (part.state.attachments ?? [])
id: Identifier.ascending("message"), const output =
role: "user", attachments.length > 0
parts: [ ? {
{ text: outputText,
type: "text", attachments,
text: `The tool ${part.tool} returned the following attachments:`, }
}, : outputText
...part.state.attachments.map((attachment) => ({
type: "file" as const,
url: attachment.url,
mediaType: attachment.mime,
filename: attachment.filename,
})),
],
})
}
assistantMessage.parts.push({ assistantMessage.parts.push({
type: ("tool-" + part.tool) as `tool-${string}`, type: ("tool-" + part.tool) as `tool-${string}`,
state: "output-available", state: "output-available",
toolCallId: part.callID, toolCallId: part.callID,
input: part.state.input, input: part.state.input,
output: part.state.time.compacted ? "[Old tool result content cleared]" : part.state.output, output,
...(differentModel ? {} : { callProviderMetadata: part.metadata }), ...(differentModel ? {} : { callProviderMetadata: part.metadata }),
}) })
} }
@ -568,7 +595,15 @@ export namespace MessageV2 {
} }
} }
return convertToModelMessages(result.filter((msg) => msg.parts.some((part) => part.type !== "step-start"))) const tools = Object.fromEntries(Array.from(toolNames).map((toolName) => [toolName, { toModelOutput }]))
return convertToModelMessages(
result.filter((msg) => msg.parts.some((part) => part.type !== "step-start")),
{
//@ts-expect-error (convertToModelMessages expects a ToolSet but only actually needs tools[name]?.toModelOutput)
tools,
},
)
} }
export const stream = fn(Identifier.schema("session"), async function* (sessionID) { export const stream = fn(Identifier.schema("session"), async function* (sessionID) {

View File

@ -722,12 +722,6 @@ export namespace SessionPrompt {
) )
return result return result
}, },
toModelOutput(result) {
return {
type: "text",
value: result.output,
}
},
}) })
} }
@ -819,12 +813,6 @@ export namespace SessionPrompt {
content: result.content, // directly return content to preserve ordering when outputting to model content: result.content, // directly return content to preserve ordering when outputting to model
} }
} }
item.toModelOutput = (result) => {
return {
type: "text",
value: result.output,
}
}
tools[key] = item tools[key] = item
} }

View File

@ -262,7 +262,7 @@ describe("session.message-v2.toModelMessage", () => {
]) ])
}) })
test("converts assistant tool completion into tool-call + tool-result messages and emits attachment message", () => { test("converts assistant tool completion into tool-call + tool-result messages with attachments", () => {
const userID = "m-user" const userID = "m-user"
const assistantID = "m-assistant" const assistantID = "m-assistant"
@ -304,7 +304,7 @@ describe("session.message-v2.toModelMessage", () => {
type: "file", type: "file",
mime: "image/png", mime: "image/png",
filename: "attachment.png", filename: "attachment.png",
url: "https://example.com/attachment.png", url: "data:image/png;base64,Zm9v",
}, },
], ],
}, },
@ -319,18 +319,6 @@ describe("session.message-v2.toModelMessage", () => {
role: "user", role: "user",
content: [{ type: "text", text: "run tool" }], content: [{ type: "text", text: "run tool" }],
}, },
{
role: "user",
content: [
{ type: "text", text: "The tool bash returned the following attachments:" },
{
type: "file",
mediaType: "image/png",
filename: "attachment.png",
data: "https://example.com/attachment.png",
},
],
},
{ {
role: "assistant", role: "assistant",
content: [ content: [
@ -352,7 +340,13 @@ describe("session.message-v2.toModelMessage", () => {
type: "tool-result", type: "tool-result",
toolCallId: "call-1", toolCallId: "call-1",
toolName: "bash", toolName: "bash",
output: { type: "text", value: "ok" }, output: {
type: "content",
value: [
{ type: "text", text: "ok" },
{ type: "media", mediaType: "image/png", data: "Zm9v" },
],
},
providerOptions: { openai: { tool: "meta" } }, providerOptions: { openai: { tool: "meta" } },
}, },
], ],