diff --git a/CHANGELOG.md b/CHANGELOG.md index 9194c464..ac5c5885 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ All notable changes to DreamCoder will be documented in this file. - Images attached from uploads, paste or the built-in browser are kept unchanged within the client limits and resized or compressed when they exceed 2000×2000 pixels or 3.75 MB; a re-encoded image is named after its new type. An image the app cannot decode is attached as it is when it is within 3.75 MB. A file that cannot be attached is named in a message, and the other files picked or dropped with it are still attached. Sent images now use compact thumbnails that still open in the gallery (#218). Thanks @henricop. - Messages that reach a session in the middle of a turn are saved with the conversation. They used to disappear when the history was reloaded, and the agent forgot them when the session was resumed (#239). - A lead session that is resumed after a restart keeps using its team's task list (#216). The lead looked for tasks in a list named after its session, so `TaskList` returned nothing and new tasks went to a different list than the one the teammates read. The lead now goes back to the team it created when the session is resumed. Thanks @WhitePepperLambSoup. +- **Context usage no longer treats screenshot base64 as text** (#230). The estimate-only path and the Messages breakdown score images and documents as 2000 tokens instead of counting their base64 as text, so a session with many browser screenshots no longer shows a nearly full context bar when the provider still reports a modest input size. An @-mentioned image or PDF is estimated from the message the API receives, so that payload is no longer counted as text and the attachment line in `/context` stays in line with the total. ## [0.8.0] — 2026-11-01 diff --git a/src/services/tokenEstimation.ts b/src/services/tokenEstimation.ts index acaef7a5..429fc92e 100644 --- a/src/services/tokenEstimation.ts +++ b/src/services/tokenEstimation.ts @@ -388,7 +388,7 @@ function roughTokenCountEstimationForContent( return totalTokens } -function roughTokenCountEstimationForBlock( +export function roughTokenCountEstimationForBlock( block: string | Anthropic.ContentBlock | Anthropic.ContentBlockParam, ): number { if (typeof block === 'string') { diff --git a/src/utils/__tests__/analyzeContext.test.ts b/src/utils/__tests__/analyzeContext.test.ts new file mode 100644 index 00000000..8f6b0b0a --- /dev/null +++ b/src/utils/__tests__/analyzeContext.test.ts @@ -0,0 +1,247 @@ +import { describe, expect, test } from 'bun:test' +import type { + AssistantMessage, + AttachmentMessage, + UserMessage, +} from '../../types/message.js' +import { roughTokenCountEstimation } from '../../services/tokenEstimation.js' +import { jsonStringify } from '../slowOperations.js' +import { + estimateTokensForContextMessages, + processAssistantMessage, + processAttachment, + processUserMessage, + type MessageBreakdown, +} from '../analyzeContext.js' + +const LARGE_BASE64 = 'A'.repeat(800_000) + +function largeImageSource() { + return { + type: 'base64' as const, + media_type: 'image/png' as const, + data: LARGE_BASE64, + } +} + +function largeDocumentSource() { + return { + type: 'base64' as const, + media_type: 'application/pdf' as const, + data: LARGE_BASE64, + } +} + +function largeImageUserMessage() { + return { + role: 'user' as const, + content: [ + { + type: 'image' as const, + source: largeImageSource(), + }, + ], + } +} + +function emptyMessageBreakdown(): MessageBreakdown { + return { + totalTokens: 0, + toolCallTokens: 0, + toolResultTokens: 0, + attachmentTokens: 0, + assistantMessageTokens: 0, + userMessageTokens: 0, + toolCallsByType: new Map(), + toolResultsByType: new Map(), + attachmentsByType: new Map(), + } +} + +function largeImageFileAttachment() { + return { + type: 'file' as const, + filename: '/tmp/screenshot.png', + displayPath: 'screenshot.png', + content: { + type: 'image' as const, + file: { + base64: LARGE_BASE64, + type: 'image/png' as const, + originalSize: 600_000, + }, + }, + } +} + +describe('estimateTokensForContextMessages', () => { + test('counts a large base64 image as 2000 tokens instead of hundreds of thousands', () => { + const tokens = estimateTokensForContextMessages([largeImageUserMessage()], []) + expect(tokens).toBe(2000) + }) + + test('counts a tool_result that holds a large base64 image as 2000 tokens', () => { + const tokens = estimateTokensForContextMessages( + [ + { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'toolu_screenshot', + content: [ + { + type: 'image', + source: largeImageSource(), + }, + ], + }, + ], + }, + ], + [], + ) + expect(tokens).toBe(2000) + }) + + test('counts a large base64 document as 2000 tokens', () => { + const tokens = estimateTokensForContextMessages( + [ + { + role: 'user', + content: [ + { + type: 'document', + source: largeDocumentSource(), + }, + ], + }, + ], + [], + ) + expect(tokens).toBe(2000) + }) + + test('counts text plus an image with the image flat estimate instead of stringify length', () => { + const text = 'Here is the screenshot from the browser tool.' + const tokens = estimateTokensForContextMessages( + [ + { + role: 'user', + content: [ + { type: 'text', text }, + { + type: 'image', + source: largeImageSource(), + }, + ], + }, + ], + [], + ) + const textTokens = roughTokenCountEstimation(text) + expect(tokens).toBe(textTokens + 2000) + }) +}) + +describe('processUserMessage', () => { + test('scores a tool_result image as 2000 tokens in the tool result breakdown', () => { + const breakdown = emptyMessageBreakdown() + const message = { + type: 'user', + message: { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'toolu_screenshot', + content: [ + { + type: 'image', + source: largeImageSource(), + }, + ], + }, + ], + }, + } as UserMessage + + processUserMessage( + message, + breakdown, + new Map([['toolu_screenshot', 'screenshot']]), + ) + + expect(breakdown.toolResultTokens).toBe(2000) + expect(breakdown.toolResultsByType.get('screenshot')).toBe(2000) + expect(breakdown.userMessageTokens).toBe(0) + }) +}) + +describe('processAssistantMessage', () => { + test('scores an image block as 2000 tokens instead of the base64 length', () => { + const breakdown = emptyMessageBreakdown() + const message = { + type: 'assistant', + message: { + role: 'assistant', + content: [ + { + type: 'image', + source: largeImageSource(), + }, + ], + }, + } as AssistantMessage + + processAssistantMessage(message, breakdown) + + expect(breakdown.assistantMessageTokens).toBe(2000) + expect(breakdown.toolCallTokens).toBe(0) + }) + + test('scores a tool call with the per-block estimator', () => { + const breakdown = emptyMessageBreakdown() + const block = { + type: 'tool_use' as const, + id: 'toolu_screenshot', + name: 'screenshot', + input: { url: 'https://example.com' }, + } + const message = { + type: 'assistant', + message: { + role: 'assistant', + content: [block], + }, + } as AssistantMessage + + processAssistantMessage(message, breakdown) + + const perBlock = roughTokenCountEstimation( + block.name + jsonStringify(block.input), + ) + expect(breakdown.toolCallTokens).toBe(perBlock) + expect(breakdown.toolCallsByType.get('screenshot')).toBe(perBlock) + expect(perBlock).not.toBe(roughTokenCountEstimation(jsonStringify(block))) + }) +}) + +describe('processAttachment', () => { + test('scores a file attachment image as 2000 tokens plus the Read wrapper', () => { + const attachment = largeImageFileAttachment() + const breakdown = emptyMessageBreakdown() + const message = { + type: 'attachment', + attachment, + } as AttachmentMessage + + processAttachment(message, breakdown) + + expect(breakdown.attachmentTokens).toBe(2030) + expect(breakdown.attachmentsByType.get('file')).toBe(2030) + expect(roughTokenCountEstimation(jsonStringify(attachment))).toBeGreaterThan( + 100_000, + ) + }) +}) diff --git a/src/utils/analyzeContext.ts b/src/utils/analyzeContext.ts index e9ea89c6..b8637a3c 100644 --- a/src/utils/analyzeContext.ts +++ b/src/utils/analyzeContext.ts @@ -18,6 +18,9 @@ import { countMessagesTokensWithAPI, countTokensViaHaikuFallback, roughTokenCountEstimation, + roughTokenCountEstimationForBlock, + roughTokenCountEstimationForMessage, + roughTokenCountEstimationForMessages, } from '../services/tokenEstimation.js' import { estimateSkillFrontmatterTokens } from '../skills/loadSkillsDir.js' import { @@ -76,13 +79,29 @@ const MANUAL_COMPACT_BUFFER_NAME = 'Compact buffer' */ export const TOOL_TOKEN_COUNT_OVERHEAD = 500 +export function estimateTokensForContextMessages( + messages: Anthropic.Beta.Messages.BetaMessageParam[], + tools: Anthropic.Beta.Messages.BetaToolUnion[], +): number { + const messageTokens = roughTokenCountEstimationForMessages( + messages.map(message => ({ + type: message.role === 'assistant' ? 'assistant' : 'user', + message: { content: message.content }, + })), + ) + if (tools.length === 0) { + return messageTokens + } + return messageTokens + roughTokenCountEstimation(jsonStringify(tools)) +} + async function countTokensWithFallback( messages: Anthropic.Beta.Messages.BetaMessageParam[], tools: Anthropic.Beta.Messages.BetaToolUnion[], estimateOnly = false, ): Promise { if (estimateOnly) { - return roughTokenCountEstimation(jsonStringify({ messages, tools })) + return estimateTokensForContextMessages(messages, tools) } try { @@ -791,7 +810,7 @@ async function countCustomAgentTokens( return { agentTokens, agentDetails } } -type MessageBreakdown = { +export type MessageBreakdown = { totalTokens: number toolCallTokens: number toolResultTokens: number @@ -803,14 +822,13 @@ type MessageBreakdown = { attachmentsByType: Map } -function processAssistantMessage( +export function processAssistantMessage( msg: AssistantMessage | NormalizedAssistantMessage, breakdown: MessageBreakdown, ): void { // Process each content block individually for (const block of msg.message.content) { - const blockStr = jsonStringify(block) - const blockTokens = roughTokenCountEstimation(blockStr) + const blockTokens = roughTokenCountEstimationForBlock(block) if ('type' in block && block.type === 'tool_use') { breakdown.toolCallTokens += blockTokens @@ -826,7 +844,7 @@ function processAssistantMessage( } } -function processUserMessage( +export function processUserMessage( msg: UserMessage | NormalizedUserMessage, breakdown: MessageBreakdown, toolUseIdToName: Map, @@ -841,8 +859,7 @@ function processUserMessage( // Process each content block individually for (const block of msg.message.content) { - const blockStr = jsonStringify(block) - const blockTokens = roughTokenCountEstimation(blockStr) + const blockTokens = roughTokenCountEstimationForBlock(block) if ('type' in block && block.type === 'tool_result') { breakdown.toolResultTokens += blockTokens @@ -860,12 +877,14 @@ function processUserMessage( } } -function processAttachment( +export function processAttachment( msg: AttachmentMessage, breakdown: MessageBreakdown, ): void { - const contentStr = jsonStringify(msg.attachment) - const tokens = roughTokenCountEstimation(contentStr) + const tokens = roughTokenCountEstimationForMessage({ + type: 'attachment', + attachment: msg.attachment, + }) breakdown.attachmentTokens += tokens const attachType = msg.attachment.type || 'unknown' breakdown.attachmentsByType.set(