fix(core): budget retained image tokens

This commit is contained in:
Aiden Cline
2026-08-04 22:40:58 -05:00
parent b4ec03e5f1
commit 0c89e24aac
3 changed files with 39 additions and 2 deletions
+18 -1
View File
@@ -23,6 +23,7 @@ const DEFAULT_BUFFER = 20_000
const DEFAULT_KEEP_TOKENS = 15_000
const OUTPUT_TOKEN_MAX = 32_000
const TOOL_OUTPUT_MAX_CHARS = 2_000
const IMAGE_TOKEN_ESTIMATE = 2_000
const SUMMARY_TEMPLATE = `Output exactly the Markdown structure shown inside <template> and keep the section order unchanged. Do not include the <template> tags in your response.
<template>
## Objective
@@ -126,6 +127,22 @@ export const serializeToolContent = (content: SessionMessage.ToolStateCompleted[
)
.join("\n")
export const estimateImageTokens = (message: SessionMessage.Info) => {
if (message.type === "user")
return (message.files?.filter((file) => file.mime.toLowerCase().startsWith("image/")).length ?? 0) * IMAGE_TOKEN_ESTIMATE
if (message.type !== "assistant") return 0
return (
message.content
.flatMap((part) =>
part.type === "tool" && (part.state.status === "completed" || part.state.status === "error")
? (part.state.content ?? [])
: [],
)
.filter((content) => content.type === "file" && content.mime.toLowerCase().startsWith("image/")).length *
IMAGE_TOKEN_ESTIMATE
)
}
const serialize = (message: SessionMessage.Info) => {
if (message.type === "user") {
const files =
@@ -188,7 +205,7 @@ const select = (
let total = 0
let split = conversation.length
for (let index = conversation.length - 1; index >= 0; index--) {
const next = total + Token.estimate(conversation[index].text)
const next = total + Token.estimate(conversation[index].text) + estimateImageTokens(conversation[index].message)
if (split < conversation.length && next > tokens) break
total = next
split = index
@@ -114,6 +114,24 @@ test("compaction describes tool media without embedding base64", () => {
expect(serialized).not.toContain(base64)
})
test("compaction estimates image context without counting base64", () => {
const image = FileAttachment.make({
data: Base64.make("a".repeat(10_000)),
mime: "image/png",
source: { type: "inline" },
name: "image.png",
})
const message = SessionMessage.User.make({
id: SessionMessage.ID.create(),
type: "user",
text: "Compare these images.",
files: [image, image],
time: { created: DateTime.makeUnsafe(0) },
})
expect(SessionCompaction.estimateImageTokens(message)).toBe(4_000)
})
test("compaction prompt requires the checkpoint headings in order", () => {
const prompt = SessionCompaction.buildPrompt({ context: ["Conversation history"] })
expect(prompt.match(/^#{2,3} .+$/gm)).toEqual([
@@ -111,7 +111,9 @@ and relevant files.
The newest serialized context up to `keep.tokens` is retained separately. This
is not a byte-for-byte transcript: tool output is limited to 2000 characters,
and non-media attachments become textual descriptors. Media in the retained
context is attached to the checkpoint in the same order as its descriptors. On
context is attached to the checkpoint in the same order as its descriptors.
Tail selection budgets 2000 additional tokens per image as a provider-neutral
planning estimate; it does not count base64 request bytes as text tokens. On
later compactions, V2 updates the previous summary and carries forward its
retained recent context before selecting a new tail.