perf(core): keep unchanged messages in request shaping (#52784)

This commit is contained in:
Shoubhit Dash
2026-10-02 22:12:21 +05:30
committed by GitHub
Unverified
parent 80caae8ad1
commit 2a01ed60b1
2 changed files with 119 additions and 59 deletions
+42 -55
View File
@@ -129,48 +129,53 @@ const mimeToModality = (mime: string) => {
const unsupportedMedia = (mime: string, name: string | undefined, capabilities: Model.Capabilities) => {
const modality = mimeToModality(mime)
if (!modality || capabilities.input.some((item) => item.startsWith(modality))) return
return {
type: "text" as const,
text: `ERROR: Cannot read ${name ? `"${name}"` : modality} (this model does not support ${modality} input). Inform the user.`,
}
return `ERROR: Cannot read ${name ? `"${name}"` : modality} (this model does not support ${modality} input). Inform the user.`
}
// Remote and provider-referenced media carry no local payload and never count toward the inline budget.
const mediaBytes = (media: Media.Asset) => {
if (media.source.type === "base64") return Buffer.byteLength(media.source.data)
if (media.source.type === "bytes") return Math.ceil(media.source.data.byteLength / 3) * 4
return 0
}
/** Replaces media with the returned text; messages without replacements are returned unchanged. */
const replaceMedia = (
messages: LLMRequest["messages"],
replace: (media: { mime: string; name: string | undefined; bytes: () => number }) => string | undefined,
) =>
messages.map((message) => {
const content = message.content.map((part) => {
if (part.type === "media") {
const text = replace({ mime: part.media.mediaType, name: part.filename, bytes: () => mediaBytes(part.media) })
return text === undefined ? part : Message.text(text)
}
if (part.type !== "tool-result" || part.result.type !== "content") return part
const result = part.result
const value = result.value.map((item): Content => {
if (item.type !== "file") return item
const text = replace({ mime: item.mime, name: item.name, bytes: () => Buffer.byteLength(item.uri) })
return text === undefined ? item : { type: "text", text }
})
return value.every((item, index) => item === result.value[index])
? part
: { ...part, result: { ...result, value } }
})
return content.every((part, index) => part === message.content[index])
? message
: new Message({ ...message, content })
})
export const unsupportedParts = (messages: LLMRequest["messages"], capabilities: Model.Capabilities) =>
messages.map((message) =>
Message.make({
...message,
content: message.content.map((part) => {
if (part.type === "media") {
return unsupportedMedia(part.media.mediaType, part.filename, capabilities) ?? part
}
if (part.type !== "tool-result" || part.result.type !== "content") return part
return {
...part,
result: {
...part.result,
value: part.result.value.map((item: Content) => {
if (item.type !== "file") return item
return unsupportedMedia(item.mime, item.name, capabilities) ?? item
}),
},
}
}),
}),
)
replaceMedia(messages, (media) => unsupportedMedia(media.mime, media.name, capabilities))
export const boundImages = (messages: LLMRequest["messages"]) => {
const isImage = (mime: string) => mime.toLowerCase().startsWith("image/")
// Remote and provider-referenced media carry no local payload and never count toward the inline budget.
const size = (media: Media.Asset) => {
if (media.source.type === "base64") return Buffer.byteLength(media.source.data)
if (media.source.type === "bytes") return Math.ceil(media.source.data.byteLength / 3) * 4
return 0
}
const imageBytes = messages.reduce(
(total, message) =>
total +
message.content.reduce((sum, part) => {
if (part.type === "media" && isImage(part.media.mediaType)) return sum + size(part.media)
if (part.type === "media" && isImage(part.media.mediaType)) return sum + mediaBytes(part.media)
if (part.type !== "tool-result" || part.result.type !== "content") return sum
return (
sum +
@@ -186,29 +191,11 @@ export const boundImages = (messages: LLMRequest["messages"]) => {
if (imageBytes <= IMAGE_BYTES_TRIGGER) return messages
let removed = 0
return messages.map((message) =>
Message.make({
...message,
content: message.content.map((part) => {
if (part.type === "media" && isImage(part.media.mediaType) && imageBytes - removed > IMAGE_BYTES_TARGET) {
removed += size(part.media)
return Message.text(IMAGE_REMOVED)
}
if (part.type !== "tool-result" || part.result.type !== "content") return part
return {
...part,
result: {
...part.result,
value: part.result.value.map((item: Content) => {
if (item.type !== "file" || !isImage(item.mime) || imageBytes - removed <= IMAGE_BYTES_TARGET) return item
removed += Buffer.byteLength(item.uri)
return { type: "text" as const, text: IMAGE_REMOVED }
}),
},
}
}),
}),
)
return replaceMedia(messages, (media) => {
if (!isImage(media.mime) || imageBytes - removed <= IMAGE_BYTES_TARGET) return
removed += media.bytes()
return IMAGE_REMOVED
})
}
type Definitions = PluginHooks.Domains["session"]["context"]["tools"]
@@ -63,6 +63,73 @@ describe("SessionModelRequest.unsupportedParts", () => {
const message = Message.user({ type: "media", media: Media.base64("aGVsbG8=", "image/png") })
expect(unsupportedParts([message], capabilities(["text", "image"]))[0]?.content).toEqual(message.content)
})
test("returns the same messages when nothing is unsupported", () => {
const messages = [
Message.user([Message.text("hi"), { type: "media", media: Media.base64("aGVsbG8=", "image/png") }]),
Message.assistant("hello"),
Message.tool(
ToolResultPart.make({
id: "call_1",
name: "read",
result: {
type: "content",
value: [{ type: "file", uri: "data:image/png;base64,aGVsbG8=", mime: "image/png", name: "logo.png" }],
},
}),
),
]
const result = unsupportedParts(messages, capabilities(["text", "image"]))
expect(result).toHaveLength(messages.length)
result.forEach((message, index) => expect(message).toBe(messages[index]))
})
test("rebuilds only messages with unsupported media", () => {
const text = Message.user("plain")
const supported = Message.tool(
ToolResultPart.make({
id: "call_1",
name: "read",
result: { type: "content", value: [{ type: "text", text: "no files" }] },
}),
)
const tool = ToolResultPart.make({
id: "call_2",
name: "read",
result: {
type: "content",
value: [
{ type: "text", text: "Read" },
{ type: "file", uri: "data:application/pdf;base64,JVBERg==", mime: "application/pdf" },
],
},
})
const unsupported = Message.make({ role: "tool", content: [tool, Message.text("caption")] })
const result = unsupportedParts([text, supported, unsupported], capabilities(["text"]))
expect(result[0]).toBe(text)
expect(result[1]).toBe(supported)
expect(result[2]).not.toBe(unsupported)
expect(result[2]).toBeInstanceOf(Message)
expect(result[2]?.content[1]).toEqual(Message.text("caption"))
expect(result[2]?.content[0]).toMatchObject({
type: "tool-result",
result: {
type: "content",
value: [
{ type: "text", text: "Read" },
{ type: "text", text: "ERROR: Cannot read pdf (this model does not support pdf input). Inform the user." },
],
},
})
expect(tool.result).toMatchObject({
type: "content",
value: [
{ type: "text", text: "Read" },
{ type: "file", mime: "application/pdf" },
],
})
})
})
describe("SessionModelRequest.boundImages", () => {
@@ -73,16 +140,22 @@ describe("SessionModelRequest.boundImages", () => {
test("replaces oldest images until the retained payload reaches the target", () => {
const image = "a".repeat(9 * 1024 * 1024)
const text = Message.user("plain")
const third = Message.user({ type: "media", media: Media.base64(image, "image/png"), filename: "third.png" })
const messages = [
Message.user({ type: "media", media: Media.base64(image, "image/png"), filename: "first.png" }),
text,
Message.user({ type: "media", media: Media.base64(image, "image/png"), filename: "second.png" }),
Message.user({ type: "media", media: Media.base64(image, "image/png"), filename: "third.png" }),
third,
]
const result = boundImages(messages)
expect(result[0]?.content[0]).toMatchObject({ type: "text" })
expect(result[1]?.content[0]).toMatchObject({ type: "text" })
expect(result[2]?.content[0]).toMatchObject({ type: "media", filename: "third.png" })
expect(result[0]).toBeInstanceOf(Message)
expect(result[0]?.content).toEqual([Message.text(expect.stringContaining("image was removed"))])
expect(result[1]).toBe(text)
expect(result[2]).not.toBe(messages[2])
expect(result[2]?.content[0]).toMatchObject({ type: "text" })
expect(result[3]).toBe(third)
})
test("replaces images nested in tool results", () => {