diff --git a/package.json b/package.json index 44e2441..e786649 100644 --- a/package.json +++ b/package.json @@ -82,6 +82,7 @@ "chalk": "^5.6.2", "cli-highlight": "^2.1.11", "diff": "^9.0.0", + "gpt-tokenizer": "^3.4.0", "ink": "^7.0.3", "openai": "^6.38.0", "react": "^19.2.6", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6f0e739..a3f5120 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -22,6 +22,9 @@ importers: diff: specifier: ^9.0.0 version: 9.0.0 + gpt-tokenizer: + specifier: ^3.4.0 + version: 3.4.0 ink: specifier: ^7.0.3 version: 7.0.5(@types/react@19.2.17)(react@19.2.7) @@ -2096,6 +2099,12 @@ packages: } engines: {node: ">=10"} + gpt-tokenizer@3.4.0: + resolution: + { + integrity: sha512-wxFLnhIXTDjYebd9A9pGl3e31ZpSypbpIJSOswbgop5jLte/AsZVDvjlbEuVFlsqZixVKqbcoNmRlFDf6pz/UQ==, + } + graceful-fs@4.2.11: resolution: { @@ -4911,6 +4920,8 @@ snapshots: merge2: 1.4.1 slash: 3.0.0 + gpt-tokenizer@3.4.0: {} + graceful-fs@4.2.11: {} has-flag@4.0.0: {} diff --git a/src/agent/runtime/context-manager.ts b/src/agent/runtime/context-manager.ts index ae7416f..c63780a 100644 --- a/src/agent/runtime/context-manager.ts +++ b/src/agent/runtime/context-manager.ts @@ -1,4 +1,10 @@ import type {LLMMessage} from "../../llm/types.js"; +import { + ContextEngine, + RuleBasedContextCompressor, + type BuildContextDiagnostics, + type ContextSection, +} from "../../context/index.js"; import type { AgentContextProvider, AgentContextValue, @@ -33,52 +39,98 @@ export type ContextManagerOptions = { /** Builds model context and owns transcript mutation for assistant/tool turns. */ export class ContextManager { private readonly contextProviders: AgentContextProvider[]; + private readonly contextEngine = new ContextEngine({ + compressor: new RuleBasedContextCompressor(), + }); + private readonly runStates = new WeakMap(); /** Creates a context manager with static agent-level context providers. */ constructor(private readonly options: ContextManagerOptions) { this.contextProviders = [...(options.contextProviders ?? [])]; } - /** Builds the initial system, history, and user messages for a run. */ + /** Initializes static context sources and the mutable transcript for a run. */ async prepare(run: AgentRunState): Promise { - const values = [ - ...(await this.loadMemory(run)), - ...(run.input.context ?? []), + const sections = [ + ...(await this.loadMemory(run)).map((value) => toContextSection(value, "memory")), + ...(run.input.context ?? []).map((value) => toContextSection(value, "user")), { title: "Workspace Profile", priority: 100, content: JSON.stringify(run.workspaceProfile, null, 2), - }, + source: {kind: "workspace"}, + } satisfies ContextSection, ]; + + /** + * Context providers are deferred runtime context producers. + * + * They are useful when some context cannot be passed as static input, + * and must be generated from the current AgentRunContext at run time. + * + * Common examples: + * - Git status provider: + * injects current branch, modified files, staged files, or recent commits. + * + * - Run mode provider: + * injects mode-specific rules, such as "plan mode: do not modify files". + * + * - Diagnostics provider: + * injects recent test failures, lint errors, or command execution summaries. + * + * - External issue provider: + * injects GitHub / Linear / Jira issue details related to the current task. + * + * A provider may return either: + * - a string, which will use provider.name as the context section title + * - an AgentContextValue object, which may define its own title, priority, and content + * + * The returned value is converted into a ContextSection and then passed to + * the generic src/context pipeline together with memory, user context, and + * workspace profile sections. + */ const providers = [...this.contextProviders, ...(run.input.contextProviders ?? [])]; for (const provider of providers) { const value = await provider.build(run.context); - values.push( + sections.push( typeof value === "string" - ? {title: provider.name, content: value} - : {title: value.title ?? provider.name, ...value}, + ? { + title: provider.name, + content: value, + source: {kind: "provider", ref: provider.name}, + } + : { + title: value.title ?? provider.name, + ...value, + source: {kind: "provider", ref: provider.name}, + }, ); } - const contextText = truncateContext( - values.sort(compareContextValue).map(formatContextValue).filter(Boolean), - this.options.config.runtime.tokensLimit, - ); - this.options.observer.contextBuilt(run, estimateTokens(contextText)); - const systemPrompt = buildSystemPrompt( - run.input.systemPrompt ?? this.options.config.runtime.systemPrompt, - contextText, - ); - - run.messages.push({role: "system", content: systemPrompt}); + this.initializeRunState(run, sections); run.messages.push(...(run.input.messages ?? [])); run.messages.push({role: "user", content: run.input.prompt}); } - /** Returns the current transcript to send to the model runtime. */ + /** Builds a fresh model transcript with dynamic runtime context. */ buildModelRequest(run: AgentRunState): readonly LLMMessage[] { - return run.messages; + const state = this.getRunState(run); + const projection = this.projectTranscript(run); + const result = this.contextEngine.build({ + systemPrompt: + 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, + }); + + this.recordContextBuild(run, state, result.diagnostics); + this.options.observer.contextBuilt(run, result.tokenEstimate); + + return [{role: "system", content: result.systemPrompt}, ...projection.messages]; } /** Appends a normalized assistant response to the transcript. */ @@ -117,6 +169,80 @@ export class ContextManager { const projectMemory = await this.options.memory.loadProjectMemory?.(run); return [...(projectMemory ?? []), ...(runMemory ?? [])]; } + + private initializeRunState( + run: AgentRunState, + baseSections: readonly ContextSection[], + ): void { + const state: ContextManagerRunState = { + baseSections: [...baseSections], + contextBuilds: [], + }; + this.runStates.set(run, state); + run.context.contextBuilds = state.contextBuilds; + } + + private getRunState(run: AgentRunState): ContextManagerRunState { + const state = this.runStates.get(run); + if (state) { + return state; + } + + const fallbackState: ContextManagerRunState = { + baseSections: [], + contextBuilds: [], + }; + this.runStates.set(run, fallbackState); + run.context.contextBuilds = fallbackState.contextBuilds; + return fallbackState; + } + + private recordContextBuild( + run: AgentRunState, + state: ContextManagerRunState, + diagnostics: BuildContextDiagnostics | undefined, + ): void { + state.lastBuildDiagnostics = diagnostics; + state.contextBuilds.push({iteration: run.iteration, diagnostics}); + run.context.lastContextBuildDiagnostics = diagnostics; + run.context.contextBuilds = state.contextBuilds; + } + + private projectTranscript(run: AgentRunState): TranscriptProjection { + const latestExchangeStart = findLatestCompletedToolExchangeStart(run.messages); + const messages: LLMMessage[] = []; + const archivedToolSections: ContextSection[] = []; + + for (let index = 0; index < run.messages.length; index += 1) { + const message = run.messages[index]; + if (!message || message.role === "system") { + continue; + } + + if (isAssistantWithToolCalls(message)) { + const exchange = readToolExchange(run.messages, index); + if (exchange) { + if (index === latestExchangeStart) { + messages.push(message, ...exchange.toolMessages); + } else { + archivedToolSections.push( + ...exchange.toolMessages.map((toolMessage) => + createToolResultSection(message, toolMessage), + ), + ); + } + index = exchange.endIndex; + continue; + } + } + + if (message.role !== "tool") { + messages.push(message); + } + } + + return {messages, archivedToolSections}; + } } /** Creates the default context manager. */ @@ -124,66 +250,110 @@ export function createContextManager(options: ContextManagerOptions): ContextMan return new ContextManager(options); } -/** Estimates tokens using the runtime's coarse character-based heuristic. */ -export function estimateTokens(text: string): number { - return Math.ceil(text.length / 4); -} - -/** Combines the configured system prompt, CLI instructions, and runtime context. */ -function buildSystemPrompt( - systemPrompt: string | undefined, - contextText: string, -): string { - const prompt = systemPrompt ?? DEFAULT_SYSTEM_PROMPT; - const promptWithCliInstructions = `${prompt}\n\n${CLI_MARKDOWN_OUTPUT_INSTRUCTIONS}`; - - if (!contextText) { - return promptWithCliInstructions; +function toContextSection( + value: AgentContextValue, + source: "memory" | "user", +): ContextSection { + if (typeof value === "string") { + return {content: value, source: {kind: source}}; } - return `${promptWithCliInstructions}\n\n# Runtime Context\n${contextText}`; + return {...value, source: {kind: source}}; } -function formatContextValue(value: AgentContextValue): string { - if (typeof value === "string") { - return value.trim(); - } +type ContextManagerRunState = { + baseSections: ContextSection[]; + contextBuilds: Array<{ + iteration: number; + diagnostics?: BuildContextDiagnostics; + }>; + lastBuildDiagnostics?: BuildContextDiagnostics; +}; - const content = value.content.trim(); - if (!content) { - return ""; - } +type TranscriptProjection = { + messages: LLMMessage[]; + archivedToolSections: ContextSection[]; +}; - return value.title ? `## ${value.title}\n${content}` : content; -} +type AssistantToolMessage = Extract & { + toolCalls: NonNullable["toolCalls"]>; +}; -function compareContextValue(left: AgentContextValue, right: AgentContextValue): number { - return getContextPriority(right) - getContextPriority(left); -} +type ToolMessage = Extract; + +type ToolExchange = { + toolMessages: ToolMessage[]; + endIndex: number; +}; -function getContextPriority(value: AgentContextValue): number { - return typeof value === "string" ? 0 : (value.priority ?? 0); +function isAssistantWithToolCalls(message: LLMMessage): message is AssistantToolMessage { + return ( + message.role === "assistant" && + Array.isArray(message.toolCalls) && + message.toolCalls.length > 0 + ); } -function truncateContext(blocks: string[], tokenLimit: number): string { - const maxChars = Math.max(0, Math.floor(tokenLimit * 4 * 0.35)); - let remaining = maxChars; - const selectedBlocks: string[] = []; +function readToolExchange( + messages: readonly LLMMessage[], + assistantIndex: number, +): ToolExchange | undefined { + const assistant = messages[assistantIndex]; + if (!assistant || !isAssistantWithToolCalls(assistant)) { + return undefined; + } + + const expectedIds = new Set(assistant.toolCalls.map((toolCall) => toolCall.id)); + const toolMessages: ToolMessage[] = []; + let index = assistantIndex + 1; - for (const block of blocks) { - if (remaining <= 0) { + while (index < messages.length) { + const message = messages[index]; + if (!message || message.role !== "tool" || !expectedIds.has(message.toolCallId)) { break; } - const separatorLength = selectedBlocks.length ? 2 : 0; - const allowed = remaining - separatorLength; - if (allowed <= 0) { - break; + toolMessages.push(message); + index += 1; + + if (toolMessages.length === expectedIds.size) { + return {toolMessages, endIndex: index - 1}; } + } - selectedBlocks.push(block.length > allowed ? block.slice(0, allowed) : block); - remaining -= Math.min(block.length, allowed) + separatorLength; + return undefined; +} + +function findLatestCompletedToolExchangeStart( + messages: readonly LLMMessage[], +): number | undefined { + for (let index = messages.length - 1; index >= 0; index -= 1) { + if (readToolExchange(messages, index)) { + return index; + } } - return selectedBlocks.join("\n\n"); + return undefined; +} + +function createToolResultSection( + assistant: AssistantToolMessage, + toolMessage: ToolMessage, +): ContextSection { + return { + id: `tool-result:${toolMessage.toolCallId}`, + replaceKey: `tool-result:${toolMessage.toolCallId}`, + title: `Tool Result: ${toolMessage.name}`, + priority: 40, + source: {kind: "tool", ref: toolMessage.toolCallId}, + content: [ + `Tool: ${toolMessage.name}`, + `Call ID: ${toolMessage.toolCallId}`, + assistant.content ? `Assistant: ${assistant.content}` : undefined, + "Result:", + toolMessage.content, + ] + .filter((line): line is string => Boolean(line)) + .join("\n"), + }; } diff --git a/src/agent/types.ts b/src/agent/types.ts index 72ee15a..088531d 100644 --- a/src/agent/types.ts +++ b/src/agent/types.ts @@ -7,6 +7,7 @@ import type { VerificationConfig, } from "../config/index.js"; import type {PixelleEvent} from "../events/index.js"; +import type {BuildContextDiagnostics} from "../context/index.js"; import type {BaseLLMClient} from "../llm/index.js"; import type { LLMGenerateInput, @@ -99,6 +100,11 @@ export type AgentRunContext = { fileWriter?: ToolFileWriter; traceStore?: TraceStore; workspaceProfile?: WorkspaceProfile; + lastContextBuildDiagnostics?: BuildContextDiagnostics; + contextBuilds?: Array<{ + iteration: number; + diagnostics?: BuildContextDiagnostics; + }>; }; /** Model request enriched with agent trace data. */ diff --git a/src/context/context-budget.ts b/src/context/context-budget.ts new file mode 100644 index 0000000..62179f2 --- /dev/null +++ b/src/context/context-budget.ts @@ -0,0 +1,73 @@ +import type {BuildContextInput, ContextBudget} from "./types.js"; + +/** Strategy used to derive runtime context budget from build input. */ +export interface ContextBudgetPolicy { + createBudget(input: BuildContextInput): ContextBudget; +} + +export type DefaultContextBudgetPolicyOptions = { + defaultMaxContextTokens?: number; + reservedOutputTokens?: number; + /** @deprecated kept only to populate legacy char diagnostics. */ + runtimeContextRatio?: number; +}; + +const DEFAULT_MAX_CONTEXT_TOKENS = 128_000; +const DEFAULT_RESERVED_OUTPUT_TOKENS = 4_000; + +/** 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; + + constructor(options: DefaultContextBudgetPolicyOptions = {}) { + this.defaultMaxContextTokens = positiveIntegerOrDefault( + options.defaultMaxContextTokens, + DEFAULT_MAX_CONTEXT_TOKENS, + ); + this.reservedOutputTokens = nonNegativeIntegerOrDefault( + options.reservedOutputTokens, + DEFAULT_RESERVED_OUTPUT_TOKENS, + ); + this.runtimeContextRatio = options.runtimeContextRatio ?? 0.35; + } + + createBudget(input: BuildContextInput): ContextBudget { + const maxContextTokens = + input.tokenLimit > 0 ? Math.floor(input.tokenLimit) : this.defaultMaxContextTokens; + const reservedOutputTokens = Math.min( + this.reservedOutputTokens, + Math.floor(maxContextTokens * 0.2), + ); + const maxInputTokens = Math.max(0, maxContextTokens - reservedOutputTokens); + + return { + tokenLimit: input.tokenLimit, + maxContextTokens, + reservedOutputTokens, + maxInputTokens, + runtimeContextRatio: this.runtimeContextRatio, + maxContextChars: Math.max(0, maxInputTokens * 4), + }; + } +} + +function positiveIntegerOrDefault(value: number | undefined, fallback: number): number { + if (value === undefined || !Number.isFinite(value) || value <= 0) { + return fallback; + } + + return Math.floor(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/context-builder.ts b/src/context/context-builder.ts new file mode 100644 index 0000000..6df4489 --- /dev/null +++ b/src/context/context-builder.ts @@ -0,0 +1,9 @@ +import {ContextEngine} from "./context-engine.js"; +import type {BuildContextInput, BuildContextResult} from "./types.js"; + +const defaultContextEngine = new ContextEngine(); + +/** Builds the final system prompt and model-visible runtime context. */ +export function buildRuntimeContext(input: BuildContextInput): BuildContextResult { + return defaultContextEngine.build(input); +} diff --git a/src/context/context-compression-pipeline.ts b/src/context/context-compression-pipeline.ts new file mode 100644 index 0000000..9bac7fb --- /dev/null +++ b/src/context/context-compression-pipeline.ts @@ -0,0 +1,80 @@ +import { + ContextCompressionResultFactory, + NoopContextCompressor, + type ContextCompressor, +} from "./context-compressor.js"; +import {formatContextSection} from "./context-formatter.js"; +import {createDefaultTokenEstimator, type TokenEstimator} from "./token-estimator.js"; +import type {ContextBudget, ContextCompressionResult, ContextSection} from "./types.js"; + +export type ContextCompressionPipelineOptions = { + compressor?: ContextCompressor; + thresholdRatio?: number; + resultFactory?: ContextCompressionResultFactory; + tokenEstimator?: TokenEstimator; +}; + +export type ContextCompressionPipelineResult = { + sections: ContextSection[]; + results: ContextCompressionResult[]; + estimatedContextChars: number; + estimatedContextTokens: number; + compressionLimitTokens: number; + thresholdRatio: number; + triggered: boolean; +}; + +/** Decides when to apply context compression and records compression diagnostics. */ +export class ContextCompressionPipeline { + private readonly compressor: ContextCompressor; + private readonly thresholdRatio: 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(); + } + + compress( + sections: readonly ContextSection[], + budget: ContextBudget, + ): ContextCompressionPipelineResult { + const formattedContextText = formatContextSections(sections); + const estimatedContextChars = formattedContextText.length; + const estimatedContextTokens = this.tokenEstimator.countText(formattedContextText); + const compressionLimitTokens = Math.floor( + budget.maxInputTokens * 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, + }), + ); + + return { + sections: results.map((result) => result.section), + results, + estimatedContextChars, + estimatedContextTokens, + compressionLimitTokens, + thresholdRatio: this.thresholdRatio, + triggered, + }; + } +} + +function formatContextSections(sections: readonly ContextSection[]): string { + return sections + .map((section) => formatContextSection(section)) + .filter((text) => text.length > 0) + .join("\n\n"); +} diff --git a/src/context/context-compressor.ts b/src/context/context-compressor.ts new file mode 100644 index 0000000..90ddc5d --- /dev/null +++ b/src/context/context-compressor.ts @@ -0,0 +1,643 @@ +import {createDefaultTokenEstimator, type TokenEstimator} from "./token-estimator.js"; +import type {ContextBudget, ContextCompressionResult, ContextSection} from "./types.js"; + +export type ContextCompressionMetadata = { + strategy?: string; + maxSectionChars?: number; + maxSectionTokens?: number; + omittedChars?: number; + savedChars?: number; + originalTokens?: number; + compressedTokens?: number; + savedTokens?: number; + tokenCompressionRatio?: number; +}; + +type HeadTailParts = { + head: string; + tail: string; + marker: string; + content: string; + omittedChars: number; +}; + +/** Extension point for pre-truncation context compression. */ +export interface ContextCompressor { + compress(section: ContextSection, budget: ContextBudget): ContextCompressionResult; +} + +/** Creates normalized compression results for compressors and pipelines. */ +export class ContextCompressionResultFactory { + private readonly tokenEstimator: TokenEstimator; + + constructor(tokenEstimator: TokenEstimator = createDefaultTokenEstimator()) { + this.tokenEstimator = tokenEstimator; + } + + unchanged( + section: ContextSection, + reason: string, + metadata: ContextCompressionMetadata = {}, + ): ContextCompressionResult { + return this.create(section, section, false, reason, metadata); + } + + compressed( + section: ContextSection, + originalSection: ContextSection, + reason: string, + metadata: ContextCompressionMetadata = {}, + ): ContextCompressionResult { + return this.create(section, originalSection, true, reason, metadata); + } + + skipped( + section: ContextSection, + reason: string, + metadata: ContextCompressionMetadata = {}, + ): ContextCompressionResult { + return this.create(section, section, false, reason, metadata); + } + + private create( + section: ContextSection, + originalSection: ContextSection, + compressed: boolean, + reason: string, + metadata: ContextCompressionMetadata, + ): ContextCompressionResult { + const savedChars = + metadata.savedChars ?? + Math.max(0, originalSection.content.length - section.content.length); + const omittedChars = metadata.omittedChars ?? savedChars; + const compressionRatio = + originalSection.content.length > 0 + ? section.content.length / originalSection.content.length + : 1; + const originalTokens = + metadata.originalTokens ?? this.tokenEstimator.countText(originalSection.content); + const compressedTokens = + metadata.compressedTokens ?? this.tokenEstimator.countText(section.content); + const savedTokens = + metadata.savedTokens ?? Math.max(0, originalTokens - compressedTokens); + const tokenCompressionRatio = + metadata.tokenCompressionRatio ?? + (originalTokens > 0 ? compressedTokens / originalTokens : 1); + + return { + section, + originalSection, + compressed, + originalChars: originalSection.content.length, + compressedChars: section.content.length, + omittedChars, + savedChars, + compressionRatio, + originalTokens, + compressedTokens, + savedTokens, + tokenCompressionRatio, + ...metadata, + reason, + }; + } +} + +/** 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; + maxSectionChars?: number; + minSectionChars?: number; + maxSectionRatio?: number; + headTokenRatio?: number; + tailTokenRatio?: number; + headChars?: number; + tailChars?: number; + preserveLineBoundaries?: boolean; + tokenEstimator?: TokenEstimator; +}; + +/** Conservative compressor for oversized tool and file context sections. */ +export class RuleBasedContextCompressor implements ContextCompressor { + private readonly maxSectionTokens: number; + private readonly minSectionTokens: number; + private readonly maxSectionChars: number; + private readonly minSectionChars: number; + private readonly maxSectionRatio: number; + private readonly headTokenRatio: number; + private readonly tailTokenRatio: number; + private readonly headChars: number; + private readonly tailChars: number; + private readonly preserveLineBoundaries: boolean; + private readonly resultFactory: ContextCompressionResultFactory; + private readonly tokenEstimator: TokenEstimator; + + constructor( + options: RuleBasedContextCompressorOptions = {}, + resultFactory?: ContextCompressionResultFactory, + ) { + this.tokenEstimator = options.tokenEstimator ?? createDefaultTokenEstimator(); + const maxSectionTokens = this.positiveIntegerOrDefault( + options.maxSectionTokens, + 8_000, + ); + const minSectionTokens = this.positiveIntegerOrDefault( + options.minSectionTokens, + 1_200, + ); + const maxSectionChars = this.positiveIntegerOrDefault(options.maxSectionChars, 8_000); + const minSectionChars = this.positiveIntegerOrDefault(options.minSectionChars, 1_200); + + this.maxSectionTokens = maxSectionTokens; + this.minSectionTokens = Math.min(minSectionTokens, maxSectionTokens); + this.maxSectionChars = maxSectionChars; + this.minSectionChars = Math.min(minSectionChars, maxSectionChars); + this.maxSectionRatio = this.clampNumber(options.maxSectionRatio ?? 0.25, 0.05, 1); + this.headTokenRatio = this.clampNumber(options.headTokenRatio ?? 0.67, 0, 1); + this.tailTokenRatio = this.clampNumber(options.tailTokenRatio ?? 0.33, 0, 1); + this.headChars = this.nonNegativeIntegerOrDefault(options.headChars, 4_000); + this.tailChars = this.nonNegativeIntegerOrDefault(options.tailChars, 2_000); + this.preserveLineBoundaries = options.preserveLineBoundaries ?? true; + this.resultFactory = + resultFactory ?? new ContextCompressionResultFactory(this.tokenEstimator); + } + + compress(section: ContextSection, budget: ContextBudget): ContextCompressionResult { + if (!isCompressibleSection(section)) { + return this.resultFactory.unchanged(section, "Section source is not compressible."); + } + + switch (section.source?.kind) { + case "tool": + return this.compressToolSection(section, budget); + case "file": + return this.compressFileSection(section, budget); + default: + return this.resultFactory.unchanged( + section, + "Section source is not compressible.", + ); + } + } + + private compressToolSection( + section: ContextSection, + budget: ContextBudget, + ): ContextCompressionResult { + return this.compressHeadTail(section, this.resolveMaxSectionTokens(budget)); + } + + private compressFileSection( + section: ContextSection, + budget: ContextBudget, + ): ContextCompressionResult { + return this.compressHeadTail(section, this.resolveMaxSectionTokens(budget)); + } + + private resolveMaxSectionTokens(budget: ContextBudget): number { + const budgetBasedLimit = Math.floor(budget.maxInputTokens * this.maxSectionRatio); + + return this.clampInteger( + Math.min(this.maxSectionTokens, budgetBasedLimit), + this.minSectionTokens, + this.maxSectionTokens, + ); + } + + private compressHeadTail( + section: ContextSection, + maxSectionTokens: number, + ): ContextCompressionResult { + const strategy = "rule-based-head-tail"; + const originalTokens = this.tokenEstimator.countText(section.content); + const maxSectionChars = this.clampInteger( + Math.min(this.maxSectionChars, maxSectionTokens * 4), + this.minSectionChars, + this.maxSectionChars, + ); + + if (!section.content.trim()) { + return this.resultFactory.unchanged( + section, + "Section content is empty after trimming.", + { + strategy, + maxSectionChars, + maxSectionTokens, + originalTokens, + compressedTokens: originalTokens, + }, + ); + } + + if (originalTokens <= maxSectionTokens) { + return this.resultFactory.unchanged( + section, + `Section is within the dynamic section token limit (${originalTokens} <= ${maxSectionTokens} tokens).`, + { + strategy, + maxSectionChars, + maxSectionTokens, + originalTokens, + compressedTokens: originalTokens, + }, + ); + } + + const parts = this.buildHeadTailContent( + section.content, + maxSectionChars, + maxSectionTokens, + strategy, + ); + const compressedTokens = this.tokenEstimator.countText(parts.content); + + if ( + !parts.content.trim() || + parts.content.length >= section.content.length || + compressedTokens >= originalTokens + ) { + return this.resultFactory.unchanged( + section, + "Rule-based compression would not reduce this section.", + { + strategy, + maxSectionChars, + maxSectionTokens, + originalTokens, + compressedTokens: originalTokens, + }, + ); + } + + const compressedSection = {...section, content: parts.content}; + const savedChars = section.content.length - parts.content.length; + const savedTokens = originalTokens - compressedTokens; + + return this.resultFactory.compressed( + compressedSection, + section, + `Section exceeded dynamic section token limit: ${originalTokens} tokens -> ${compressedTokens} tokens, omitted ${parts.omittedChars} original chars, saved ${savedChars} chars and ${savedTokens} tokens using rule-based-head-tail compression.`, + { + strategy, + maxSectionChars, + maxSectionTokens, + omittedChars: parts.omittedChars, + savedChars, + originalTokens, + compressedTokens, + savedTokens, + tokenCompressionRatio: originalTokens > 0 ? compressedTokens / originalTokens : 1, + }, + ); + } + + private buildHeadTailContent( + content: string, + maxSectionChars: number, + maxSectionTokens: number, + strategy: string, + ): HeadTailParts { + let marker = this.createOmissionMarker(content.length, strategy); + let availableContentChars = Math.max(0, maxSectionChars - marker.length); + let budgets = this.allocateHeadTailBudgets(availableContentChars); + let tokenBudgets = this.allocateHeadTailTokenBudgets(maxSectionTokens, marker); + + for (let attempts = 0; attempts < 12; attempts += 1) { + const head = + budgets.headChars > 0 && tokenBudgets.headTokens > 0 + ? this.truncateHeadToTokens( + this.sliceHead(content, budgets.headChars), + tokenBudgets.headTokens, + ) + : ""; + const tail = + budgets.tailChars > 0 && tokenBudgets.tailTokens > 0 + ? this.truncateTailToTokens( + this.sliceTail(content, budgets.tailChars), + tokenBudgets.tailTokens, + ) + : ""; + const omittedChars = Math.max(0, content.length - head.length - tail.length); + marker = this.createFittingOmissionMarker(omittedChars, strategy, maxSectionChars); + const compressedContent = `${head}${marker}${tail}`; + + if ( + compressedContent.length <= maxSectionChars && + this.tokenEstimator.countText(compressedContent) <= maxSectionTokens + ) { + return {head, tail, marker, content: compressedContent, omittedChars}; + } + + const overflow = compressedContent.length - maxSectionChars; + budgets = this.reduceHeadTailBudgets(budgets, Math.max(1, overflow)); + tokenBudgets = this.reduceHeadTailTokenBudgets(tokenBudgets, 1); + availableContentChars = Math.max(0, availableContentChars - overflow); + if ( + budgets.headChars + budgets.tailChars <= 0 && + tokenBudgets.headTokens + tokenBudgets.tailTokens <= 0 && + marker.length <= maxSectionChars && + this.tokenEstimator.countText(marker) <= maxSectionTokens + ) { + return { + head: "", + tail: "", + marker, + content: marker, + omittedChars: content.length, + }; + } + } + + const markerOnly = this.createFittingOmissionMarker( + content.length, + strategy, + maxSectionChars, + ); + return { + head: "", + tail: "", + marker: markerOnly, + content: markerOnly, + omittedChars: content.length, + }; + } + + private allocateHeadTailTokenBudgets( + maxSectionTokens: number, + marker: string, + ): { + headTokens: number; + tailTokens: number; + } { + const availableTokens = Math.max( + 0, + maxSectionTokens - this.tokenEstimator.countText(marker), + ); + if (availableTokens <= 0) { + return {headTokens: 0, tailTokens: 0}; + } + + const ratioTotal = this.headTokenRatio + this.tailTokenRatio; + const headRatio = ratioTotal > 0 ? this.headTokenRatio / ratioTotal : 0.67; + const headTokens = Math.floor(availableTokens * headRatio); + + return { + headTokens, + tailTokens: Math.max(0, availableTokens - headTokens), + }; + } + + private allocateHeadTailBudgets(availableContentChars: number): { + headChars: number; + tailChars: number; + } { + if (availableContentChars <= 0) { + return {headChars: 0, tailChars: 0}; + } + + let configuredHeadChars = this.headChars; + let configuredTailChars = this.tailChars; + + if (configuredHeadChars + configuredTailChars === 0) { + configuredHeadChars = Math.ceil(availableContentChars * 0.67); + configuredTailChars = availableContentChars - configuredHeadChars; + } + + const totalConfiguredChars = configuredHeadChars + configuredTailChars; + if (totalConfiguredChars <= availableContentChars) { + return { + headChars: configuredHeadChars, + tailChars: configuredTailChars, + }; + } + + const headChars = + configuredHeadChars === 0 + ? 0 + : Math.floor( + (availableContentChars * configuredHeadChars) / totalConfiguredChars, + ); + + return { + headChars, + tailChars: Math.max(0, availableContentChars - headChars), + }; + } + + private sliceHead(content: string, maxChars: number): string { + return this.preserveLineBoundaries + ? this.sliceHeadPreservingLines(content, maxChars) + : content.slice(0, maxChars); + } + + private sliceTail(content: string, maxChars: number): string { + return this.preserveLineBoundaries + ? this.sliceTailPreservingLines(content, maxChars) + : content.slice(-maxChars); + } + + private truncateHeadToTokens(content: string, maxTokens: number): string { + if (!content || maxTokens <= 0) { + return ""; + } + + if (this.tokenEstimator.countText(content) <= maxTokens) { + return content; + } + + let low = 0; + let high = content.length; + let best = ""; + + while (low <= high) { + const mid = Math.floor((low + high) / 2); + const candidate = content.slice(0, mid); + if (this.tokenEstimator.countText(candidate) <= maxTokens) { + best = candidate; + low = mid + 1; + } else { + high = mid - 1; + } + } + + return this.preserveLineBoundaries + ? this.sliceHeadPreservingLines(best, best.length) + : best; + } + + private truncateTailToTokens(content: string, maxTokens: number): string { + if (!content || maxTokens <= 0) { + return ""; + } + + if (this.tokenEstimator.countText(content) <= maxTokens) { + return content; + } + + let low = 0; + let high = content.length; + let best = ""; + + while (low <= high) { + const mid = Math.floor((low + high) / 2); + const candidate = content.slice(content.length - mid); + if (this.tokenEstimator.countText(candidate) <= maxTokens) { + best = candidate; + low = mid + 1; + } else { + high = mid - 1; + } + } + + return this.preserveLineBoundaries + ? this.sliceTailPreservingLines(best, best.length) + : best; + } + + private createOmissionMarker(omittedChars: number, strategy: string): string { + return `\n\n[...${omittedChars} chars omitted by ${strategy} compressor...]\n\n`; + } + + private createFittingOmissionMarker( + omittedChars: number, + strategy: string, + maxChars: number, + ): string { + const fullMarker = this.createOmissionMarker(omittedChars, strategy); + if (fullMarker.length <= maxChars) { + return fullMarker; + } + + const compactMarker = `\n\n[...${omittedChars} chars omitted...]\n\n`; + if (compactMarker.length <= maxChars) { + return compactMarker; + } + + const minimalMarker = "[...omitted...]"; + if (minimalMarker.length <= maxChars) { + return minimalMarker; + } + + return ".".repeat(Math.max(0, maxChars)); + } + + private sliceHeadPreservingLines(content: string, maxChars: number): string { + const sliced = content.slice(0, Math.max(1, maxChars)); + const lastNewline = sliced.lastIndexOf("\n"); + + if (lastNewline > 0) { + const linePreserved = sliced.slice(0, lastNewline + 1); + if (linePreserved.length >= Math.ceil(maxChars * 0.5)) { + return linePreserved; + } + } + + return sliced; + } + + private sliceTailPreservingLines(content: string, maxChars: number): string { + const sliced = content.slice(-Math.max(1, maxChars)); + const firstNewline = sliced.indexOf("\n"); + + if (firstNewline >= 0 && firstNewline < sliced.length - 1) { + const linePreserved = sliced.slice(firstNewline + 1); + if (linePreserved.length >= Math.ceil(maxChars * 0.5)) { + return linePreserved; + } + } + + return sliced; + } + + private reduceHeadTailBudgets( + budgets: {headChars: number; tailChars: number}, + overflow: number, + ): {headChars: number; tailChars: number} { + let remainingOverflow = Math.max(1, overflow); + const tailReduction = Math.min(budgets.tailChars, remainingOverflow); + remainingOverflow -= tailReduction; + + return { + headChars: Math.max(0, budgets.headChars - remainingOverflow), + tailChars: budgets.tailChars - tailReduction, + }; + } + + private reduceHeadTailTokenBudgets( + budgets: {headTokens: number; tailTokens: number}, + overflow: number, + ): {headTokens: number; tailTokens: number} { + let remainingOverflow = Math.max(1, overflow); + const tailReduction = Math.min(budgets.tailTokens, remainingOverflow); + remainingOverflow -= tailReduction; + + return { + headTokens: Math.max(0, budgets.headTokens - remainingOverflow), + tailTokens: budgets.tailTokens - tailReduction, + }; + } + + private positiveIntegerOrDefault(value: number | undefined, fallback: number): number { + if (value === undefined || !Number.isFinite(value) || value <= 0) { + return fallback; + } + + return Math.floor(value); + } + + private nonNegativeIntegerOrDefault( + value: number | undefined, + fallback: number, + ): number { + if (value === undefined || !Number.isFinite(value) || value < 0) { + return fallback; + } + + return Math.floor(value); + } + + private clampInteger(value: number, min: number, max: number): number { + return Math.floor(this.clampNumber(value, min, max)); + } + + private clampNumber(value: number, min: number, max: number): number { + if (!Number.isFinite(value)) { + return min; + } + + return Math.min(max, Math.max(min, value)); + } +} + +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); + } + + return factory.unchanged(section, reason ?? "", metadata); +} diff --git a/src/context/context-engine.ts b/src/context/context-engine.ts new file mode 100644 index 0000000..25fb09a --- /dev/null +++ b/src/context/context-engine.ts @@ -0,0 +1,80 @@ +import {ContextCompressionPipeline} from "./context-compression-pipeline.js"; +import {DefaultContextBudgetPolicy} from "./context-budget.js"; +import {ContextRegistry} from "./context-registry.js"; +import {ContextTruncator} from "./context-truncator.js"; +import {DefaultContextPriorityPolicy} from "./priority-policy.js"; +import {SystemPromptAssembler} from "./system-prompt-assembler.js"; +import {createDefaultTokenEstimator, type TokenEstimator} from "./token-estimator.js"; +import type {ContextBudgetPolicy} from "./context-budget.js"; +import type {ContextPriorityPolicy} from "./priority-policy.js"; +import type { + BuildContextInput, + BuildContextResult, + ContextEngineOptions, +} from "./types.js"; + +/** Class-based context engine that owns the full runtime context pipeline. */ +export class ContextEngine { + private readonly priorityPolicy: ContextPriorityPolicy; + private readonly budgetPolicy: ContextBudgetPolicy; + private readonly compressionPipeline: ContextCompressionPipeline; + private readonly truncator: ContextTruncator; + private readonly assembler: SystemPromptAssembler; + private readonly tokenEstimator: TokenEstimator; + + constructor(options: ContextEngineOptions = {}) { + this.tokenEstimator = options.tokenEstimator ?? createDefaultTokenEstimator(); + this.priorityPolicy = options.priorityPolicy ?? new DefaultContextPriorityPolicy(); + this.budgetPolicy = options.budgetPolicy ?? new DefaultContextBudgetPolicy(); + this.compressionPipeline = + options.compressionPipeline ?? + new ContextCompressionPipeline({ + compressor: options.compressor, + thresholdRatio: options.compressionThresholdRatio, + tokenEstimator: this.tokenEstimator, + }); + this.truncator = + options.truncator ?? new ContextTruncator({tokenEstimator: this.tokenEstimator}); + this.assembler = options.assembler ?? new SystemPromptAssembler(); + } + + build(input: BuildContextInput): BuildContextResult { + const registry = new ContextRegistry() + .addMany(input.sections) + .normalize() + .dedupe() + .sort(this.priorityPolicy); + const budget = this.budgetPolicy.createBudget(input); + const sections = registry.getAll(); + const compression = this.compressionPipeline.compress(sections, budget); + const truncation = this.truncator.truncate(compression.sections, budget); + const systemPrompt = this.assembler.assemble({ + systemPrompt: input.systemPrompt, + outputInstructions: input.outputInstructions, + contextText: truncation.contextText, + }); + const contextTextTokens = this.tokenEstimator.countText(truncation.contextText); + const systemPromptTokens = this.tokenEstimator.countText(systemPrompt); + + return { + systemPrompt, + contextText: truncation.contextText, + tokenEstimate: systemPromptTokens, + includedSections: truncation.includedSections, + partialSections: truncation.partialSections, + droppedSections: truncation.droppedSections, + sectionUsages: truncation.sectionUsages, + diagnostics: { + budget, + estimatedContextChars: compression.estimatedContextChars, + estimatedContextTokens: compression.estimatedContextTokens, + compressionThresholdRatio: compression.thresholdRatio, + compressionTriggered: compression.triggered, + compressionLimitTokens: compression.compressionLimitTokens, + compressionResults: compression.results, + contextTextTokens, + systemPromptTokens, + }, + }; + } +} diff --git a/src/context/context-formatter.ts b/src/context/context-formatter.ts new file mode 100644 index 0000000..c9c7aff --- /dev/null +++ b/src/context/context-formatter.ts @@ -0,0 +1,21 @@ +import type {ContextSection} from "./types.js"; +import {DefaultContextPriorityPolicy} from "./priority-policy.js"; + +/** Returns a formatted section string, or an empty string for blank content. */ +export function formatContextSection(section: ContextSection): string { + const content = section.content.trim(); + + if (!content) { + return ""; + } + + 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/context-registry.ts b/src/context/context-registry.ts new file mode 100644 index 0000000..e21795f --- /dev/null +++ b/src/context/context-registry.ts @@ -0,0 +1,64 @@ +import type {ContextPriorityPolicy} from "./priority-policy.js"; +import type {ContextSection} from "./types.js"; + +/** Mutable registry for normalizing, deduping, and ordering context sections. */ +export class ContextRegistry { + private sections: ContextSection[] = []; + + add(section: ContextSection): this { + this.sections.push(section); + return this; + } + + addMany(sections: readonly ContextSection[]): this { + this.sections.push(...sections); + return this; + } + + normalize(): this { + this.sections = this.sections + .map((section) => ({...section, content: section.content.trim()})) + .filter((section) => section.content.length > 0); + return this; + } + + dedupe(): this { + const seen = new Set(); + const deduped: ContextSection[] = []; + + for (let index = this.sections.length - 1; index >= 0; index -= 1) { + const section = this.sections[index]; + if (!section) { + continue; + } + const key = section.replaceKey ?? section.id; + + if (key) { + if (seen.has(key)) { + continue; + } + seen.add(key); + } + + deduped.push(section); + } + + this.sections = deduped.reverse(); + return this; + } + + sort(policy: ContextPriorityPolicy): this { + this.sections = this.sections + .map((section, index) => ({section, index})) + .sort((left, right) => { + const priorityComparison = policy.compare(left.section, right.section); + return priorityComparison || left.index - right.index; + }) + .map((entry) => entry.section); + return this; + } + + getAll(): ContextSection[] { + return [...this.sections]; + } +} diff --git a/src/context/context-truncator.ts b/src/context/context-truncator.ts new file mode 100644 index 0000000..7ba066d --- /dev/null +++ b/src/context/context-truncator.ts @@ -0,0 +1,221 @@ +import {DefaultContextBudgetPolicy} from "./context-budget.js"; +import {formatContextSection} from "./context-formatter.js"; +import {createDefaultTokenEstimator, type TokenEstimator} from "./token-estimator.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/index.ts b/src/context/index.ts new file mode 100644 index 0000000..f39ae53 --- /dev/null +++ b/src/context/index.ts @@ -0,0 +1,54 @@ +export {ContextEngine} from "./context-engine.js"; +export {buildRuntimeContext} from "./context-builder.js"; +export {ContextCompressionPipeline} from "./context-compression-pipeline.js"; +export type { + ContextCompressionPipelineOptions, + ContextCompressionPipelineResult, +} from "./context-compression-pipeline.js"; +export {DefaultContextBudgetPolicy} from "./context-budget.js"; +export type {ContextBudgetPolicy} from "./context-budget.js"; +export { + ContextCompressionResultFactory, + createCompressionResult, + isCompressibleSection, + NoopContextCompressor, + RuleBasedContextCompressor, +} from "./context-compressor.js"; +export type { + ContextCompressor, + ContextCompressionMetadata, + RuleBasedContextCompressorOptions, +} from "./context-compressor.js"; +export {compareContextSection, formatContextSection} from "./context-formatter.js"; +export {ContextRegistry} from "./context-registry.js"; +export { + ContextTruncator, + truncateContext, + truncateTextToTokens, +} from "./context-truncator.js"; +export type { + FormattedContextSection, + TruncateContextResult, +} from "./context-truncator.js"; +export {DefaultContextPriorityPolicy} from "./priority-policy.js"; +export type {ContextPriorityPolicy} from "./priority-policy.js"; +export {SystemPromptAssembler} from "./system-prompt-assembler.js"; +export { + ApproxTokenEstimator, + createDefaultTokenEstimator, + estimateTokens, + GptTokenEstimator, +} from "./token-estimator.js"; +export type {TokenCountableMessage, TokenEstimator} from "./token-estimator.js"; +export type { + BuildContextInput, + BuildContextDiagnostics, + BuildContextResult, + ContextBudget, + ContextCompressionResult, + ContextEngineOptions, + ContextSection, + ContextSectionUsage, + ContextSectionUsageStatus, + ContextSource, +} from "./types.js"; diff --git a/src/context/priority-policy.ts b/src/context/priority-policy.ts new file mode 100644 index 0000000..c17e956 --- /dev/null +++ b/src/context/priority-policy.ts @@ -0,0 +1,39 @@ +import type {ContextSection, ContextSource} from "./types.js"; + +/** Strategy used to order context sections. */ +export interface ContextPriorityPolicy { + priorityOf(section: ContextSection): number; + compare(left: ContextSection, right: ContextSection): number; +} + +/** Default priority policy with explicit priority taking precedence. */ +export class DefaultContextPriorityPolicy implements ContextPriorityPolicy { + priorityOf(section: ContextSection): number { + return section.priority ?? priorityForSource(section.source); + } + + compare(left: ContextSection, right: ContextSection): number { + return this.priorityOf(right) - this.priorityOf(left); + } +} + +function priorityForSource(source: ContextSource | undefined): number { + switch (source?.kind) { + case "workspace": + return 100; + case "user": + return 80; + case "memory": + return 60; + case "provider": + return 50; + case "tool": + return 40; + case "file": + return 30; + case "system": + return 20; + default: + return 0; + } +} diff --git a/src/context/system-prompt-assembler.ts b/src/context/system-prompt-assembler.ts new file mode 100644 index 0000000..91b9b9d --- /dev/null +++ b/src/context/system-prompt-assembler.ts @@ -0,0 +1,23 @@ +/** Assembles the final system prompt from base prompt, output rules, and context. */ +export class SystemPromptAssembler { + assemble(input: { + systemPrompt?: string; + outputInstructions?: string; + contextText: string; + }): string { + const promptParts = [input.systemPrompt, input.outputInstructions].filter( + (part): part is string => Boolean(part), + ); + const prompt = promptParts.join("\n\n"); + + if (!input.contextText) { + return prompt; + } + + if (!prompt) { + return `# Runtime Context\n${input.contextText}`; + } + + return `${prompt}\n\n# Runtime Context\n${input.contextText}`; + } +} diff --git a/src/context/token-estimator.ts b/src/context/token-estimator.ts new file mode 100644 index 0000000..eaf661a --- /dev/null +++ b/src/context/token-estimator.ts @@ -0,0 +1,71 @@ +import {countTokens} from "gpt-tokenizer"; + +export type TokenCountableMessage = { + role: string; + content?: string | null; +}; + +/** Counts model tokens for runtime context budgeting. */ +export interface TokenEstimator { + countText(text: string): number; + countMessages?(messages: readonly TokenCountableMessage[]): number; +} + +/** Fallback estimator used when an exact tokenizer is unavailable. */ +export class ApproxTokenEstimator implements TokenEstimator { + countText(text: string): number { + if (!text) { + return 0; + } + + return Math.max(1, Math.ceil(text.length / 4)); + } + + countMessages(messages: readonly TokenCountableMessage[]): number { + return messages.reduce( + (total, message) => + total + this.countText(`${message.role}\n${message.content ?? ""}`), + 0, + ); + } +} + +/** Token estimator backed by gpt-tokenizer with a safe approximation fallback. */ +export class GptTokenEstimator implements TokenEstimator { + private readonly fallback = new ApproxTokenEstimator(); + + countText(text: string): number { + if (!text) { + return 0; + } + + try { + return countTokens(text); + } catch { + return this.fallback.countText(text); + } + } + + countMessages(messages: readonly TokenCountableMessage[]): number { + return messages.reduce( + (total, message) => + total + this.countText(`${message.role}\n${message.content ?? ""}`), + 0, + ); + } +} + +export function createDefaultTokenEstimator(): TokenEstimator { + try { + return new GptTokenEstimator(); + } catch { + return new ApproxTokenEstimator(); + } +} + +const defaultTokenEstimator = createDefaultTokenEstimator(); + +/** Compatibility helper for older call sites. */ +export function estimateTokens(text: string): number { + return defaultTokenEstimator.countText(text); +} diff --git a/src/context/types.ts b/src/context/types.ts new file mode 100644 index 0000000..8e1f5b7 --- /dev/null +++ b/src/context/types.ts @@ -0,0 +1,120 @@ +import type {ContextCompressionPipeline} from "./context-compression-pipeline.js"; +import type {ContextBudgetPolicy} from "./context-budget.js"; +import type {ContextCompressor} from "./context-compressor.js"; +import type {ContextTruncator} from "./context-truncator.js"; +import type {ContextPriorityPolicy} from "./priority-policy.js"; +import type {SystemPromptAssembler} from "./system-prompt-assembler.js"; +import type {TokenEstimator} from "./token-estimator.js"; + +/** Source metadata for a context section. */ +export type ContextSource = + | {kind: "system"; ref?: string} + | {kind: "workspace"; ref?: string} + | {kind: "memory"; ref?: string} + | {kind: "provider"; ref?: string} + | {kind: "user"; ref?: string} + | {kind: "tool"; ref?: string} + | {kind: "file"; ref?: string}; + +/** A model-visible runtime context block. */ +export type ContextSection = { + id?: string; + replaceKey?: string; + title?: string; + content: string; + priority?: number; + source?: ContextSource; +}; + +/** Input accepted by the generic runtime context builder. */ +export type BuildContextInput = { + systemPrompt?: string; + outputInstructions?: string; + sections: readonly ContextSection[]; + tokenLimit: number; +}; + +/** Result returned after formatting and truncating runtime context. */ +export type BuildContextResult = { + systemPrompt: string; + contextText: string; + tokenEstimate: number; + includedSections: ContextSection[]; + partialSections: ContextSection[]; + droppedSections: ContextSection[]; + sectionUsages: ContextSectionUsage[]; + diagnostics?: BuildContextDiagnostics; +}; + +/** Runtime context budget derived from the model token limit. */ +export type ContextBudget = { + tokenLimit: number; + maxContextTokens: number; + reservedOutputTokens: number; + maxInputTokens: number; + /** @deprecated kept for compatibility and diagnostics only. */ + runtimeContextRatio?: number; + /** @deprecated char budget is a diagnostic fallback only. */ + maxContextChars?: number; +}; + +/** Explicit section-level truncation status. */ +export type ContextSectionUsageStatus = "included" | "partial" | "dropped"; + +/** Truncation and injection details for a single context section. */ +export type ContextSectionUsage = { + section: ContextSection; + status: ContextSectionUsageStatus; + originalLength: number; + includedLength: number; + formattedLength: number; + originalTokens?: number; + includedTokens?: number; + formattedTokens?: number; + reason?: string; +}; + +/** Result returned by a context compressor. */ +export type ContextCompressionResult = { + section: ContextSection; + originalSection: ContextSection; + compressed: boolean; + originalChars: number; + compressedChars: number; + strategy?: string; + omittedChars?: number; + savedChars?: number; + compressionRatio?: number; + maxSectionChars?: number; + originalTokens?: number; + compressedTokens?: number; + savedTokens?: number; + tokenCompressionRatio?: number; + maxSectionTokens?: number; + reason?: string; +}; + +/** Diagnostics produced while building runtime context. */ +export type BuildContextDiagnostics = { + budget: ContextBudget; + estimatedContextChars: number; + estimatedContextTokens: number; + compressionThresholdRatio: number; + compressionTriggered: boolean; + compressionLimitTokens: number; + compressionResults: ContextCompressionResult[]; + contextTextTokens: number; + systemPromptTokens: number; +}; + +/** Optional strategy overrides for the class-based context engine. */ +export type ContextEngineOptions = { + priorityPolicy?: ContextPriorityPolicy; + budgetPolicy?: ContextBudgetPolicy; + compressionPipeline?: ContextCompressionPipeline; + compressor?: ContextCompressor; + compressionThresholdRatio?: number; + truncator?: ContextTruncator; + assembler?: SystemPromptAssembler; + tokenEstimator?: TokenEstimator; +}; diff --git a/src/index.ts b/src/index.ts index 0ab4efc..e39a4fa 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,5 +1,6 @@ export * from "./agent/index.js"; export * from "./config/index.js"; +export * from "./context/index.js"; export * from "./events/index.js"; export * from "./llm/index.js"; export * from "./runtime/index.js"; diff --git a/tests/agent/runtime-loop.test.ts b/tests/agent/runtime-loop.test.ts index 3b2c0ca..1bcafdb 100644 --- a/tests/agent/runtime-loop.test.ts +++ b/tests/agent/runtime-loop.test.ts @@ -207,6 +207,80 @@ describe("Agent runtime loop", () => { }); }); + it("rebuilds runtime context each model request and archives older tool results", async () => { + const workspaceRoot = await createWorkspace(); + const registry = new ToolRegistry(); + const beforeModelBuildCounts: number[] = []; + const longOutput = "0123456789".repeat(1_500); + + registry.register({ + definition: { + name: "large_output", + description: "Returns a large output.", + parameters: z.object({label: z.string()}), + }, + execute: (input: {label: string}) => + okToolResult("Large output.", { + label: input.label, + output: input.label === "first" ? longOutput : "short output", + }), + }); + + const llm = new QueueLLMClient([ + { + content: "First call.", + toolCalls: [ + {id: "call-first", name: "large_output", arguments: {label: "first"}}, + ], + }, + { + content: "Second call.", + toolCalls: [ + {id: "call-second", name: "large_output", arguments: {label: "second"}}, + ], + }, + {content: "Done.", toolCalls: []}, + ]); + + const config = createConfig(workspaceRoot); + config.runtime.tokensLimit = 2_000; + config.runtime.maxIterations = 4; + + const result = await new Agent({ + config, + llm, + toolRegistry: registry, + workspaceScanner: createWorkspaceScanner(workspaceRoot), + middleware: [ + { + beforeModel: (_request, context) => { + beforeModelBuildCounts.push(context.contextBuilds?.length ?? 0); + }, + }, + ], + }).run({prompt: "Use large output."}); + + const thirdRequest = llm.requests[2]; + const thirdSystemPrompt = thirdRequest?.messages[0]?.content ?? ""; + + expect(result.stopReason).toBe("completed"); + expect(llm.requests).toHaveLength(3); + expect(beforeModelBuildCounts).toEqual([1, 2, 3]); + expect(thirdSystemPrompt).toContain("## Tool Result: large_output"); + expect(thirdSystemPrompt).toContain("Call ID: call-first"); + expect(thirdSystemPrompt).toContain("chars omitted"); + expect( + thirdRequest?.messages.some( + (message) => message.role === "tool" && message.toolCallId === "call-first", + ), + ).toBe(false); + expect(thirdRequest?.messages.at(-1)).toMatchObject({ + role: "tool", + toolCallId: "call-second", + name: "large_output", + }); + }); + it("emits one failed tool event for execution failures through the ToolRunner adapter", async () => { const workspaceRoot = await createWorkspace(); const registry = new ToolRegistry(); diff --git a/tests/agent/system-prompt.test.ts b/tests/agent/system-prompt.test.ts index 7bef48b..ed5e342 100644 --- a/tests/agent/system-prompt.test.ts +++ b/tests/agent/system-prompt.test.ts @@ -1,6 +1,8 @@ import {describe, expect, it} from "vitest"; import {Agent} from "../../src/agent/index.js"; +import {estimateTokens} from "../../src/context/index.js"; +import {EventBus, type PixelleEvent} from "../../src/events/index.js"; import {BaseLLMClient} from "../../src/llm/index.js"; import type {LLMGenerateInput, LLMResponse} from "../../src/llm/types.js"; import type {WorkspaceProfile} from "../../src/runtime/index.js"; @@ -65,4 +67,94 @@ describe("ContextManager system prompt", () => { expect(prompt).toContain("Do not use Markdown tables"); expect(prompt).toContain("always include a language identifier"); }); + + it("keeps runtime context inputs compatible after context core extraction", async () => { + const workspaceRoot = process.cwd(); + const llm = new CapturingLLMClient(); + const eventBus = new EventBus(); + const events: PixelleEvent[] = []; + const profile: WorkspaceProfile = { + root: workspaceRoot, + packageManager: "pnpm", + scripts: {test: "vitest run"}, + projectFiles: ["package.json"], + detectedFrameworks: ["typescript"], + }; + + eventBus.subscribe((event) => events.push(event)); + + await new Agent({ + config: { + runtime: { + systemPrompt: "Config prompt.", + workspaceDir: workspaceRoot, + maxIterations: 1, + maxRepairAttempts: 0, + tokensLimit: 1000, + rollbackOnFailure: false, + }, + permissions: { + readFile: true, + writeFile: false, + network: false, + shell: false, + }, + verification: { + enabled: false, + commands: [], + }, + trace: { + enabled: false, + directory: workspaceRoot, + }, + }, + llm, + eventBus, + memory: { + loadProjectMemory: () => [{title: "Project Memory", content: "project facts"}], + loadRunMemory: () => ["run facts"], + }, + contextProviders: [ + { + name: "Agent Provider", + build: () => "agent provider facts", + }, + ], + workspaceScanner: { + async scan(): Promise { + return profile; + }, + }, + }).run({ + prompt: "test", + systemPrompt: "Input prompt.", + context: [{title: "User Context", content: "user facts", priority: 50}], + contextProviders: [ + { + name: "Input Provider", + build: () => ({content: "input provider facts"}), + }, + ], + }); + + const prompt = llm.request?.messages[0]?.content ?? ""; + const contextBuilt = events.find((event) => event.type === "runtime.context_built"); + const runtimeContext = prompt.split("# Runtime Context\n")[1] ?? ""; + + expect(prompt).toContain("Input prompt."); + expect(prompt).not.toContain("Config prompt."); + expect(prompt).toContain("Do not use Markdown tables"); + expect(prompt).toContain("# Runtime Context"); + expect(runtimeContext).toContain("## Workspace Profile"); + expect(runtimeContext).toContain('"packageManager": "pnpm"'); + expect(runtimeContext).toContain("## User Context\nuser facts"); + expect(runtimeContext).toContain("## Project Memory\nproject facts"); + expect(runtimeContext).toContain("run facts"); + expect(runtimeContext).toContain("## Agent Provider\nagent provider facts"); + expect(runtimeContext).toContain("## Input Provider\ninput provider facts"); + expect(contextBuilt).toMatchObject({ + type: "runtime.context_built", + tokenEstimate: estimateTokens(prompt), + }); + }); }); diff --git a/tests/context-builder.test.ts b/tests/context-builder.test.ts new file mode 100644 index 0000000..4ce3084 --- /dev/null +++ b/tests/context-builder.test.ts @@ -0,0 +1,758 @@ +import {describe, expect, it} from "vitest"; + +import { + ApproxTokenEstimator, + buildRuntimeContext, + ContextCompressionPipeline, + ContextCompressionResultFactory, + ContextEngine, + ContextRegistry, + DefaultContextBudgetPolicy, + estimateTokens, + GptTokenEstimator, + NoopContextCompressor, + RuleBasedContextCompressor, + truncateTextToTokens, + type ContextBudget, + type ContextCompressionResult, + type ContextCompressor, + type ContextSection, + type TokenEstimator, +} from "../src/context/index.js"; + +class CharTokenEstimator implements TokenEstimator { + countText(text: string): number { + return text.length; + } +} + +const charTokenEstimator = new CharTokenEstimator(); + +class ThrowingCompressor implements ContextCompressor { + compress(): ContextCompressionResult { + throw new Error("Compressor should not run."); + } +} + +class MarkingCompressor implements ContextCompressor { + compress(section: ContextSection, _budget: ContextBudget): ContextCompressionResult { + const compressedSection = { + ...section, + content: section.content.slice( + 0, + Math.max(1, Math.floor(section.content.length / 2)), + ), + }; + + return { + section: compressedSection, + originalSection: section, + compressed: true, + originalChars: section.content.length, + compressedChars: compressedSection.content.length, + originalTokens: section.content.length, + compressedTokens: compressedSection.content.length, + omittedChars: section.content.length - compressedSection.content.length, + savedChars: section.content.length - compressedSection.content.length, + savedTokens: section.content.length - compressedSection.content.length, + compressionRatio: compressedSection.content.length / section.content.length, + tokenCompressionRatio: compressedSection.content.length / section.content.length, + reason: "Marked by test compressor.", + }; + } +} + +describe("buildRuntimeContext", () => { + it("provides approximate and gpt-backed token estimators", () => { + const approx = new ApproxTokenEstimator(); + const gpt = new GptTokenEstimator(); + + expect(approx.countText("hello world")).toBeGreaterThan(0); + expect(gpt.countText("hello world")).toBeGreaterThan(0); + expect(estimateTokens("hello world")).toBeGreaterThan(0); + }); + + it("matches the default ContextEngine output", () => { + const input = { + systemPrompt: "Base.", + outputInstructions: "Use Markdown.", + sections: [{title: "Runtime", content: "details"}], + tokenLimit: 100, + }; + + expect(buildRuntimeContext(input)).toEqual(new ContextEngine().build(input)); + }); + + it("sorts context sections by descending priority", () => { + const result = new ContextEngine().build({ + systemPrompt: "Base.", + sections: [ + {id: "low", content: "low", priority: 1}, + {id: "high", content: "high", priority: 10}, + {id: "default", content: "default"}, + ], + tokenLimit: 100, + }); + + expect(result.contextText).toBe("high\n\nlow\n\ndefault"); + expect(result.includedSections.map((section) => section.id)).toEqual([ + "high", + "low", + "default", + ]); + }); + + it("uses source priority when explicit priority is absent", () => { + const result = new ContextEngine().build({ + systemPrompt: "Base.", + sections: [ + {id: "file", content: "file", source: {kind: "file"}}, + {id: "workspace", content: "workspace", source: {kind: "workspace"}}, + {id: "memory", content: "memory", source: {kind: "memory"}}, + ], + tokenLimit: 100, + }); + + expect(result.contextText).toBe("workspace\n\nmemory\n\nfile"); + }); + + it("skips empty section content", () => { + const result = buildRuntimeContext({ + systemPrompt: "Base.", + sections: [ + {id: "empty", title: "Empty", content: " ", priority: 100}, + {id: "filled", content: "filled"}, + ], + tokenLimit: 100, + }); + + 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([]); + }); + + it("formats titled sections with markdown headings", () => { + const result = buildRuntimeContext({ + systemPrompt: "Base.", + sections: [{title: "Notes", content: " keep this "}], + tokenLimit: 100, + }); + + expect(result.contextText).toBe("## Notes\nkeep this"); + }); + + it("builds a token-first runtime context budget policy", () => { + expect( + new DefaultContextBudgetPolicy().createBudget({ + sections: [], + tokenLimit: 100, + }), + ).toEqual({ + tokenLimit: 100, + maxContextTokens: 100, + reservedOutputTokens: 20, + maxInputTokens: 80, + runtimeContextRatio: 0.35, + maxContextChars: 320, + }); + }); + + it("truncates over-budget runtime context and records included, partial, and dropped sections", () => { + const result = new ContextEngine({ + budgetPolicy: new DefaultContextBudgetPolicy({reservedOutputTokens: 0}), + tokenEstimator: charTokenEstimator, + }).build({ + systemPrompt: "Base.", + sections: [ + {id: "first", content: "abcdefghij", priority: 10}, + {id: "second", content: "klmnopqrst", priority: 9}, + ], + tokenLimit: 4, + }); + + 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([ + "first", + "second", + ]); + expect(result.sectionUsages.map((usage) => usage.status)).toEqual([ + "partial", + "dropped", + ]); + expect(result.diagnostics?.contextTextTokens).toBe(4); + expect(result.tokenEstimate).toBe(charTokenEstimator.countText(result.systemPrompt)); + }); + + it("keeps fully injected sections separate from partial and dropped sections", () => { + const result = new ContextEngine({ + budgetPolicy: new DefaultContextBudgetPolicy({reservedOutputTokens: 0}), + tokenEstimator: charTokenEstimator, + }).build({ + systemPrompt: "Base.", + sections: [ + {id: "first", content: "abc", priority: 10}, + {id: "second", content: "defghij", priority: 9}, + {id: "third", content: "klm", priority: 8}, + ], + tokenLimit: 7, + }); + + 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([ + "second", + "third", + ]); + expect(result.sectionUsages).toEqual([ + expect.objectContaining({ + section: expect.objectContaining({id: "first"}), + status: "included", + reason: "Section fits within the remaining context budget.", + }), + expect.objectContaining({ + section: expect.objectContaining({id: "second"}), + status: "partial", + reason: + "Section exceeded the remaining context budget and was partially included.", + }), + expect.objectContaining({ + section: expect.objectContaining({id: "third"}), + status: "dropped", + reason: "No remaining context budget for this section.", + }), + ]); + }); + + it("builds system prompt with output instructions and runtime context", () => { + const result = buildRuntimeContext({ + systemPrompt: "Base prompt.", + outputInstructions: "Use Markdown.", + sections: [{title: "Runtime", content: "details"}], + tokenLimit: 100, + }); + + expect(result.systemPrompt).toBe( + "Base prompt.\n\nUse Markdown.\n\n# Runtime Context\n## Runtime\ndetails", + ); + }); + + it("dedupes sections by replaceKey or id with the later section winning", () => { + const registry = new ContextRegistry() + .addMany([ + {id: "same", content: "old"}, + {id: "same", content: "new"}, + {replaceKey: "profile", content: "old profile"}, + {replaceKey: "profile", content: "new profile"}, + {content: "kept"}, + ]) + .normalize() + .dedupe(); + + expect(registry.getAll().map((section) => section.content)).toEqual([ + "new", + "new profile", + "kept", + ]); + }); + + 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({ + compressor, + thresholdRatio: 0.85, + tokenEstimator: charTokenEstimator, + }); + const compression = pipeline.compress( + [{id: "short", content: "short"}], + new DefaultContextBudgetPolicy().createBudget({ + sections: [], + tokenLimit: 100, + }), + ); + const result = new ContextEngine({ + compressionPipeline: pipeline, + }).build({ + systemPrompt: "Base.", + sections: [{id: "short", content: "short"}], + tokenLimit: 100, + }); + + expect(compression).toMatchObject({ + estimatedContextChars: 5, + estimatedContextTokens: 5, + compressionLimitTokens: 68, + thresholdRatio: 0.85, + triggered: false, + }); + expect(compression.sections).toEqual([{id: "short", content: "short"}]); + expect(compression.results).toEqual([ + expect.objectContaining({ + compressed: false, + reason: "Compression was not triggered.", + }), + ]); + expect(result.diagnostics).toMatchObject({ + compressionThresholdRatio: 0.85, + compressionTriggered: false, + estimatedContextChars: 5, + estimatedContextTokens: 5, + compressionLimitTokens: 68, + }); + expect(result.diagnostics?.compressionResults).toEqual([ + expect.objectContaining({ + section: expect.objectContaining({id: "short"}), + compressed: false, + originalChars: 5, + compressedChars: 5, + reason: "Compression was not triggered.", + }), + ]); + }); + + it("triggers compression above the configured threshold and exposes diagnostics", () => { + const pipeline = new ContextCompressionPipeline({ + compressor: new MarkingCompressor(), + thresholdRatio: 0.5, + tokenEstimator: charTokenEstimator, + }); + const result = new ContextEngine({ + compressionPipeline: pipeline, + budgetPolicy: new DefaultContextBudgetPolicy({reservedOutputTokens: 0}), + tokenEstimator: charTokenEstimator, + }).build({ + systemPrompt: "Base prompt.", + outputInstructions: "Use Markdown.", + sections: [{id: "long", content: "abcdefghijklmnop", source: {kind: "tool"}}], + tokenLimit: 8, + }); + + expect(result.contextText).toBe("abcdefgh"); + expect(result.tokenEstimate).toBe(charTokenEstimator.countText(result.systemPrompt)); + expect(result.diagnostics).toMatchObject({ + budget: { + tokenLimit: 8, + maxContextTokens: 8, + reservedOutputTokens: 0, + maxInputTokens: 8, + runtimeContextRatio: 0.35, + maxContextChars: 32, + }, + estimatedContextChars: 16, + estimatedContextTokens: 16, + compressionThresholdRatio: 0.5, + compressionTriggered: true, + compressionLimitTokens: 4, + contextTextTokens: charTokenEstimator.countText(result.contextText), + systemPromptTokens: charTokenEstimator.countText(result.systemPrompt), + }); + expect(result.diagnostics?.compressionResults).toEqual([ + expect.objectContaining({ + originalSection: expect.objectContaining({id: "long"}), + section: expect.objectContaining({id: "long", content: "abcdefgh"}), + compressed: true, + originalChars: 16, + compressedChars: 8, + originalTokens: 16, + compressedTokens: 8, + omittedChars: 8, + savedChars: 8, + savedTokens: 8, + compressionRatio: 0.5, + tokenCompressionRatio: 0.5, + reason: "Marked by test compressor.", + }), + ]); + }); + + it("uses a default noop compression pipeline when no compressor is provided", () => { + const compression = new ContextCompressionPipeline({ + thresholdRatio: 0, + tokenEstimator: charTokenEstimator, + }).compress( + [{id: "plain", content: "plain"}], + new DefaultContextBudgetPolicy().createBudget({ + sections: [], + tokenLimit: 1, + }), + ); + + expect(compression.triggered).toBe(true); + expect(compression.sections).toEqual([{id: "plain", content: "plain"}]); + expect(compression.results).toEqual([ + expect.objectContaining({ + compressed: false, + reason: "No compression applied.", + }), + ]); + }); + + it("truncates text to a token limit without exceeding the limit", () => { + const truncated = truncateTextToTokens( + "line-one\nline-two\nline-three", + 10, + charTokenEstimator, + ); + + expect(charTokenEstimator.countText(truncated)).toBeLessThanOrEqual(10); + expect(truncated).toBe("line-one\n"); + }); + + it("creates compression results through the shared factory", () => { + const factory = new ContextCompressionResultFactory(); + const original = {content: "abcdefghij"}; + const compressed = {content: "abc"}; + + expect(factory.skipped(original, "skip")).toMatchObject({ + section: original, + originalSection: original, + compressed: false, + reason: "skip", + }); + expect(factory.compressed(compressed, original, "compressed")).toMatchObject({ + section: compressed, + originalSection: original, + compressed: true, + originalChars: 10, + compressedChars: 3, + savedChars: 7, + omittedChars: 7, + compressionRatio: 0.3, + reason: "compressed", + }); + }); + + it("rule-based compressor only compresses long tool and file sections", () => { + const compressor = new RuleBasedContextCompressor({ + maxSectionTokens: 80, + minSectionTokens: 80, + maxSectionChars: 80, + headChars: 20, + tailChars: 10, + tokenEstimator: charTokenEstimator, + }); + const longContent = "abcdefghijklmnopqrstuvwxyz".repeat(10); + const kinds = ["user", "memory", "workspace", "provider", "system"] as const; + + for (const kind of kinds) { + const section = {content: longContent, source: {kind}}; + expect(compressor.compress(section, budgetFor(section))).toMatchObject({ + section, + originalSection: section, + compressed: false, + reason: "Section source is not compressible.", + }); + } + + for (const kind of ["tool", "file"] as const) { + const section = {content: longContent, source: {kind}}; + const result = compressor.compress(section, budgetFor(section)); + + expect(result.compressed).toBe(true); + expect(result.section.content).toContain("abcdefghij"); + expect(result.section.content).toContain("uvwxyz"); + expect(result.section.content).toMatch( + /\[\.\.\.\d+ chars omitted by rule-based-head-tail compressor\.\.\.\]/, + ); + expect(result.section.content.length).toBeLessThan(longContent.length); + expect(result.section.content.length).toBeLessThanOrEqual(80); + expect(result).toMatchObject({ + strategy: "rule-based-head-tail", + maxSectionTokens: 80, + maxSectionChars: 80, + }); + expect(result.omittedChars).toBeGreaterThan(0); + expect(result.savedChars).toBeGreaterThan(0); + expect(result.compressionRatio).toBeLessThan(1); + } + }); + + it("reports omitted original chars separately from saved chars", () => { + const compressor = new RuleBasedContextCompressor({ + maxSectionTokens: 90, + minSectionTokens: 90, + maxSectionChars: 90, + minSectionChars: 90, + maxSectionRatio: 1, + headChars: 20, + tailChars: 20, + preserveLineBoundaries: false, + tokenEstimator: charTokenEstimator, + }); + const section = { + content: "0123456789".repeat(30), + source: {kind: "tool" as const}, + }; + const result = compressor.compress(section, rawBudget(90)); + const markerMatch = /\[\.\.\.(\d+) chars omitted/.exec(result.section.content); + + expect(result.compressed).toBe(true); + expect(markerMatch?.[1]).toBe(String(result.omittedChars)); + expect(result.omittedChars).toBeGreaterThan(result.savedChars ?? 0); + expect(result.savedChars).toBe(result.originalChars - result.compressedChars); + expect(result.section.content.length).toBeLessThanOrEqual( + result.maxSectionChars ?? 0, + ); + }); + + it("supports explicit zero head or tail retention", () => { + const section = { + content: "HEAD-".repeat(20) + "TAIL-".repeat(20), + source: {kind: "tool" as const}, + }; + const tailOnly = new RuleBasedContextCompressor({ + maxSectionTokens: 80, + minSectionTokens: 80, + maxSectionChars: 80, + minSectionChars: 80, + maxSectionRatio: 1, + headChars: 0, + tailChars: 20, + preserveLineBoundaries: false, + tokenEstimator: charTokenEstimator, + }).compress(section, rawBudget(80)); + const headOnly = new RuleBasedContextCompressor({ + maxSectionTokens: 80, + minSectionTokens: 80, + maxSectionChars: 80, + minSectionChars: 80, + maxSectionRatio: 1, + headChars: 20, + tailChars: 0, + preserveLineBoundaries: false, + tokenEstimator: charTokenEstimator, + }).compress(section, rawBudget(80)); + + expect(tailOnly.section.content.startsWith("\n\n[...")).toBe(true); + expect(tailOnly.section.content).toContain("TAIL-"); + expect(headOnly.section.content.startsWith("HEAD-")).toBe(true); + expect(headOnly.section.content.endsWith("compressor...]\n\n")).toBe(true); + }); + + it("uses default head-heavy allocation when head and tail are both zero", () => { + const section = { + content: "abcdefghijklmnopqrstuvwxyz".repeat(20), + source: {kind: "file" as const}, + }; + const result = new RuleBasedContextCompressor({ + maxSectionTokens: 100, + minSectionTokens: 100, + maxSectionChars: 100, + minSectionChars: 100, + maxSectionRatio: 1, + headChars: 0, + tailChars: 0, + preserveLineBoundaries: false, + tokenEstimator: charTokenEstimator, + }).compress(section, rawBudget(100)); + const markerStart = result.section.content.indexOf("[..."); + const markerEnd = result.section.content.indexOf("]\n\n"); + const head = result.section.content.slice(0, markerStart); + const tail = result.section.content.slice(markerEnd + 3); + + expect(head.length).toBeGreaterThan(tail.length); + expect(result.section.content.length).toBeLessThanOrEqual(100); + }); + + it("uses ContextBudget to resolve dynamic per-section compression limits", () => { + const compressor = new RuleBasedContextCompressor({ + maxSectionTokens: 500, + minSectionTokens: 80, + maxSectionChars: 500, + minSectionChars: 80, + maxSectionRatio: 0.25, + headChars: 60, + tailChars: 40, + preserveLineBoundaries: false, + tokenEstimator: charTokenEstimator, + }); + const section = { + content: "0123456789".repeat(100), + source: {kind: "tool" as const}, + }; + const smallBudgetResult = compressor.compress(section, rawBudget(400)); + const largeBudgetResult = compressor.compress(section, rawBudget(20_000)); + + expect(smallBudgetResult.compressed).toBe(true); + expect(smallBudgetResult.maxSectionTokens).toBe(100); + expect(smallBudgetResult.compressedTokens).toBeLessThanOrEqual(100); + expect(largeBudgetResult.compressed).toBe(true); + expect(largeBudgetResult.maxSectionTokens).toBe(500); + expect(largeBudgetResult.section.content.length).toBeGreaterThan( + smallBudgetResult.section.content.length, + ); + }); + + it("does not compress tool or file sections below the dynamic limit", () => { + const compressor = new RuleBasedContextCompressor({ + maxSectionTokens: 500, + minSectionTokens: 1, + maxSectionChars: 500, + minSectionChars: 1, + maxSectionRatio: 1, + tokenEstimator: charTokenEstimator, + }); + const section = { + content: "short tool output", + source: {kind: "tool" as const}, + }; + const result = compressor.compress(section, budgetFor(section)); + + expect(result).toMatchObject({ + section, + originalSection: section, + compressed: false, + reason: "Section is within the dynamic section token limit (17 <= 100 tokens).", + strategy: "rule-based-head-tail", + maxSectionTokens: 100, + }); + }); + + it("preserves line boundaries for head and tail slices when configured", () => { + const compressor = new RuleBasedContextCompressor({ + maxSectionTokens: 140, + minSectionTokens: 140, + maxSectionChars: 140, + minSectionChars: 140, + maxSectionRatio: 1, + headChars: 40, + tailChars: 40, + headTokenRatio: 0.5, + tailTokenRatio: 0.5, + preserveLineBoundaries: true, + tokenEstimator: charTokenEstimator, + }); + const section = { + content: [ + "head-line-1", + "head-line-2", + "head-line-3", + "middle-line".repeat(12), + "tail-line-1", + "tail-line-2", + ].join("\n"), + source: {kind: "file" as const}, + }; + const result = compressor.compress(section, rawBudget(140)); + + expect(result.compressed).toBe(true); + expect(result.section.content).toContain("head-line-1\nhead-line-2\n"); + expect(result.section.content).not.toContain("middle-linemiddle-line"); + expect(result.section.content).toContain("tail-line-1\ntail-line-2"); + }); + + it("falls back to plain line slicing when line preservation would keep too little", () => { + const compressor = new RuleBasedContextCompressor({ + maxSectionTokens: 100, + minSectionTokens: 100, + maxSectionChars: 100, + minSectionChars: 100, + maxSectionRatio: 1, + headChars: 10, + tailChars: 10, + headTokenRatio: 0.5, + tailTokenRatio: 0.5, + preserveLineBoundaries: true, + tokenEstimator: charTokenEstimator, + }); + const section = { + content: `ab\n${"c".repeat(120)}\n${"d".repeat(120)}\nij`, + source: {kind: "tool" as const}, + }; + const result = compressor.compress(section, rawBudget(100)); + + expect(result.section.content).toContain("ab\nccccc"); + expect(result.section.content).toContain("ddddd\nij"); + }); + + it("normalizes invalid rule-based compressor options without throwing", () => { + const compressor = new RuleBasedContextCompressor({ + maxSectionChars: 0, + minSectionChars: 99_999, + maxSectionTokens: 0, + minSectionTokens: 99_999, + maxSectionRatio: 10, + headChars: -1, + tailChars: -1, + tokenEstimator: charTokenEstimator, + }); + const section = { + content: "x".repeat(20_000), + source: {kind: "tool" as const}, + }; + const result = compressor.compress(section, rawBudget(100)); + + expect(result.compressed).toBe(true); + expect(result.section.content.trim().length).toBeGreaterThan(0); + expect(result.maxSectionTokens).toBe(8_000); + }); + + it("keeps omission marker details visible through compressor output", () => { + const section = { + content: "abcdefghijklmnopqrstuvwxyz".repeat(20), + source: {kind: "tool" as const}, + }; + const result = new RuleBasedContextCompressor({ + maxSectionTokens: 100, + minSectionTokens: 100, + maxSectionChars: 100, + minSectionChars: 100, + maxSectionRatio: 1, + preserveLineBoundaries: false, + tokenEstimator: charTokenEstimator, + }).compress(section, rawBudget(100)); + + expect(result.section.content).toMatch( + new RegExp(`\\[\\.\\.\\.${result.omittedChars} chars omitted`), + ); + }); +}); + +function budgetFor(section: ContextSection): ContextBudget { + return new DefaultContextBudgetPolicy({reservedOutputTokens: 0}).createBudget({ + sections: [section], + tokenLimit: 100, + }); +} + +function rawBudget(maxInputTokens: number): ContextBudget { + return { + tokenLimit: maxInputTokens, + maxContextTokens: maxInputTokens, + reservedOutputTokens: 0, + maxInputTokens, + runtimeContextRatio: 0.35, + maxContextChars: maxInputTokens * 4, + }; +}