From 734dac8b730c6f57471380d01be8acd7d90a5112 Mon Sep 17 00:00:00 2001 From: TianYi <2739022972@qq.com> Date: Tue, 23 Jun 2026 00:33:06 +0800 Subject: [PATCH 1/2] feat: context-compression and budget --- src/agent/agent.ts | 1 + src/agent/runtime/context-manager.ts | 132 ++++++- .../context-formatter.ts | 9 - .../system-prompt-assembler.ts | 4 +- src/context/budget/context-budget.ts | 96 ++++- src/context/budget/priority-policy.ts | 6 +- .../context-compression-pipeline.ts | 131 ++++++- src/context/compression/context-compressor.ts | 45 +-- src/context/compression/context-truncator.ts | 224 ----------- src/context/engine/context-engine.ts | 144 ++++++- src/context/engine/context-registry.ts | 3 + src/context/index.ts | 22 +- src/context/truncate/context-truncator.ts | 353 ++++++++++++++++++ src/context/types.ts | 73 +++- tests/agent/runtime-loop.test.ts | 6 +- tests/agent/system-prompt.test.ts | 4 +- tests/context-builder.test.ts | 260 ++++++++----- 17 files changed, 1079 insertions(+), 434 deletions(-) rename src/context/{formatting => assembler}/context-formatter.ts (53%) rename src/context/{formatting => assembler}/system-prompt-assembler.ts (82%) delete mode 100644 src/context/compression/context-truncator.ts create mode 100644 src/context/truncate/context-truncator.ts diff --git a/src/agent/agent.ts b/src/agent/agent.ts index b22f9cd..de2846f 100644 --- a/src/agent/agent.ts +++ b/src/agent/agent.ts @@ -147,6 +147,7 @@ export class Agent { memory: this.memory, observer: this.observer, contextProviders: options.contextProviders, + toolSchemasProvider: () => this.tools.schemas(), }); this.verification = options.verification ?? diff --git a/src/agent/runtime/context-manager.ts b/src/agent/runtime/context-manager.ts index c63780a..4fb5763 100644 --- a/src/agent/runtime/context-manager.ts +++ b/src/agent/runtime/context-manager.ts @@ -1,4 +1,4 @@ -import type {LLMMessage} from "../../llm/types.js"; +import type {LLMMessage, LLMTool, LLMToolCall} from "../../llm/types.js"; import { ContextEngine, RuleBasedContextCompressor, @@ -34,6 +34,8 @@ export type ContextManagerOptions = { observer: AgentObserver; /** Agent-level context providers invoked for every run. */ contextProviders?: readonly AgentContextProvider[]; + /** Returns the tool schemas that will be sent with model requests. */ + toolSchemasProvider?: () => readonly LLMTool[]; }; /** Builds model context and owns transcript mutation for assistant/tool turns. */ @@ -117,20 +119,26 @@ export class ContextManager { buildModelRequest(run: AgentRunState): readonly LLMMessage[] { const state = this.getRunState(run); const projection = this.projectTranscript(run); + const toolSchemas = this.options.toolSchemasProvider?.() ?? []; const result = this.contextEngine.build({ - systemPrompt: + baseSystemPrompt: run.input.systemPrompt ?? this.options.config.runtime.systemPrompt ?? DEFAULT_SYSTEM_PROMPT, outputInstructions: CLI_MARKDOWN_OUTPUT_INSTRUCTIONS, sections: [...state.baseSections, ...projection.archivedToolSections], tokenLimit: this.options.config.runtime.tokensLimit, + conversationMessages: projection.messages, + toolSchemas, }); this.recordContextBuild(run, state, result.diagnostics); this.options.observer.contextBuilt(run, result.tokenEstimate); - return [{role: "system", content: result.systemPrompt}, ...projection.messages]; + return [ + {role: "system", content: result.assembledSystemPrompt}, + ...projection.messages, + ]; } /** Appends a normalized assistant response to the transcript. */ @@ -226,9 +234,10 @@ export class ContextManager { messages.push(message, ...exchange.toolMessages); } else { archivedToolSections.push( - ...exchange.toolMessages.map((toolMessage) => - createToolResultSection(message, toolMessage), - ), + ...createArchivedToolSections(message, exchange.toolMessages, { + exchangeAge: + latestExchangeStart === undefined ? 0 : latestExchangeStart - index, + }), ); } index = exchange.endIndex; @@ -336,9 +345,62 @@ function findLatestCompletedToolExchangeStart( return undefined; } +function createArchivedToolSections( + assistant: AssistantToolMessage, + toolMessages: readonly ToolMessage[], + metadata: {exchangeAge: number}, +): ContextSection[] { + return toolMessages.flatMap((toolMessage) => { + const toolCall = assistant.toolCalls.find( + (candidate) => candidate.id === toolMessage.toolCallId, + ); + + return [ + createToolArgsSection(toolCall, toolMessage, metadata), + createToolResultSection(assistant, toolMessage, metadata), + ]; + }); +} + +function createToolArgsSection( + toolCall: LLMToolCall | undefined, + toolMessage: ToolMessage, + metadata: {exchangeAge: number}, +): ContextSection { + const argumentsText = JSON.stringify( + redactToolArguments(toolMessage.name, toolCall?.arguments ?? {}), + null, + 2, + ); + + return { + id: `tool-args:${toolMessage.toolCallId}`, + replaceKey: `tool-args:${toolMessage.toolCallId}`, + title: `Tool Args: ${toolMessage.name}`, + priority: 45, + source: {kind: "tool", ref: toolMessage.toolCallId}, + compressible: false, + truncation: "none", + metadata: { + toolName: toolMessage.name, + toolCallId: toolMessage.toolCallId, + toolPart: "args", + exchangeAge: metadata.exchangeAge, + safeToCompress: false, + }, + content: [ + `Tool: ${toolMessage.name}`, + `Call ID: ${toolMessage.toolCallId}`, + "Arguments:", + argumentsText, + ].join("\n"), + }; +} + function createToolResultSection( assistant: AssistantToolMessage, toolMessage: ToolMessage, + metadata: {exchangeAge: number}, ): ContextSection { return { id: `tool-result:${toolMessage.toolCallId}`, @@ -346,6 +408,16 @@ function createToolResultSection( title: `Tool Result: ${toolMessage.name}`, priority: 40, source: {kind: "tool", ref: toolMessage.toolCallId}, + compressible: true, + truncation: "middle", + metadata: { + toolName: toolMessage.name, + toolCallId: toolMessage.toolCallId, + toolPart: "result", + toolResultKind: classifyToolResultKind(toolMessage.name, toolMessage.content), + exchangeAge: metadata.exchangeAge, + safeToCompress: isSafeToolResultToCompress(toolMessage.name), + }, content: [ `Tool: ${toolMessage.name}`, `Call ID: ${toolMessage.toolCallId}`, @@ -357,3 +429,51 @@ function createToolResultSection( .join("\n"), }; } + +function redactToolArguments( + toolName: string, + args: Record, +): Record { + if (toolName !== "write_file") { + return args; + } + + const content = args.content; + if (typeof content !== "string" || content.length === 0) { + return args; + } + + return { + ...args, + content: ``, + }; +} + +function classifyToolResultKind(toolName: string, content: string): string { + if (toolName === "bash") { + return content.toLowerCase().includes("test") ? "test" : "bash"; + } + + if (toolName === "grep") { + return "grep"; + } + + if (toolName === "read_file") { + return "read_file"; + } + + return "old_tool_result"; +} + +function isSafeToolResultToCompress(toolName: string): boolean { + return ( + toolName === "bash" || + toolName === "grep" || + toolName === "read_file" || + toolName === "glob" || + toolName === "web_fetch" || + classifyToolResultKind(toolName, "") === "old_tool_result" + ); +} diff --git a/src/context/formatting/context-formatter.ts b/src/context/assembler/context-formatter.ts similarity index 53% rename from src/context/formatting/context-formatter.ts rename to src/context/assembler/context-formatter.ts index 7c74a7f..b18be71 100644 --- a/src/context/formatting/context-formatter.ts +++ b/src/context/assembler/context-formatter.ts @@ -1,4 +1,3 @@ -import {DefaultContextPriorityPolicy} from "../budget/priority-policy.js"; import type {ContextSection} from "../types.js"; /** Returns a formatted section string, or an empty string for blank content. */ @@ -11,11 +10,3 @@ export function formatContextSection(section: ContextSection): string { return section.title ? `## ${section.title}\n${content}` : content; } - -/** Sorts context sections from highest to lowest priority. */ -export function compareContextSection( - left: ContextSection, - right: ContextSection, -): number { - return new DefaultContextPriorityPolicy().compare(left, right); -} diff --git a/src/context/formatting/system-prompt-assembler.ts b/src/context/assembler/system-prompt-assembler.ts similarity index 82% rename from src/context/formatting/system-prompt-assembler.ts rename to src/context/assembler/system-prompt-assembler.ts index 91b9b9d..ad5ba36 100644 --- a/src/context/formatting/system-prompt-assembler.ts +++ b/src/context/assembler/system-prompt-assembler.ts @@ -1,11 +1,11 @@ /** Assembles the final system prompt from base prompt, output rules, and context. */ export class SystemPromptAssembler { assemble(input: { - systemPrompt?: string; + baseSystemPrompt?: string; outputInstructions?: string; contextText: string; }): string { - const promptParts = [input.systemPrompt, input.outputInstructions].filter( + const promptParts = [input.baseSystemPrompt, input.outputInstructions].filter( (part): part is string => Boolean(part), ); const prompt = promptParts.join("\n\n"); diff --git a/src/context/budget/context-budget.ts b/src/context/budget/context-budget.ts index d75635e..fb93453 100644 --- a/src/context/budget/context-budget.ts +++ b/src/context/budget/context-budget.ts @@ -1,25 +1,26 @@ +import {createDefaultTokenEstimator, type TokenEstimator} from "./token-estimator.js"; import type {BuildContextInput, ContextBudget} from "../types.js"; /** Strategy used to derive runtime context budget from build input. */ export interface ContextBudgetPolicy { - createBudget(input: BuildContextInput): ContextBudget; + createBudget(input: BuildContextInput, tokenEstimator?: TokenEstimator): ContextBudget; } export type DefaultContextBudgetPolicyOptions = { defaultMaxContextTokens?: number; reservedOutputTokens?: number; - /** @deprecated kept only to populate legacy char diagnostics. */ - runtimeContextRatio?: number; + safetyMarginTokens?: number; }; const DEFAULT_MAX_CONTEXT_TOKENS = 128_000; const DEFAULT_RESERVED_OUTPUT_TOKENS = 4_000; +const DEFAULT_SAFETY_MARGIN_TOKENS = 0; /** Default budget policy that derives an input token budget from the model limit. */ export class DefaultContextBudgetPolicy implements ContextBudgetPolicy { private readonly defaultMaxContextTokens: number; private readonly reservedOutputTokens: number; - private readonly runtimeContextRatio: number; + private readonly safetyMarginTokens: number; constructor(options: DefaultContextBudgetPolicyOptions = {}) { this.defaultMaxContextTokens = positiveIntegerOrDefault( @@ -30,29 +31,98 @@ export class DefaultContextBudgetPolicy implements ContextBudgetPolicy { options.reservedOutputTokens, DEFAULT_RESERVED_OUTPUT_TOKENS, ); - this.runtimeContextRatio = options.runtimeContextRatio ?? 0.35; + this.safetyMarginTokens = nonNegativeIntegerOrDefault( + options.safetyMarginTokens, + DEFAULT_SAFETY_MARGIN_TOKENS, + ); } - createBudget(input: BuildContextInput): ContextBudget { - const maxContextTokens = + createBudget( + input: BuildContextInput, + tokenEstimator: TokenEstimator = createDefaultTokenEstimator(), + ): ContextBudget { + const modelContextWindow = input.tokenLimit > 0 ? Math.floor(input.tokenLimit) : this.defaultMaxContextTokens; const reservedOutputTokens = Math.min( this.reservedOutputTokens, - Math.floor(maxContextTokens * 0.2), + Math.floor(modelContextWindow * 0.2), + ); + const inputTokenBudget = Math.max(0, modelContextWindow - reservedOutputTokens); + // Count non-runtime prompt components before allocating runtime context. + const systemPromptTokens = tokenEstimator.countText(input.baseSystemPrompt ?? ""); + const outputInstructionTokens = tokenEstimator.countText( + input.outputInstructions ?? "", + ); + const toolSchemaTokens = + input.toolSchemaTokens ?? + estimateSerializableTokens(input.toolSchemas ?? [], tokenEstimator); + const conversationTokens = + input.conversationTokens ?? + estimateConversationTokens(input.conversationMessages ?? [], tokenEstimator); + const safetyMarginTokens = nonNegativeIntegerOrDefault( + input.safetyMarginTokens, + this.safetyMarginTokens, + ); + // Runtime context receives only the input tokens left after fixed prompt costs. + const runtimeContextTokens = Math.max( + 0, + inputTokenBudget - + systemPromptTokens - + outputInstructionTokens - + toolSchemaTokens - + conversationTokens - + safetyMarginTokens, ); - const maxInputTokens = Math.max(0, maxContextTokens - reservedOutputTokens); return { tokenLimit: input.tokenLimit, - maxContextTokens, + modelContextWindow, reservedOutputTokens, - maxInputTokens, - runtimeContextRatio: this.runtimeContextRatio, - maxContextChars: Math.max(0, maxInputTokens * 4), + systemPromptTokens, + outputInstructionTokens, + toolSchemaTokens, + conversationTokens, + safetyMarginTokens, + runtimeContextTokens, + finalInputTokens: undefined, + remainingInputTokens: undefined, }; } } +function estimateConversationTokens( + messages: NonNullable, + tokenEstimator: TokenEstimator, +): number { + if (!messages.length) { + return 0; + } + + return ( + tokenEstimator.countMessages?.(messages) ?? + messages.reduce( + (total, message) => + total + tokenEstimator.countText(`${message.role}\n${message.content ?? ""}`), + 0, + ) + ); +} + +function estimateSerializableTokens( + value: unknown, + tokenEstimator: TokenEstimator, +): number { + if (Array.isArray(value) && value.length === 0) { + return 0; + } + + try { + return tokenEstimator.countText(JSON.stringify(value)); + } catch { + return tokenEstimator.countText(String(value)); + } +} + function positiveIntegerOrDefault(value: number | undefined, fallback: number): number { if (value === undefined || !Number.isFinite(value) || value <= 0) { return fallback; diff --git a/src/context/budget/priority-policy.ts b/src/context/budget/priority-policy.ts index 3f3cf23..64be501 100644 --- a/src/context/budget/priority-policy.ts +++ b/src/context/budget/priority-policy.ts @@ -9,7 +9,11 @@ export interface ContextPriorityPolicy { /** Default priority policy with explicit priority taking precedence. */ export class DefaultContextPriorityPolicy implements ContextPriorityPolicy { priorityOf(section: ContextSection): number { - return section.priority ?? priorityForSource(section.source); + const basePriority = section.priority ?? priorityForSource(section.source); + const requiredBoost = section.required ? 10_000 : 0; + const pinnedBoost = section.pinned ? 1_000 : 0; + + return basePriority + requiredBoost + pinnedBoost; } compare(left: ContextSection, right: ContextSection): number { diff --git a/src/context/compression/context-compression-pipeline.ts b/src/context/compression/context-compression-pipeline.ts index 935e624..0254f8c 100644 --- a/src/context/compression/context-compression-pipeline.ts +++ b/src/context/compression/context-compression-pipeline.ts @@ -1,18 +1,19 @@ import { ContextCompressionResultFactory, - NoopContextCompressor, + RuleBasedContextCompressor, type ContextCompressor, } from "./context-compressor.js"; import { createDefaultTokenEstimator, type TokenEstimator, } from "../budget/token-estimator.js"; -import {formatContextSection} from "../formatting/context-formatter.js"; +import {formatContextSection} from "../assembler/context-formatter.js"; import type {ContextBudget, ContextCompressionResult, ContextSection} from "../types.js"; export type ContextCompressionPipelineOptions = { compressor?: ContextCompressor; thresholdRatio?: number; + minSectionTokens?: number; resultFactory?: ContextCompressionResultFactory; tokenEstimator?: TokenEstimator; }; @@ -31,37 +32,53 @@ export type ContextCompressionPipelineResult = { export class ContextCompressionPipeline { private readonly compressor: ContextCompressor; private readonly thresholdRatio: number; + private readonly minSectionTokens: number; private readonly resultFactory: ContextCompressionResultFactory; private readonly tokenEstimator: TokenEstimator; constructor(options: ContextCompressionPipelineOptions = {}) { - this.resultFactory = options.resultFactory ?? new ContextCompressionResultFactory(); - this.compressor = options.compressor ?? new NoopContextCompressor(this.resultFactory); - this.thresholdRatio = options.thresholdRatio ?? 0.85; this.tokenEstimator = options.tokenEstimator ?? createDefaultTokenEstimator(); + this.resultFactory = + options.resultFactory ?? new ContextCompressionResultFactory(this.tokenEstimator); + // Default to real rule-based compression; callers can inject a custom strategy. + this.compressor = + options.compressor ?? + new RuleBasedContextCompressor( + {tokenEstimator: this.tokenEstimator}, + this.resultFactory, + ); + this.thresholdRatio = clampNumber(options.thresholdRatio ?? 0.85, 0, 1); + this.minSectionTokens = nonNegativeIntegerOrDefault(options.minSectionTokens, 200); } compress( sections: readonly ContextSection[], budget: ContextBudget, ): ContextCompressionPipelineResult { + // Estimate the full formatted context before deciding whether compression is needed. const formattedContextText = formatContextSections(sections); const estimatedContextChars = formattedContextText.length; const estimatedContextTokens = this.tokenEstimator.countText(formattedContextText); const compressionLimitTokens = Math.floor( - budget.maxInputTokens * this.thresholdRatio, + budget.runtimeContextTokens * this.thresholdRatio, ); const triggered = estimatedContextTokens > compressionLimitTokens; - const results = triggered - ? sections.map((section) => this.compressor.compress(section, budget)) - : sections.map((section) => - this.resultFactory.skipped(section, "Compression was not triggered.", { - originalTokens: this.tokenEstimator.countText(section.content), - compressedTokens: this.tokenEstimator.countText(section.content), - savedTokens: 0, - tokenCompressionRatio: section.content ? 1 : 0, - }), - ); + const results = sections.map((section) => { + const originalTokens = this.tokenEstimator.countText(section.content); + // Policy skips are recorded as compression results for diagnostics. + const skipReason = this.skipReason(section, originalTokens, triggered); + + if (skipReason) { + return this.resultFactory.skipped(section, skipReason, { + originalTokens, + compressedTokens: originalTokens, + savedTokens: 0, + tokenCompressionRatio: section.content ? 1 : 0, + }); + } + + return this.compressor.compress(section, budget); + }); return { sections: results.map((result) => result.section), @@ -73,6 +90,36 @@ export class ContextCompressionPipeline { triggered, }; } + + private skipReason( + section: ContextSection, + originalTokens: number, + triggered: boolean, + ): string | undefined { + // Global threshold prevents unnecessary work when context already fits comfortably. + if (!triggered) { + return "Compression was not triggered because runtime context is below the configured total threshold."; + } + + if (section.compressible === false) { + return "Section opted out of compression."; + } + + if (section.required && section.truncation === "none") { + return "Required section disallows truncation and is protected from compression."; + } + + if (originalTokens < this.minSectionTokens) { + return `Section is below the minimum compression threshold (${originalTokens} < ${this.minSectionTokens} tokens).`; + } + + // Only compress source types whose content can be safely summarized or abbreviated. + if (!isSafeToCompress(section)) { + return "Section type is not safe to compress."; + } + + return undefined; + } } function formatContextSections(sections: readonly ContextSection[]): string { @@ -81,3 +128,55 @@ function formatContextSections(sections: readonly ContextSection[]): string { .filter((text) => text.length > 0) .join("\n\n"); } + +function isSafeToCompress(section: ContextSection): boolean { + if (section.metadata?.safeToCompress === true) { + return true; + } + + if (section.metadata?.safeToCompress === false) { + return false; + } + + if (section.source?.kind === "file") { + return true; + } + + if (section.source?.kind !== "tool") { + return false; + } + + const toolPart = section.metadata?.toolPart; + if (toolPart && toolPart !== "result") { + return false; + } + + const kind = section.metadata?.toolResultKind; + return ( + kind === "bash" || + kind === "test" || + kind === "grep" || + kind === "read_file" || + kind === "old_tool_result" || + kind === undefined + ); +} + +function clampNumber(value: number, min: number, max: number): number { + if (!Number.isFinite(value)) { + return min; + } + + return Math.min(max, Math.max(min, value)); +} + +function nonNegativeIntegerOrDefault( + value: number | undefined, + fallback: number, +): number { + if (value === undefined || !Number.isFinite(value) || value < 0) { + return fallback; + } + + return Math.floor(value); +} diff --git a/src/context/compression/context-compressor.ts b/src/context/compression/context-compressor.ts index c6fe3a2..49a0194 100644 --- a/src/context/compression/context-compressor.ts +++ b/src/context/compression/context-compressor.ts @@ -106,19 +106,6 @@ export class ContextCompressionResultFactory { } } -/** Default compressor that leaves sections unchanged. */ -export class NoopContextCompressor implements ContextCompressor { - private readonly resultFactory: ContextCompressionResultFactory; - - constructor(resultFactory = new ContextCompressionResultFactory()) { - this.resultFactory = resultFactory; - } - - compress(section: ContextSection, _budget: ContextBudget): ContextCompressionResult { - return this.resultFactory.unchanged(section, "No compression applied."); - } -} - export type RuleBasedContextCompressorOptions = { maxSectionTokens?: number; minSectionTokens?: number; @@ -179,6 +166,17 @@ export class RuleBasedContextCompressor implements ContextCompressor { } compress(section: ContextSection, budget: ContextBudget): ContextCompressionResult { + if (section.compressible === false) { + return this.resultFactory.unchanged(section, "Section opted out of compression."); + } + + if (section.required && section.truncation === "none") { + return this.resultFactory.unchanged( + section, + "Required section disallows truncation and is protected from compression.", + ); + } + if (!isCompressibleSection(section)) { return this.resultFactory.unchanged(section, "Section source is not compressible."); } @@ -211,7 +209,9 @@ export class RuleBasedContextCompressor implements ContextCompressor { } private resolveMaxSectionTokens(budget: ContextBudget): number { - const budgetBasedLimit = Math.floor(budget.maxInputTokens * this.maxSectionRatio); + const budgetBasedLimit = Math.floor( + budget.runtimeContextTokens * this.maxSectionRatio, + ); return this.clampInteger( Math.min(this.maxSectionTokens, budgetBasedLimit), @@ -627,20 +627,9 @@ export class RuleBasedContextCompressor implements ContextCompressor { } export function isCompressibleSection(section: ContextSection): boolean { - return section.source?.kind === "tool" || section.source?.kind === "file"; -} - -export function createCompressionResult( - section: ContextSection, - originalSection: ContextSection, - compressed: boolean, - reason?: string, - metadata: ContextCompressionMetadata = {}, -): ContextCompressionResult { - const factory = new ContextCompressionResultFactory(); - if (compressed) { - return factory.compressed(section, originalSection, reason ?? "", metadata); + if (section.compressible === false) { + return false; } - return factory.unchanged(section, reason ?? "", metadata); + return section.source?.kind === "tool" || section.source?.kind === "file"; } diff --git a/src/context/compression/context-truncator.ts b/src/context/compression/context-truncator.ts deleted file mode 100644 index dcb5927..0000000 --- a/src/context/compression/context-truncator.ts +++ /dev/null @@ -1,224 +0,0 @@ -import {DefaultContextBudgetPolicy} from "../budget/context-budget.js"; -import { - createDefaultTokenEstimator, - type TokenEstimator, -} from "../budget/token-estimator.js"; -import {formatContextSection} from "../formatting/context-formatter.js"; -import type {ContextBudget, ContextSection, ContextSectionUsage} from "../types.js"; - -export type FormattedContextSection = { - section: ContextSection; - text: string; -}; - -export type TruncateContextResult = { - contextText: string; - includedSections: ContextSection[]; - partialSections: ContextSection[]; - droppedSections: ContextSection[]; - sectionUsages: ContextSectionUsage[]; -}; - -/** Truncates formatted context blocks to the runtime context budget. */ -export class ContextTruncator { - private readonly tokenEstimator: TokenEstimator; - - constructor(options: {tokenEstimator?: TokenEstimator} = {}) { - this.tokenEstimator = options.tokenEstimator ?? createDefaultTokenEstimator(); - } - - truncate( - sections: readonly ContextSection[], - budget: ContextBudget, - ): TruncateContextResult { - return this.truncateFormatted( - sections - .map((section) => ({section, text: formatContextSection(section)})) - .filter((block) => block.text.length > 0), - budget, - ); - } - - truncateFormatted( - blocks: readonly FormattedContextSection[], - budget: ContextBudget, - ): TruncateContextResult { - let remaining = budget.maxInputTokens; - const selectedBlocks: string[] = []; - const includedSections: ContextSection[] = []; - const partialSections: ContextSection[] = []; - const droppedSections: ContextSection[] = []; - const sectionUsages: ContextSectionUsage[] = []; - - for (const block of blocks) { - const separatorText = selectedBlocks.length ? "\n\n" : ""; - const separatorTokens = this.tokenEstimator.countText(separatorText); - const allowed = remaining - separatorTokens; - const blockTokens = this.tokenEstimator.countText(block.text); - - if (remaining <= 0 || allowed <= 0) { - droppedSections.push(block.section); - sectionUsages.push( - this.createUsage( - block, - "dropped", - 0, - blockTokens, - "No remaining context budget for this section.", - ), - ); - remaining = 0; - continue; - } - - if (blockTokens > allowed) { - const partialText = truncateTextToTokens( - block.text, - allowed, - this.tokenEstimator, - ); - const partialTokens = this.tokenEstimator.countText(partialText); - - if (!partialText) { - droppedSections.push(block.section); - sectionUsages.push( - this.createUsage( - block, - "dropped", - 0, - blockTokens, - "No remaining context budget for this section.", - ), - ); - remaining = 0; - continue; - } - - selectedBlocks.push(partialText); - partialSections.push(block.section); - droppedSections.push(block.section); - sectionUsages.push( - this.createUsage( - block, - "partial", - partialText.length, - blockTokens, - "Section exceeded the remaining context budget and was partially included.", - partialTokens, - ), - ); - remaining = 0; - continue; - } - - selectedBlocks.push(block.text); - includedSections.push(block.section); - sectionUsages.push( - this.createUsage( - block, - "included", - block.text.length, - blockTokens, - "Section fits within the remaining context budget.", - blockTokens, - ), - ); - remaining -= blockTokens + separatorTokens; - } - - return { - contextText: selectedBlocks.join("\n\n"), - includedSections, - partialSections, - droppedSections, - sectionUsages, - }; - } - - private createUsage( - block: FormattedContextSection, - status: ContextSectionUsage["status"], - includedLength: number, - formattedTokens: number, - reason: string, - includedTokens = 0, - ): ContextSectionUsage { - return { - section: block.section, - status, - originalLength: block.section.content.length, - includedLength, - formattedLength: block.text.length, - originalTokens: this.tokenEstimator.countText(block.section.content), - includedTokens, - formattedTokens, - reason, - }; - } -} - -/** Compatibility function for callers that still pass formatted blocks and a token limit. */ -export function truncateContext( - blocks: readonly FormattedContextSection[], - tokenLimit: number, -): TruncateContextResult { - return new ContextTruncator().truncateFormatted( - blocks, - new DefaultContextBudgetPolicy().createBudget({ - sections: blocks.map((block) => block.section), - tokenLimit, - }), - ); -} - -export function truncateTextToTokens( - text: string, - maxTokens: number, - tokenEstimator: TokenEstimator = createDefaultTokenEstimator(), -): string { - if (!text || maxTokens <= 0) { - return ""; - } - - if (tokenEstimator.countText(text) <= maxTokens) { - return text; - } - - let low = 0; - let high = text.length; - let best = ""; - - while (low <= high) { - const mid = Math.floor((low + high) / 2); - const candidate = text.slice(0, mid); - if (tokenEstimator.countText(candidate) <= maxTokens) { - best = candidate; - low = mid + 1; - } else { - high = mid - 1; - } - } - - if (!best) { - return ""; - } - - const linePreserved = trimToLineBoundary(best); - if ( - linePreserved.length >= Math.ceil(best.length * 0.5) && - tokenEstimator.countText(linePreserved) <= maxTokens - ) { - return linePreserved; - } - - return best; -} - -function trimToLineBoundary(text: string): string { - const lastNewline = text.lastIndexOf("\n"); - if (lastNewline > 0) { - return text.slice(0, lastNewline + 1); - } - - return text; -} diff --git a/src/context/engine/context-engine.ts b/src/context/engine/context-engine.ts index e40bbcb..05a5eb6 100644 --- a/src/context/engine/context-engine.ts +++ b/src/context/engine/context-engine.ts @@ -5,18 +5,35 @@ import { type TokenEstimator, } from "../budget/token-estimator.js"; import {ContextCompressionPipeline} from "../compression/context-compression-pipeline.js"; -import {ContextTruncator} from "../compression/context-truncator.js"; -import {SystemPromptAssembler} from "../formatting/system-prompt-assembler.js"; +import {ContextTruncator} from "../truncate/context-truncator.js"; +import {SystemPromptAssembler} from "../assembler/system-prompt-assembler.js"; import type {ContextBudgetPolicy} from "../budget/context-budget.js"; import type {ContextPriorityPolicy} from "../budget/priority-policy.js"; import type { BuildContextInput, BuildContextResult, ContextEngineOptions, + ContextTokenUsageDiagnostics, } from "../types.js"; import {ContextRegistry} from "./context-registry.js"; -/** Class-based context engine that owns the full runtime context pipeline. */ +/** Thrown when the final assembled prompt exceeds the model input budget. */ +export class ContextBudgetExceededError extends Error { + constructor( + message: string, + readonly tokenUsage: ContextTokenUsageDiagnostics, + ) { + super(message); + this.name = "ContextBudgetExceededError"; + } +} + +/** + * Class-based context engine that owns the full runtime context pipeline. + * + * Pipeline: + * sections -> registry -> budget -> compression -> truncation -> assembly -> validation + */ export class ContextEngine { private readonly priorityPolicy: ContextPriorityPolicy; private readonly budgetPolicy: ContextBudgetPolicy; @@ -26,9 +43,14 @@ export class ContextEngine { private readonly tokenEstimator: TokenEstimator; constructor(options: ContextEngineOptions = {}) { + // Shared token estimator used by budget, compression, truncation and diagnostics. this.tokenEstimator = options.tokenEstimator ?? createDefaultTokenEstimator(); + + // Pluggable policies keep context ordering and budget calculation replaceable. this.priorityPolicy = options.priorityPolicy ?? new DefaultContextPriorityPolicy(); this.budgetPolicy = options.budgetPolicy ?? new DefaultContextBudgetPolicy(); + + // Compress oversized or low-value context before truncation. this.compressionPipeline = options.compressionPipeline ?? new ContextCompressionPipeline({ @@ -36,47 +58,143 @@ export class ContextEngine { thresholdRatio: options.compressionThresholdRatio, tokenEstimator: this.tokenEstimator, }); + + // Enforces the final runtime context budget. this.truncator = options.truncator ?? new ContextTruncator({tokenEstimator: this.tokenEstimator}); + + // Assembles system prompt, output instructions and runtime context. this.assembler = options.assembler ?? new SystemPromptAssembler(); } build(input: BuildContextInput): BuildContextResult { + // Normalize, dedupe and sort sections before budget-sensitive processing. const registry = new ContextRegistry() .addMany(input.sections) .normalize() .dedupe() .sort(this.priorityPolicy); - const budget = this.budgetPolicy.createBudget(input); + + // Create a token budget that accounts for prompt, tools, history and output reserve. + const budget = this.budgetPolicy.createBudget(input, this.tokenEstimator); const sections = registry.getAll(); + + // Compress first so truncation works on smaller, higher-value sections. const compression = this.compressionPipeline.compress(sections, budget); + + // Truncate compressed sections to the runtime context budget. const truncation = this.truncator.truncate(compression.sections, budget); - const systemPrompt = this.assembler.assemble({ - systemPrompt: input.systemPrompt, + + // Build the assembled system message content sent to the model. + const assembledSystemPrompt = this.assembler.assemble({ + baseSystemPrompt: input.baseSystemPrompt, outputInstructions: input.outputInstructions, contextText: truncation.contextText, }); + + // Re-count final text after compression/truncation/assembly for validation. const contextTextTokens = this.tokenEstimator.countText(truncation.contextText); - const systemPromptTokens = this.tokenEstimator.countText(systemPrompt); + const assembledSystemPromptTokens = + this.tokenEstimator.countText(assembledSystemPrompt); + + // Final model input also includes conversation history and tool schemas. + const finalInputTokens = + assembledSystemPromptTokens + budget.conversationTokens + budget.toolSchemaTokens; + + // Input limit excludes the output tokens reserved for model generation. + const inputTokenLimit = Math.max( + 0, + budget.modelContextWindow - budget.reservedOutputTokens, + ); + + const remainingInputTokens = inputTokenLimit - finalInputTokens; + + // Attach final runtime numbers to the budget diagnostics. + const finalBudget = { + ...budget, + finalInputTokens, + remainingInputTokens, + }; + + // Token usage breakdown for debugging and observability. + const tokenUsage = { + modelContextWindow: budget.modelContextWindow, + reservedOutputTokens: budget.reservedOutputTokens, + systemPromptTokens: budget.systemPromptTokens, + outputInstructionTokens: budget.outputInstructionTokens, + toolSchemaTokens: budget.toolSchemaTokens, + conversationTokens: budget.conversationTokens, + safetyMarginTokens: budget.safetyMarginTokens, + runtimeContextTokens: budget.runtimeContextTokens, + contextTextTokens, + finalInputTokens, + remainingInputTokens, + }; + + // Fail fast if the final assembled request still exceeds the model window. + if (remainingInputTokens < 0) { + throw new ContextBudgetExceededError( + `Final input exceeds the model context window by ${Math.abs( + remainingInputTokens, + )} tokens after reserving output tokens.`, + tokenUsage, + ); + } return { - systemPrompt, + assembledSystemPrompt, contextText: truncation.contextText, - tokenEstimate: systemPromptTokens, + tokenEstimate: assembledSystemPromptTokens, + + // Section-level inclusion results after truncation. includedSections: truncation.includedSections, - partialSections: truncation.partialSections, - droppedSections: truncation.droppedSections, + partiallyIncludedSections: truncation.partiallyIncludedSections, + omittedSections: truncation.omittedSections, sectionUsages: truncation.sectionUsages, + diagnostics: { - budget, + budget: finalBudget, + tokenUsage, + + // Compression-level diagnostics. estimatedContextChars: compression.estimatedContextChars, estimatedContextTokens: compression.estimatedContextTokens, compressionThresholdRatio: compression.thresholdRatio, compressionTriggered: compression.triggered, compressionLimitTokens: compression.compressionLimitTokens, compressionResults: compression.results, + + // Final prompt-level diagnostics. contextTextTokens, - systemPromptTokens, + systemPromptTokens: assembledSystemPromptTokens, + finalInputTokens, + remainingInputTokens, + + // Unified decision log for compression and truncation. + sectionDecisions: [ + ...compression.results.map((result) => ({ + section: result.originalSection, + stage: "compression" as const, + action: result.compressed ? ("compressed" as const) : ("skipped" as const), + reason: result.reason ?? "", + originalTokens: result.originalTokens, + outputTokens: result.compressedTokens, + savedTokens: result.savedTokens, + })), + ...truncation.sectionUsages.map((usage) => ({ + section: usage.section, + stage: "truncation" as const, + action: + usage.status === "included" + ? ("included" as const) + : usage.status === "partial" + ? ("partial" as const) + : ("omitted" as const), + reason: usage.reason ?? "", + originalTokens: usage.formattedTokens, + outputTokens: usage.includedTokens, + })), + ], }, }; } diff --git a/src/context/engine/context-registry.ts b/src/context/engine/context-registry.ts index 258e4b0..b395088 100644 --- a/src/context/engine/context-registry.ts +++ b/src/context/engine/context-registry.ts @@ -16,6 +16,7 @@ export class ContextRegistry { } normalize(): this { + // Trim blank padding early so downstream token estimates match inserted text. this.sections = this.sections .map((section) => ({...section, content: section.content.trim()})) .filter((section) => section.content.length > 0); @@ -26,6 +27,7 @@ export class ContextRegistry { const seen = new Set(); const deduped: ContextSection[] = []; + // Walk backward so the newest section wins for the same id or replaceKey. for (let index = this.sections.length - 1; index >= 0; index -= 1) { const section = this.sections[index]; if (!section) { @@ -48,6 +50,7 @@ export class ContextRegistry { } sort(policy: ContextPriorityPolicy): this { + // Stable sort keeps original order when two sections have equal priority. this.sections = this.sections .map((section, index) => ({section, index})) .sort((left, right) => { diff --git a/src/context/index.ts b/src/context/index.ts index 1c8dde1..b374d5c 100644 --- a/src/context/index.ts +++ b/src/context/index.ts @@ -1,4 +1,4 @@ -export {ContextEngine} from "./engine/context-engine.js"; +export {ContextBudgetExceededError, ContextEngine} from "./engine/context-engine.js"; export {buildRuntimeContext} from "./engine/context-builder.js"; export {ContextRegistry} from "./engine/context-registry.js"; export {ContextCompressionPipeline} from "./compression/context-compression-pipeline.js"; @@ -10,9 +10,7 @@ export {DefaultContextBudgetPolicy} from "./budget/context-budget.js"; export type {ContextBudgetPolicy} from "./budget/context-budget.js"; export { ContextCompressionResultFactory, - createCompressionResult, isCompressibleSection, - NoopContextCompressor, RuleBasedContextCompressor, } from "./compression/context-compressor.js"; export type { @@ -20,22 +18,15 @@ export type { ContextCompressionMetadata, RuleBasedContextCompressorOptions, } from "./compression/context-compressor.js"; -export { - compareContextSection, - formatContextSection, -} from "./formatting/context-formatter.js"; -export { - ContextTruncator, - truncateContext, - truncateTextToTokens, -} from "./compression/context-truncator.js"; +export {formatContextSection} from "./assembler/context-formatter.js"; +export {ContextTruncator, truncateTextToTokens} from "./truncate/context-truncator.js"; export type { FormattedContextSection, TruncateContextResult, -} from "./compression/context-truncator.js"; +} from "./truncate/context-truncator.js"; export {DefaultContextPriorityPolicy} from "./budget/priority-policy.js"; export type {ContextPriorityPolicy} from "./budget/priority-policy.js"; -export {SystemPromptAssembler} from "./formatting/system-prompt-assembler.js"; +export {SystemPromptAssembler} from "./assembler/system-prompt-assembler.js"; export { ApproxTokenEstimator, createDefaultTokenEstimator, @@ -51,7 +42,10 @@ export type { ContextCompressionResult, ContextEngineOptions, ContextSection, + ContextSectionDecision, + ContextSectionTruncation, ContextSectionUsage, ContextSectionUsageStatus, ContextSource, + ContextTokenUsageDiagnostics, } from "./types.js"; diff --git a/src/context/truncate/context-truncator.ts b/src/context/truncate/context-truncator.ts new file mode 100644 index 0000000..251eca5 --- /dev/null +++ b/src/context/truncate/context-truncator.ts @@ -0,0 +1,353 @@ +import { + createDefaultTokenEstimator, + type TokenEstimator, +} from "../budget/token-estimator.js"; +import {formatContextSection} from "../assembler/context-formatter.js"; +import type {ContextBudget, ContextSection, ContextSectionUsage} from "../types.js"; + +export type FormattedContextSection = { + section: ContextSection; + text: string; +}; + +export type TruncateContextResult = { + contextText: string; + includedSections: ContextSection[]; + partiallyIncludedSections: ContextSection[]; + omittedSections: ContextSection[]; + sectionUsages: ContextSectionUsage[]; +}; + +/** Truncates formatted context blocks to the runtime context budget. */ +export class ContextTruncator { + private readonly tokenEstimator: TokenEstimator; + + constructor(options: {tokenEstimator?: TokenEstimator} = {}) { + this.tokenEstimator = options.tokenEstimator ?? createDefaultTokenEstimator(); + } + + truncate( + sections: readonly ContextSection[], + budget: ContextBudget, + ): TruncateContextResult { + // Format once, then truncate the exact text that will be inserted. + return this.truncateFormatted( + sections + .map((section) => ({section, text: formatContextSection(section)})) + .filter((block) => block.text.length > 0), + budget, + ); + } + + truncateFormatted( + blocks: readonly FormattedContextSection[], + budget: ContextBudget, + ): TruncateContextResult { + // Only runtime context can spend this budget; prompt/history/tool costs are pre-deducted. + let remaining = budget.runtimeContextTokens; + const selectedBlocks: string[] = []; + const includedSections: ContextSection[] = []; + const partiallyIncludedSections: ContextSection[] = []; + const omittedSections: ContextSection[] = []; + const sectionUsages: ContextSectionUsage[] = []; + + for (const block of blocks) { + const separatorText = selectedBlocks.length ? "\n\n" : ""; + const separatorTokens = this.tokenEstimator.countText(separatorText); + const allowed = remaining - separatorTokens; + const blockTokens = this.tokenEstimator.countText(block.text); + + if (remaining <= 0 || allowed <= 0) { + // Required sections are allowed through; final validation catches true overflow. + if (block.section.required) { + selectedBlocks.push(block.text); + includedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "included", + block.text.length, + blockTokens, + "Required section was included even though it exceeded the remaining runtime context budget.", + blockTokens, + ), + ); + remaining -= blockTokens + separatorTokens; + continue; + } + + omittedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "omitted", + 0, + blockTokens, + "No remaining context budget for this section.", + ), + ); + remaining = 0; + continue; + } + + if (blockTokens > allowed) { + // Non-truncatable sections are either included whole if required or omitted. + if (block.section.truncation === "none") { + if (block.section.required) { + selectedBlocks.push(block.text); + includedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "included", + block.text.length, + blockTokens, + "Required section disallows truncation and was included in full.", + blockTokens, + ), + ); + remaining -= blockTokens + separatorTokens; + continue; + } + + omittedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "omitted", + 0, + blockTokens, + "Section exceeded the remaining context budget and disallows truncation.", + ), + ); + remaining = Math.max(0, remaining); + continue; + } + + // Strategy controls which part of an oversized section is preserved. + const partialText = truncateTextToTokens( + block.text, + allowed, + block.section.truncation, + this.tokenEstimator, + ); + const partialTokens = this.tokenEstimator.countText(partialText); + + if (!partialText) { + if (block.section.required) { + selectedBlocks.push(block.text); + includedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "included", + block.text.length, + blockTokens, + "Required section could not be partially truncated and was included in full.", + blockTokens, + ), + ); + remaining -= blockTokens + separatorTokens; + continue; + } + + omittedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "omitted", + 0, + blockTokens, + "No remaining context budget for this section.", + ), + ); + remaining = 0; + continue; + } + + selectedBlocks.push(partialText); + partiallyIncludedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "partial", + partialText.length, + blockTokens, + "Section exceeded the remaining context budget and was partially included.", + partialTokens, + ), + ); + remaining = 0; + continue; + } + + selectedBlocks.push(block.text); + includedSections.push(block.section); + sectionUsages.push( + this.createUsage( + block, + "included", + block.text.length, + blockTokens, + "Section fits within the remaining context budget.", + blockTokens, + ), + ); + remaining -= blockTokens + separatorTokens; + } + + return { + contextText: selectedBlocks.join("\n\n"), + includedSections, + partiallyIncludedSections, + omittedSections, + sectionUsages, + }; + } + + private createUsage( + block: FormattedContextSection, + status: ContextSectionUsage["status"], + includedLength: number, + formattedTokens: number, + reason: string, + includedTokens = 0, + ): ContextSectionUsage { + return { + section: block.section, + status, + originalLength: block.section.content.length, + includedLength, + formattedLength: block.text.length, + originalTokens: this.tokenEstimator.countText(block.section.content), + includedTokens, + formattedTokens, + reason, + }; + } +} + +export function truncateTextToTokens( + text: string, + maxTokens: number, + strategyOrEstimator: ContextSection["truncation"] | TokenEstimator = "head", + tokenEstimator: TokenEstimator = createDefaultTokenEstimator(), +): string { + const strategy = typeof strategyOrEstimator === "string" ? strategyOrEstimator : "head"; + const estimator = + typeof strategyOrEstimator === "string" ? tokenEstimator : strategyOrEstimator; + + if (!text || maxTokens <= 0) { + return ""; + } + + if (estimator.countText(text) <= maxTokens) { + return text; + } + + if (strategy === "tail") { + return truncateTailTextToTokens(text, maxTokens, estimator); + } + + if (strategy === "middle" || strategy === "semantic") { + return truncateMiddleTextToTokens(text, maxTokens, estimator); + } + + return truncateHeadTextToTokens(text, maxTokens, estimator); +} + +function truncateHeadTextToTokens( + text: string, + maxTokens: number, + tokenEstimator: TokenEstimator, +): string { + let low = 0; + let high = text.length; + let best = ""; + + while (low <= high) { + const mid = Math.floor((low + high) / 2); + const candidate = text.slice(0, mid); + if (tokenEstimator.countText(candidate) <= maxTokens) { + best = candidate; + low = mid + 1; + } else { + high = mid - 1; + } + } + + if (!best) { + return ""; + } + + const linePreserved = trimToLineBoundary(best); + if ( + linePreserved.length >= Math.ceil(best.length * 0.5) && + tokenEstimator.countText(linePreserved) <= maxTokens + ) { + return linePreserved; + } + + return best; +} + +function truncateTailTextToTokens( + text: string, + maxTokens: number, + tokenEstimator: TokenEstimator, +): string { + let low = 0; + let high = text.length; + let best = ""; + + while (low <= high) { + const mid = Math.floor((low + high) / 2); + const candidate = text.slice(text.length - mid); + if (tokenEstimator.countText(candidate) <= maxTokens) { + best = candidate; + low = mid + 1; + } else { + high = mid - 1; + } + } + + return best; +} + +function truncateMiddleTextToTokens( + text: string, + maxTokens: number, + tokenEstimator: TokenEstimator, +): string { + if (maxTokens <= 0) { + return ""; + } + + const marker = "\n[...context omitted...]\n"; + const markerTokens = tokenEstimator.countText(marker); + if (markerTokens >= maxTokens) { + return truncateHeadTextToTokens(text, maxTokens, tokenEstimator); + } + + const sideBudget = Math.floor((maxTokens - markerTokens) / 2); + const head = truncateHeadTextToTokens(text, sideBudget, tokenEstimator); + const tail = truncateTailTextToTokens( + text, + maxTokens - markerTokens - tokenEstimator.countText(head), + tokenEstimator, + ); + const candidate = `${head}${marker}${tail}`; + + return tokenEstimator.countText(candidate) <= maxTokens + ? candidate + : truncateHeadTextToTokens(text, maxTokens, tokenEstimator); +} + +function trimToLineBoundary(text: string): string { + const lastNewline = text.lastIndexOf("\n"); + if (lastNewline > 0) { + return text.slice(0, lastNewline + 1); + } + + return text; +} diff --git a/src/context/types.ts b/src/context/types.ts index 7167c21..cc3ba57 100644 --- a/src/context/types.ts +++ b/src/context/types.ts @@ -1,10 +1,10 @@ import type {ContextBudgetPolicy} from "./budget/context-budget.js"; import type {ContextPriorityPolicy} from "./budget/priority-policy.js"; -import type {TokenEstimator} from "./budget/token-estimator.js"; +import type {TokenCountableMessage, TokenEstimator} from "./budget/token-estimator.js"; import type {ContextCompressionPipeline} from "./compression/context-compression-pipeline.js"; import type {ContextCompressor} from "./compression/context-compressor.js"; -import type {ContextTruncator} from "./compression/context-truncator.js"; -import type {SystemPromptAssembler} from "./formatting/system-prompt-assembler.js"; +import type {ContextTruncator} from "./truncate/context-truncator.js"; +import type {SystemPromptAssembler} from "./assembler/system-prompt-assembler.js"; /** Source metadata for a context section. */ export type ContextSource = @@ -24,24 +24,38 @@ export type ContextSection = { content: string; priority?: number; source?: ContextSource; + required?: boolean; + pinned?: boolean; + compressible?: boolean; + truncation?: ContextSectionTruncation; + metadata?: Record; }; +export type ContextSectionTruncation = "none" | "head" | "tail" | "middle" | "semantic"; + /** Input accepted by the generic runtime context builder. */ export type BuildContextInput = { - systemPrompt?: string; + // Base system prompt before output instructions and runtime context are appended. + baseSystemPrompt?: string; outputInstructions?: string; sections: readonly ContextSection[]; tokenLimit: number; + conversationMessages?: readonly TokenCountableMessage[]; + toolSchemas?: readonly unknown[]; + conversationTokens?: number; + toolSchemaTokens?: number; + safetyMarginTokens?: number; }; /** Result returned after formatting and truncating runtime context. */ export type BuildContextResult = { - systemPrompt: string; + // Final system message content after assembling base prompt, output rules, and runtime context. + assembledSystemPrompt: string; contextText: string; tokenEstimate: number; includedSections: ContextSection[]; - partialSections: ContextSection[]; - droppedSections: ContextSection[]; + partiallyIncludedSections: ContextSection[]; + omittedSections: ContextSection[]; sectionUsages: ContextSectionUsage[]; diagnostics?: BuildContextDiagnostics; }; @@ -49,17 +63,20 @@ export type BuildContextResult = { /** Runtime context budget derived from the model token limit. */ export type ContextBudget = { tokenLimit: number; - maxContextTokens: number; + modelContextWindow: number; reservedOutputTokens: number; - maxInputTokens: number; - /** @deprecated kept for compatibility and diagnostics only. */ - runtimeContextRatio?: number; - /** @deprecated char budget is a diagnostic fallback only. */ - maxContextChars?: number; + systemPromptTokens: number; + outputInstructionTokens: number; + toolSchemaTokens: number; + conversationTokens: number; + safetyMarginTokens: number; + runtimeContextTokens: number; + finalInputTokens?: number; + remainingInputTokens?: number; }; /** Explicit section-level truncation status. */ -export type ContextSectionUsageStatus = "included" | "partial" | "dropped"; +export type ContextSectionUsageStatus = "included" | "partial" | "omitted"; /** Truncation and injection details for a single context section. */ export type ContextSectionUsage = { @@ -94,9 +111,34 @@ export type ContextCompressionResult = { reason?: string; }; +export type ContextSectionDecision = { + section: ContextSection; + stage: "compression" | "truncation"; + action: "included" | "partial" | "omitted" | "compressed" | "skipped"; + reason: string; + originalTokens?: number; + outputTokens?: number; + savedTokens?: number; +}; + +export type ContextTokenUsageDiagnostics = { + modelContextWindow: number; + reservedOutputTokens: number; + systemPromptTokens: number; + outputInstructionTokens: number; + toolSchemaTokens: number; + conversationTokens: number; + safetyMarginTokens: number; + runtimeContextTokens: number; + contextTextTokens: number; + finalInputTokens: number; + remainingInputTokens: number; +}; + /** Diagnostics produced while building runtime context. */ export type BuildContextDiagnostics = { budget: ContextBudget; + tokenUsage: ContextTokenUsageDiagnostics; estimatedContextChars: number; estimatedContextTokens: number; compressionThresholdRatio: number; @@ -105,6 +147,9 @@ export type BuildContextDiagnostics = { compressionResults: ContextCompressionResult[]; contextTextTokens: number; systemPromptTokens: number; + finalInputTokens: number; + remainingInputTokens: number; + sectionDecisions: ContextSectionDecision[]; }; /** Optional strategy overrides for the class-based context engine. */ diff --git a/tests/agent/runtime-loop.test.ts b/tests/agent/runtime-loop.test.ts index 1bcafdb..1f19b78 100644 --- a/tests/agent/runtime-loop.test.ts +++ b/tests/agent/runtime-loop.test.ts @@ -211,7 +211,7 @@ describe("Agent runtime loop", () => { const workspaceRoot = await createWorkspace(); const registry = new ToolRegistry(); const beforeModelBuildCounts: number[] = []; - const longOutput = "0123456789".repeat(1_500); + const longOutput = "0123456789".repeat(15_000); registry.register({ definition: { @@ -243,7 +243,7 @@ describe("Agent runtime loop", () => { ]); const config = createConfig(workspaceRoot); - config.runtime.tokensLimit = 2_000; + config.runtime.tokensLimit = 60_000; config.runtime.maxIterations = 4; const result = await new Agent({ @@ -266,8 +266,10 @@ describe("Agent runtime loop", () => { expect(result.stopReason).toBe("completed"); expect(llm.requests).toHaveLength(3); expect(beforeModelBuildCounts).toEqual([1, 2, 3]); + expect(thirdSystemPrompt).toContain("## Tool Args: large_output"); expect(thirdSystemPrompt).toContain("## Tool Result: large_output"); expect(thirdSystemPrompt).toContain("Call ID: call-first"); + expect(thirdSystemPrompt).toContain('"label": "first"'); expect(thirdSystemPrompt).toContain("chars omitted"); expect( thirdRequest?.messages.some( diff --git a/tests/agent/system-prompt.test.ts b/tests/agent/system-prompt.test.ts index ed5e342..2bed742 100644 --- a/tests/agent/system-prompt.test.ts +++ b/tests/agent/system-prompt.test.ts @@ -35,7 +35,7 @@ describe("ContextManager system prompt", () => { workspaceDir: workspaceRoot, maxIterations: 1, maxRepairAttempts: 0, - tokensLimit: 1000, + tokensLimit: 32_000, rollbackOnFailure: false, }, permissions: { @@ -90,7 +90,7 @@ describe("ContextManager system prompt", () => { workspaceDir: workspaceRoot, maxIterations: 1, maxRepairAttempts: 0, - tokensLimit: 1000, + tokensLimit: 32_000, rollbackOnFailure: false, }, permissions: { diff --git a/tests/context-builder.test.ts b/tests/context-builder.test.ts index 4ce3084..1336a2d 100644 --- a/tests/context-builder.test.ts +++ b/tests/context-builder.test.ts @@ -5,15 +5,16 @@ import { buildRuntimeContext, ContextCompressionPipeline, ContextCompressionResultFactory, + ContextBudgetExceededError, ContextEngine, ContextRegistry, DefaultContextBudgetPolicy, estimateTokens, GptTokenEstimator, - NoopContextCompressor, RuleBasedContextCompressor, truncateTextToTokens, type ContextBudget, + type ContextBudgetPolicy, type ContextCompressionResult, type ContextCompressor, type ContextSection, @@ -74,7 +75,7 @@ describe("buildRuntimeContext", () => { it("matches the default ContextEngine output", () => { const input = { - systemPrompt: "Base.", + baseSystemPrompt: "Base.", outputInstructions: "Use Markdown.", sections: [{title: "Runtime", content: "details"}], tokenLimit: 100, @@ -85,7 +86,7 @@ describe("buildRuntimeContext", () => { it("sorts context sections by descending priority", () => { const result = new ContextEngine().build({ - systemPrompt: "Base.", + baseSystemPrompt: "Base.", sections: [ {id: "low", content: "low", priority: 1}, {id: "high", content: "high", priority: 10}, @@ -104,7 +105,7 @@ describe("buildRuntimeContext", () => { it("uses source priority when explicit priority is absent", () => { const result = new ContextEngine().build({ - systemPrompt: "Base.", + baseSystemPrompt: "Base.", sections: [ {id: "file", content: "file", source: {kind: "file"}}, {id: "workspace", content: "workspace", source: {kind: "workspace"}}, @@ -118,7 +119,7 @@ describe("buildRuntimeContext", () => { it("skips empty section content", () => { const result = buildRuntimeContext({ - systemPrompt: "Base.", + baseSystemPrompt: "Base.", sections: [ {id: "empty", title: "Empty", content: " ", priority: 100}, {id: "filled", content: "filled"}, @@ -129,12 +130,12 @@ describe("buildRuntimeContext", () => { expect(result.contextText).toBe("filled"); expect(result.includedSections.map((section) => section.id)).toEqual(["filled"]); expect(result.sectionUsages.map((usage) => usage.section.id)).toEqual(["filled"]); - expect(result.droppedSections).toEqual([]); + expect(result.omittedSections).toEqual([]); }); it("formats titled sections with markdown headings", () => { const result = buildRuntimeContext({ - systemPrompt: "Base.", + baseSystemPrompt: "Base.", sections: [{title: "Notes", content: " keep this "}], tokenLimit: 100, }); @@ -148,65 +149,96 @@ describe("buildRuntimeContext", () => { sections: [], tokenLimit: 100, }), - ).toEqual({ + ).toMatchObject({ tokenLimit: 100, - maxContextTokens: 100, + modelContextWindow: 100, reservedOutputTokens: 20, - maxInputTokens: 80, - runtimeContextRatio: 0.35, - maxContextChars: 320, + runtimeContextTokens: 80, + systemPromptTokens: 0, + outputInstructionTokens: 0, + toolSchemaTokens: 0, + conversationTokens: 0, + safetyMarginTokens: 0, }); }); - it("truncates over-budget runtime context and records included, partial, and dropped sections", () => { + it("deducts prompt, conversation, tool schema, and safety tokens from runtime context budget", () => { + const budget = new DefaultContextBudgetPolicy({ + reservedOutputTokens: 10, + }).createBudget( + { + baseSystemPrompt: "system", + outputInstructions: "output", + sections: [], + tokenLimit: 100, + conversationMessages: [{role: "user", content: "conversation"}], + toolSchemas: [{name: "tool"}], + safetyMarginTokens: 3, + }, + charTokenEstimator, + ); + + expect(budget).toMatchObject({ + modelContextWindow: 100, + reservedOutputTokens: 10, + systemPromptTokens: 6, + outputInstructionTokens: 6, + conversationTokens: 17, + toolSchemaTokens: 17, + safetyMarginTokens: 3, + runtimeContextTokens: 41, + }); + }); + + it("truncates over-budget runtime context and records included, partial, and omitted sections", () => { const result = new ContextEngine({ - budgetPolicy: new DefaultContextBudgetPolicy({reservedOutputTokens: 0}), + budgetPolicy: fixedRuntimeBudgetPolicy(4), tokenEstimator: charTokenEstimator, }).build({ - systemPrompt: "Base.", + baseSystemPrompt: "Base.", sections: [ {id: "first", content: "abcdefghij", priority: 10}, {id: "second", content: "klmnopqrst", priority: 9}, ], - tokenLimit: 4, + tokenLimit: 100, }); expect(result.contextText).toBe("abcd"); expect(result.includedSections).toEqual([]); - expect(result.partialSections.map((section) => section.id)).toEqual(["first"]); - expect(result.droppedSections.map((section) => section.id)).toEqual([ + expect(result.partiallyIncludedSections.map((section) => section.id)).toEqual([ "first", - "second", ]); + expect(result.omittedSections.map((section) => section.id)).toEqual(["second"]); expect(result.sectionUsages.map((usage) => usage.status)).toEqual([ "partial", - "dropped", + "omitted", ]); expect(result.diagnostics?.contextTextTokens).toBe(4); - expect(result.tokenEstimate).toBe(charTokenEstimator.countText(result.systemPrompt)); + expect(result.tokenEstimate).toBe( + charTokenEstimator.countText(result.assembledSystemPrompt), + ); }); - it("keeps fully injected sections separate from partial and dropped sections", () => { + it("keeps fully injected sections separate from partial and omitted sections", () => { const result = new ContextEngine({ - budgetPolicy: new DefaultContextBudgetPolicy({reservedOutputTokens: 0}), + budgetPolicy: fixedRuntimeBudgetPolicy(7), tokenEstimator: charTokenEstimator, }).build({ - systemPrompt: "Base.", + baseSystemPrompt: "Base.", sections: [ {id: "first", content: "abc", priority: 10}, {id: "second", content: "defghij", priority: 9}, {id: "third", content: "klm", priority: 8}, ], - tokenLimit: 7, + tokenLimit: 100, }); expect(result.contextText).toBe("abc\n\nde"); expect(result.includedSections.map((section) => section.id)).toEqual(["first"]); - expect(result.partialSections.map((section) => section.id)).toEqual(["second"]); - expect(result.droppedSections.map((section) => section.id)).toEqual([ + expect(result.partiallyIncludedSections.map((section) => section.id)).toEqual([ "second", - "third", ]); + expect(result.omittedSections.map((section) => section.id)).toEqual(["third"]); expect(result.sectionUsages).toEqual([ expect.objectContaining({ section: expect.objectContaining({id: "first"}), @@ -221,7 +253,7 @@ describe("buildRuntimeContext", () => { }), expect.objectContaining({ section: expect.objectContaining({id: "third"}), - status: "dropped", + status: "omitted", reason: "No remaining context budget for this section.", }), ]); @@ -229,17 +261,37 @@ describe("buildRuntimeContext", () => { it("builds system prompt with output instructions and runtime context", () => { const result = buildRuntimeContext({ - systemPrompt: "Base prompt.", + baseSystemPrompt: "Base prompt.", outputInstructions: "Use Markdown.", sections: [{title: "Runtime", content: "details"}], tokenLimit: 100, }); - expect(result.systemPrompt).toBe( + expect(result.assembledSystemPrompt).toBe( "Base prompt.\n\nUse Markdown.\n\n# Runtime Context\n## Runtime\ndetails", ); }); + it("throws when protected required context makes the final prompt exceed the model window", () => { + expect(() => + new ContextEngine({ + budgetPolicy: fixedRuntimeBudgetPolicy(0), + tokenEstimator: charTokenEstimator, + }).build({ + baseSystemPrompt: "Base.", + sections: [ + { + id: "required", + content: "required content", + required: true, + truncation: "none", + }, + ], + tokenLimit: 20, + }), + ).toThrow(ContextBudgetExceededError); + }); + it("dedupes sections by replaceKey or id with the later section winning", () => { const registry = new ContextRegistry() .addMany([ @@ -259,34 +311,6 @@ describe("buildRuntimeContext", () => { ]); }); - it("does not alter ordinary sections with the default compressor", () => { - const section = { - id: "plain", - title: "Plain", - content: "content", - source: {kind: "user" as const}, - }; - const result = new NoopContextCompressor().compress( - section, - new DefaultContextBudgetPolicy().createBudget({ - sections: [section], - tokenLimit: 100, - }), - ); - - expect(result).toMatchObject({ - section, - originalSection: section, - compressed: false, - originalChars: section.content.length, - compressedChars: section.content.length, - omittedChars: 0, - savedChars: 0, - compressionRatio: 1, - reason: "No compression applied.", - }); - }); - it("does not trigger compression below the configured threshold", () => { const compressor = new ThrowingCompressor(); const pipeline = new ContextCompressionPipeline({ @@ -304,7 +328,7 @@ describe("buildRuntimeContext", () => { const result = new ContextEngine({ compressionPipeline: pipeline, }).build({ - systemPrompt: "Base.", + baseSystemPrompt: "Base.", sections: [{id: "short", content: "short"}], tokenLimit: 100, }); @@ -320,7 +344,8 @@ describe("buildRuntimeContext", () => { expect(compression.results).toEqual([ expect.objectContaining({ compressed: false, - reason: "Compression was not triggered.", + reason: + "Compression was not triggered because runtime context is below the configured total threshold.", }), ]); expect(result.diagnostics).toMatchObject({ @@ -328,7 +353,7 @@ describe("buildRuntimeContext", () => { compressionTriggered: false, estimatedContextChars: 5, estimatedContextTokens: 5, - compressionLimitTokens: 68, + compressionLimitTokens: 66, }); expect(result.diagnostics?.compressionResults).toEqual([ expect.objectContaining({ @@ -336,15 +361,54 @@ describe("buildRuntimeContext", () => { compressed: false, originalChars: 5, compressedChars: 5, - reason: "Compression was not triggered.", + reason: + "Compression was not triggered because runtime context is below the configured total threshold.", }), ]); }); + it("skips compression for sections that are protected or below strategy thresholds", () => { + const pipeline = new ContextCompressionPipeline({ + compressor: new ThrowingCompressor(), + thresholdRatio: 0, + minSectionTokens: 10, + tokenEstimator: charTokenEstimator, + }); + const sections: ContextSection[] = [ + { + id: "required", + content: "required long content", + required: true, + truncation: "none", + source: {kind: "tool"}, + }, + { + id: "opt-out", + content: "opt out long content", + compressible: false, + source: {kind: "tool"}, + }, + { + id: "short", + content: "short", + source: {kind: "file"}, + }, + ]; + + const compression = pipeline.compress(sections, rawBudget(100)); + + expect(compression.results.map((result) => result.reason)).toEqual([ + "Required section disallows truncation and is protected from compression.", + "Section opted out of compression.", + "Section is below the minimum compression threshold (5 < 10 tokens).", + ]); + }); + it("triggers compression above the configured threshold and exposes diagnostics", () => { const pipeline = new ContextCompressionPipeline({ compressor: new MarkingCompressor(), thresholdRatio: 0.5, + minSectionTokens: 1, tokenEstimator: charTokenEstimator, }); const result = new ContextEngine({ @@ -352,30 +416,30 @@ describe("buildRuntimeContext", () => { budgetPolicy: new DefaultContextBudgetPolicy({reservedOutputTokens: 0}), tokenEstimator: charTokenEstimator, }).build({ - systemPrompt: "Base prompt.", + baseSystemPrompt: "Base prompt.", outputInstructions: "Use Markdown.", sections: [{id: "long", content: "abcdefghijklmnop", source: {kind: "tool"}}], - tokenLimit: 8, + tokenLimit: 55, }); expect(result.contextText).toBe("abcdefgh"); - expect(result.tokenEstimate).toBe(charTokenEstimator.countText(result.systemPrompt)); + expect(result.tokenEstimate).toBe( + charTokenEstimator.countText(result.assembledSystemPrompt), + ); expect(result.diagnostics).toMatchObject({ budget: { - tokenLimit: 8, - maxContextTokens: 8, + tokenLimit: 55, + modelContextWindow: 55, reservedOutputTokens: 0, - maxInputTokens: 8, - runtimeContextRatio: 0.35, - maxContextChars: 32, + runtimeContextTokens: 30, }, estimatedContextChars: 16, estimatedContextTokens: 16, compressionThresholdRatio: 0.5, compressionTriggered: true, - compressionLimitTokens: 4, + compressionLimitTokens: 15, contextTextTokens: charTokenEstimator.countText(result.contextText), - systemPromptTokens: charTokenEstimator.countText(result.systemPrompt), + systemPromptTokens: charTokenEstimator.countText(result.assembledSystemPrompt), }); expect(result.diagnostics?.compressionResults).toEqual([ expect.objectContaining({ @@ -396,24 +460,24 @@ describe("buildRuntimeContext", () => { ]); }); - it("uses a default noop compression pipeline when no compressor is provided", () => { + it("uses rule-based compression by default when no compressor is provided", () => { + const section = { + id: "plain", + content: "0123456789".repeat(300), + source: {kind: "file" as const}, + }; const compression = new ContextCompressionPipeline({ thresholdRatio: 0, + minSectionTokens: 1, tokenEstimator: charTokenEstimator, - }).compress( - [{id: "plain", content: "plain"}], - new DefaultContextBudgetPolicy().createBudget({ - sections: [], - tokenLimit: 1, - }), - ); + }).compress([section], rawBudget(4_000)); expect(compression.triggered).toBe(true); - expect(compression.sections).toEqual([{id: "plain", content: "plain"}]); + expect(compression.sections[0]?.content.length).toBeLessThan(section.content.length); expect(compression.results).toEqual([ expect.objectContaining({ - compressed: false, - reason: "No compression applied.", + compressed: true, + strategy: "rule-based-head-tail", }), ]); }); @@ -746,13 +810,29 @@ function budgetFor(section: ContextSection): ContextBudget { }); } -function rawBudget(maxInputTokens: number): ContextBudget { +function fixedRuntimeBudgetPolicy(runtimeContextTokens: number): ContextBudgetPolicy { + const basePolicy = new DefaultContextBudgetPolicy({reservedOutputTokens: 0}); + + return { + createBudget(input, tokenEstimator) { + return { + ...basePolicy.createBudget(input, tokenEstimator), + runtimeContextTokens, + }; + }, + }; +} + +function rawBudget(runtimeContextTokens: number): ContextBudget { return { - tokenLimit: maxInputTokens, - maxContextTokens: maxInputTokens, + tokenLimit: runtimeContextTokens, + modelContextWindow: runtimeContextTokens, reservedOutputTokens: 0, - maxInputTokens, - runtimeContextRatio: 0.35, - maxContextChars: maxInputTokens * 4, + systemPromptTokens: 0, + outputInstructionTokens: 0, + toolSchemaTokens: 0, + conversationTokens: 0, + safetyMarginTokens: 0, + runtimeContextTokens, }; } From 1145972a4eba475716318b0a63c93d79ef6f6765 Mon Sep 17 00:00:00 2001 From: TianYi <2739022972@qq.com> Date: Tue, 23 Jun 2026 01:05:46 +0800 Subject: [PATCH 2/2] feat: add docs --- src/context/budget/context-budget.ts | 66 +++++++++++++++++++++++++-- src/context/budget/token-estimator.ts | 31 +++++++++---- 2 files changed, 84 insertions(+), 13 deletions(-) diff --git a/src/context/budget/context-budget.ts b/src/context/budget/context-budget.ts index fb93453..45f6b53 100644 --- a/src/context/budget/context-budget.ts +++ b/src/context/budget/context-budget.ts @@ -1,14 +1,24 @@ import {createDefaultTokenEstimator, type TokenEstimator} from "./token-estimator.js"; import type {BuildContextInput, ContextBudget} from "../types.js"; -/** Strategy used to derive runtime context budget from build input. */ +/** + * Budget policy interface. + * + * Its job is to calculate how many tokens can be used by runtime context + * after reserving space for system prompt, tools, conversation history and output. + */ export interface ContextBudgetPolicy { createBudget(input: BuildContextInput, tokenEstimator?: TokenEstimator): ContextBudget; } export type DefaultContextBudgetPolicyOptions = { + /** Fallback model context window when input.tokenLimit is not provided. */ defaultMaxContextTokens?: number; + + /** Tokens reserved for model output generation. */ reservedOutputTokens?: number; + + /** Extra buffer to avoid hitting the model limit exactly. */ safetyMarginTokens?: number; }; @@ -16,21 +26,38 @@ const DEFAULT_MAX_CONTEXT_TOKENS = 128_000; const DEFAULT_RESERVED_OUTPUT_TOKENS = 4_000; const DEFAULT_SAFETY_MARGIN_TOKENS = 0; -/** Default budget policy that derives an input token budget from the model limit. */ +/** + * Default budget policy. + * + * It derives the runtime context budget from: + * + * model context window + * - reserved output tokens + * - base system prompt tokens + * - output instruction tokens + * - tool schema tokens + * - conversation history tokens + * - safety margin tokens + */ export class DefaultContextBudgetPolicy implements ContextBudgetPolicy { private readonly defaultMaxContextTokens: number; private readonly reservedOutputTokens: number; private readonly safetyMarginTokens: number; constructor(options: DefaultContextBudgetPolicyOptions = {}) { + // Normalize and validate default model context window. this.defaultMaxContextTokens = positiveIntegerOrDefault( options.defaultMaxContextTokens, DEFAULT_MAX_CONTEXT_TOKENS, ); + + // Normalize and validate output token reserve. this.reservedOutputTokens = nonNegativeIntegerOrDefault( options.reservedOutputTokens, DEFAULT_RESERVED_OUTPUT_TOKENS, ); + + // Normalize and validate safety margin. this.safetyMarginTokens = nonNegativeIntegerOrDefault( options.safetyMarginTokens, DEFAULT_SAFETY_MARGIN_TOKENS, @@ -41,29 +68,43 @@ export class DefaultContextBudgetPolicy implements ContextBudgetPolicy { input: BuildContextInput, tokenEstimator: TokenEstimator = createDefaultTokenEstimator(), ): ContextBudget { + // Use input.tokenLimit if provided; otherwise fall back to the default model window. const modelContextWindow = input.tokenLimit > 0 ? Math.floor(input.tokenLimit) : this.defaultMaxContextTokens; + + // Reserve output tokens, but cap it to 20% of the model window for small models. const reservedOutputTokens = Math.min( this.reservedOutputTokens, Math.floor(modelContextWindow * 0.2), ); + + // This is the total input-side budget after reserving model output space. const inputTokenBudget = Math.max(0, modelContextWindow - reservedOutputTokens); - // Count non-runtime prompt components before allocating runtime context. + + // Count fixed prompt components before allocating budget to runtime context. const systemPromptTokens = tokenEstimator.countText(input.baseSystemPrompt ?? ""); + const outputInstructionTokens = tokenEstimator.countText( input.outputInstructions ?? "", ); + + // Tool schemas also consume input tokens. Prefer caller-provided estimate if present. const toolSchemaTokens = input.toolSchemaTokens ?? estimateSerializableTokens(input.toolSchemas ?? [], tokenEstimator); + + // Conversation history also consumes input tokens. Prefer caller-provided estimate if present. const conversationTokens = input.conversationTokens ?? estimateConversationTokens(input.conversationMessages ?? [], tokenEstimator); + + // Safety margin can be overridden per build input. const safetyMarginTokens = nonNegativeIntegerOrDefault( input.safetyMarginTokens, this.safetyMarginTokens, ); - // Runtime context receives only the input tokens left after fixed prompt costs. + + // Runtime context gets only the remaining input budget. const runtimeContextTokens = Math.max( 0, inputTokenBudget - @@ -74,6 +115,8 @@ export class DefaultContextBudgetPolicy implements ContextBudgetPolicy { safetyMarginTokens, ); + // finalInputTokens and remainingInputTokens are filled later by ContextEngine + // after compression, truncation and prompt assembly are complete. return { tokenLimit: input.tokenLimit, modelContextWindow, @@ -90,6 +133,12 @@ export class DefaultContextBudgetPolicy implements ContextBudgetPolicy { } } +/** + * Estimate token usage of conversation messages. + * + * If the estimator supports countMessages, use it. + * Otherwise, fall back to counting each role + content pair as plain text. + */ function estimateConversationTokens( messages: NonNullable, tokenEstimator: TokenEstimator, @@ -108,6 +157,9 @@ function estimateConversationTokens( ); } +/** + * Estimate tokens for serializable objects, mainly tool schemas. + */ function estimateSerializableTokens( value: unknown, tokenEstimator: TokenEstimator, @@ -123,6 +175,9 @@ function estimateSerializableTokens( } } +/** + * Return a positive integer, or fallback if the value is invalid. + */ function positiveIntegerOrDefault(value: number | undefined, fallback: number): number { if (value === undefined || !Number.isFinite(value) || value <= 0) { return fallback; @@ -131,6 +186,9 @@ function positiveIntegerOrDefault(value: number | undefined, fallback: number): return Math.floor(value); } +/** + * Return a non-negative integer, or fallback if the value is invalid. + */ function nonNegativeIntegerOrDefault( value: number | undefined, fallback: number, diff --git a/src/context/budget/token-estimator.ts b/src/context/budget/token-estimator.ts index eaf661a..1bc8429 100644 --- a/src/context/budget/token-estimator.ts +++ b/src/context/budget/token-estimator.ts @@ -2,9 +2,15 @@ import {countTokens} from "gpt-tokenizer"; export type TokenCountableMessage = { role: string; - content?: string | null; + content?: unknown; + name?: string; + tool_call_id?: string; + tool_calls?: unknown[]; }; +const APPROX_CHARS_PER_TOKEN = 3; +const MESSAGE_OVERHEAD_TOKENS = 4; + /** Counts model tokens for runtime context budgeting. */ export interface TokenEstimator { countText(text: string): number; @@ -18,13 +24,16 @@ export class ApproxTokenEstimator implements TokenEstimator { return 0; } - return Math.max(1, Math.ceil(text.length / 4)); + const cjkChars = text.match(/[\u4e00-\u9fff]/g)?.length ?? 0; + const nonCjkChars = text.length - cjkChars; + + return Math.max(1, Math.ceil(cjkChars + nonCjkChars / APPROX_CHARS_PER_TOKEN)); } countMessages(messages: readonly TokenCountableMessage[]): number { return messages.reduce( (total, message) => - total + this.countText(`${message.role}\n${message.content ?? ""}`), + total + MESSAGE_OVERHEAD_TOKENS + this.countText(serializeMessage(message)), 0, ); } @@ -49,18 +58,14 @@ export class GptTokenEstimator implements TokenEstimator { countMessages(messages: readonly TokenCountableMessage[]): number { return messages.reduce( (total, message) => - total + this.countText(`${message.role}\n${message.content ?? ""}`), + total + MESSAGE_OVERHEAD_TOKENS + this.countText(serializeMessage(message)), 0, ); } } export function createDefaultTokenEstimator(): TokenEstimator { - try { - return new GptTokenEstimator(); - } catch { - return new ApproxTokenEstimator(); - } + return new GptTokenEstimator(); } const defaultTokenEstimator = createDefaultTokenEstimator(); @@ -69,3 +74,11 @@ const defaultTokenEstimator = createDefaultTokenEstimator(); export function estimateTokens(text: string): number { return defaultTokenEstimator.countText(text); } + +function serializeMessage(message: TokenCountableMessage): string { + try { + return JSON.stringify(message) ?? ""; + } catch { + return `${message.role}\n${String(message.content ?? "")}`; + } +}