From c0e4006f1e71e20180ec94663f2395003c2d077d Mon Sep 17 00:00:00 2001 From: Ruiming Zhao Date: Sat, 15 Aug 2026 11:34:35 -0700 Subject: [PATCH] Count the fixed request overhead in context calibration so compression doesn't underestimate --- src/agent.ts | 5 +++ src/context.ts | 31 +++++++++++--- tests/context-calibration.test.ts | 67 +++++++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 6 deletions(-) create mode 100644 tests/context-calibration.test.ts diff --git a/src/agent.ts b/src/agent.ts index 064492c..c5accda 100644 --- a/src/agent.ts +++ b/src/agent.ts @@ -63,6 +63,11 @@ export class Agent { for (const t of this.tools) { if (t instanceof SubAgentTool) t.parentAgent = this } + + // the API bills for the system prompt and tool schemas on every call; + // count them as fixed overhead so the compressor's calibration doesn't + // underestimate once the conversation gets shrunk (see context.ts) + this.context.setFixedOverhead(this.system + '\n' + JSON.stringify(this.toolSchemas())) } private fullMessages(): ChatMessage[] { diff --git a/src/context.ts b/src/context.ts index ce44b7b..2aade4f 100644 --- a/src/context.ts +++ b/src/context.ts @@ -44,6 +44,19 @@ export class ContextManager { */ private ratio = 1 + /** + * Fixed per-request token overhead: the system prompt plus the tool + * schemas. The API bills for these on every call, but they're not part of + * `messages`, so a naive estimate can't see them. Counting them explicitly + * keeps `ratio` a pure chars-per-token rate. Without this, compression + * would silently distort the calibration: after layer 2/3 shrink the + * message list, the fixed overhead makes up a larger share of the real + * usage, the ratio (calibrated on an unshrunk conversation) overcounts it, + * and `measure()` underestimates — pushing the compression thresholds + * later than intended, close to or past the provider's real limit. + */ + private fixedTokens = 0 + constructor(maxTokens = 128_000) { this.maxTokens = maxTokens // layer thresholds (fraction of maxTokens) @@ -52,25 +65,31 @@ export class ContextManager { this.collapseAt = Math.floor(maxTokens * 0.9) // 90% -> hard collapse } + /** Register the fixed per-request text (system prompt + tool schemas). */ + setFixedOverhead(text: string): void { + this.fixedTokens = approxTokens(text) + } + /** * Calibrate the estimator against the real prompt_tokens the API just * reported for `messages`. The char/3 guess is systematically off — CJK - * text runs ~1 token per char, and the estimate can't see the fixed - * system-prompt + tool-schema overhead the API bills for. Scaling by - * real/estimated absorbs both. Survives compression: the ratio is a - * chars-to-tokens rate, still valid after messages shrink. + * text runs ~1 token per char. The fixed overhead is counted explicitly + * (see setFixedOverhead), so the ratio stays a pure chars-per-token rate + * and survives compression: shrinking the messages doesn't change what + * the ratio means. */ observe(realPromptTokens: number, messages: ChatMessage[]): void { if (realPromptTokens <= 0) return - const est = estimateTokens(messages) + const est = estimateTokens(messages) + this.fixedTokens if (est > 0) this.ratio = realPromptTokens / est } /** Best-available token count: char estimate scaled by observed reality. */ measure(messages: ChatMessage[]): number { - return Math.round(estimateTokens(messages) * this.ratio) + return Math.round((estimateTokens(messages) + this.fixedTokens) * this.ratio) } + /** Apply compression layers as needed (mutates `messages` in place). */ async maybeCompress(messages: ChatMessage[], llm?: LLMClient): Promise { let current = this.measure(messages) diff --git a/tests/context-calibration.test.ts b/tests/context-calibration.test.ts new file mode 100644 index 0000000..e57ff9a --- /dev/null +++ b/tests/context-calibration.test.ts @@ -0,0 +1,67 @@ +/** + * Calibration tests: the fixed per-request overhead (system prompt + tool + * schemas) must be counted explicitly, otherwise compression silently + * distorts the estimate and pushes the compression thresholds too late. + */ + +import assert from 'node:assert/strict' +import { test } from 'node:test' + +import { Agent } from '../src/agent.js' +import { ContextManager, estimateTokens } from '../src/context.js' +import { ScriptedLLM, type ChatMessage } from '../src/llm.js' + +const msg = (role: ChatMessage['role'], content: string): ChatMessage => + ({ role, content }) as ChatMessage + +test('fixed overhead is counted in observe and measure', () => { + const cm = new ContextManager(100_000) + cm.setFixedOverhead('system ' + 'x'.repeat(600)) // ~202 tokens of fixed text + + const messages = [msg('user', 'hello')] + // naive estimate: 5 chars / 3 ≈ 1 token; with overhead ≈ 203 + cm.observe(203, messages) + + assert.equal(cm.measure(messages), 203) +}) + +test('measure stays honest after the conversation shrinks', () => { + const cm = new ContextManager(100_000) + cm.setFixedOverhead('system ' + 'x'.repeat(600)) // ~202 tokens fixed + + // a long conversation: 20 messages × ~153 chars each ≈ 1000 estimate tokens + const big: ChatMessage[] = Array.from({ length: 20 }, (_, i) => + msg('user', 'q' + i + ' ' + 'y'.repeat(150)), + ) + const estBig = estimateTokens(big) + assert.ok(estBig > 950 && estBig < 1100, 'sanity: big estimate ~1000, got ' + estBig) + + // the API bills ~1000 + 202 overhead; calibrate against that + const realBig = estBig + 202 + cm.observe(realBig, big) + + // now the conversation got compressed to a tiny summary + const small = [msg('user', '[Context compressed]\nshort summary here')] + + // with fixed overhead counted: estimate = ~10 + 202 = 212, measure ≈ 212 + // without it: ratio ≈ (1000+202)/1000 ≈ 1.2 applied to ~10 → ~12, a ~95% underestimate + const measured = cm.measure(small) + assert.ok( + measured >= 190 && measured <= 230, + 'measure should reflect the fixed overhead, got ' + measured + ' (naive would be ~12)', + ) +}) + +test('without fixed overhead the estimator behaves as before', () => { + const cm = new ContextManager(100_000) + const messages = [msg('user', 'a'.repeat(300))] // ~100 estimate tokens + cm.observe(150, messages) // real is 1.5× the estimate + assert.equal(cm.measure(messages), 150) +}) + +test('agent registers system prompt and tool schemas as fixed overhead', () => { + const agent = new Agent({ llm: new ScriptedLLM([]) }) + // empty history still measures the fixed overhead, not zero + const measured = agent.context.measure(agent.messages) + assert.ok(measured > 0, 'empty history should measure the fixed overhead, got ' + measured) +})