fix: ensure images are properly returned as tool results

This commit is contained in:
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[] {
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) {
if (msg.parts.length === 0) continue
@@ -505,31 +539,24 @@ export namespace MessageV2 {
type: "step-start",
})
if (part.type === "tool") {
toolNames.add(part.tool)
if (part.state.status === "completed") {
if (part.state.attachments?.length) {
result.push({
id: Identifier.ascending("message"),
role: "user",
parts: [
{
type: "text",
text: `The tool ${part.tool} returned the following attachments:`,
},
...part.state.attachments.map((attachment) => ({
type: "file" as const,
url: attachment.url,
mediaType: attachment.mime,
filename: attachment.filename,
})),
],
})
}
const outputText = part.state.time.compacted ? "[Old tool result content cleared]" : part.state.output
const attachments = part.state.time.compacted ? [] : (part.state.attachments ?? [])
const output =
attachments.length > 0
? {
text: outputText,
attachments,
}
: outputText
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,
...(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) {

View File

@@ -722,12 +722,6 @@ export namespace SessionPrompt {
)
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
}
}
item.toModelOutput = (result) => {
return {
type: "text",
value: result.output,
}
}
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 assistantID = "m-assistant"
@@ -304,7 +304,7 @@ describe("session.message-v2.toModelMessage", () => {
type: "file",
mime: "image/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",
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",
content: [
@@ -352,7 +340,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", mediaType: "image/png", data: "Zm9v" },
],
},
providerOptions: { openai: { tool: "meta" } },
},
],