diff --git a/packages/core/src/context/DESIGN.md b/packages/core/src/context/DESIGN.md index 80ad9772a1..37c76e9be2 100644 --- a/packages/core/src/context/DESIGN.md +++ b/packages/core/src/context/DESIGN.md @@ -2,80 +2,174 @@ ## 1. Introduction & Motivation -This document provides a high-level orientation to the Context Management system within `@google/gemini-cli-core`. +This document provides a high-level orientation to the Context Management system +within `@google/gemini-cli-core`. -Previously, context management in the CLI was decentralized, synchronous, and relied on fixed-function, destructive mutations of the raw Gemini `Content[]` history. Because all context management was local, this approach made it nearly impossible to reason about the global impact of any specific change. For example, should we distill tool outputs, or mask them? Or maybe it's contextual? What about other processors like the snapshotter, should they see masked results? Distilled results? What about new approaches to context management, how do they fit into the solution we've already built. The old approach to context management made it nearly challenging to even attempt to answer any one of these questions, let alone to try and answer all of them. +Previously, context management in the CLI was decentralized, synchronous, and +relied on fixed-function, destructive mutations of the raw Gemini `Content[]` +history. Because all context management was local, this approach made it nearly +impossible to reason about the global impact of any specific change. For +example, should we distill tool outputs, or mask them? Or maybe it's contextual? +What about other processors like the snapshotter, should they see masked +results? Distilled results? What about new approaches to context management, how +do they fit into the solution we've already built. The old approach to context +management made it nearly challenging to even attempt to answer any one of these +questions, let alone to try and answer all of them. -To address these issues, we went back to the drawing board to create an explicit Context Manager. As opposed to our old approach, the new Context Manager V0 is a robust, event-driven, pluggable system. It introduces a non-destructive Episodic Intermediate Representation (IR) and an asynchronous processing pipeline, allowing the CLI to run expensive LLM summarization tasks in the background and opportunistically project an optimized view of the history only when budget constraints require it. +To address these issues, we went back to the drawing board to create an explicit +Context Manager. As opposed to our old approach, the new Context Manager V0 is a +robust, event-driven, pluggable system. It introduces a non-destructive Episodic +Intermediate Representation (IR) and an asynchronous processing pipeline, +allowing the CLI to run expensive LLM summarization tasks in the background and +opportunistically project an optimized view of the history only when budget +constraints require it. --- ## 2. Chief Innovations & Salient Features -The architecture is built upon seven core principles that distinguish it from the legacy system: +The architecture is built upon seven core principles that distinguish it from +the legacy system: -1. **Centralized Budgeting:** The `ContextManager` is the sole source of truth for the token budget. It makes the final, just-in-time decision about what gets projected to the LLM. -2. **Statelessness via IR:** Raw history is never mutated or deleted. Instead, it is translated into an Intermediate Representation (IR). Context reduction is achieved by attaching compressed `Variant`s to the IR graph. The original text is always recoverable. -3. **Asynchronicity:** Designed around a `ContextEventBus`. Heavy context operations (like LLM-powered summarization) run as detached background tasks without blocking the main agent loop. -4. **Configurability:** Driven by a typed JSON "Sidecar" configuration. Token ceilings, fallback strategies, and processing pipelines are entirely data-driven. -5. **Pluggability:** `ContextProcessor`s are isolated plugins with typed schemas. They are registered via Dependency Injection and can be arranged into arbitrary pipelines. -6. **Debuggability:** A built-in `ContextTracer` tracks every step of the pipeline, providing full audit trails of exactly when, why, and how a message was altered. -7. **Testability:** Global state has been eliminated. The system uses strict Dependency Injection (`SidecarRegistry`, `ContextEnvironment`, `ContextEventBus`), making every layer easily unit-testable. -8. **Orthogonality via Targets:** Processors do not implicitly scan the entire history graph. The `ContextManager` computes exact Deltas (e.g., new nodes just added, or specific nodes that just aged out of the retained buffer). Processors are sandboxed by the `EpisodeEditor` to only iterate over and mutate these specific `targetNodes`, ensuring surgical and highly efficient reductions. +1. **Centralized Budgeting:** The `ContextManager` is the sole source of truth + for the token budget. It makes the final, just-in-time decision about what + gets projected to the LLM. +2. **Statelessness via IR:** Raw history is never mutated or deleted. Instead, + it is translated into an Intermediate Representation (IR). Context reduction + is achieved by attaching compressed `Variant`s to the IR graph. The original + text is always recoverable. +3. **Asynchronicity:** Designed around a `ContextEventBus`. Heavy context + operations (like LLM-powered summarization) run as detached background tasks + without blocking the main agent loop. +4. **Configurability:** Driven by a typed JSON "Sidecar" configuration. Token + ceilings, fallback strategies, and processing pipelines are entirely + data-driven. +5. **Pluggability:** `ContextProcessor`s are isolated plugins with typed + schemas. They are registered via Dependency Injection and can be arranged + into arbitrary pipelines. +6. **Debuggability:** A built-in `ContextTracer` tracks every step of the + pipeline, providing full audit trails of exactly when, why, and how a + message was altered. +7. **Testability:** Global state has been eliminated. The system uses strict + Dependency Injection (`SidecarRegistry`, `ContextEnvironment`, + `ContextEventBus`), making every layer easily unit-testable. +8. **Orthogonality via Targets:** Processors do not implicitly scan the entire + history graph. The `ContextManager` computes exact Deltas (e.g., new nodes + just added, or specific nodes that just aged out of the retained buffer). + Processors are sandboxed by the `EpisodeEditor` to only iterate over and + mutate these specific `targetNodes`, ensuring surgical and highly efficient + reductions. --- ## 3. The Major Pieces: Roles & Responsibilities ### The Brain: `ContextManager` -The central coordinator. It owns the "Pristine History" (the ground-truth Episodic IR graph). Its primary responsibility is exposing `projectCompressedHistory()`, which flattens the IR graph into a standard `Content[]` array strictly adhering to the configured token budget. + +The central coordinator. It owns the "Pristine History" (the ground-truth +Episodic IR graph). Its primary responsibility is exposing +`projectCompressedHistory()`, which flattens the IR graph into a standard +`Content[]` array strictly adhering to the configured token budget. ### The Data Model: Episodic Intermediate Representation (IR) -Instead of a flat array of messages, interactions are grouped into `Episode`s. An Episode represents a single turn: a User Prompt, followed by the Agent's Thoughts and Tool Executions (Steps), concluding with a Yield. -* **`IrNode`:** The base unit (e.g., `ToolExecution`, `AgentThought`). -* **`Variant`:** Compressed alternatives to the raw node (e.g., `SummaryVariant`, `MaskedVariant`, `SnapshotVariant`). -* **`IrMetadata`:** An audit trail attached to every node, tracking token counts and the chronological list of `transformations` applied by processors. + +Instead of a flat array of messages, interactions are grouped into `Episode`s. +An Episode represents a single turn: a User Prompt, followed by the Agent's +Thoughts and Tool Executions (Steps), concluding with a Yield. + +- **`IrNode`:** The base unit (e.g., `ToolExecution`, `AgentThought`). +- **`Variant`:** Compressed alternatives to the raw node (e.g., + `SummaryVariant`, `MaskedVariant`, `SnapshotVariant`). +- **`IrMetadata`:** An audit trail attached to every node, tracking token counts + and the chronological list of `transformations` applied by processors. ### The Engine: `PipelineOrchestrator` & Sidecar -The orchestrator reads the `SidecarConfig`. It manages the lifecycle of the pipelines, registering triggers and executing processors in order. It dictates whether a pipeline blocks the main thread or runs in the background. + +The orchestrator reads the `SidecarConfig`. It manages the lifecycle of the +pipelines, registering triggers and executing processors in order. It dictates +whether a pipeline blocks the main thread or runs in the background. ### The Workers: `ContextProcessor`s -Small, highly-focused classes that implement context reduction strategies. They do not mutate the graph directly; instead, they are given an `EpisodeEditor` which provides a safe, scoped API to attach `Variant`s and append metadata. -* *Examples:* `ToolMaskingProcessor`, `NodeDistillationProcessor`, `BlobDegradationProcessor`. + +Small, highly-focused classes that implement context reduction strategies. They +do not mutate the graph directly; instead, they are given an `EpisodeEditor` +which provides a safe, scoped API to attach `Variant`s and append metadata. + +- _Examples:_ `ToolMaskingProcessor`, `NodeDistillationProcessor`, + `BlobDegradationProcessor`. ### The Glue: `ContextEventBus` -A Pub/Sub bus that decouples the components. It enables the `HistoryObserver` to notify the system of new messages, and allows background processors to notify the `ContextManager` when a new compressed variant is ready to be used. + +A Pub/Sub bus that decouples the components. It enables the `HistoryObserver` to +notify the system of new messages, and allows background processors to notify +the `ContextManager` when a new compressed variant is ready to be used. --- ## 4. How They Interact: The Life of a Message -To understand how these pieces fit together, let's walk through the lifecycle of a single interaction as it moves through the context system. +To understand how these pieces fit together, let's walk through the lifecycle of +a single interaction as it moves through the context system. ### Phase 1: Ingestion & Translation -1. **Action:** The user sends a prompt, and the agent responds with a tool call. These raw messages are appended to the standard `AgentChatHistory`. + +1. **Action:** The user sends a prompt, and the agent responds with a tool + call. These raw messages are appended to the standard `AgentChatHistory`. 2. **Observation:** The `HistoryObserver` detects the new messages. -3. **Translation:** The observer passes the raw `Content[]` to the `IrMapper`. The mapper groups the prompt and the tool execution into a single, structured `Episode`. -4. **Registration:** The new `Episode` is added to the `ContextManager`'s pristine graph. +3. **Translation:** The observer passes the raw `Content[]` to the `IrMapper`. + The mapper groups the prompt and the tool execution into a single, + structured `Episode`. +4. **Registration:** The new `Episode` is added to the `ContextManager`'s + pristine graph. ### Phase 2: Triggering the Pipelines -1. **Delta Generation:** The `ContextManager` receives the updated pristine graph. It diffs it against the previous state and extracts a Delta—the exact Set of new `IrNode` IDs. -2. **Event Emission:** The `ContextManager` fires a `ChunkReceivedEvent` (with the Delta targets) over the `ContextEventBus`. -3. **Orchestration:** The `PipelineOrchestrator` hears the event and evaluates its configured `PipelineDef`s. It finds a pipeline with the trigger `on_turn`. -4. **Execution:** The Orchestrator creates an `EpisodeEditor` heavily sandboxed to *only* allow access to the targeted Delta nodes, and begins running the processors in that pipeline sequentially. + +1. **Delta Generation:** The `ContextManager` receives the updated pristine + graph. It diffs it against the previous state and extracts a Delta—the exact + Set of new `IrNode` IDs. +2. **Event Emission:** The `ContextManager` fires a `ChunkReceivedEvent` (with + the Delta targets) over the `ContextEventBus`. +3. **Orchestration:** The `PipelineOrchestrator` hears the event and evaluates + its configured `PipelineDef`s. It finds a pipeline with the trigger + `on_turn`. +4. **Execution:** The Orchestrator creates an `EpisodeEditor` heavily sandboxed + to _only_ allow access to the targeted Delta nodes, and begins running the + processors in that pipeline sequentially. ### Phase 3: Processing & Safe Editing -1. **Processing:** A processor (e.g., `ToolMaskingProcessor`) receives the `EpisodeEditor`. It iterates over `editor.targets` (ignoring the rest of the historical graph). It identifies a massive JSON payload in one of the new tool executions. -2. **Editing:** Instead of deleting the JSON, it calls `editor.editEpisode()`. It creates a `MaskedVariant` containing a string summary of the JSON. If it had attempted to edit a node outside its target Delta, the editor would have thrown an error. -3. **Auditing:** The editor automatically appends a record to the node's `IrMetadata.transformations` indicating that the `ToolMaskingProcessor` applied a `MASKED` action. + +1. **Processing:** A processor (e.g., `ToolMaskingProcessor`) receives the + `EpisodeEditor`. It iterates over `editor.targets` (ignoring the rest of the + historical graph). It identifies a massive JSON payload in one of the new + tool executions. +2. **Editing:** Instead of deleting the JSON, it calls `editor.editEpisode()`. + It creates a `MaskedVariant` containing a string summary of the JSON. If it + had attempted to edit a node outside its target Delta, the editor would have + thrown an error. +3. **Auditing:** The editor automatically appends a record to the node's + `IrMetadata.transformations` indicating that the `ToolMaskingProcessor` + applied a `MASKED` action. ### Phase 4: Async Resolution -1. **Completion:** The background pipeline finishes. The orchestrator fires a `VariantReadyEvent` over the bus. -2. **Integration:** The `ContextManager` receives the event and securely attaches the `MaskedVariant` to the correct `Episode` in the pristine graph. (If the pipeline was synchronous/blocking, this happens immediately). + +1. **Completion:** The background pipeline finishes. The orchestrator fires a + `VariantReadyEvent` over the bus. +2. **Integration:** The `ContextManager` receives the event and securely + attaches the `MaskedVariant` to the correct `Episode` in the pristine graph. + (If the pipeline was synchronous/blocking, this happens immediately). ### Phase 5: Just-In-Time Projection -1. **Request:** The agent is ready to send the next prompt to Gemini. The core routing logic calls `contextManager.projectCompressedHistory()`. -2. **Budget Evaluation:** The `IrProjector` calculates the current total tokens of the pristine graph and compares it to the `SidecarConfig` budget. -3. **Variant Selection:** If the graph exceeds the budget, the projector looks for available `Variant`s. It sees the newly attached `MaskedVariant` and calculates the token deficit recovered by using it. -4. **Flattening:** The `graphUtils` safely swap the raw node for the `MaskedVariant` in a temporary view, and flatten the Episodic IR back into a raw Gemini `Content[]` array. -5. **Delivery:** The optimized, budget-compliant array is sent to the LLM. The underlying pristine graph remains completely untouched and available for future reference or alternative projections. + +1. **Request:** The agent is ready to send the next prompt to Gemini. The core + routing logic calls `contextManager.projectCompressedHistory()`. +2. **Budget Evaluation:** The `IrProjector` calculates the current total tokens + of the pristine graph and compares it to the `SidecarConfig` budget. +3. **Variant Selection:** If the graph exceeds the budget, the projector looks + for available `Variant`s. It sees the newly attached `MaskedVariant` and + calculates the token deficit recovered by using it. +4. **Flattening:** The `graphUtils` safely swap the raw node for the + `MaskedVariant` in a temporary view, and flatten the Episodic IR back into a + raw Gemini `Content[]` array. +5. **Delivery:** The optimized, budget-compliant array is sent to the LLM. The + underlying pristine graph remains completely untouched and available for + future reference or alternative projections. diff --git a/packages/core/src/context/contextManager.golden.test.ts b/packages/core/src/context/contextManager.golden.test.ts index 3982672561..9a1da81981 100644 --- a/packages/core/src/context/contextManager.golden.test.ts +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -17,7 +17,10 @@ import type { ContextManager } from './contextManager.js'; import type { Content } from '@google/genai'; import type { Episode } from './ir/types.js'; import type { SidecarConfig } from './sidecar/types.js'; -import { createMockContextConfig, setupContextComponentTest } from './testing/contextTestUtils.js'; +import { + createMockContextConfig, + setupContextComponentTest, +} from './testing/contextTestUtils.js'; expect.addSnapshotSerializer({ test: (val) => @@ -41,7 +44,9 @@ describe('ContextManager Golden Tests', () => { let contextManager: ContextManager; beforeEach(() => { - contextManager = setupContextComponentTest(createMockContextConfig()).contextManager; + contextManager = setupContextComponentTest( + createMockContextConfig(), + ).contextManager; }); const createLargeHistory = (): Content[] => [ @@ -78,12 +83,18 @@ describe('ContextManager Golden Tests', () => { const history = createLargeHistory(); // Use the actual public methods or carefully type the internal state for testing // To seed the manager purely for testing without invoking generateContent, we bypass the pipeline: - const managerAsAny = contextManager as unknown as { - pristineEpisodes: Episode[]; - env: { irMapper: { toIr(h: unknown, t: unknown): Episode[] }, tokenCalculator: unknown } + const managerAsAny = contextManager as unknown as { + pristineEpisodes: Episode[]; + env: { + irMapper: { toIr(h: unknown, t: unknown): Episode[] }; + tokenCalculator: unknown; + }; }; - managerAsAny.pristineEpisodes = managerAsAny.env.irMapper.toIr(history, managerAsAny.env.tokenCalculator); - + managerAsAny.pristineEpisodes = managerAsAny.env.irMapper.toIr( + history, + managerAsAny.env.tokenCalculator, + ); + const result = await contextManager.projectCompressedHistory(); expect(result).toMatchSnapshot(); }); @@ -92,10 +103,11 @@ describe('ContextManager Golden Tests', () => { const history = createLargeHistory(); const config = createMockContextConfig(); - const { chatHistory, contextManager: localManager } = setupContextComponentTest(config, { + const { chatHistory, contextManager: localManager } = + setupContextComponentTest(config, { budget: { retainedTokens: 100000, maxTokens: 150000 }, pipelines: [], - } as unknown as SidecarConfig); + } as unknown as SidecarConfig); chatHistory.set(history); @@ -105,4 +117,3 @@ describe('ContextManager Golden Tests', () => { expect(result.length).toEqual(history.length + 1); }); }); - diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index 6fa3b400ee..5f4e5dea05 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -54,14 +54,14 @@ export class ContextManager { const existingIds = new Set(this.currentNodes.map((n) => n.id)); const addedNodes = event.nodes.filter((n) => !existingIds.has(n.id)); if (addedNodes.length > 0) { - this.currentNodes = [...this.currentNodes, ...addedNodes]; + this.currentNodes = [...this.currentNodes, ...addedNodes]; } this.evaluateTriggers(event.newNodes); }); this.eventBus.onVariantReady((event) => { - // In V2, async workers write back patches. + // In V2, async workers write back patches. // The old variant dict logic is replaced by the orchestrator applying patches directly. // For now we log it. this.tracer.logEvent( @@ -90,13 +90,15 @@ export class ContextManager { if (!this.sidecar.budget) return; if (newNodes.size > 0) { - this.eventBus.emitChunkReceived({ - nodes: this.currentNodes, - targetNodeIds: newNodes - }); + this.eventBus.emitChunkReceived({ + nodes: this.currentNodes, + targetNodeIds: newNodes, + }); } - const currentTokens = this.env.tokenCalculator.calculateConcreteListTokens(this.currentNodes); + const currentTokens = this.env.tokenCalculator.calculateConcreteListTokens( + this.currentNodes, + ); if (currentTokens > this.sidecar.budget.retainedTokens) { const agedOutNodes = new Set(); @@ -104,7 +106,9 @@ export class ContextManager { // Walk backwards finding nodes that fall out of the retained budget for (let i = this.currentNodes.length - 1; i >= 0; i--) { const node = this.currentNodes[i]; - rollingTokens += this.env.tokenCalculator.calculateConcreteListTokens([node]); + rollingTokens += this.env.tokenCalculator.calculateConcreteListTokens([ + node, + ]); if (rollingTokens > this.sidecar.budget.retainedTokens) { agedOutNodes.add(node.id); } diff --git a/packages/core/src/context/ir/behaviorRegistry.ts b/packages/core/src/context/ir/behaviorRegistry.ts index 0148a662de..6729e3b1ad 100644 --- a/packages/core/src/context/ir/behaviorRegistry.ts +++ b/packages/core/src/context/ir/behaviorRegistry.ts @@ -15,10 +15,10 @@ export interface IrSerializationWriter { export interface IrNodeBehavior { readonly type: T['type']; - + /** Serializes the node into the Gemini Content structure. */ serialize(node: T, writer: IrSerializationWriter): void; - + /** * Generates a structural representation of the node for the purpose * of estimating its token cost. diff --git a/packages/core/src/context/ir/builtinBehaviors.ts b/packages/core/src/context/ir/builtinBehaviors.ts index 66555f40ca..c5ec5f8f89 100644 --- a/packages/core/src/context/ir/builtinBehaviors.ts +++ b/packages/core/src/context/ir/builtinBehaviors.ts @@ -4,7 +4,10 @@ * SPDX-License-Identifier: Apache-2.0 */ import type { Part } from '@google/genai'; -import type { IrNodeBehavior, IrNodeBehaviorRegistry } from './behaviorRegistry.js'; +import type { + IrNodeBehavior, + IrNodeBehaviorRegistry, +} from './behaviorRegistry.js'; import type { UserPrompt, AgentThought, @@ -29,7 +32,9 @@ export const UserPromptBehavior: IrNodeBehavior = { parts.push({ inlineData: { mimeType: sp.mimeType, data: sp.data } }); break; case 'file_data': - parts.push({ fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri } }); + parts.push({ + fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri }, + }); break; case 'raw_part': parts.push(sp.part); @@ -46,7 +51,7 @@ export const UserPromptBehavior: IrNodeBehavior = { writer.flushModelParts(); writer.appendContent({ role: 'user', parts }); } - } + }, }; export const AgentThoughtBehavior: IrNodeBehavior = { @@ -56,7 +61,7 @@ export const AgentThoughtBehavior: IrNodeBehavior = { }, serialize(thought, writer) { writer.appendModelPart({ text: thought.text }); - } + }, }; export const ToolExecutionBehavior: IrNodeBehavior = { @@ -64,7 +69,16 @@ export const ToolExecutionBehavior: IrNodeBehavior = { getEstimatableParts(tool) { return [ { functionCall: { id: tool.id, name: tool.toolName, args: tool.intent } }, - { functionResponse: { id: tool.id, name: tool.toolName, response: typeof tool.observation === 'string' ? { message: tool.observation } : tool.observation } } + { + functionResponse: { + id: tool.id, + name: tool.toolName, + response: + typeof tool.observation === 'string' + ? { message: tool.observation } + : tool.observation, + }, + }, ]; }, serialize(tool, writer) { @@ -72,15 +86,30 @@ export const ToolExecutionBehavior: IrNodeBehavior = { writer.appendModelPart(parts[0]); writer.flushModelParts(); writer.appendUserPart(parts[1]); - } + }, }; export const MaskedToolBehavior: IrNodeBehavior = { type: 'MASKED_TOOL', getEstimatableParts(tool) { return [ - { functionCall: { id: tool.id, name: tool.toolName, args: tool.intent ?? {} } }, - { functionResponse: { id: tool.id, name: tool.toolName, response: typeof tool.observation === 'string' ? { message: tool.observation } : (tool.observation ?? {}) } } + { + functionCall: { + id: tool.id, + name: tool.toolName, + args: tool.intent ?? {}, + }, + }, + { + functionResponse: { + id: tool.id, + name: tool.toolName, + response: + typeof tool.observation === 'string' + ? { message: tool.observation } + : (tool.observation ?? {}), + }, + }, ]; }, serialize(tool, writer) { @@ -88,7 +117,7 @@ export const MaskedToolBehavior: IrNodeBehavior = { writer.appendModelPart(parts[0]); writer.flushModelParts(); writer.appendUserPart(parts[1]); - } + }, }; export const AgentYieldBehavior: IrNodeBehavior = { @@ -99,33 +128,39 @@ export const AgentYieldBehavior: IrNodeBehavior = { serialize(yieldNode, writer) { writer.appendModelPart({ text: yieldNode.text }); writer.flushModelParts(); - } + }, }; export const SystemEventBehavior: IrNodeBehavior = { type: 'SYSTEM_EVENT', - getEstimatableParts() { return []; }, + getEstimatableParts() { + return []; + }, serialize(node, writer) { writer.flushModelParts(); - } + }, }; export const SnapshotBehavior: IrNodeBehavior = { type: 'SNAPSHOT', - getEstimatableParts(node) { return [{ text: node.text }]; }, + getEstimatableParts(node) { + return [{ text: node.text }]; + }, serialize(node, writer) { writer.flushModelParts(); writer.appendUserPart({ text: node.text }); - } + }, }; export const RollingSummaryBehavior: IrNodeBehavior = { type: 'ROLLING_SUMMARY', - getEstimatableParts(node) { return [{ text: node.text }]; }, + getEstimatableParts(node) { + return [{ text: node.text }]; + }, serialize(node, writer) { writer.flushModelParts(); writer.appendUserPart({ text: node.text }); - } + }, }; export function registerBuiltInBehaviors(registry: IrNodeBehaviorRegistry) { diff --git a/packages/core/src/context/ir/fromIr.ts b/packages/core/src/context/ir/fromIr.ts index 4f0ccfcd79..fa3afb3053 100644 --- a/packages/core/src/context/ir/fromIr.ts +++ b/packages/core/src/context/ir/fromIr.ts @@ -5,7 +5,10 @@ */ import type { Content, Part } from '@google/genai'; import type { ConcreteNode } from './types.js'; -import type { IrSerializationWriter, IrNodeBehaviorRegistry } from './behaviorRegistry.js'; +import type { + IrSerializationWriter, + IrNodeBehaviorRegistry, +} from './behaviorRegistry.js'; class IrSerializer implements IrSerializationWriter { private history: Content[] = []; @@ -38,7 +41,10 @@ class IrSerializer implements IrSerializationWriter { } } -export function fromIr(nodes: readonly ConcreteNode[], registry: IrNodeBehaviorRegistry): Content[] { +export function fromIr( + nodes: readonly ConcreteNode[], + registry: IrNodeBehaviorRegistry, +): Content[] { const writer = new IrSerializer(); for (const node of nodes) { const behavior = registry.get(node.type); diff --git a/packages/core/src/context/ir/graphUtils.ts b/packages/core/src/context/ir/graphUtils.ts index de10a70783..455dede8b6 100644 --- a/packages/core/src/context/ir/graphUtils.ts +++ b/packages/core/src/context/ir/graphUtils.ts @@ -4,7 +4,19 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { Episode, Task, IrNode, AgentThought, ToolExecution, MaskedTool, UserPrompt, AgentYield, SystemEvent, Snapshot, RollingSummary } from './types.js'; +import type { + Episode, + Task, + IrNode, + AgentThought, + ToolExecution, + MaskedTool, + UserPrompt, + AgentYield, + SystemEvent, + Snapshot, + RollingSummary, +} from './types.js'; export function isEpisode(node: IrNode): node is Episode { return node.type === 'EPISODE'; @@ -42,4 +54,4 @@ export function isSnapshot(node: IrNode): node is Snapshot { } export function isRollingSummary(node: IrNode): node is RollingSummary { return node.type === 'ROLLING_SUMMARY'; -} \ No newline at end of file +} diff --git a/packages/core/src/context/ir/projector.ts b/packages/core/src/context/ir/projector.ts index 881bf6ed1a..f3925e2c9f 100644 --- a/packages/core/src/context/ir/projector.ts +++ b/packages/core/src/context/ir/projector.ts @@ -36,7 +36,8 @@ export class IrProjector { } const maxTokens = sidecar.budget.maxTokens; - const currentTokens = env.tokenCalculator.calculateConcreteListTokens(nodes); + const currentTokens = + env.tokenCalculator.calculateConcreteListTokens(nodes); // V0: Always protect the first node (System Prompt) and the last turn if (nodes.length > 0) { @@ -74,7 +75,9 @@ export class IrProjector { // Start from newest and count backwards for (let i = nodes.length - 1; i >= 0; i--) { const node = nodes[i]; - const nodeTokens = env.tokenCalculator.calculateConcreteListTokens([node]); + const nodeTokens = env.tokenCalculator.calculateConcreteListTokens([ + node, + ]); rollingTokens += nodeTokens; if (rollingTokens > sidecar.budget.retainedTokens) { agedOutNodes.add(node.id); diff --git a/packages/core/src/context/ir/toIr.ts b/packages/core/src/context/ir/toIr.ts index 8129908244..55deb08e98 100644 --- a/packages/core/src/context/ir/toIr.ts +++ b/packages/core/src/context/ir/toIr.ts @@ -17,7 +17,10 @@ import type { import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; // We remove the global nodeIdentityMap and instead rely on one passed from IrMapper -export function getStableId(obj: object, nodeIdentityMap: WeakMap): string { +export function getStableId( + obj: object, + nodeIdentityMap: WeakMap, +): string { let id = nodeIdentityMap.get(obj); if (!id) { id = randomUUID(); @@ -42,14 +45,12 @@ function isCompleteEpisode(ep: Partial): ep is Episode { export function toIr( history: readonly Content[], tokenCalculator: ContextTokenCalculator, - nodeIdentityMap: WeakMap + nodeIdentityMap: WeakMap, ): Episode[] { const episodes: Episode[] = []; let currentEpisode: Partial | null = null; const pendingCallParts: Map = new Map(); - - const finalizeEpisode = () => { if (currentEpisode && isCompleteEpisode(currentEpisode)) { episodes.push(currentEpisode); @@ -72,7 +73,7 @@ export function toIr( currentEpisode, pendingCallParts, tokenCalculator, - nodeIdentityMap + nodeIdentityMap, ); } @@ -85,8 +86,8 @@ export function toIr( msg, currentEpisode, pendingCallParts, - nodeIdentityMap - ); + nodeIdentityMap, + ); } } @@ -103,7 +104,7 @@ function parseToolResponses( currentEpisode: Partial | null, pendingCallParts: Map, tokenCalculator: ContextTokenCalculator, - nodeIdentityMap: WeakMap + nodeIdentityMap: WeakMap, ): Partial { if (!currentEpisode) { currentEpisode = { @@ -151,7 +152,7 @@ function parseToolResponses( function parseUserParts( msg: Content, - nodeIdentityMap: WeakMap + nodeIdentityMap: WeakMap, ): Partial { const semanticParts: SemanticPart[] = []; const parts = msg.parts || []; @@ -191,7 +192,7 @@ function parseModelParts( msg: Content, currentEpisode: Partial | null, pendingCallParts: Map, - nodeIdentityMap: WeakMap + nodeIdentityMap: WeakMap, ): Partial { if (!currentEpisode) { currentEpisode = { diff --git a/packages/core/src/context/processors/blobDegradationProcessor.test.ts b/packages/core/src/context/processors/blobDegradationProcessor.test.ts index 9ebb0bfbc0..c18f61c5e4 100644 --- a/packages/core/src/context/processors/blobDegradationProcessor.test.ts +++ b/packages/core/src/context/processors/blobDegradationProcessor.test.ts @@ -32,7 +32,7 @@ describe('BlobDegradationProcessor', () => { ]; const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, { - semanticParts: parts + semanticParts: parts, }) as UserPrompt; const targets = [prompt]; @@ -41,10 +41,10 @@ describe('BlobDegradationProcessor', () => { expect(result.length).toBe(1); const modifiedPrompt = result[0] as UserPrompt; - + expect(modifiedPrompt.id).not.toBe(prompt.id); expect(modifiedPrompt.semanticParts.length).toBe(3); - + // Text parts should be untouched expect(modifiedPrompt.semanticParts[0]).toEqual(parts[0]); expect(modifiedPrompt.semanticParts[2]).toEqual(parts[2]); @@ -53,7 +53,9 @@ describe('BlobDegradationProcessor', () => { const degradedPart = modifiedPrompt.semanticParts[1]; expect(degradedPart.type).toBe('text'); assert(degradedPart.type === 'text'); - expect(degradedPart.text).toContain('[Multi-Modal Blob (image/png, 0.00MB) degraded to text'); + expect(degradedPart.text).toContain( + '[Multi-Modal Blob (image/png, 0.00MB) degraded to text', + ); }); it('should degrade all blobs unconditionally', async () => { @@ -64,14 +66,14 @@ describe('BlobDegradationProcessor', () => { // Tokens for fileData = 258. // Degraded text = "[File Reference (video/mp4) degraded to text to preserve context window. Original URI: gs://test1]" // Degraded text length ~100 characters. - // Since charsPerToken=1, degraded text = 100 tokens. + // Since charsPerToken=1, degraded text = 100 tokens. // Tokens saved = 258 - 100 = 158. This is > 0, so it WILL degrade it! const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, { semanticParts: [ { type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test1' }, { type: 'file_data', mimeType: 'video/mp4', fileUri: 'gs://test2' }, - ] + ], }) as UserPrompt; const targets = [prompt]; @@ -80,7 +82,7 @@ describe('BlobDegradationProcessor', () => { const modifiedPrompt = result[0] as UserPrompt; expect(modifiedPrompt.semanticParts.length).toBe(2); - + // Both parts should be degraded expect(modifiedPrompt.semanticParts[0].type).toBe('text'); expect(modifiedPrompt.semanticParts[1].type).toBe('text'); diff --git a/packages/core/src/context/processors/blobDegradationProcessor.ts b/packages/core/src/context/processors/blobDegradationProcessor.ts index 6294cc7559..ff8320b06c 100644 --- a/packages/core/src/context/processors/blobDegradationProcessor.ts +++ b/packages/core/src/context/processors/blobDegradationProcessor.ts @@ -70,12 +70,15 @@ export class BlobDegradationProcessor implements ContextProcessor { let newText = ''; let tokensSaved = 0; - switch (part.type) { + switch (part.type) { case 'inline_data': { await ensureDir(); const ext = part.mimeType.split('/')[1] || 'bin'; const fileName = `blob_${Date.now()}_${this.env.idGenerator.generateId()}.${ext}`; - const filePath = this.env.fileSystem.join(blobOutputsDir, fileName); + const filePath = this.env.fileSystem.join( + blobOutputsDir, + fileName, + ); const buffer = Buffer.from(part.data, 'base64'); await this.env.fileSystem.writeFile(filePath, buffer); @@ -83,34 +86,45 @@ export class BlobDegradationProcessor implements ContextProcessor { const mb = (buffer.byteLength / 1024 / 1024).toFixed(2); newText = `[Multi-Modal Blob (${part.mimeType}, ${mb}MB) degraded to text to preserve context window. Saved to: ${filePath}]`; - const oldTokens = this.env.tokenCalculator.estimateTokensForParts([ - { inlineData: { mimeType: part.mimeType, data: part.data } }, - ]); - const newTokens = this.env.tokenCalculator.estimateTokensForParts([ - { text: newText }, - ]); + const oldTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { + inlineData: { mimeType: part.mimeType, data: part.data }, + }, + ]); + const newTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { text: newText }, + ]); tokensSaved = oldTokens - newTokens; break; } case 'file_data': { newText = `[File Reference (${part.mimeType}) degraded to text to preserve context window. Original URI: ${part.fileUri}]`; - const oldTokens = this.env.tokenCalculator.estimateTokensForParts([ - { fileData: { mimeType: part.mimeType, fileUri: part.fileUri } }, - ]); - const newTokens = this.env.tokenCalculator.estimateTokensForParts([ - { text: newText }, - ]); + const oldTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { + fileData: { + mimeType: part.mimeType, + fileUri: part.fileUri, + }, + }, + ]); + const newTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { text: newText }, + ]); tokensSaved = oldTokens - newTokens; break; } case 'raw_part': { newText = `[Unknown Part degraded to text to preserve context window.]`; - const oldTokens = this.env.tokenCalculator.estimateTokensForParts([ - part.part, - ]); - const newTokens = this.env.tokenCalculator.estimateTokensForParts([ - { text: newText }, - ]); + const oldTokens = + this.env.tokenCalculator.estimateTokensForParts([part.part]); + const newTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { text: newText }, + ]); tokensSaved = oldTokens - newTokens; break; } diff --git a/packages/core/src/context/processors/historyTruncationProcessor.ts b/packages/core/src/context/processors/historyTruncationProcessor.ts index 4291c8239a..ca1631c30f 100644 --- a/packages/core/src/context/processors/historyTruncationProcessor.ts +++ b/packages/core/src/context/processors/historyTruncationProcessor.ts @@ -50,21 +50,19 @@ export class HistoryTruncationProcessor implements ContextProcessor { this.options = options; } - async process({ - targets, - }: ProcessArgs): Promise { + async process({ targets }: ProcessArgs): Promise { // Calculate how many tokens we need to remove based on the configured knob let targetTokensToRemove = 0; const strategy = this.options.target ?? 'max'; - + if (strategy === 'incremental') { - targetTokensToRemove = Infinity; + targetTokensToRemove = Infinity; } else if (strategy === 'freeNTokens') { - targetTokensToRemove = this.options.freeTokensTarget ?? 0; - if (targetTokensToRemove <= 0) return targets; + targetTokensToRemove = this.options.freeTokensTarget ?? 0; + if (targetTokensToRemove <= 0) return targets; } else if (strategy === 'max') { - // 'max' means we remove all targets without stopping early - targetTokensToRemove = Infinity; + // 'max' means we remove all targets without stopping early + targetTokensToRemove = Infinity; } let removedTokens = 0; diff --git a/packages/core/src/context/processors/nodeDistillationProcessor.test.ts b/packages/core/src/context/processors/nodeDistillationProcessor.test.ts index 485fce5a85..d3900fc35c 100644 --- a/packages/core/src/context/processors/nodeDistillationProcessor.test.ts +++ b/packages/core/src/context/processors/nodeDistillationProcessor.test.ts @@ -12,7 +12,7 @@ import { createMockEnvironment, createDummyNode, createDummyToolNode, - createMockLlmClient + createMockLlmClient, } from '../testing/contextTestUtils.js'; import type { UserPrompt, AgentThought, ToolExecution } from '../ir/types.js'; @@ -22,7 +22,7 @@ describe('NodeDistillationProcessor', () => { // Use charsPerToken=1 naturally. const env = createMockEnvironment({ - llmClient: mockLlmClient, + llmClient: mockLlmClient, }); const processor = NodeDistillationProcessor.create(env, { @@ -31,19 +31,35 @@ describe('NodeDistillationProcessor', () => { const longText = 'A'.repeat(50); // 50 chars - const prompt = createDummyNode('ep1', 'USER_PROMPT', 50, { - semanticParts: [ - { type: 'text', text: longText } - ], - }, 'prompt-id') as UserPrompt; + const prompt = createDummyNode( + 'ep1', + 'USER_PROMPT', + 50, + { + semanticParts: [{ type: 'text', text: longText }], + }, + 'prompt-id', + ) as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 50, { - text: longText, - }, 'thought-id') as AgentThought; + const thought = createDummyNode( + 'ep1', + 'AGENT_THOUGHT', + 50, + { + text: longText, + }, + 'thought-id', + ) as AgentThought; - const tool = createDummyToolNode('ep1', 5, 500, { - observation: { result: 'A'.repeat(500) }, - }, 'tool-id'); + const tool = createDummyToolNode( + 'ep1', + 5, + 500, + { + observation: { result: 'A'.repeat(500) }, + }, + 'tool-id', + ); const targets = [prompt, thought, tool]; @@ -57,7 +73,7 @@ describe('NodeDistillationProcessor', () => { expect(compressedPrompt.semanticParts[0].type).toBe('text'); assert(compressedPrompt.semanticParts[0].type === 'text'); expect(compressedPrompt.semanticParts[0].text).toBe('Mocked Summary!'); - + // 2. Agent Thought const compressedThought = result[1] as AgentThought; expect(compressedThought.id).not.toBe(thought.id); @@ -75,7 +91,7 @@ describe('NodeDistillationProcessor', () => { const mockLlmClient = createMockLlmClient(['S']); // length = 1 const env = createMockEnvironment({ - llmClient: mockLlmClient, + llmClient: mockLlmClient, }); const processor = NodeDistillationProcessor.create(env, { @@ -84,15 +100,25 @@ describe('NodeDistillationProcessor', () => { const shortText = 'Short text'; // 10 chars - const prompt = createDummyNode('ep1', 'USER_PROMPT', 10, { - semanticParts: [ - { type: 'text', text: shortText } - ], - }, 'prompt-id') as UserPrompt; + const prompt = createDummyNode( + 'ep1', + 'USER_PROMPT', + 10, + { + semanticParts: [{ type: 'text', text: shortText }], + }, + 'prompt-id', + ) as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 13, { - text: 'Short thought', - }, 'thought-id') as AgentThought; + const thought = createDummyNode( + 'ep1', + 'AGENT_THOUGHT', + 13, + { + text: 'Short thought', + }, + 'thought-id', + ) as AgentThought; const targets = [prompt, thought]; @@ -106,7 +132,7 @@ describe('NodeDistillationProcessor', () => { // 2. Agent Thought (NOT compressed) const untouchedThought = result[1] as AgentThought; - expect(untouchedThought.id).toBe(thought.id); + expect(untouchedThought.id).toBe(thought.id); // LLM should not have been called expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(0); diff --git a/packages/core/src/context/processors/nodeDistillationProcessor.ts b/packages/core/src/context/processors/nodeDistillationProcessor.ts index 8057c8024a..dedde91050 100644 --- a/packages/core/src/context/processors/nodeDistillationProcessor.ts +++ b/packages/core/src/context/processors/nodeDistillationProcessor.ts @@ -53,31 +53,32 @@ export class NodeDistillationProcessor implements ContextProcessor { contextInfo: string, ): Promise { try { - const response = await this.env.llmClient.generateContent( - { - role: LlmRole.UTILITY_COMPRESSOR, - modelConfigKey: { model: 'default' }, - promptId: this.env.promptId, - abortSignal: new AbortController().signal, - contents: [ + const response = await this.env.llmClient.generateContent({ + role: LlmRole.UTILITY_COMPRESSOR, + modelConfigKey: { model: 'default' }, + promptId: this.env.promptId, + abortSignal: new AbortController().signal, + contents: [ + { + role: 'user', + parts: [{ text }], + }, + ], + systemInstruction: { + role: 'system', + parts: [ { - role: 'user', - parts: [{ text }], + text: `You are an expert context compressor. Your job is to drastically shorten the following ${contextInfo} while preserving the absolute core semantic meaning, facts, and intent. Omit all conversational filler, pleasantries, or redundant information. Return ONLY the compressed summary.`, }, ], - systemInstruction: { - role: 'system', - parts: [ - { - text: `You are an expert context compressor. Your job is to drastically shorten the following ${contextInfo} while preserving the absolute core semantic meaning, facts, and intent. Omit all conversational filler, pleasantries, or redundant information. Return ONLY the compressed summary.`, - }, - ], - }, - } - ); + }, + }); return getResponseText(response) || text; } catch (e) { - debugLogger.warn(`NodeDistillationProcessor failed to summarize ${contextInfo}`, e); + debugLogger.warn( + `NodeDistillationProcessor failed to summarize ${contextInfo}`, + e, + ); return text; // Fallback to original text on API failure } } @@ -86,7 +87,7 @@ export class NodeDistillationProcessor implements ContextProcessor { const semanticConfig = this.options; const limitTokens = semanticConfig.nodeThresholdTokens; const thresholdChars = this.env.tokenCalculator.tokensToChars(limitTokens); - + const returnedNodes: ConcreteNode[] = []; // Scan the target working buffer and unconditionally apply the configured hyperparameter threshold @@ -101,9 +102,16 @@ export class NodeDistillationProcessor implements ContextProcessor { if (part.type !== 'text') continue; if (part.text.length > thresholdChars) { - const summary = await this.generateSummary(part.text, 'User Prompt'); - const newTokens = this.env.tokenCalculator.estimateTokensForParts([{ text: summary }]); - const oldTokens = this.env.tokenCalculator.estimateTokensForParts([{ text: part.text }]); + const summary = await this.generateSummary( + part.text, + 'User Prompt', + ); + const newTokens = this.env.tokenCalculator.estimateTokensForParts( + [{ text: summary }], + ); + const oldTokens = this.env.tokenCalculator.estimateTokensForParts( + [{ text: part.text }], + ); if (newTokens < oldTokens) { newParts[j] = { type: 'text', text: summary }; @@ -126,8 +134,13 @@ export class NodeDistillationProcessor implements ContextProcessor { case 'AGENT_THOUGHT': { if (node.text.length > thresholdChars) { - const summary = await this.generateSummary(node.text, 'Agent Thought'); - const newTokens = this.env.tokenCalculator.estimateTokensForParts([{ text: summary }]); + const summary = await this.generateSummary( + node.text, + 'Agent Thought', + ); + const newTokens = this.env.tokenCalculator.estimateTokensForParts([ + { text: summary }, + ]); const oldTokens = this.env.tokenCalculator.getTokenCost(node); if (newTokens < oldTokens) { @@ -158,20 +171,26 @@ export class NodeDistillationProcessor implements ContextProcessor { } if (stringifiedObs.length > thresholdChars) { - const summary = await this.generateSummary(stringifiedObs, node.toolName || 'unknown'); + const summary = await this.generateSummary( + stringifiedObs, + node.toolName || 'unknown', + ); const newObsObject = { summary }; - const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([ - { - functionResponse: { - name: node.toolName || 'unknown', - response: newObsObject, - id: node.id, + const newObsTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { + functionResponse: { + name: node.toolName || 'unknown', + response: newObsObject, + id: node.id, + }, }, - }, - ]); + ]); - const oldObsTokens = node.tokens?.observation ?? this.env.tokenCalculator.getTokenCost(node); + const oldObsTokens = + node.tokens?.observation ?? + this.env.tokenCalculator.getTokenCost(node); const intentTokens = node.tokens?.intent ?? 0; if (newObsTokens < oldObsTokens) { diff --git a/packages/core/src/context/processors/nodeTruncationProcessor.test.ts b/packages/core/src/context/processors/nodeTruncationProcessor.test.ts index 1573e60cf6..09764e05ea 100644 --- a/packages/core/src/context/processors/nodeTruncationProcessor.test.ts +++ b/packages/core/src/context/processors/nodeTruncationProcessor.test.ts @@ -25,17 +25,35 @@ describe('NodeTruncationProcessor', () => { const longText = 'A'.repeat(50); // 50 tokens - const prompt = createDummyNode('ep1', 'USER_PROMPT', 50, { - semanticParts: [{ type: 'text', text: longText }] - }, 'prompt-id') as UserPrompt; + const prompt = createDummyNode( + 'ep1', + 'USER_PROMPT', + 50, + { + semanticParts: [{ type: 'text', text: longText }], + }, + 'prompt-id', + ) as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 50, { - text: longText, - }, 'thought-id') as AgentThought; + const thought = createDummyNode( + 'ep1', + 'AGENT_THOUGHT', + 50, + { + text: longText, + }, + 'thought-id', + ) as AgentThought; - const yieldNode = createDummyNode('ep1', 'AGENT_YIELD', 50, { - text: longText, - }, 'yield-id') as AgentYield; + const yieldNode = createDummyNode( + 'ep1', + 'AGENT_YIELD', + 50, + { + text: longText, + }, + 'yield-id', + ) as AgentYield; const targets = [prompt, thought, yieldNode]; @@ -70,13 +88,25 @@ describe('NodeTruncationProcessor', () => { const shortText = 'Short text'; // 10 chars - const prompt = createDummyNode('ep1', 'USER_PROMPT', 10, { - semanticParts: [{ type: 'text', text: shortText }] - }, 'prompt-id') as UserPrompt; + const prompt = createDummyNode( + 'ep1', + 'USER_PROMPT', + 10, + { + semanticParts: [{ type: 'text', text: shortText }], + }, + 'prompt-id', + ) as UserPrompt; - const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 13, { - text: 'Short thought', // 13 chars - }, 'thought-id') as AgentThought; + const thought = createDummyNode( + 'ep1', + 'AGENT_THOUGHT', + 13, + { + text: 'Short thought', // 13 chars + }, + 'thought-id', + ) as AgentThought; const targets = [prompt, thought]; diff --git a/packages/core/src/context/processors/nodeTruncationProcessor.ts b/packages/core/src/context/processors/nodeTruncationProcessor.ts index 69b231683d..40014d494d 100644 --- a/packages/core/src/context/processors/nodeTruncationProcessor.ts +++ b/packages/core/src/context/processors/nodeTruncationProcessor.ts @@ -25,7 +25,8 @@ export class NodeTruncationProcessor implements ContextProcessor { properties: { maxTokensPerNode: { type: 'number', - description: 'The maximum tokens a node can have before being truncated.', + description: + 'The maximum tokens a node can have before being truncated.', }, }, required: ['maxTokensPerNode'], @@ -48,7 +49,12 @@ export class NodeTruncationProcessor implements ContextProcessor { private tryApplySquash( text: string, limitChars: number, - ): { text: string; newTokens: number; oldTokens: number; tokensSaved: number } | null { + ): { + text: string; + newTokens: number; + oldTokens: number; + tokensSaved: number; + } | null { const originalLength = text.length; if (originalLength <= limitChars) return null; @@ -60,7 +66,8 @@ export class NodeTruncationProcessor implements ContextProcessor { if (newText !== text) { // Using accurate TokenCalculator instead of simple math - const newTokens = this.env.tokenCalculator.estimateTokensForString(newText); + const newTokens = + this.env.tokenCalculator.estimateTokensForString(newText); const oldTokens = this.env.tokenCalculator.estimateTokensForString(text); const tokensSaved = oldTokens - newTokens; diff --git a/packages/core/src/context/processors/rollingSummaryProcessor.test.ts b/packages/core/src/context/processors/rollingSummaryProcessor.test.ts index 528aab97d9..86949e1582 100644 --- a/packages/core/src/context/processors/rollingSummaryProcessor.test.ts +++ b/packages/core/src/context/processors/rollingSummaryProcessor.test.ts @@ -5,32 +5,43 @@ */ import { describe, it, expect } from 'vitest'; import { RollingSummaryProcessor } from './rollingSummaryProcessor.js'; -import { createMockProcessArgs, - createMockEnvironment, createDummyNode } from '../testing/contextTestUtils.js'; +import { + createMockProcessArgs, + createMockEnvironment, + createDummyNode, +} from '../testing/contextTestUtils.js'; describe('RollingSummaryProcessor', () => { it('should initialize with correct default options', () => { const env = createMockEnvironment(); - const processor = RollingSummaryProcessor.create(env, { target: 'incremental' }); + const processor = RollingSummaryProcessor.create(env, { + target: 'incremental', + }); expect(processor.id).toBe('RollingSummaryProcessor'); }); it('should summarize older nodes when the deficit exceeds the threshold', async () => { // env.tokenCalculator uses charsPerToken=1 based on createMockEnvironment const env = createMockEnvironment(); - - // We want to free exactly 100 tokens. + + // We want to free exactly 100 tokens. // We will supply nodes that cost 50 tokens each. - const processor = RollingSummaryProcessor.create(env, { - target: 'freeNTokens', - freeTokensTarget: 100 + const processor = RollingSummaryProcessor.create(env, { + target: 'freeNTokens', + freeTokensTarget: 100, }); const text50 = 'A'.repeat(50); const targets = [ - createDummyNode('ep1', 'USER_PROMPT', 50, { semanticParts: [{ type: 'text', text: text50 }] }, 'id1'), + createDummyNode( + 'ep1', + 'USER_PROMPT', + 50, + { semanticParts: [{ type: 'text', text: text50 }] }, + 'id1', + ), createDummyNode('ep1', 'AGENT_THOUGHT', 50, { text: text50 }, 'id2'), - createDummyNode('ep1', 'AGENT_YIELD', 50, { text: text50 }, 'id3'), + createDummyNode('ep1', 'AGENT_YIELD', 50, { text: text50 }, 'id3'), ]; const result = await processor.process(createMockProcessArgs(targets)); @@ -46,21 +57,27 @@ describe('RollingSummaryProcessor', () => { it('should preserve targets if deficit does not trigger summary', async () => { const env = createMockEnvironment(); - + // We want to free 100 tokens, but our nodes will only cost 10 tokens each. - const processor = RollingSummaryProcessor.create(env, { - target: 'freeNTokens', - freeTokensTarget: 100 + const processor = RollingSummaryProcessor.create(env, { + target: 'freeNTokens', + freeTokensTarget: 100, }); const text10 = 'A'.repeat(10); const targets = [ - createDummyNode('ep1', 'USER_PROMPT', 10, { semanticParts: [{ type: 'text', text: text10 }] }, 'id1'), + createDummyNode( + 'ep1', + 'USER_PROMPT', + 10, + { semanticParts: [{ type: 'text', text: text10 }] }, + 'id1', + ), createDummyNode('ep1', 'AGENT_THOUGHT', 10, { text: text10 }, 'id2'), ]; const result = await processor.process(createMockProcessArgs(targets)); - + // Deficit accumulator reaches 10. This is < 100 limit, and total summarizable nodes < 2 anyway. expect(result.length).toBe(2); expect(result[0].type).toBe('USER_PROMPT'); diff --git a/packages/core/src/context/processors/rollingSummaryProcessor.ts b/packages/core/src/context/processors/rollingSummaryProcessor.ts index b131281440..f87b4c4a2f 100644 --- a/packages/core/src/context/processors/rollingSummaryProcessor.ts +++ b/packages/core/src/context/processors/rollingSummaryProcessor.ts @@ -3,7 +3,11 @@ * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -import type { ContextProcessor, ProcessArgs, BackstopTargetOptions } from '../pipeline.js'; +import type { + ContextProcessor, + ProcessArgs, + BackstopTargetOptions, +} from '../pipeline.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import type { ConcreteNode, RollingSummary } from '../ir/types.js'; import { SnapshotGenerator } from '../utils/snapshotGenerator.js'; @@ -37,7 +41,10 @@ export class RollingSummaryProcessor implements ContextProcessor { private readonly env: ContextEnvironment; private readonly generator: SnapshotGenerator; - constructor(env: ContextEnvironment, options: RollingSummaryProcessorOptions) { + constructor( + env: ContextEnvironment, + options: RollingSummaryProcessorOptions, + ) { this.env = env; this.options = options; this.generator = new SnapshotGenerator(env); @@ -50,14 +57,14 @@ export class RollingSummaryProcessor implements ContextProcessor { let targetTokensToRemove = 0; if (strategy === 'incremental') { - // A rolling summary should target a small chunk. For now, since state isn't passed, - // we'll default to a fixed threshold, like 10000 tokens, to avoid eating the whole history. - // Ideally, the orchestrator should pass `tokensToRemove` explicitly. - targetTokensToRemove = 10000; + // A rolling summary should target a small chunk. For now, since state isn't passed, + // we'll default to a fixed threshold, like 10000 tokens, to avoid eating the whole history. + // Ideally, the orchestrator should pass `tokensToRemove` explicitly. + targetTokensToRemove = 10000; } else if (strategy === 'freeNTokens') { - targetTokensToRemove = this.options.freeTokensTarget ?? Infinity; + targetTokensToRemove = this.options.freeTokensTarget ?? Infinity; } else if (strategy === 'max') { - targetTokensToRemove = Infinity; + targetTokensToRemove = Infinity; } if (targetTokensToRemove <= 0) return targets; @@ -68,10 +75,10 @@ export class RollingSummaryProcessor implements ContextProcessor { // Scan oldest to newest to find the oldest block that exceeds the token requirement for (const node of targets) { if (node.id === targets[0].id && node.type === 'USER_PROMPT') { - // Keep system prompt if it's the very first node - continue; + // Keep system prompt if it's the very first node + continue; } - + nodesToSummarize.push(node); deficitAccumulator += this.env.tokenCalculator.getTokenCost(node); @@ -81,34 +88,38 @@ export class RollingSummaryProcessor implements ContextProcessor { if (nodesToSummarize.length < 2) return targets; // Not enough context to summarize try { - // Synthesize the rolling summary synchronously - const snapshotText = await this.generator.synthesizeSnapshot(nodesToSummarize, this.options.systemInstruction); - const newId = this.env.idGenerator.generateId(); - - const summaryNode: RollingSummary = { - id: newId, - logicalParentId: newId, - type: 'ROLLING_SUMMARY', - timestamp: Date.now(), - text: snapshotText, - }; + // Synthesize the rolling summary synchronously + const snapshotText = await this.generator.synthesizeSnapshot( + nodesToSummarize, + this.options.systemInstruction, + ); + const newId = this.env.idGenerator.generateId(); - const consumedIds = nodesToSummarize.map(n => n.id); - const returnedNodes = targets.filter(t => !consumedIds.includes(t.id)); - const firstRemovedIdx = targets.findIndex(t => consumedIds.includes(t.id)); - - if (firstRemovedIdx !== -1) { - const idx = Math.max(0, firstRemovedIdx); - returnedNodes.splice(idx, 0, summaryNode); - } else { - returnedNodes.unshift(summaryNode); - } + const summaryNode: RollingSummary = { + id: newId, + logicalParentId: newId, + type: 'ROLLING_SUMMARY', + timestamp: Date.now(), + text: snapshotText, + }; - return returnedNodes; + const consumedIds = nodesToSummarize.map((n) => n.id); + const returnedNodes = targets.filter((t) => !consumedIds.includes(t.id)); + const firstRemovedIdx = targets.findIndex((t) => + consumedIds.includes(t.id), + ); + if (firstRemovedIdx !== -1) { + const idx = Math.max(0, firstRemovedIdx); + returnedNodes.splice(idx, 0, summaryNode); + } else { + returnedNodes.unshift(summaryNode); + } + + return returnedNodes; } catch (e) { - debugLogger.error('RollingSummaryProcessor failed sync backstop', e); - return targets; + debugLogger.error('RollingSummaryProcessor failed sync backstop', e); + return targets; } } } diff --git a/packages/core/src/context/processors/stateSnapshotProcessor.test.ts b/packages/core/src/context/processors/stateSnapshotProcessor.test.ts index 414e0f035d..31103534cc 100644 --- a/packages/core/src/context/processors/stateSnapshotProcessor.test.ts +++ b/packages/core/src/context/processors/stateSnapshotProcessor.test.ts @@ -15,7 +15,9 @@ import type { InboxSnapshotImpl } from '../sidecar/inbox.js'; describe('StateSnapshotProcessor', () => { it('should ignore if budget is satisfied', async () => { const env = createMockEnvironment(); - const processor = StateSnapshotProcessor.create(env, { target: 'incremental' }); + const processor = StateSnapshotProcessor.create(env, { + target: 'incremental', + }); const targets = [createDummyNode('ep1', 'USER_PROMPT')]; const result = await processor.process(createMockProcessArgs(targets)); expect(result).toBe(targets); // Strict equality @@ -23,12 +25,14 @@ describe('StateSnapshotProcessor', () => { it('should apply a valid snapshot from the Inbox (Fast Path)', async () => { const env = createMockEnvironment(); - const processor = StateSnapshotProcessor.create(env, { target: 'incremental' }); + const processor = StateSnapshotProcessor.create(env, { + target: 'incremental', + }); const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A'); const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B'); const nodeC = createDummyNode('ep2', 'USER_PROMPT', 50, {}, 'node-C'); - + const targets = [nodeA, nodeB, nodeC]; // The background worker created a snapshot of A and B @@ -41,8 +45,8 @@ describe('StateSnapshotProcessor', () => { consumedIds: ['node-A', 'node-B'], newText: '', type: 'point-in-time', - } - } + }, + }, ]; const processArgs = createMockProcessArgs(targets, [], messages); @@ -54,12 +58,16 @@ describe('StateSnapshotProcessor', () => { expect(result[1].id).toBe('node-C'); // Should consume the message - expect((processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1')).toBe(true); + expect( + (processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1'), + ).toBe(true); }); it('should reject a snapshot if the nodes were modified/deleted (Cache Invalidated)', async () => { const env = createMockEnvironment(); - const processor = StateSnapshotProcessor.create(env, { target: 'incremental' }); + const processor = StateSnapshotProcessor.create(env, { + target: 'incremental', + }); // Make deficit 0 so we don't fall through to the sync backstop and fail the test that way // node-A is MISSING (user deleted it) @@ -74,8 +82,8 @@ describe('StateSnapshotProcessor', () => { payload: { consumedIds: ['node-A', 'node-B'], newText: '', - } - } + }, + }, ]; const processArgs = createMockProcessArgs(targets, [], messages); @@ -84,7 +92,9 @@ describe('StateSnapshotProcessor', () => { // Because deficit is 0, and Inbox was rejected, nothing should change expect(result.length).toBe(1); expect(result[0].id).toBe('node-B'); - expect((processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1')).toBe(false); + expect( + (processArgs.inbox as InboxSnapshotImpl).getConsumedIds().has('msg-1'), + ).toBe(false); }); it('should fall back to sync backstop if inbox is empty', async () => { diff --git a/packages/core/src/context/processors/stateSnapshotProcessor.ts b/packages/core/src/context/processors/stateSnapshotProcessor.ts index 8794c7c47a..1c6c5d6f5a 100644 --- a/packages/core/src/context/processors/stateSnapshotProcessor.ts +++ b/packages/core/src/context/processors/stateSnapshotProcessor.ts @@ -3,7 +3,11 @@ * Copyright 2026 Google LLC * SPDX-License-Identifier: Apache-2.0 */ -import type { ContextProcessor, ProcessArgs, BackstopTargetOptions } from '../pipeline.js'; +import type { + ContextProcessor, + ProcessArgs, + BackstopTargetOptions, +} from '../pipeline.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import type { ConcreteNode, Snapshot } from '../ir/types.js'; import { SnapshotGenerator } from '../utils/snapshotGenerator.js'; @@ -46,36 +50,48 @@ export class StateSnapshotProcessor implements ContextProcessor { } // --- ContextProcessor Interface (Sync Backstop / Cache Application) --- - async process({ targets, inbox }: ProcessArgs): Promise { + async process({ + targets, + inbox, + }: ProcessArgs): Promise { if (targets.length === 0) { return targets; } // Determine what mode we are looking for: 'incremental' -> 'point-in-time', 'max' -> 'accumulate' const strategy = this.options.target ?? 'max'; - const expectedType = strategy === 'incremental' ? 'point-in-time' : 'accumulate'; + const expectedType = + strategy === 'incremental' ? 'point-in-time' : 'accumulate'; // 1. Check Inbox for a completed Snapshot (The Fast Path) - const proposedSnapshots = inbox.getMessages<{ newText: string; consumedIds: string[]; type: string }>('PROPOSED_SNAPSHOT'); - + const proposedSnapshots = inbox.getMessages<{ + newText: string; + consumedIds: string[]; + type: string; + }>('PROPOSED_SNAPSHOT'); + if (proposedSnapshots.length > 0) { // Filter for the snapshot type that matches our processor mode - const matchingSnapshots = proposedSnapshots.filter(s => s.payload.type === expectedType); + const matchingSnapshots = proposedSnapshots.filter( + (s) => s.payload.type === expectedType, + ); // Sort by newest timestamp first (we want the most accumulated snapshot) - const sorted = [...matchingSnapshots].sort((a, b) => b.timestamp - a.timestamp); - + const sorted = [...matchingSnapshots].sort( + (a, b) => b.timestamp - a.timestamp, + ); + for (const proposed of sorted) { const { consumedIds, newText } = proposed.payload; - + // Verify all consumed IDs still exist sequentially in targets - const targetIds = new Set(targets.map(t => t.id)); - const isValid = consumedIds.every(id => targetIds.has(id)); - + const targetIds = new Set(targets.map((t) => t.id)); + const isValid = consumedIds.every((id) => targetIds.has(id)); + if (isValid) { // If valid, apply it! const newId = this.env.idGenerator.generateId(); - + const snapshotNode: Snapshot = { id: newId, logicalParentId: newId, @@ -85,14 +101,18 @@ export class StateSnapshotProcessor implements ContextProcessor { }; // Remove the consumed nodes and insert the snapshot at the earliest index - const returnedNodes = targets.filter(t => !consumedIds.includes(t.id)); - const firstRemovedIdx = targets.findIndex(t => consumedIds.includes(t.id)); - + const returnedNodes = targets.filter( + (t) => !consumedIds.includes(t.id), + ); + const firstRemovedIdx = targets.findIndex((t) => + consumedIds.includes(t.id), + ); + if (firstRemovedIdx !== -1) { - const idx = Math.max(0, firstRemovedIdx); - returnedNodes.splice(idx, 0, snapshotNode); + const idx = Math.max(0, firstRemovedIdx); + returnedNodes.splice(idx, 0, snapshotNode); } else { - returnedNodes.unshift(snapshotNode); + returnedNodes.unshift(snapshotNode); } inbox.consume(proposed.id); @@ -105,11 +125,11 @@ export class StateSnapshotProcessor implements ContextProcessor { let targetTokensToRemove = 0; if (strategy === 'incremental') { - targetTokensToRemove = Infinity; // incremental implies removing as much as possible if no state is passed + targetTokensToRemove = Infinity; // incremental implies removing as much as possible if no state is passed } else if (strategy === 'freeNTokens') { - targetTokensToRemove = this.options.freeTokensTarget ?? Infinity; + targetTokensToRemove = this.options.freeTokensTarget ?? Infinity; } else if (strategy === 'max') { - targetTokensToRemove = Infinity; + targetTokensToRemove = Infinity; } let deficitAccumulator = 0; @@ -118,11 +138,11 @@ export class StateSnapshotProcessor implements ContextProcessor { // Scan oldest to newest for (const node of targets) { if (node.id === targets[0].id && node.type === 'USER_PROMPT') { - // Keep system prompt if it's the very first node - // In a real system, system prompt is protected, but we double check - continue; + // Keep system prompt if it's the very first node + // In a real system, system prompt is protected, but we double check + continue; } - + nodesToSummarize.push(node); deficitAccumulator += this.env.tokenCalculator.getTokenCost(node); @@ -132,32 +152,36 @@ export class StateSnapshotProcessor implements ContextProcessor { if (nodesToSummarize.length < 2) return targets; // Not enough context try { - const snapshotText = await this.generator.synthesizeSnapshot(nodesToSummarize, this.options.systemInstruction); - const newId = this.env.idGenerator.generateId(); - const snapshotNode: Snapshot = { - id: newId, - logicalParentId: newId, - type: 'SNAPSHOT', - timestamp: Date.now(), - text: snapshotText, - }; + const snapshotText = await this.generator.synthesizeSnapshot( + nodesToSummarize, + this.options.systemInstruction, + ); + const newId = this.env.idGenerator.generateId(); + const snapshotNode: Snapshot = { + id: newId, + logicalParentId: newId, + type: 'SNAPSHOT', + timestamp: Date.now(), + text: snapshotText, + }; - const consumedIds = nodesToSummarize.map(n => n.id); - const returnedNodes = targets.filter(t => !consumedIds.includes(t.id)); - const firstRemovedIdx = targets.findIndex(t => consumedIds.includes(t.id)); - - if (firstRemovedIdx !== -1) { - const idx = Math.max(0, firstRemovedIdx); - returnedNodes.splice(idx, 0, snapshotNode); - } else { - returnedNodes.unshift(snapshotNode); - } + const consumedIds = nodesToSummarize.map((n) => n.id); + const returnedNodes = targets.filter((t) => !consumedIds.includes(t.id)); + const firstRemovedIdx = targets.findIndex((t) => + consumedIds.includes(t.id), + ); - return returnedNodes; + if (firstRemovedIdx !== -1) { + const idx = Math.max(0, firstRemovedIdx); + returnedNodes.splice(idx, 0, snapshotNode); + } else { + returnedNodes.unshift(snapshotNode); + } + return returnedNodes; } catch (e) { - debugLogger.error('StateSnapshotProcessor failed sync backstop', e); - return targets; + debugLogger.error('StateSnapshotProcessor failed sync backstop', e); + return targets; } } } diff --git a/packages/core/src/context/processors/stateSnapshotWorker.test.ts b/packages/core/src/context/processors/stateSnapshotWorker.test.ts index 6abe4c42a1..97b198a595 100644 --- a/packages/core/src/context/processors/stateSnapshotWorker.test.ts +++ b/packages/core/src/context/processors/stateSnapshotWorker.test.ts @@ -21,7 +21,7 @@ describe('StateSnapshotWorker', () => { const nodeA = createDummyNode('ep1', 'USER_PROMPT', 50, {}, 'node-A'); const nodeB = createDummyNode('ep1', 'AGENT_THOUGHT', 60, {}, 'node-B'); - + const targets = [nodeA, nodeB]; const inbox = new InboxSnapshotImpl([]); @@ -38,7 +38,7 @@ describe('StateSnapshotWorker', () => { consumedIds: ['node-A', 'node-B'], type: 'point-in-time', }), - env.idGenerator + env.idGenerator, ); }); @@ -62,8 +62,8 @@ describe('StateSnapshotWorker', () => { consumedIds: ['node-A', 'node-B'], newText: '', type: 'accumulate', - } - } + }, + }, ]); await worker.execute({ targets, inbox }); @@ -80,7 +80,7 @@ describe('StateSnapshotWorker', () => { consumedIds: ['node-A', 'node-B', 'node-C'], // Aggregated! type: 'accumulate', }), - env.idGenerator + env.idGenerator, ); // Verify the LLM was called with the old snapshot prepended @@ -91,11 +91,11 @@ describe('StateSnapshotWorker', () => { parts: expect.arrayContaining([ expect.objectContaining({ text: expect.stringContaining(''), - }) - ]) - }) - ]) - }) + }), + ]), + }), + ]), + }), ); }); diff --git a/packages/core/src/context/processors/stateSnapshotWorker.ts b/packages/core/src/context/processors/stateSnapshotWorker.ts index bfafd582d6..012debbba6 100644 --- a/packages/core/src/context/processors/stateSnapshotWorker.ts +++ b/packages/core/src/context/processors/stateSnapshotWorker.ts @@ -48,7 +48,13 @@ export class StateSnapshotWorker implements ContextWorker { this.generator = new SnapshotGenerator(env); } - async execute({ targets, inbox }: { targets: readonly ConcreteNode[]; inbox: InboxSnapshot }): Promise { + async execute({ + targets, + inbox, + }: { + targets: readonly ConcreteNode[]; + inbox: InboxSnapshot; + }): Promise { if (targets.length === 0) return; try { @@ -58,16 +64,24 @@ export class StateSnapshotWorker implements ContextWorker { if (workerType === 'accumulate') { // Look for the most recent unconsumed accumulate snapshot in the inbox - const proposedSnapshots = inbox.getMessages<{ newText: string; consumedIds: string[]; type: string }>('PROPOSED_SNAPSHOT'); - const accumulateSnapshots = proposedSnapshots.filter(s => s.payload.type === 'accumulate'); - + const proposedSnapshots = inbox.getMessages<{ + newText: string; + consumedIds: string[]; + type: string; + }>('PROPOSED_SNAPSHOT'); + const accumulateSnapshots = proposedSnapshots.filter( + (s) => s.payload.type === 'accumulate', + ); + if (accumulateSnapshots.length > 0) { // Sort to find the most recent - const latest = [...accumulateSnapshots].sort((a, b) => b.timestamp - a.timestamp)[0]; - + const latest = [...accumulateSnapshots].sort( + (a, b) => b.timestamp - a.timestamp, + )[0]; + // Consume the old draft so the inbox doesn't fill up with stale drafts inbox.consume(latest.id); - // And we must persist its consumption back to the live inbox immediately, + // And we must persist its consumption back to the live inbox immediately, // because we are effectively "taking" it from the shelf to modify. this.env.inbox.drainConsumed(new Set([latest.id])); @@ -91,15 +105,21 @@ export class StateSnapshotWorker implements ContextWorker { this.options.systemInstruction, ); - const newConsumedIds = [...previousConsumedIds, ...targets.map((t) => t.id)]; + const newConsumedIds = [ + ...previousConsumedIds, + ...targets.map((t) => t.id), + ]; // In V2, workers communicate their work to the inbox, and the processor picks it up. - this.env.inbox.publish('PROPOSED_SNAPSHOT', { - newText: snapshotText, - consumedIds: newConsumedIds, - type: workerType, - }, this.env.idGenerator); - + this.env.inbox.publish( + 'PROPOSED_SNAPSHOT', + { + newText: snapshotText, + consumedIds: newConsumedIds, + type: workerType, + }, + this.env.idGenerator, + ); } catch (e) { debugLogger.error('StateSnapshotWorker failed to generate snapshot', e); } diff --git a/packages/core/src/context/processors/toolMaskingProcessor.test.ts b/packages/core/src/context/processors/toolMaskingProcessor.test.ts index 7c85f1bca5..e7689ce128 100644 --- a/packages/core/src/context/processors/toolMaskingProcessor.test.ts +++ b/packages/core/src/context/processors/toolMaskingProcessor.test.ts @@ -35,12 +35,12 @@ describe('ToolMaskingProcessor', () => { expect(result.length).toBe(1); const masked = result[0] as ToolExecution; - + // It should have generated a new ID because it modified it expect(masked.id).not.toBe(toolStep.id); // It should have masked the observation - const obs = masked.observation as { result: string, metadata: string }; + const obs = masked.observation as { result: string; metadata: string }; expect(obs.result).toContain(''); expect(obs.metadata).toBe('short'); // Untouched }); @@ -55,8 +55,9 @@ describe('ToolMaskingProcessor', () => { const toolStep = createDummyToolNode('ep1', 10, 10, { toolName: 'activate_skill', observation: { - result: 'this is a really long string that normally would get masked but wont because of the tool name', - } + result: + 'this is a really long string that normally would get masked but wont because of the tool name', + }, }); const result = await processor.process(createMockProcessArgs([toolStep])); diff --git a/packages/core/src/context/processors/toolMaskingProcessor.ts b/packages/core/src/context/processors/toolMaskingProcessor.ts index 924526ad5f..6a0f5050e7 100644 --- a/packages/core/src/context/processors/toolMaskingProcessor.ts +++ b/packages/core/src/context/processors/toolMaskingProcessor.ts @@ -214,34 +214,39 @@ export class ToolMaskingProcessor implements ContextProcessor { ? (intentRes.masked as Record) : undefined; // Handle observation explicitly as string vs object - const maskedObs = typeof obsRes.masked === 'string' - ? { message: obsRes.masked } as Record - : isMaskableRecord(obsRes.masked) - ? (obsRes.masked as Record) - : undefined; + const maskedObs = + typeof obsRes.masked === 'string' + ? ({ message: obsRes.masked } as Record) + : isMaskableRecord(obsRes.masked) + ? (obsRes.masked as Record) + : undefined; - const newIntentTokens = this.env.tokenCalculator.estimateTokensForParts([ - { - functionCall: { - name: toolName || 'unknown', - args: maskedIntent, - id: callId, + const newIntentTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { + functionCall: { + name: toolName || 'unknown', + args: maskedIntent, + id: callId, + }, }, - }, - ]); - + ]); + let obsPart: Record = {}; if (maskedObs) { - obsPart = { - functionResponse: { - name: toolName || 'unknown', - response: maskedObs, - id: callId - } - }; + obsPart = { + functionResponse: { + name: toolName || 'unknown', + response: maskedObs, + id: callId, + }, + }; } - - const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([obsPart as Part]); + + const newObsTokens = + this.env.tokenCalculator.estimateTokensForParts([ + obsPart as Part, + ]); const tokensSaved = this.env.tokenCalculator.getTokenCost(node) - diff --git a/packages/core/src/context/sidecar/builtins.ts b/packages/core/src/context/sidecar/builtins.ts index 024f706bf7..b3f552bd1d 100644 --- a/packages/core/src/context/sidecar/builtins.ts +++ b/packages/core/src/context/sidecar/builtins.ts @@ -4,14 +4,35 @@ * SPDX-License-Identifier: Apache-2.0 */ import type { SidecarRegistry } from './registry.js'; -import { HistoryTruncationProcessor, type HistoryTruncationProcessorOptions } from '../processors/historyTruncationProcessor.js'; +import { + HistoryTruncationProcessor, + type HistoryTruncationProcessorOptions, +} from '../processors/historyTruncationProcessor.js'; import { BlobDegradationProcessor } from '../processors/blobDegradationProcessor.js'; -import { NodeTruncationProcessor, type NodeTruncationProcessorOptions } from '../processors/nodeTruncationProcessor.js'; -import { NodeDistillationProcessor, type NodeDistillationProcessorOptions } from '../processors/nodeDistillationProcessor.js'; -import { ToolMaskingProcessor, type ToolMaskingProcessorOptions } from '../processors/toolMaskingProcessor.js'; -import { StateSnapshotProcessor, type StateSnapshotProcessorOptions } from '../processors/stateSnapshotProcessor.js'; -import { StateSnapshotWorker, type StateSnapshotWorkerOptions } from '../processors/stateSnapshotWorker.js'; -import { RollingSummaryProcessor, type RollingSummaryProcessorOptions } from '../processors/rollingSummaryProcessor.js'; +import { + NodeTruncationProcessor, + type NodeTruncationProcessorOptions, +} from '../processors/nodeTruncationProcessor.js'; +import { + NodeDistillationProcessor, + type NodeDistillationProcessorOptions, +} from '../processors/nodeDistillationProcessor.js'; +import { + ToolMaskingProcessor, + type ToolMaskingProcessorOptions, +} from '../processors/toolMaskingProcessor.js'; +import { + StateSnapshotProcessor, + type StateSnapshotProcessorOptions, +} from '../processors/stateSnapshotProcessor.js'; +import { + StateSnapshotWorker, + type StateSnapshotWorkerOptions, +} from '../processors/stateSnapshotWorker.js'; +import { + RollingSummaryProcessor, + type RollingSummaryProcessorOptions, +} from '../processors/rollingSummaryProcessor.js'; export function registerBuiltInProcessors(registry: SidecarRegistry) { registry.registerProcessor>({ diff --git a/packages/core/src/context/sidecar/environmentImpl.test.ts b/packages/core/src/context/sidecar/environmentImpl.test.ts index e990b106b2..4636c5ccd2 100644 --- a/packages/core/src/context/sidecar/environmentImpl.test.ts +++ b/packages/core/src/context/sidecar/environmentImpl.test.ts @@ -10,7 +10,6 @@ import { ContextEventBus } from '../eventBus.js'; import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js'; - import { createMockLlmClient } from '../testing/contextTestUtils.js'; describe('ContextEnvironmentImpl', () => { diff --git a/packages/core/src/context/sidecar/environmentImpl.ts b/packages/core/src/context/sidecar/environmentImpl.ts index 6a7f0aa8d9..55a828021d 100644 --- a/packages/core/src/context/sidecar/environmentImpl.ts +++ b/packages/core/src/context/sidecar/environmentImpl.ts @@ -42,7 +42,10 @@ export class ContextEnvironmentImpl implements ContextEnvironment { ) { this.behaviorRegistry = new IrNodeBehaviorRegistry(); registerBuiltInBehaviors(this.behaviorRegistry); - this.tokenCalculator = new ContextTokenCalculator(this.charsPerToken, this.behaviorRegistry); + this.tokenCalculator = new ContextTokenCalculator( + this.charsPerToken, + this.behaviorRegistry, + ); this.fileSystem = fileSystem || new NodeFileSystem(); this.idGenerator = idGenerator || new NodeIdGenerator(); this.inbox = new LiveInbox(); diff --git a/packages/core/src/context/sidecar/inbox.ts b/packages/core/src/context/sidecar/inbox.ts index 1b96dc835b..b9aded9d6a 100644 --- a/packages/core/src/context/sidecar/inbox.ts +++ b/packages/core/src/context/sidecar/inbox.ts @@ -8,7 +8,11 @@ import type { InboxMessage, InboxSnapshot } from '../pipeline.js'; export class LiveInbox { private messages: InboxMessage[] = []; - publish(topic: string, payload: T, idGenerator: { generateId(): string }): void { + publish( + topic: string, + payload: T, + idGenerator: { generateId(): string }, + ): void { this.messages.push({ id: idGenerator.generateId(), topic, diff --git a/packages/core/src/context/sidecar/orchestrator.test.ts b/packages/core/src/context/sidecar/orchestrator.test.ts index 041f257954..70215353d7 100644 --- a/packages/core/src/context/sidecar/orchestrator.test.ts +++ b/packages/core/src/context/sidecar/orchestrator.test.ts @@ -13,10 +13,15 @@ import { createDummyNode, } from '../testing/contextTestUtils.js'; import type { ContextEnvironment } from './environment.js'; -import type { ContextProcessor, ContextWorker } from '../pipeline.js'; +import type { + ContextProcessor, + ContextWorker, + InboxSnapshot, + ProcessArgs, +} from '../pipeline.js'; import type { PipelineDef, ProcessorConfig, SidecarConfig } from './types.js'; import type { ContextEventBus } from '../eventBus.js'; -import type { UserPrompt } from '../ir/types.js'; +import type { ConcreteNode, UserPrompt } from '../ir/types.js'; // A realistic mock processor that modifies the text of the first target node class ModifyingProcessor implements ContextProcessor { @@ -27,13 +32,16 @@ class ModifyingProcessor implements ContextProcessor { readonly name = 'ModifyingProcessor'; readonly id = 'ModifyingProcessor'; readonly options = {}; - async process(args: import('../pipeline.js').ProcessArgs) { + async process(args: ProcessArgs) { const newTargets = [...args.targets]; if (newTargets.length > 0 && newTargets[0].type === 'USER_PROMPT') { const prompt = newTargets[0]; const newParts = [...prompt.semanticParts]; if (newParts.length > 0 && newParts[0].type === 'text') { - newParts[0] = { ...newParts[0], text: newParts[0].text + ' [modified]' }; + newParts[0] = { + ...newParts[0], + text: newParts[0].text + ' [modified]', + }; } newTargets[0] = { ...prompt, semanticParts: newParts }; } @@ -50,7 +58,7 @@ class ThrowingProcessor implements ContextProcessor { readonly name = 'Throwing'; readonly id = 'Throwing'; readonly options = {}; - async process(): Promise> { + async process(): Promise { throw new Error('Processor failed intentionally'); } } @@ -69,8 +77,8 @@ class MockWorker implements ContextWorker { wasExecuted = false; async execute(args: { - targets: ReadonlyArray; - inbox: import('../pipeline.js').InboxSnapshot; + targets: readonly ConcreteNode[]; + inbox: InboxSnapshot; }) { this.wasExecuted = true; if (args.targets.length > 0 && args.targets[0].type === 'USER_PROMPT') { @@ -114,23 +122,29 @@ describe('PipelineOrchestrator (Component)', () => { registry.clear(); }); - const createConfig = (pipelines: PipelineDef[], workers?: Array<{ workerId: string }>): SidecarConfig => ({ + const createConfig = ( + pipelines: PipelineDef[], + workers?: Array<{ workerId: string }>, + ): SidecarConfig => ({ budget: { maxTokens: 100, retainedTokens: 50 }, pipelines, workers, }); it('instantiates processors and workers from the registry on initialization', () => { - const config = createConfig([ - { - name: 'SyncPipe', - execution: 'blocking', - triggers: ['new_message'], - processors: [ - { processorId: 'ModifyingProcessor' } as unknown as ProcessorConfig, - ], - }, - ], [{ workerId: 'MockWorker' }]); + const config = createConfig( + [ + { + name: 'SyncPipe', + execution: 'blocking', + triggers: ['new_message'], + processors: [ + { processorId: 'ModifyingProcessor' } as unknown as ProcessorConfig, + ], + }, + ], + [{ workerId: 'MockWorker' }], + ); const orchestrator = new PipelineOrchestrator( config, @@ -139,11 +153,15 @@ describe('PipelineOrchestrator (Component)', () => { env.tracer, registry, ); - + // eslint-disable-next-line @typescript-eslint/no-explicit-any - expect((orchestrator as any).instantiatedProcessors.has('ModifyingProcessor')).toBe(true); + expect( + (orchestrator as any).instantiatedProcessors.has('ModifyingProcessor'), + ).toBe(true); // eslint-disable-next-line @typescript-eslint/no-explicit-any - expect((orchestrator as any).instantiatedWorkers.has('MockWorker')).toBe(true); + expect((orchestrator as any).instantiatedWorkers.has('MockWorker')).toBe( + true, + ); }); it('throws an error if a config requests an unknown processor', () => { @@ -183,9 +201,17 @@ describe('PipelineOrchestrator (Component)', () => { registry, ); - const nodes = [createDummyNode('not-protected-ep', 'USER_PROMPT', 100, { - semanticParts: [{ type: 'text', text: 'original text' }] - }, 'not-protected-id')]; + const nodes = [ + createDummyNode( + 'not-protected-ep', + 'USER_PROMPT', + 100, + { + semanticParts: [{ type: 'text', text: 'original text' }], + }, + 'not-protected-id', + ), + ]; const result = await orchestrator.executeTriggerSync( 'new_message', @@ -196,8 +222,13 @@ describe('PipelineOrchestrator (Component)', () => { expect(result).toHaveLength(1); const modifiedPrompt = result[0] as UserPrompt; - assert(modifiedPrompt.semanticParts[0].type === 'text', 'Expected a text part'); - expect(modifiedPrompt.semanticParts[0].text).toBe('original text [modified]'); + assert( + modifiedPrompt.semanticParts[0].type === 'text', + 'Expected a text part', + ); + expect(modifiedPrompt.semanticParts[0].text).toBe( + 'original text [modified]', + ); }); it('gracefully handles and swallows processor errors in synchronous routes', async () => { @@ -219,7 +250,15 @@ describe('PipelineOrchestrator (Component)', () => { registry, ); - const nodes = [createDummyNode('not-protected-ep', 'USER_PROMPT', 100, undefined, 'not-protected-id')]; + const nodes = [ + createDummyNode( + 'not-protected-ep', + 'USER_PROMPT', + 100, + undefined, + 'not-protected-id', + ), + ]; // It should not throw! It should swallow the error and return the unmodified array. const result = await orchestrator.executeTriggerSync( @@ -269,20 +308,36 @@ describe('PipelineOrchestrator (Component)', () => { it('automatically dispatches workers when matching EventBus events occur', async () => { const config = createConfig([], [{ workerId: 'MockWorker' }]); - const orchestrator = new PipelineOrchestrator(config, env, eventBus, env.tracer, registry); - + const orchestrator = new PipelineOrchestrator( + config, + env, + eventBus, + env.tracer, + registry, + ); + // eslint-disable-next-line @typescript-eslint/no-explicit-any - const workerInstance = (orchestrator as any).instantiatedWorkers.get('MockWorker') as MockWorker; + const workerInstance = (orchestrator as any).instantiatedWorkers.get( + 'MockWorker', + ) as MockWorker; expect(workerInstance.wasExecuted).toBe(false); - const nodes = [createDummyNode('not-protected-ep', 'USER_PROMPT', 100, { - semanticParts: [{ type: 'text', text: 'worker trigger text' }] - }, 'not-protected-id')]; + const nodes = [ + createDummyNode( + 'not-protected-ep', + 'USER_PROMPT', + 100, + { + semanticParts: [{ type: 'text', text: 'worker trigger text' }], + }, + 'not-protected-id', + ), + ]; // Emit the new_message chunk which maps to onNodesAdded for workers eventBus.emitChunkReceived({ nodes, - targetNodeIds: new Set(nodes.map(n => n.id)), + targetNodeIds: new Set(nodes.map((n) => n.id)), }); // Worker execute is fire and forget, so we yield to the event loop @@ -290,4 +345,4 @@ describe('PipelineOrchestrator (Component)', () => { expect(workerInstance.wasExecuted).toBe(true); }); -}); \ No newline at end of file +}); diff --git a/packages/core/src/context/sidecar/orchestrator.ts b/packages/core/src/context/sidecar/orchestrator.ts index 8775b351d3..38afe1520b 100644 --- a/packages/core/src/context/sidecar/orchestrator.ts +++ b/packages/core/src/context/sidecar/orchestrator.ts @@ -8,7 +8,8 @@ import type { ConcreteNode } from '../ir/types.js'; import type { ContextProcessor, ContextWorker, - ContextWorkingBuffer } from '../pipeline.js'; + ContextWorkingBuffer, +} from '../pipeline.js'; import type { SidecarConfig, PipelineDef, PipelineTrigger } from './types.js'; import type { ContextEnvironment, @@ -22,10 +23,8 @@ import { InboxSnapshotImpl } from './inbox.js'; class ContextWorkingBufferImpl implements ContextWorkingBuffer { private readonly nodesMap: Map; - constructor( - readonly nodes: readonly ConcreteNode[], - ) { - this.nodesMap = new Map(nodes.map(n => [n.id, n])); + constructor(readonly nodes: readonly ConcreteNode[]) { + this.nodesMap = new Map(nodes.map((n) => [n.id, n])); } getPristineNode(id: string): ConcreteNode | undefined { @@ -77,8 +76,7 @@ export class PipelineOrchestrator { return ( triggerTargets.has(node.id) && !protectedLogicalIds.has(node.id) && - (!node.logicalParentId || - !protectedLogicalIds.has(node.logicalParentId)) + (!node.logicalParentId || !protectedLogicalIds.has(node.logicalParentId)) ); } @@ -87,10 +85,7 @@ export class PipelineOrchestrator { for (const procDef of pipeline.processors) { if (!this.instantiatedProcessors.has(procDef.processorId)) { const factory = this.registry.getProcessor(procDef.processorId); - const instance = factory.create( - this.env, - procDef.options || {}, - ); + const instance = factory.create(this.env, procDef.options || {}); this.instantiatedProcessors.set(procDef.processorId, instance); } } @@ -102,10 +97,7 @@ export class PipelineOrchestrator { for (const workerDef of this.config.workers) { if (!this.instantiatedWorkers.has(workerDef.workerId)) { const factory = this.registry.getWorker(workerDef.workerId); - const instance = factory.create( - this.env, - workerDef.options || {}, - ); + const instance = factory.create(this.env, workerDef.options || {}); this.instantiatedWorkers.set(workerDef.workerId, instance); } } @@ -150,7 +142,9 @@ export class PipelineOrchestrator { const inboxSnapshot = new InboxSnapshotImpl( this.env.inbox.getMessages() || [], ); - const targets = event.nodes.filter(n => event.targetNodeIds.has(n.id)); + const targets = event.nodes.filter((n) => + event.targetNodeIds.has(n.id), + ); // Fire and forget worker.execute({ targets, inbox: inboxSnapshot }).catch((e) => { debugLogger.error(`Worker ${worker.name} failed onNodesAdded:`, e); @@ -166,10 +160,15 @@ export class PipelineOrchestrator { const inboxSnapshot = new InboxSnapshotImpl( this.env.inbox.getMessages() || [], ); - const targets = event.nodes.filter(n => event.targetNodeIds.has(n.id)); + const targets = event.nodes.filter((n) => + event.targetNodeIds.has(n.id), + ); // Fire and forget worker.execute({ targets, inbox: inboxSnapshot }).catch((e) => { - debugLogger.error(`Worker ${worker.name} failed onNodesAgedOut:`, e); + debugLogger.error( + `Worker ${worker.name} failed onNodesAgedOut:`, + e, + ); }); } } diff --git a/packages/core/src/context/sidecar/profiles.ts b/packages/core/src/context/sidecar/profiles.ts index da98ce4193..eefee516f4 100644 --- a/packages/core/src/context/sidecar/profiles.ts +++ b/packages/core/src/context/sidecar/profiles.ts @@ -56,9 +56,9 @@ export const defaultSidecarProfile: SidecarConfig = { triggers: ['gc_backstop'], execution: 'blocking', processors: [ - { - processorId: 'StateSnapshotProcessor', - options: { target: 'max' } + { + processorId: 'StateSnapshotProcessor', + options: { target: 'max' }, }, ], }, diff --git a/packages/core/src/context/sidecar/registry.test.ts b/packages/core/src/context/sidecar/registry.test.ts index f808d17f14..34613c963d 100644 --- a/packages/core/src/context/sidecar/registry.test.ts +++ b/packages/core/src/context/sidecar/registry.test.ts @@ -14,7 +14,7 @@ describe('SidecarRegistry', () => { const processorDef: ContextProcessorDef = { id: 'TestProcessor', schema: { type: 'object' }, - create: () => ({} as ContextProcessor), + create: () => ({}) as ContextProcessor, }; registry.registerProcessor(processorDef); @@ -27,7 +27,7 @@ describe('SidecarRegistry', () => { const workerDef: ContextWorkerDef = { id: 'TestWorker', schema: { type: 'object' }, - create: () => ({} as ContextWorker), + create: () => ({}) as ContextWorker, }; registry.registerWorker(workerDef); @@ -37,12 +37,16 @@ describe('SidecarRegistry', () => { it('should throw an error when retrieving unregistered processors', () => { const registry = new SidecarRegistry(); - expect(() => registry.getProcessor('Unknown')).toThrow('Context Processor [Unknown] is not registered.'); + expect(() => registry.getProcessor('Unknown')).toThrow( + 'Context Processor [Unknown] is not registered.', + ); }); it('should throw an error when retrieving unregistered workers', () => { const registry = new SidecarRegistry(); - expect(() => registry.getWorker('Unknown')).toThrow('Context Worker [Unknown] is not registered.'); + expect(() => registry.getWorker('Unknown')).toThrow( + 'Context Worker [Unknown] is not registered.', + ); }); it('should return combined schemas', () => { @@ -50,24 +54,32 @@ describe('SidecarRegistry', () => { registry.registerProcessor({ id: 'TestProcessor', schema: { title: 'processorSchema' }, - create: () => ({} as ContextProcessor), + create: () => ({}) as ContextProcessor, }); registry.registerWorker({ id: 'TestWorker', schema: { title: 'workerSchema' }, - create: () => ({} as ContextWorker), + create: () => ({}) as ContextWorker, }); - const schemas = registry.getSchemas() as Array<{title?: string}>; + const schemas = registry.getSchemas() as Array<{ title?: string }>; expect(schemas.length).toBe(2); - expect(schemas.find(s => s.title === 'processorSchema')).toBeDefined(); - expect(schemas.find(s => s.title === 'workerSchema')).toBeDefined(); + expect(schemas.find((s) => s.title === 'processorSchema')).toBeDefined(); + expect(schemas.find((s) => s.title === 'workerSchema')).toBeDefined(); }); it('should safely clear the registry', () => { const registry = new SidecarRegistry(); - registry.registerProcessor({ id: 'TestProcessor', schema: {}, create: () => ({} as ContextProcessor) }); - registry.registerWorker({ id: 'TestWorker', schema: {}, create: () => ({} as ContextWorker) }); + registry.registerProcessor({ + id: 'TestProcessor', + schema: {}, + create: () => ({}) as ContextProcessor, + }); + registry.registerWorker({ + id: 'TestWorker', + schema: {}, + create: () => ({}) as ContextWorker, + }); registry.clear(); diff --git a/packages/core/src/context/sidecar/registry.ts b/packages/core/src/context/sidecar/registry.ts index 1b95877b6e..92bbb866ac 100644 --- a/packages/core/src/context/sidecar/registry.ts +++ b/packages/core/src/context/sidecar/registry.ts @@ -65,4 +65,4 @@ export class SidecarRegistry { this.processors.clear(); this.workers.clear(); } -} \ No newline at end of file +} diff --git a/packages/core/src/context/sidecar/testProfile.ts b/packages/core/src/context/sidecar/testProfile.ts index 68641b7d76..6d25a262af 100644 --- a/packages/core/src/context/sidecar/testProfile.ts +++ b/packages/core/src/context/sidecar/testProfile.ts @@ -15,9 +15,7 @@ export const testTruncateProfile: SidecarConfig = { name: 'Emergency Backstop (Truncate Only)', triggers: ['gc_backstop', 'retained_exceeded'], execution: 'blocking', - processors: [ - { processorId: 'HistoryTruncationProcessor', options: {} }, - ], + processors: [{ processorId: 'HistoryTruncationProcessor', options: {} }], }, ], }; diff --git a/packages/core/src/context/system-tests/SimulationHarness.ts b/packages/core/src/context/system-tests/SimulationHarness.ts index ba170ca967..e772bdded4 100644 --- a/packages/core/src/context/system-tests/SimulationHarness.ts +++ b/packages/core/src/context/system-tests/SimulationHarness.ts @@ -90,7 +90,7 @@ export class SimulationHarness { this.env, this.tracer, this.orchestrator, - this.chatHistory + this.chatHistory, ); } @@ -127,17 +127,13 @@ export class SimulationHarness { currentView = await orchestrator.executeTriggerSync( 'gc_backstop', currentView, - new Set(currentView.map(e => e.id)), + new Set(currentView.map((e) => e.id)), new Set(), ); // Inject the truncated view back into the graph for (let i = 0; i < currentView.length; i++) { const ep = currentView[i]; - if ( - !this.contextManager - .getNodes() - .find((c) => c.id === ep.id) - ) { + if (!this.contextManager.getNodes().find((c) => c.id === ep.id)) { this.eventBus.emitVariantReady({ targetId: ep.id, variantId: 'v-emergency', @@ -145,7 +141,9 @@ export class SimulationHarness { type: 'MASKED_TOOL', id: 'mock-id', tokens: { intent: 0, observation: 0 }, - intent: {}, observation: {}, toolName: 'tool', + intent: {}, + observation: {}, + toolName: 'tool', }, }); } diff --git a/packages/core/src/context/system-tests/lifecycle.golden.test.ts b/packages/core/src/context/system-tests/lifecycle.golden.test.ts index 0a3437ae06..bdab334aa6 100644 --- a/packages/core/src/context/system-tests/lifecycle.golden.test.ts +++ b/packages/core/src/context/system-tests/lifecycle.golden.test.ts @@ -9,7 +9,6 @@ import { SimulationHarness } from './SimulationHarness.js'; import { createMockLlmClient } from '../testing/contextTestUtils.js'; import type { SidecarConfig } from '../sidecar/types.js'; - expect.addSnapshotSerializer({ test: (val) => typeof val === 'string' && @@ -59,7 +58,9 @@ describe('System Lifecycle Golden Tests', () => { ], }); - const mockLlmClient = createMockLlmClient(['']); + const mockLlmClient = createMockLlmClient([ + '', + ]); it('Scenario 1: Organic Growth with Huge Tool Output & Images', async () => { const harness = await SimulationHarness.create( diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index 6715e66d4e..b6e5ff41cf 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -23,20 +23,24 @@ import type { Config } from '../../config/config.js'; import type { BaseLlmClient } from '../../core/baseLlmClient.js'; import type { Content, GenerateContentResponse } from '@google/genai'; import { InboxSnapshotImpl } from '../sidecar/inbox.js'; -import type { ContextWorkingBuffer, InboxMessage, ProcessArgs } from '../pipeline.js'; - +import type { + ContextWorkingBuffer, + InboxMessage, + ProcessArgs, +} from '../pipeline.js'; /** * Creates a valid mock GenerateContentResponse with the provided text. * Used to avoid having to manually construct the deeply nested candidate/content/part structure. */ -export const createMockGenerateContentResponse = (text: string): GenerateContentResponse => +export const createMockGenerateContentResponse = ( + text: string, +): GenerateContentResponse => // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion ({ candidates: [{ content: { role: 'model', parts: [{ text }] }, index: 0 }], }) as GenerateContentResponse; - export function createDummyNode( logicalParentId: string, type: 'USER_PROMPT' | 'SYSTEM_EVENT' | 'AGENT_THOUGHT' | 'AGENT_YIELD', @@ -101,13 +105,17 @@ export interface MockLlmClient extends BaseLlmClient { generateContent: Mock; } -export function createMockLlmClient(responses?: Array): MockLlmClient { +export function createMockLlmClient( + responses?: Array, +): MockLlmClient { const generateContentMock = vi.fn(); - + if (responses && responses.length > 0) { for (const response of responses) { if (typeof response === 'string') { - generateContentMock.mockResolvedValueOnce(createMockGenerateContentResponse(response)); + generateContentMock.mockResolvedValueOnce( + createMockGenerateContentResponse(response), + ); } else { generateContentMock.mockResolvedValueOnce(response); } @@ -115,13 +123,17 @@ export function createMockLlmClient(responses?: Array, ): ContextEnvironment { const llmClient = createMockLlmClient(['Mock LLM summary response']); - - const tracer = new ContextTracer({ targetDir: '/tmp', sessionId: 'mock-session' }); + + const tracer = new ContextTracer({ + targetDir: '/tmp', + sessionId: 'mock-session', + }); const eventBus = new ContextEventBus(); - + const env = new ContextEnvironmentImpl( llmClient, 'mock-session', @@ -148,7 +163,7 @@ export function createMockEnvironment( 1, eventBus, new InMemoryFileSystem(), - new DeterministicIdGenerator('mock-uuid-') + new DeterministicIdGenerator('mock-uuid-'), ); if (overrides) { @@ -188,10 +203,16 @@ export class FakeContextWorkingBuffer implements ContextWorkingBuffer { } } -export function createMockProcessArgs(targets: ConcreteNode[], bufferNodes: ConcreteNode[] = [], inboxMessages: InboxMessage[] = []): ProcessArgs { +export function createMockProcessArgs( + targets: ConcreteNode[], + bufferNodes: ConcreteNode[] = [], + inboxMessages: InboxMessage[] = [], +): ProcessArgs { return { targets, - buffer: new FakeContextWorkingBuffer(bufferNodes.length ? bufferNodes : targets), + buffer: new FakeContextWorkingBuffer( + bufferNodes.length ? bufferNodes : targets, + ), inbox: new InboxSnapshotImpl(inboxMessages), }; } @@ -253,7 +274,7 @@ export function createMockContextConfig( export function setupContextComponentTest( config: Config, sidecarOverride?: SidecarConfig, -): {chatHistory: AgentChatHistory, contextManager: ContextManager} { +): { chatHistory: AgentChatHistory; contextManager: ContextManager } { const chatHistory = new AgentChatHistory(); const registry = new SidecarRegistry(); registerBuiltInProcessors(registry); @@ -273,13 +294,13 @@ export function setupContextComponentTest( 1, eventBus, ); - + const orchestrator = new PipelineOrchestrator( sidecar, env, eventBus, tracer, - registry + registry, ); const contextManager = new ContextManager( @@ -287,7 +308,7 @@ export function setupContextComponentTest( env, tracer, orchestrator, - chatHistory + chatHistory, ); // The async worker is now internally managed by ContextManager diff --git a/packages/core/src/context/typed-context-ir.md b/packages/core/src/context/typed-context-ir.md index 0abf4b8b16..fc11091e2f 100644 --- a/packages/core/src/context/typed-context-ir.md +++ b/packages/core/src/context/typed-context-ir.md @@ -1,60 +1,90 @@ # Context Manager: The Pure Functional "Nodes of Theseus" IR -This document outlines the architectural transition from the V0 Mutating Editor pattern to the V1 Pure Functional, Immutable Episodic IR, designed to scale into a multi-agent, async state transformation system. +This document outlines the architectural transition from the V0 Mutating Editor +pattern to the V1 Pure Functional, Immutable Episodic IR, designed to scale into +a multi-agent, async state transformation system. ## 1. Core Philosophy: The Nodes of Theseus -The primary constraint of deep immutable trees is the cascading cost of cloning parent nodes when a leaf node changes. To solve this, we decouple the structural hierarchy of the context from the actual data sent to the LLM. +The primary constraint of deep immutable trees is the cascading cost of cloning +parent nodes when a leaf node changes. To solve this, we decouple the structural +hierarchy of the context from the actual data sent to the LLM. The IR is divided into two distinct domains: -1. **Logical Nodes:** Structural boundaries that define the hierarchy (e.g., `Task`, `Episode`). These nodes **do not render** to the LLM. They exist to group related interactions and provide semantic meaning. -2. **Concrete Nodes:** The atomic, renderable pieces of data (e.g., `UserPrompt`, `ToolExecution`, `Snapshot`, `RollingSummary`). These are the actual "planks" of the nodes. -Because Concrete Nodes carry a reference to their Logical Parent (e.g., `episodeId`), they can be stored and processed as a **Flat List**. +1. **Logical Nodes:** Structural boundaries that define the hierarchy (e.g., + `Task`, `Episode`). These nodes **do not render** to the LLM. They exist to + group related interactions and provide semantic meaning. +2. **Concrete Nodes:** The atomic, renderable pieces of data (e.g., + `UserPrompt`, `ToolExecution`, `Snapshot`, `RollingSummary`). These are the + actual "planks" of the nodes. + +Because Concrete Nodes carry a reference to their Logical Parent (e.g., +`episodeId`), they can be stored and processed as a **Flat List**. ## 2. The Autonomous `ContextWorkingBuffer` -The "Nodes" is no longer a dumb array; it is encapsulated in a rich `ContextWorkingBuffer` entity. +The "Nodes" is no longer a dumb array; it is encapsulated in a rich +`ContextWorkingBuffer` entity. ### Encapsulation of History -The Buffer manages its own audit trail and lineage. If a processor needs the pristine, unaltered data of a deeply compressed node (e.g., a Snapshotter summarizing masked tools), it queries the Buffer directly: + +The Buffer manages its own audit trail and lineage. If a processor needs the +pristine, unaltered data of a deeply compressed node (e.g., a Snapshotter +summarizing masked tools), it queries the Buffer directly: `buffer.getPristineNode(id)` ### Linear Temporal Progression (The Conveyor Belt) -Processors do not vote or compete. Context degradation is a linear temporal progression defined by triggers: -1. **Frontbuffer Trim:** E.g., Tool Masking replaces raw tools immediately. -2. **Backbuffer Normalize:** E.g., Summarization replaces aging nodes in the background. -3. **GC Backstop:** E.g., Truncation brutally destroys nodes only when the absolute budget is breached. -When a pipeline triggers, the Orchestrator runs its processors, gathers their `ContextPatch`es, and applies them to the Buffer immediately. The state simply advances. +Processors do not vote or compete. Context degradation is a linear temporal +progression defined by triggers: + +1. **Frontbuffer Trim:** E.g., Tool Masking replaces raw tools immediately. +2. **Backbuffer Normalize:** E.g., Summarization replaces aging nodes in the + background. +3. **GC Backstop:** E.g., Truncation brutally destroys nodes only when the + absolute budget is breached. + +When a pipeline triggers, the Orchestrator runs its processors, gathers their +`ContextPatch`es, and applies them to the Buffer immediately. The state simply +advances. ## 3. Type-Safe Async Coordination (The `ContextInbox`) -To solve the async/sync barrier (where a slow background worker generates a summary that a fast synchronous emergency backstop needs instantly), we introduce the `ContextInbox`. +To solve the async/sync barrier (where a slow background worker generates a +summary that a fast synchronous emergency backstop needs instantly), we +introduce the `ContextInbox`. -This is a strictly-typed messaging system. A worker dispatches a `SNAPSHOT_READY` message to the Inbox. The backstop peeks at the Inbox, instantly retrieving the pre-computed summary and applying it. +This is a strictly-typed messaging system. A worker dispatches a +`SNAPSHOT_READY` message to the Inbox. The backstop peeks at the Inbox, +instantly retrieving the pre-computed summary and applying it. ## 4. The Processor Contract -Processors are purely functional map/filter operations. They evaluate a list of unprotected targets and return the exact list of nodes they intend to substitute. They do **not** generate manual `ContextPatch` objects or manage `IrMetadata`. +Processors are purely functional map/filter operations. They evaluate a list of +unprotected targets and return the exact list of nodes they intend to +substitute. They do **not** generate manual `ContextPatch` objects or manage +`IrMetadata`. ```typescript -export type InboxMessage = +export type InboxMessage = | { type: 'SNAPSHOT_READY'; snapshot: Snapshot; abstractsIds: string[] } | { type: 'BACKGROUND_SUMMARY'; summary: RollingSummary; targetId: string }; export interface ContextInbox { dispatch(message: InboxMessage): void; - peek(type: T): Extract | undefined; + peek( + type: T, + ): Extract | undefined; } export interface ContextWorkingBuffer { /** The current active (projected) flat list of ConcreteNodes. */ readonly nodes: ReadonlyArray; - + /** Retrieves the historical, pristine version of a node (before any masks/summaries). */ getPristineNode(id: string): ConcreteNode | undefined; - + /** Retrieves the full audit lineage of a specific node ID. */ getLineage(id: string): ReadonlyArray; } @@ -62,17 +92,17 @@ export interface ContextWorkingBuffer { export interface ProcessArgs { /** The rich buffer containing current nodes and their history. */ readonly buffer: ContextWorkingBuffer; - - /** + + /** * The specific unprotected, mutable nodes the pipeline is allowed to operate on. * The Orchestrator strictly filters out ANY protected nodes (like active tasks) before calling. * Processors can assume all targets passed here are legally theirs to mutate or drop. */ readonly targets: ReadonlyArray; - + /** The token budget and accounting state. */ readonly state: ContextAccountingState; - + /** Type-safe messaging system for async/sync coordination. */ readonly inbox: ContextInbox; } @@ -80,9 +110,9 @@ export interface ProcessArgs { export interface ContextProcessor { readonly id: string; readonly name: string; - - /** - * A pure function. Returns the new state of the `targets`. + + /** + * A pure function. Returns the new state of the `targets`. * If an ID from `targets` is missing in the return array, the Orchestrator deletes it. * If a new synthetic node is in the return array, the Orchestrator inserts it. * The Orchestrator automatically appends audit `IrMetadata` to any changes. @@ -93,7 +123,8 @@ export interface ContextProcessor { ## 5. The Node Taxonomy (`IrNodeType`) -The `IrNodeType` union explicitly defines all valid nodes. Synthetic nodes (like `Snapshot`) are first-class citizens. +The `IrNodeType` union explicitly defines all valid nodes. Synthetic nodes (like +`Snapshot`) are first-class citizens. ```typescript export type IrNodeType = @@ -107,9 +138,9 @@ export type IrNodeType = | 'AGENT_THOUGHT' | 'TOOL_EXECUTION' | 'AGENT_YIELD' - + // Synthetic Concrete Nodes | 'SNAPSHOT' | 'ROLLING_SUMMARY' | 'MASKED_TOOL'; -``` \ No newline at end of file +``` diff --git a/packages/core/src/context/utils/contextTokenCalculator.ts b/packages/core/src/context/utils/contextTokenCalculator.ts index 2d0126e5db..6dfbc67035 100644 --- a/packages/core/src/context/utils/contextTokenCalculator.ts +++ b/packages/core/src/context/utils/contextTokenCalculator.ts @@ -14,13 +14,12 @@ import type { IrNodeBehaviorRegistry } from '../ir/behaviorRegistry.js'; * by the Gemini API. We use this as a baseline heuristic for inlineData/fileData. */ - export class ContextTokenCalculator { private readonly tokenCache = new Map(); constructor( private readonly charsPerToken: number, - private readonly registry: IrNodeBehaviorRegistry + private readonly registry: IrNodeBehaviorRegistry, ) {} /** @@ -68,7 +67,7 @@ export class ContextTokenCalculator { calculateConcreteListTokens(nodes: readonly ConcreteNode[]): number { let tokens = 0; for (const node of nodes) { - tokens += this.getTokenCost(node); + tokens += this.getTokenCost(node); } return tokens; }