From c1b06fec0dd57652950dfc641158a8e06d29b9a7 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 6 Apr 2026 19:48:44 +0000 Subject: [PATCH] token counting service --- packages/core/src/context/historyObserver.ts | 2 +- packages/core/src/context/ir/mapper.ts | 11 +-- packages/core/src/context/ir/toIr.ts | 43 +++++----- .../semanticCompressionProcessor.ts | 26 +++--- .../processors/toolMaskingProcessor.ts | 21 +---- .../core/src/context/sidecar/environment.ts | 4 +- .../src/context/sidecar/environmentImpl.ts | 8 +- .../src/context/testing/contextTestUtils.ts | 2 + .../context/utils/contextTokenCalculator.ts | 79 ++++++++++++------- 9 files changed, 99 insertions(+), 97 deletions(-) diff --git a/packages/core/src/context/historyObserver.ts b/packages/core/src/context/historyObserver.ts index 560dd11ebe..f46e9f5a0f 100644 --- a/packages/core/src/context/historyObserver.ts +++ b/packages/core/src/context/historyObserver.ts @@ -36,7 +36,7 @@ export class HistoryObserver { this.unsubscribeHistory = this.chatHistory.subscribe((_event: HistoryEvent) => { // Rebuild the pristine IR graph from the full source history on every change. - const pristineEpisodes = IrMapper.toIr(this.chatHistory.get()); + const pristineEpisodes = IrMapper.toIr(this.chatHistory.get(), this.sidecar.tokenCalculator); this.tracer.logEvent('HistoryObserver', 'Rebuilt pristine graph from chat history update', { episodeCount: pristineEpisodes.length }); this.onIrRebuilt(pristineEpisodes); diff --git a/packages/core/src/context/ir/mapper.ts b/packages/core/src/context/ir/mapper.ts index 5a57bdf046..4b0a34f222 100644 --- a/packages/core/src/context/ir/mapper.ts +++ b/packages/core/src/context/ir/mapper.ts @@ -6,20 +6,17 @@ import type { Content } from '@google/genai'; import type { Episode } from './types.js'; -import { toIr, setMapperConfig } from './toIr.js'; +import { toIr } from './toIr.js'; import { fromIr } from './fromIr.js'; +import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; export class IrMapper { - static setConfig(cfg: { charsPerToken?: number }) { - setMapperConfig(cfg); - } - /** * Translates a flat Gemini Content[] array into our rich Episodic Intermediate Representation. * Groups adjacent function calls and responses into unified ToolExecution nodes. */ - static toIr(history: readonly Content[]): Episode[] { - return toIr(history); + static toIr(history: readonly Content[], tokenCalculator: ContextTokenCalculator): Episode[] { + return toIr(history, tokenCalculator); } /** diff --git a/packages/core/src/context/ir/toIr.ts b/packages/core/src/context/ir/toIr.ts index a7f4b28ce1..e4c956f060 100644 --- a/packages/core/src/context/ir/toIr.ts +++ b/packages/core/src/context/ir/toIr.ts @@ -16,7 +16,7 @@ import type { UserPrompt, SystemEvent, } from './types.js'; -import { estimateContextTokenCountSync } from '../utils/contextTokenCalculator.js'; +import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; // WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references const nodeIdentityMap = new WeakMap(); @@ -30,26 +30,20 @@ export function getStableId(obj: object): string { return id; } -export let charsPerTokenConfig: { charsPerToken?: number } | undefined; - -export function setMapperConfig(cfg: { charsPerToken?: number }) { - charsPerTokenConfig = cfg; -} - -export function createMetadata(parts: Part[]): IrMetadata { - const tokens = estimateContextTokenCountSync(parts, 0, charsPerTokenConfig); - return { - originalTokens: tokens, - currentTokens: tokens, - transformations: [], - }; -} - -export function toIr(history: readonly Content[]): Episode[] { +export function toIr(history: readonly Content[], tokenCalculator: ContextTokenCalculator): Episode[] { const episodes: Episode[] = []; let currentEpisode: Partial | null = null; const pendingCallParts: Map = new Map(); + const createMetadata = (parts: Part[]): IrMetadata => { + const tokens = tokenCalculator.estimateTokensForParts(parts, 0); + return { + originalTokens: tokens, + currentTokens: tokens, + transformations: [], + }; + }; + const finalizeEpisode = () => { if (currentEpisode && currentEpisode.trigger) { episodes.push(currentEpisode as unknown as Episode); // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion @@ -67,15 +61,15 @@ export function toIr(history: readonly Content[]): Episode[] { ); if (hasToolResponses) { - currentEpisode = parseToolResponses(msg, currentEpisode, pendingCallParts); + currentEpisode = parseToolResponses(msg, currentEpisode, pendingCallParts, tokenCalculator, createMetadata); } if (hasUserParts) { finalizeEpisode(); - currentEpisode = parseUserParts(msg); + currentEpisode = parseUserParts(msg, createMetadata); } } else if (msg.role === 'model') { - currentEpisode = parseModelParts(msg, currentEpisode, pendingCallParts); + currentEpisode = parseModelParts(msg, currentEpisode, pendingCallParts, createMetadata); } } @@ -91,6 +85,8 @@ function parseToolResponses( msg: Content, currentEpisode: Partial | null, pendingCallParts: Map, + tokenCalculator: ContextTokenCalculator, + createMetadata: (parts: Part[]) => IrMetadata ): Partial { if (!currentEpisode) { currentEpisode = { @@ -113,9 +109,9 @@ function parseToolResponses( const matchingCall = pendingCallParts.get(callId); const intentTokens = matchingCall - ? estimateContextTokenCountSync([matchingCall]) + ? tokenCalculator.estimateTokensForParts([matchingCall]) : 0; - const obsTokens = estimateContextTokenCountSync([part]); + const obsTokens = tokenCalculator.estimateTokensForParts([part]); const step: ToolExecution = { id: getStableId(part), @@ -150,7 +146,7 @@ function parseToolResponses( return currentEpisode; } -function parseUserParts(msg: Content): Partial { +function parseUserParts(msg: Content, createMetadata: (parts: Part[]) => IrMetadata): Partial { const semanticParts: SemanticPart[] = []; for (const p of msg.parts!) { if (p.text !== undefined) @@ -192,6 +188,7 @@ function parseModelParts( msg: Content, currentEpisode: Partial | null, pendingCallParts: Map, + createMetadata: (parts: Part[]) => IrMetadata ): Partial { if (!currentEpisode) { currentEpisode = { diff --git a/packages/core/src/context/processors/semanticCompressionProcessor.ts b/packages/core/src/context/processors/semanticCompressionProcessor.ts index 7c4a56546a..d38fe1d1f0 100644 --- a/packages/core/src/context/processors/semanticCompressionProcessor.ts +++ b/packages/core/src/context/processors/semanticCompressionProcessor.ts @@ -10,7 +10,7 @@ import type { ContextEnvironment } from '../sidecar/environment.js'; import { debugLogger } from '../../utils/debugLogger.js'; import { LlmRole } from '../../telemetry/types.js'; import { getResponseText } from '../../utils/partUtils.js'; -import { estimateTokenCountSync } from '../../utils/tokenCalculation.js'; +import { estimateContextTokenCountSync } from '../utils/contextTokenCalculator.js'; export class SemanticCompressionProcessor implements ContextProcessor { readonly name = 'SemanticCompression'; @@ -30,22 +30,14 @@ export class SemanticCompressionProcessor implements ContextProcessor { episodes: Episode[], state: ContextAccountingState, ): Promise { - require('fs').appendFileSync( - '/tmp/debug2.json', - 'SEMANTIC PROCESS: First episode ID: ' + - episodes[0]?.id + - '\nProtected IDs: ' + - Array.from(state.protectedEpisodeIds).join(', ') + - '\n', - ); // If the budget is satisfied, or semantic compression isn't enabled if (state.isBudgetSatisfied) { return episodes; } const semanticConfig = this.options; - // We estimate 4 chars per token for truncation logic - const thresholdChars = semanticConfig.nodeThresholdTokens * 4; + const limitTokens = semanticConfig.nodeThresholdTokens; + const thresholdChars = limitTokens * this.env.charsPerToken; this.modelToUse = 'gemini-2.5-flash'; let currentDeficit = state.deficitTokens; @@ -70,8 +62,8 @@ export class SemanticCompressionProcessor implements ContextProcessor { part.text, 'User Prompt', ); - const newTokens = estimateTokenCountSync([{ text: summary }]); - const oldTokens = estimateTokenCountSync([{ text: part.text }]); + const newTokens = estimateContextTokenCountSync([{ text: summary }], 0, { charsPerToken: this.env.charsPerToken }); + const oldTokens = estimateContextTokenCountSync([{ text: part.text }], 0, { charsPerToken: this.env.charsPerToken }); if (newTokens < oldTokens) { part.presentation = { text: summary, tokens: newTokens }; @@ -96,8 +88,8 @@ export class SemanticCompressionProcessor implements ContextProcessor { step.text, 'Agent Thought', ); - const newTokens = estimateTokenCountSync([{ text: summary }]); - const oldTokens = estimateTokenCountSync([{ text: step.text }]); + const newTokens = estimateContextTokenCountSync([{ text: summary }], 0, { charsPerToken: this.env.charsPerToken }); + const oldTokens = estimateContextTokenCountSync([{ text: step.text }], 0, { charsPerToken: this.env.charsPerToken }); if (newTokens < oldTokens) { step.presentation = { text: summary, tokens: newTokens }; @@ -138,7 +130,7 @@ export class SemanticCompressionProcessor implements ContextProcessor { // Wrap the summary in an object so the Gemini API accepts it as a valid functionResponse.response const newObsObject = { summary }; - const newObsTokens = estimateTokenCountSync([ + const newObsTokens = estimateContextTokenCountSync([ { functionResponse: { name: step.toolName, @@ -146,7 +138,7 @@ export class SemanticCompressionProcessor implements ContextProcessor { id: step.id, }, }, - ]); + ], 0, { charsPerToken: this.env.charsPerToken }); const oldObsTokens = step.presentation?.tokens.observation ?? step.tokens.observation; diff --git a/packages/core/src/context/processors/toolMaskingProcessor.ts b/packages/core/src/context/processors/toolMaskingProcessor.ts index e254600d8c..c50788fb83 100644 --- a/packages/core/src/context/processors/toolMaskingProcessor.ts +++ b/packages/core/src/context/processors/toolMaskingProcessor.ts @@ -6,7 +6,7 @@ import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; -import { estimateTokenCountSync } from '../../utils/tokenCalculation.js'; +import { estimateContextTokenCountSync } from '../utils/contextTokenCalculator.js'; import { sanitizeFilenamePart } from '../../utils/fileUtils.js'; import * as fsPromises from 'node:fs/promises'; import path from 'node:path'; @@ -50,7 +50,7 @@ export class ToolMaskingProcessor implements ContextProcessor { const newEpisodes = [...episodes]; let currentDeficit = state.deficitTokens; - const limitChars = maskingConfig.stringLengthThresholdTokens * 4; + const limitChars = maskingConfig.stringLengthThresholdTokens * this.env.charsPerToken; let toolOutputsDir = path.join( this.env.projectTempDir, @@ -125,23 +125,6 @@ export class ToolMaskingProcessor implements ContextProcessor { nodeType: string, ): Promise<{ masked: any; changed: boolean }> => { if (typeof obj === 'string') { - require('fs').appendFileSync( - '/tmp/debug.json', - 'STRING FOUND. length: ' + - obj.length + - ' limitChars: ' + - limitChars + - '\n', - ); - if (obj.length > 1000) - console.log( - 'Found string of length:', - obj.length, - 'limitChars is:', - limitChars, - 'isAlreadyMasked:', - this.isAlreadyMasked(obj), - ); if (obj.length > limitChars && !this.isAlreadyMasked(obj)) { const newString = await handleMasking( obj, diff --git a/packages/core/src/context/sidecar/environment.ts b/packages/core/src/context/sidecar/environment.ts index f1e0f99b23..a113fe707f 100644 --- a/packages/core/src/context/sidecar/environment.ts +++ b/packages/core/src/context/sidecar/environment.ts @@ -6,6 +6,7 @@ import type { BaseLlmClient } from '../../core/baseLlmClient.js'; import type { ContextTracer } from '../tracer.js'; import type { ContextEventBus } from '../eventBus.js'; +import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; export type { ContextTracer, ContextEventBus }; export interface ContextEnvironment { @@ -16,6 +17,7 @@ readonly projectTempDir: string; readonly tracer: ContextTracer; readonly charsPerToken: number; + readonly tokenCalculator: ContextTokenCalculator; - readonly eventBus: ContextEventBus; + eventBus: ContextEventBus; } diff --git a/packages/core/src/context/sidecar/environmentImpl.ts b/packages/core/src/context/sidecar/environmentImpl.ts index fabd089cfc..52bf9dce0b 100644 --- a/packages/core/src/context/sidecar/environmentImpl.ts +++ b/packages/core/src/context/sidecar/environmentImpl.ts @@ -10,7 +10,11 @@ import type { ContextEnvironment } from './environment.js'; import type { ContextEventBus } from '../eventBus.js'; +import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; + export class ContextEnvironmentImpl implements ContextEnvironment { + public readonly tokenCalculator: ContextTokenCalculator; + constructor( public readonly llmClient: BaseLlmClient, public readonly sessionId: string, @@ -20,5 +24,7 @@ export class ContextEnvironmentImpl implements ContextEnvironment { public readonly tracer: ContextTracer, public readonly charsPerToken: number, public readonly eventBus: ContextEventBus, - ) {} + ) { + this.tokenCalculator = new ContextTokenCalculator(this.charsPerToken); + } } diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index 9001bc3f4b..6a855c783f 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -26,6 +26,7 @@ export function createMockEnvironment(): ContextEnvironment { eventBus: new ContextEventBus(), tracer: new ContextTracer('/tmp', 'mock-session'), charsPerToken: 1, + tokenCalculator: new ContextTokenCalculator(1), }; } @@ -89,6 +90,7 @@ import { ContextTracer } from '../tracer.js'; import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js'; import { SidecarLoader } from '../sidecar/SidecarLoader.js'; import { ContextEventBus } from '../eventBus.js'; +import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; import type { BaseLlmClient } from 'src/core/baseLlmClient.js'; export function setupContextComponentTest(config: Config) { diff --git a/packages/core/src/context/utils/contextTokenCalculator.ts b/packages/core/src/context/utils/contextTokenCalculator.ts index 1ed5fb452e..7590965775 100644 --- a/packages/core/src/context/utils/contextTokenCalculator.ts +++ b/packages/core/src/context/utils/contextTokenCalculator.ts @@ -8,37 +8,60 @@ import type { Part } from '@google/genai'; import { estimateTokenCountSync as baseEstimate } from '../../utils/tokenCalculation.js'; import type { Episode } from '../ir/types.js'; -export function calculateEpisodeListTokens(episodes: Episode[]): number { - let tokens = 0; - for (const ep of episodes) { - if (ep.trigger) tokens += ep.trigger.metadata.currentTokens; - for (const step of ep.steps) { - tokens += step.metadata.currentTokens; - } - if (ep.yield) tokens += ep.yield.metadata.currentTokens; - } - return tokens; -} +export class ContextTokenCalculator { + constructor(private readonly charsPerToken: number) {} -export function estimateContextTokenCountSync( - parts: Part[], - depth: number = 0, - config?: { charsPerToken?: number }, -): number { - if (config?.charsPerToken !== undefined && config.charsPerToken !== 4) { - let totalTokens = 0; - for (const part of parts) { - if (typeof part.text === 'string') { - totalTokens += Math.ceil(part.text.length / config.charsPerToken); - } else { - totalTokens += Math.ceil( - JSON.stringify(part).length / config.charsPerToken, - ); + /** + * Fast, simple heuristic estimation for a raw string. + */ + estimateTokensForString(text: string): number { + return Math.ceil(text.length / this.charsPerToken); + } + + /** + * Fast, simple heuristic conversion from tokens to expected character length. + * Useful for calculating truncation thresholds. + */ + tokensToChars(tokens: number): number { + return tokens * this.charsPerToken; + } + + /** + * Calculates the total token count for a complete Episodic IR graph. + * This is fast because it relies on pre-computed metadata where available. + */ + calculateEpisodeListTokens(episodes: Episode[]): number { + let tokens = 0; + for (const ep of episodes) { + if (ep.trigger) tokens += ep.trigger.metadata.currentTokens; + for (const step of ep.steps) { + tokens += step.metadata.currentTokens; } + if (ep.yield) tokens += ep.yield.metadata.currentTokens; } - return totalTokens; + return tokens; } - // The baseEstimate no longer accepts config because we forked it! - return baseEstimate(parts, depth); + /** + * Slower, precise estimation for a Gemini Content/Part graph. + * Deeply inspects the nested structure and uses the base tokenization math. + */ + estimateTokensForParts(parts: Part[], depth: number = 0): number { + if (this.charsPerToken !== 4) { + let totalTokens = 0; + for (const part of parts) { + if (typeof part.text === 'string') { + totalTokens += Math.ceil(part.text.length / this.charsPerToken); + } else { + totalTokens += Math.ceil( + JSON.stringify(part).length / this.charsPerToken, + ); + } + } + return totalTokens; + } + + // The baseEstimate no longer accepts config because we forked it! + return baseEstimate(parts, depth); + } }