From f77db3577ea5959209e6cff316fc06e0d5949766 Mon Sep 17 00:00:00 2001 From: Oscar <319719771+Oscar-Williams@users.noreply.github.com> Date: Mon, 2 Nov 2026 00:30:11 +0800 Subject: [PATCH 1/2] fix(context): count screenshot images as 2000 tokens in usage estimates When estimateOnly is true, context usage serialized the whole message payload and treated base64 screenshots as text. Route that path and the Messages breakdown through the existing per block estimators instead. Closes #230 --- CHANGELOG.md | 4 + src/services/tokenEstimation.ts | 2 +- src/utils/__tests__/analyzeContext.test.ts | 146 +++++++++++++++++++++ src/utils/analyzeContext.ts | 26 +++- 4 files changed, 172 insertions(+), 6 deletions(-) create mode 100644 src/utils/__tests__/analyzeContext.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index e1c28e56..b1b62944 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to DreamCoder will be documented in this file. ## [Unreleased] +### Fixed + +- **Context usage no longer treats screenshot base64 as text** (#230). The estimate-only path and the Messages breakdown now use the same per-block estimator that already scores images and documents as 2000 tokens, so a session with many browser screenshots no longer shows a nearly full context bar when the provider still reports a modest input size. + ## [0.8.0] — 2026-11-01 The first stable release of the Electron desktop app. These notes cover everything since 0.7.0, including 0.8.0-beta.1 and 0.8.0-beta.2. 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..b68b6603 --- /dev/null +++ b/src/utils/__tests__/analyzeContext.test.ts @@ -0,0 +1,146 @@ +import { describe, expect, test } from 'bun:test' +import { + roughTokenCountEstimation, + roughTokenCountEstimationForBlock, +} from '../../services/tokenEstimation.js' +import { estimateTokensForContextMessages } 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(), + }, + ], + } +} + +describe('estimateTokensForContextMessages', () => { + test('counts a large base64 image near 2000 tokens instead of hundreds of thousands', () => { + const tokens = estimateTokensForContextMessages([largeImageUserMessage()], []) + expect(tokens).toBeGreaterThanOrEqual(1_500) + expect(tokens).toBeLessThanOrEqual(2_500) + expect(tokens).toBeLessThan(100_000) + }) + + test('counts a tool_result that holds a large base64 image near 2000 tokens', () => { + const tokens = estimateTokensForContextMessages( + [ + { + role: 'user', + content: [ + { + type: 'tool_result', + tool_use_id: 'toolu_screenshot', + content: [ + { + type: 'image', + source: largeImageSource(), + }, + ], + }, + ], + }, + ], + [], + ) + expect(tokens).toBeGreaterThanOrEqual(1_500) + expect(tokens).toBeLessThanOrEqual(2_500) + expect(tokens).toBeLessThan(100_000) + }) + + 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) + expect(tokens).toBeLessThan(100_000) + }) +}) + +describe('roughTokenCountEstimationForBlock', () => { + test('scores an image block as 2000 tokens', () => { + expect( + roughTokenCountEstimationForBlock({ + type: 'image', + source: largeImageSource(), + }), + ).toBe(2000) + }) + + test('scores a document block as 2000 tokens', () => { + expect( + roughTokenCountEstimationForBlock({ + type: 'document', + source: largeDocumentSource(), + }), + ).toBe(2000) + }) + + test('scores a tool_result that holds a large image as 2000 tokens', () => { + expect( + roughTokenCountEstimationForBlock({ + type: 'tool_result', + tool_use_id: 'toolu_screenshot', + content: [ + { + type: 'image', + source: largeImageSource(), + }, + ], + }), + ).toBe(2000) + }) +}) diff --git a/src/utils/analyzeContext.ts b/src/utils/analyzeContext.ts index e9ea89c6..41263312 100644 --- a/src/utils/analyzeContext.ts +++ b/src/utils/analyzeContext.ts @@ -18,6 +18,8 @@ import { countMessagesTokensWithAPI, countTokensViaHaikuFallback, roughTokenCountEstimation, + roughTokenCountEstimationForBlock, + roughTokenCountEstimationForMessages, } from '../services/tokenEstimation.js' import { estimateSkillFrontmatterTokens } from '../skills/loadSkillsDir.js' import { @@ -76,13 +78,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 { @@ -809,8 +827,7 @@ function processAssistantMessage( ): 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 @@ -841,8 +858,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 From 6bf29ec999c0841ba46a65bca738050ed695643c Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sun, 8 Nov 2026 08:41:31 +0000 Subject: [PATCH 2/2] fix(context): score file attachments with the image-aware estimator The Messages breakdown tests now fail if tool results or assistant blocks are stringified, and @-mentioned images use the same normalized estimate as the API instead of counting base64 as text. Co-authored-by: Oscar --- CHANGELOG.md | 2 +- src/utils/__tests__/analyzeContext.test.ts | 173 ++++++++++++++++----- src/utils/analyzeContext.ts | 15 +- 3 files changed, 147 insertions(+), 43 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b1b62944..edf34da6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ All notable changes to DreamCoder will be documented in this file. ### Fixed -- **Context usage no longer treats screenshot base64 as text** (#230). The estimate-only path and the Messages breakdown now use the same per-block estimator that already scores images and documents as 2000 tokens, so a session with many browser screenshots no longer shows a nearly full context bar when the provider still reports a modest input size. +- **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/utils/__tests__/analyzeContext.test.ts b/src/utils/__tests__/analyzeContext.test.ts index b68b6603..8f6b0b0a 100644 --- a/src/utils/__tests__/analyzeContext.test.ts +++ b/src/utils/__tests__/analyzeContext.test.ts @@ -1,9 +1,18 @@ 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 { - roughTokenCountEstimation, - roughTokenCountEstimationForBlock, -} from '../../services/tokenEstimation.js' -import { estimateTokensForContextMessages } from '../analyzeContext.js' + estimateTokensForContextMessages, + processAssistantMessage, + processAttachment, + processUserMessage, + type MessageBreakdown, +} from '../analyzeContext.js' const LARGE_BASE64 = 'A'.repeat(800_000) @@ -35,15 +44,43 @@ function largeImageUserMessage() { } } +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 near 2000 tokens instead of hundreds of thousands', () => { + test('counts a large base64 image as 2000 tokens instead of hundreds of thousands', () => { const tokens = estimateTokensForContextMessages([largeImageUserMessage()], []) - expect(tokens).toBeGreaterThanOrEqual(1_500) - expect(tokens).toBeLessThanOrEqual(2_500) - expect(tokens).toBeLessThan(100_000) + expect(tokens).toBe(2000) }) - test('counts a tool_result that holds a large base64 image near 2000 tokens', () => { + test('counts a tool_result that holds a large base64 image as 2000 tokens', () => { const tokens = estimateTokensForContextMessages( [ { @@ -64,9 +101,7 @@ describe('estimateTokensForContextMessages', () => { ], [], ) - expect(tokens).toBeGreaterThanOrEqual(1_500) - expect(tokens).toBeLessThanOrEqual(2_500) - expect(tokens).toBeLessThan(100_000) + expect(tokens).toBe(2000) }) test('counts a large base64 document as 2000 tokens', () => { @@ -106,41 +141,107 @@ describe('estimateTokensForContextMessages', () => { ) const textTokens = roughTokenCountEstimation(text) expect(tokens).toBe(textTokens + 2000) - expect(tokens).toBeLessThan(100_000) }) }) -describe('roughTokenCountEstimationForBlock', () => { - test('scores an image block as 2000 tokens', () => { - expect( - roughTokenCountEstimationForBlock({ - type: 'image', - source: largeImageSource(), - }), - ).toBe(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 - test('scores a document block as 2000 tokens', () => { - expect( - roughTokenCountEstimationForBlock({ - type: 'document', - source: largeDocumentSource(), - }), - ).toBe(2000) + 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) }) +}) - test('scores a tool_result that holds a large image as 2000 tokens', () => { - expect( - roughTokenCountEstimationForBlock({ - type: 'tool_result', - tool_use_id: 'toolu_screenshot', +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(), }, ], - }), - ).toBe(2000) + }, + } 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 41263312..b8637a3c 100644 --- a/src/utils/analyzeContext.ts +++ b/src/utils/analyzeContext.ts @@ -19,6 +19,7 @@ import { countTokensViaHaikuFallback, roughTokenCountEstimation, roughTokenCountEstimationForBlock, + roughTokenCountEstimationForMessage, roughTokenCountEstimationForMessages, } from '../services/tokenEstimation.js' import { estimateSkillFrontmatterTokens } from '../skills/loadSkillsDir.js' @@ -809,7 +810,7 @@ async function countCustomAgentTokens( return { agentTokens, agentDetails } } -type MessageBreakdown = { +export type MessageBreakdown = { totalTokens: number toolCallTokens: number toolResultTokens: number @@ -821,7 +822,7 @@ type MessageBreakdown = { attachmentsByType: Map } -function processAssistantMessage( +export function processAssistantMessage( msg: AssistantMessage | NormalizedAssistantMessage, breakdown: MessageBreakdown, ): void { @@ -843,7 +844,7 @@ function processAssistantMessage( } } -function processUserMessage( +export function processUserMessage( msg: UserMessage | NormalizedUserMessage, breakdown: MessageBreakdown, toolUseIdToName: Map, @@ -876,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(