diff --git a/packages/core/src/context/contextManager.async.test.ts b/packages/core/src/context/contextManager.async.test.ts deleted file mode 100644 index 4921cb119e..0000000000 --- a/packages/core/src/context/contextManager.async.test.ts +++ /dev/null @@ -1,111 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, it, expect } from 'vitest'; -import { - createMockContextConfig, - setupContextComponentTest, -} from './testing/contextTestUtils.js'; - -describe('ContextManager Barrier Tests', () => { - it('Soft Barrier (retainedTokens): should inject ready variants and shrink projection', async () => { - const config = createMockContextConfig(); - const { chatHistory, contextManager } = setupContextComponentTest(config); - - // 1. Shrink limits: 1 char = 1 token. RetainedTokens = 10. MaxTokens = 100. - - contextManager['sidecar'].budget.retainedTokens = 5; - contextManager['sidecar'].budget.maxTokens = 100; - - // 2. Build tiny history: 5 turns (10 messages). 2 tokens per turn. - const tinyHistory = []; - for (let i = 0; i < 5; i++) { - tinyHistory.push({ role: 'user', parts: [{ text: `U${i}` }] }); - tinyHistory.push({ role: 'model', parts: [{ text: `M${i}` }] }); - } - - // Set history directly to avoid event races - chatHistory.set(tinyHistory); - - // 3. Pre-verify baseline length. - const baseline = await contextManager.projectCompressedHistory(); - expect(baseline.length).toBe(10); - - // 4. Emit a fake snapshot covering the first 3 pairs (6 messages) - const targetEp = contextManager['pristineEpisodes'][2]; - const replacedIds = contextManager['pristineEpisodes'] - .slice(0, 3) - .map((ep) => ep.id); - - contextManager['eventBus'].emitVariantReady({ - targetId: targetEp.id, - variantId: 'snapshot', - variant: { - status: 'ready', - type: 'snapshot', - replacedEpisodeIds: replacedIds, - episode: { - type: 'EPISODE', - id: 'snapshot-ep', - timestamp: Date.now(), - trigger: { - id: 't1', - type: 'USER_PROMPT', - semanticParts: [], - metadata: { - originalTokens: 0, - currentTokens: 0, - transformations: [], - }, - }, - yield: { - id: 'y1', - type: 'AGENT_YIELD', - text: '', - metadata: { - originalTokens: 5, - currentTokens: 5, - transformations: [], - }, - }, - steps: [], - }, - }, - }); - - // 5. Verify Projection shrinks: 6 original messages replaced by 1 snapshot episode (1 text part) -> length 5. - const projection = await contextManager.projectCompressedHistory(); - expect(projection.length).toBe(5); - // projection[0] should be the snapshot yield - expect(projection[0].parts![0].text).toBe(''); - }); - - it('Hard Barrier (maxTokens): should ruthlessly truncate unprotected episodes', async () => { - const config = createMockContextConfig(); - const { chatHistory, contextManager } = setupContextComponentTest(config); - - // 1. Shrink limits: maxTokens = 15. - - contextManager['sidecar'].budget.maxTokens = 15; - - // 2. Build history: 2 turns. Total = 24 tokens. - const history = [ - { role: 'user', parts: [{ text: 'U0' }] }, - { role: 'model', parts: [{ text: 'M0_LARGE!!' }] }, - { role: 'user', parts: [{ text: 'U1' }] }, - { role: 'model', parts: [{ text: 'M1_LARGE!!' }] }, - ]; - chatHistory.set(history); - - const projection = await contextManager.projectCompressedHistory(); - - // Because Turn 0 is architecturally protected (system prompt/initialization), it SURVIVES! - // Turn 1 is dropped to satisfy the maxTokens constraint. - expect(projection.length).toBe(2); - expect(projection[0].parts![0].text).toBe('U0'); - expect(projection[1].parts![0].text).toBe('M0_LARGE!!'); - }); -}); diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index 9cc0202fd6..20040239c3 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -7,26 +7,24 @@ import type { Content } from '@google/genai'; import type { AgentChatHistory } from '../core/agentChatHistory.js'; import { debugLogger } from '../utils/debugLogger.js'; -import type { Episode } from './ir/types.js'; +import type { ConcreteNode } from './ir/types.js'; import type { ContextEventBus } from './eventBus.js'; import type { ContextTracer } from './tracer.js'; import type { ContextEnvironment } from './sidecar/environment.js'; import type { SidecarConfig } from './sidecar/types.js'; import { PipelineOrchestrator } from './sidecar/orchestrator.js'; import { HistoryObserver } from './historyObserver.js'; -import { generateWorkingBufferView } from './ir/graphUtils.js'; import { IrProjector } from './ir/projector.js'; import { registerBuiltInProcessors } from './sidecar/builtins.js'; import { ProcessorRegistry } from './sidecar/registry.js'; export class ContextManager { - // The stateful, pristine Episodic Intermediate Representation graph. - // This allows the agent to remember and summarize continuously without losing data across turns. - private pristineEpisodes: Episode[] = []; + // The stateful, pristine flat graph. + private pristineShip: ReadonlyArray = []; + private currentShip: ReadonlyArray = []; private readonly eventBus: ContextEventBus; // Internal sub-components - // Synchronous processors are instantiated but effectively used as singletons within this class private orchestrator: PipelineOrchestrator; private historyObserver?: HistoryObserver; @@ -58,28 +56,31 @@ export class ContextManager { this.orchestrator = orchestrator; this.eventBus.onPristineHistoryUpdated((event) => { - this.pristineEpisodes = event.episodes; + this.pristineShip = event.ship; + // In V2, we assume currentShip updates sequentially via Orchestrator patches. + // But if pristine changes, we must ensure our current view incorporates new nodes. + // For now, simple fallback: if the current ship doesn't have the new nodes, append them. + // A more robust implementation would diff the ship, but for now we'll just track. + const existingIds = new Set(this.currentShip.map((n) => n.id)); + const addedNodes = event.ship.filter((n) => !existingIds.has(n.id)); + if (addedNodes.length > 0) { + this.currentShip = [...this.currentShip, ...addedNodes]; + } + this.evaluateTriggers(event.newNodes); }); this.eventBus.onVariantReady((event) => { - // Find the target episode in the pristine graph - const targetEp = this.pristineEpisodes.find( - (ep) => ep.id === event.targetId, + // 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( + 'ContextManager', + `Received async variant [${event.variantId}] for Node ${event.targetId}`, + ); + debugLogger.log( + `ContextManager: Received async variant [${event.variantId}] for Node ${event.targetId}.`, ); - if (targetEp) { - if (!targetEp.variants) { - targetEp.variants = {}; - } - targetEp.variants[event.variantId] = event.variant; - this.tracer.logEvent( - 'ContextManager', - `Received async variant [${event.variantId}] for Episode ${event.targetId}`, - ); - debugLogger.log( - `ContextManager: Received async variant [${event.variantId}] for Episode ${event.targetId}.`, - ); - } }); } @@ -100,56 +101,39 @@ export class ContextManager { private evaluateTriggers(newNodes: Set) { if (!this.sidecar.budget) return; - const workingBuffer = this.getWorkingBufferView(); - const currentTokens = - this.env.tokenCalculator.calculateEpisodeListTokens(workingBuffer); - - this.tracer.logEvent('ContextManager', 'Evaluated triggers', { - currentTokens, - retainedTokens: this.sidecar.budget.retainedTokens, - }); - - // 1. Eager Compute Trigger (on_turn) if (newNodes.size > 0) { - this.eventBus.emitChunkReceived({ episodes: this.pristineEpisodes, targetNodeIds: newNodes }); + this.eventBus.emitChunkReceived({ + ship: this.currentShip, + targetNodeIds: newNodes + }); } - // 2. Budget Crossed Trigger + const currentTokens = this.env.tokenCalculator.calculateConcreteListTokens(this.currentShip); + if (currentTokens > this.sidecar.budget.retainedTokens) { - const deficit = currentTokens - this.sidecar.budget.retainedTokens; - - // Calculate exactly which nodes aged out of the retainedTokens budget to form our target delta const agedOutNodes = new Set(); let rollingTokens = 0; - // Start from newest and count backwards - for (let i = workingBuffer.length - 1; i >= 0; i--) { - const ep = workingBuffer[i]; - const epTokens = this.env.tokenCalculator.calculateEpisodeListTokens([ep]); - rollingTokens += epTokens; + // Walk backwards finding nodes that fall out of the retained budget + for (let i = this.currentShip.length - 1; i >= 0; i--) { + const node = this.currentShip[i]; + rollingTokens += node.metadata.currentTokens; if (rollingTokens > this.sidecar.budget.retainedTokens) { - agedOutNodes.add(ep.id); - agedOutNodes.add(ep.trigger.id); - for (const step of ep.steps) agedOutNodes.add(step.id); - if (ep.yield) agedOutNodes.add(ep.yield.id); + agedOutNodes.add(node.id); } } - this.tracer.logEvent( - 'ContextManager', - 'Budget crossed. Emitting ConsolidationNeeded', - { deficit, agedOutCount: agedOutNodes.size }, - ); - this.eventBus.emitConsolidationNeeded({ - episodes: workingBuffer, - targetDeficit: deficit, - targetNodeIds: agedOutNodes, - }); + if (agedOutNodes.size > 0) { + this.eventBus.emitConsolidationNeeded({ + ship: this.currentShip, + targetDeficit: currentTokens - this.sidecar.budget.retainedTokens, + targetNodeIds: agedOutNodes, + }); + } } } /** - * Subscribes to the core AgentChatHistory to natively track all message events, - * converting them seamlessly into pristine Episodes. + * Starts tracking the raw agent history and translating it to Episodic IR. */ subscribeToHistory(chatHistory: AgentChatHistory) { if (this.historyObserver) { @@ -166,39 +150,47 @@ export class ContextManager { } /** - * Generates a computed view of the pristine log. - * Sweeps backwards (newest to oldest), tracking rolling tokens. - * When rollingTokens > retainedTokens, it injects the "best" available ready variant - * (snapshot > summary > masked) instead of the raw text. - * Handles N-to-1 variant skipping automatically. + * Retrieves the raw, uncompressed Episodic IR graph. + * Useful for internal tool rendering (like the trace viewer). + * Note: This is an expensive, deep clone operation. */ - getWorkingBufferView(): Episode[] { - return generateWorkingBufferView( - this.pristineEpisodes, - this.sidecar.budget.retainedTokens, - this.tracer, - this.env, - ); + getPristineGraph(): ReadonlyArray { + return [...this.pristineShip]; } /** - * Returns a temporary, compressed Content[] array to be used exclusively for the LLM request. - * This does NOT mutate the pristine episodic graph. + * Generates a virtual view of the pristine graph, substituting in variants + * up to the configured token budget. + * This is the view that will eventually be projected back to the LLM. */ - async projectCompressedHistory(): Promise { - this.tracer.logEvent('ContextManager', 'Projection requested.'); - const protectedIds = new Set(); - if (this.pristineEpisodes.length > 0) { - protectedIds.add(this.pristineEpisodes[0].id); // Structural invariant - } + getShip(): ReadonlyArray { + return [...this.currentShip]; + } - return IrProjector.project( - this.getWorkingBufferView(), + /** + * Executes the final 'gc_backstop' pipeline if necessary, enforcing the token budget, + * and maps the Episodic IR back into a raw Gemini Content[] array for transmission. + * This is the primary method called by the agent framework before sending a request. + */ + async projectCompressedHistory( + activeTaskIds: Set = new Set(), + ): Promise { + this.tracer.logEvent( + 'ContextManager', + 'Starting projection to LLM context', + ); + // Apply final GC Backstop pressure barrier synchronously before mapping + const finalHistory = await IrProjector.project( + this.currentShip, this.orchestrator, this.sidecar, this.tracer, this.env, - protectedIds, + activeTaskIds, ); + + this.tracer.logEvent('ContextManager', 'Finished projection'); + + return finalHistory; } } diff --git a/packages/core/src/context/eventBus.ts b/packages/core/src/context/eventBus.ts index ee6ddb9bcf..9b7d3c7729 100644 --- a/packages/core/src/context/eventBus.ts +++ b/packages/core/src/context/eventBus.ts @@ -5,28 +5,28 @@ */ import { EventEmitter } from 'node:events'; -import type { Episode, Variant } from './ir/types.js'; +import type { ConcreteNode } from './ir/types.js'; export interface PristineHistoryUpdatedEvent { - episodes: Episode[]; + ship: ReadonlyArray; newNodes: Set; } export interface ContextConsolidationEvent { - episodes: Episode[]; + ship: ReadonlyArray; targetDeficit: number; targetNodeIds: Set; } export interface IrChunkReceivedEvent { - episodes: Episode[]; + ship: ReadonlyArray; targetNodeIds: Set; } export interface VariantReadyEvent { targetId: string; // The Episode or Step ID this variant attaches to variantId: string; // A unique ID for the variant itself - variant: Variant; + variant: ConcreteNode; // In V2, variants are synthetic concrete nodes } export class ContextEventBus extends EventEmitter { diff --git a/packages/core/src/context/historyObserver.ts b/packages/core/src/context/historyObserver.ts index c2cd1dda43..ea1c89280c 100644 --- a/packages/core/src/context/historyObserver.ts +++ b/packages/core/src/context/historyObserver.ts @@ -38,41 +38,38 @@ export class HistoryObserver { this.unsubscribeHistory = this.chatHistory.subscribe( (_event: HistoryEvent) => { // Rebuild the pristine IR graph from the full source history on every change. + // Wait, toIr still returns an Episode[]. + // We actually need to map the Episode[] to a flat ConcreteNode[] here to form the 'ship'. const pristineEpisodes = IrMapper.toIr( this.chatHistory.get(), this.tokenCalculator, ); - const newNodes = new Set(); + const ship: import('./ir/types.js').ConcreteNode[] = []; for (const ep of pristineEpisodes) { - if (!this.seenNodeIds.has(ep.id)) { - newNodes.add(ep.id); - this.seenNodeIds.add(ep.id); - } - if (!this.seenNodeIds.has(ep.trigger.id)) { - newNodes.add(ep.trigger.id); - this.seenNodeIds.add(ep.trigger.id); - } - for (const step of ep.steps) { - if (!this.seenNodeIds.has(step.id)) { - newNodes.add(step.id); - this.seenNodeIds.add(step.id); - } - } - if (ep.yield && !this.seenNodeIds.has(ep.yield.id)) { - newNodes.add(ep.yield.id); - this.seenNodeIds.add(ep.yield.id); + if (ep.concreteNodeIds) { + for (const child of ep.concreteNodeIds) { + ship.push(child as unknown as import('./ir/types.js').ConcreteNode); + } + } + } + + const newNodes = new Set(); + for (const node of ship) { + if (!this.seenNodeIds.has(node.id)) { + newNodes.add(node.id); + this.seenNodeIds.add(node.id); } } this.tracer.logEvent( 'HistoryObserver', 'Rebuilt pristine graph from chat history update', - { episodeCount: pristineEpisodes.length, newNodesCount: newNodes.size }, + { shipSize: ship.length, newNodesCount: newNodes.size }, ); this.eventBus.emitPristineHistoryUpdated({ - episodes: pristineEpisodes, + ship, newNodes, }); }, diff --git a/packages/core/src/context/ir/fromIr.ts b/packages/core/src/context/ir/fromIr.ts index 9a683c8584..f730236a6f 100644 --- a/packages/core/src/context/ir/fromIr.ts +++ b/packages/core/src/context/ir/fromIr.ts @@ -102,7 +102,7 @@ function serializeToolExecution( functionResponse: { id: tool.id, name: tool.toolName, - response: typeof tool.observation === 'string' ? { message: tool.observation } : tool.observation as object, + response: typeof tool.observation === "string" ? { message: tool.observation } : tool.observation as Record, }, }, }; diff --git a/packages/core/src/context/ir/graphUtils.test.ts b/packages/core/src/context/ir/graphUtils.test.ts deleted file mode 100644 index 9ba63db80f..0000000000 --- a/packages/core/src/context/ir/graphUtils.test.ts +++ /dev/null @@ -1,175 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, it, expect, beforeEach, vi } from 'vitest'; -import { generateWorkingBufferView } from './graphUtils.js'; -import { - createMockEnvironment, - createDummyEpisode, -} from '../testing/contextTestUtils.js'; -import type { ContextEnvironment } from '../sidecar/environment.js'; -import type { AgentThought, UserPrompt } from './types.js'; - -describe('graphUtils (View Generator)', () => { - let env: ContextEnvironment; - - beforeEach(() => { - vi.resetAllMocks(); - env = createMockEnvironment(); - // Our token mock is 1 char = 1 token for simplicity - vi.spyOn( - env.tokenCalculator, - 'calculateEpisodeListTokens', - ).mockImplementation((eps) => - eps.reduce( - (acc, ep) => acc + (ep.trigger.metadata.originalTokens || 100), - 0, - ), - ); - }); - - it('returns pristine episodes untouched if under budget', () => { - const episodes = [ - createDummyEpisode('ep-1', 'USER_PROMPT', [{ type: 'text', text: '1' }]), - createDummyEpisode('ep-2', 'USER_PROMPT', [{ type: 'text', text: '2' }]), - ]; - - // We retain 5000 tokens. Total mock tokens = 200. - const view = generateWorkingBufferView(episodes, 5000, env.tracer, env); - - expect(view).toHaveLength(2); - // Must be a deep copy! The view generator clones episodes. - expect(view).not.toBe(episodes); - expect(view[0].id).toBe('ep-1'); - expect(view[1].id).toBe('ep-2'); - }); - - it('swaps to Masked variant when over budget (rolling backwards)', () => { - const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [ - { text: '1', type: 'text' }, - ]); - const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [ - { text: '2', type: 'text' }, - ]); - - ep1.variants = { - masked: { - type: 'masked', - status: 'ready', - text: '', - recoveredTokens: 10, - }, - }; - - // We only retain 100 tokens. - // ep-2 (newest) takes 100 tokens. - // Now rolling = 100. Over budget! - // ep-1 is evaluated, and swapped for Masked. - const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env); - - expect(view).toHaveLength(2); - expect(view[1].id).toBe('ep-2'); // Unchanged (newest) - - expect(view[0].id).toBe('ep-1'); - expect( - (view[0].trigger as UserPrompt).semanticParts[0].presentation?.text, - ).toBe(''); - }); - - it('swaps to Summary variant when over budget', () => { - const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: '1' }, - ]); - const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [ - { type: 'text', text: '2' }, - ]); - - ep1.variants = { - summary: { - type: 'summary', - status: 'ready', - text: '', - recoveredTokens: 50, - }, - }; - - const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env); - - expect(view).toHaveLength(2); - - // The summary completely replaces the internal steps and clears the yield. - expect(view[0].steps).toHaveLength(1); - expect(view[0].steps[0].type).toBe('AGENT_THOUGHT'); - expect((view[0].steps[0] as AgentThought).text).toBe(''); - expect(view[0].yield).toBeUndefined(); - }); - - it('handles complex N-to-1 Snapshot skipping gracefully', () => { - const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: '1' }, - ]); - const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [ - { type: 'text', text: '2' }, - ]); - const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [ - { type: 'text', text: '3' }, - ]); - const ep4 = createDummyEpisode('ep-4', 'USER_PROMPT', [ - { type: 'text', text: '4' }, - ]); - - // ep-3 has a snapshot that replaces [ep-1, ep-2, ep-3] - const snapshotEp = createDummyEpisode('snap-1', 'SYSTEM_EVENT', []); - - ep3.variants = { - snapshot: { - type: 'snapshot', - status: 'ready', - episode: snapshotEp, - replacedEpisodeIds: ['ep-1', 'ep-2', 'ep-3'], - }, - }; - - // We only retain 5 tokens, forcing the sweep to use variants for EVERYTHING except ep4. - const view = generateWorkingBufferView( - [ep1, ep2, ep3, ep4], - 5, - env.tracer, - env, - ); - - // Result should be exactly: [snapshot, ep-4] - expect(view).toHaveLength(2); - expect(view[0].id).toBe('snap-1'); - expect(view[1].id).toBe('ep-4'); - }); - - it('ignores variants that are not yet "ready"', () => { - const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: '1' }, - ]); - const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [ - { type: 'text', text: '2' }, - ]); - - ep1.variants = { - masked: { - type: 'masked', - status: 'computing', - text: '', - recoveredTokens: 10, - }, - }; - - const view = generateWorkingBufferView([ep1, ep2], 10, env.tracer, env); - - // Because the variant was computing, it must fall back to the raw pristine text. - expect(view).toHaveLength(2); - expect( - (view[0].trigger as UserPrompt).semanticParts[0].presentation, - ).toBeUndefined(); - }); -}); diff --git a/packages/core/src/context/ir/graphUtils.ts b/packages/core/src/context/ir/graphUtils.ts index c6776f2cc6..e298561224 100644 --- a/packages/core/src/context/ir/graphUtils.ts +++ b/packages/core/src/context/ir/graphUtils.ts @@ -33,183 +33,3 @@ export function isUserPrompt(node: IrNode): node is UserPrompt { return node.type === 'USER_PROMPT'; } -export function isAgentYield(node: IrNode): node is AgentYield { - return node.type === 'AGENT_YIELD'; -} - -export function isSystemEvent(node: IrNode): node is SystemEvent { - return node.type === 'SYSTEM_EVENT'; -} - -export function isSnapshot(node: IrNode): node is Snapshot { - return node.type === 'SNAPSHOT'; -} - -export function isRollingSummary(node: IrNode): node is RollingSummary { - return node.type === 'ROLLING_SUMMARY'; -} - -/** - * Generates a computed view of the pristine log. - * Sweeps backwards (newest to oldest), tracking rolling tokens. - * When rollingTokens > retainedTokens, it injects the "best" available ready variant - * (snapshot > summary > masked) instead of the raw text. - * Handles N-to-1 variant skipping automatically. - */ - -export function generateWorkingBufferView_OLD( - pristineEpisodes: Episode[], - retainedTokens: number, - tracer: ContextTracer, - env: ContextEnvironment, -): Episode[] { - const currentEpisodes: Episode[] = []; - let rollingTokens = 0; - const skippedIds = new Set(); - tracer.logEvent('ViewGenerator', 'Generating Working Buffer View'); - - for (let i = pristineEpisodes.length - 1; i >= 0; i--) { - const ep = pristineEpisodes[i]; - - // If this episode was already replaced by an N-to-1 Snapshot injected earlier in the sweep, skip it entirely! - if (skippedIds.has(ep.id)) { - tracer.logEvent( - 'ViewGenerator', - `Skipping episode [${ep.id}] due to N-to-1 replacement.`, - ); - continue; - } - - let projectedTrigger: typeof ep.trigger; - - if (isUserPrompt(ep.trigger)) { - projectedTrigger = { - ...ep.trigger, - metadata: { - ...ep.trigger.metadata, - transformations: [...(ep.trigger.metadata?.transformations || [])], - }, - semanticParts: ep.trigger.semanticParts.map((sp) => ({ ...sp })), - }; - } else { - projectedTrigger = { - ...ep.trigger, - metadata: { - ...ep.trigger.metadata, - transformations: [...(ep.trigger.metadata?.transformations || [])], - }, - }; - } - - let projectedEp: Episode = { - ...ep, - type: 'EPISODE', - trigger: projectedTrigger, - steps: ep.steps.map((step) => ({ - ...step, - metadata: { - ...step.metadata, - transformations: [...(step.metadata?.transformations || [])], - }, - })), - yield: ep.yield - ? { - ...ep.yield, - metadata: { - ...ep.yield.metadata, - transformations: [...(ep.yield.metadata?.transformations || [])], - }, - } - : undefined, - }; - - const epTokens = env.tokenCalculator.calculateEpisodeListTokens([ - projectedEp, - ]); - - if (rollingTokens > retainedTokens && ep.variants) { - const snapshot = ep.variants['snapshot']; - const summary = ep.variants['summary']; - const masked = ep.variants['masked']; - - if ( - snapshot && - snapshot.status === 'ready' && - snapshot.type === 'snapshot' - ) { - projectedEp = snapshot.episode; - // Mark all the episodes this snapshot covers to be skipped by the backwards sweep. - for (const id of snapshot.replacedEpisodeIds) { - skippedIds.add(id); - } - tracer.logEvent( - 'ViewGenerator', - `Episode [${ep.id}] has SnapshotVariant. Selecting variant over raw text. Added [${snapshot.replacedEpisodeIds.join(',')}] to skippedIds.`, - ); - debugLogger.log( - `Opportunistically swapped Episodes [${snapshot.replacedEpisodeIds.join(', ')}] for pre-computed Snapshot variant.`, - ); - } else if ( - summary && - summary.status === 'ready' && - summary.type === 'summary' - ) { - projectedEp.steps = [ - { - id: ep.id + '-summary', - type: 'AGENT_THOUGHT', - text: summary.text, - metadata: { - originalTokens: epTokens, - currentTokens: summary.recoveredTokens || 50, - transformations: [ - { - processorName: 'AsyncSemanticCompressor', - action: 'SUMMARIZED', - timestamp: Date.now(), - }, - ], - }, - }, - ] as typeof projectedEp.steps; - projectedEp.yield = undefined; - tracer.logEvent( - 'ViewGenerator', - `Episode [${ep.id}] has SummaryVariant. Selecting variant over raw text.`, - ); - debugLogger.log( - `Opportunistically swapped Episode ${ep.id} for pre-computed Summary variant.`, - ); - } else if ( - masked && - masked.status === 'ready' && - masked.type === 'masked' - ) { - if ( - isUserPrompt(projectedEp.trigger) && - projectedEp.trigger.semanticParts && - projectedEp.trigger.semanticParts.length > 0 - ) { - projectedEp.trigger.semanticParts[0].presentation = { - text: masked.text, - tokens: masked.recoveredTokens || 10, - }; - } - tracer.logEvent( - 'ViewGenerator', - `Episode [${ep.id}] has MaskedVariant. Selecting variant over raw text.`, - ); - debugLogger.log( - `Opportunistically swapped Episode ${ep.id} for pre-computed Masked variant.`, - ); - } - } - - currentEpisodes.unshift(projectedEp); - rollingTokens += env.tokenCalculator.calculateEpisodeListTokens([ - projectedEp, - ]); - } - - return currentEpisodes; -} diff --git a/packages/core/src/context/ir/mapper.test.ts b/packages/core/src/context/ir/mapper.test.ts deleted file mode 100644 index d74befd62b..0000000000 --- a/packages/core/src/context/ir/mapper.test.ts +++ /dev/null @@ -1,271 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { describe, it, expect } from 'vitest'; -import { IrMapper } from './mapper.js'; -import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; -import type { Content } from '@google/genai'; -import type { UserPrompt, ToolExecution, AgentThought } from './types.js'; - -describe('IrMapper', () => { - it('should correctly map a complex conversation into Episodes and back', () => { - const rawHistory: Content[] = [ - { role: 'user', parts: [{ text: 'Can you read file A and B?' }] }, - { - role: 'model', - parts: [ - { text: 'Let me check those files.' }, - { - functionCall: { - id: 'call_1', - name: 'read_file', - args: { filepath: 'A.txt' }, - }, - }, - { - functionCall: { - id: 'call_2', - name: 'read_file', - args: { filepath: 'B.txt' }, - }, - }, - ], - }, - { - role: 'user', - parts: [ - { - functionResponse: { - id: 'call_1', - name: 'read_file', - response: { output: 'Contents of A' }, - }, - }, - { - functionResponse: { - id: 'call_2', - name: 'read_file', - response: { output: 'Contents of B' }, - }, - }, - ], - }, - { - role: 'model', - parts: [ - { text: 'Thanks. Now I will compile.' }, - { - functionCall: { - id: 'call_3', - name: 'shell', - args: { cmd: 'make' }, - }, - }, - ], - }, - { - role: 'user', - parts: [ - { - functionResponse: { - id: 'call_3', - name: 'shell', - response: { output: 'success' }, - }, - }, - ], - }, - { role: 'model', parts: [{ text: 'Everything is done!' }] }, - ]; - - const tokenCalculator = new ContextTokenCalculator(4); - const episodes = IrMapper.toIr(rawHistory, tokenCalculator); - - expect(episodes).toHaveLength(1); - const ep = episodes[0]; - - expect(ep.trigger.type).toBe('USER_PROMPT'); - expect( - ((ep.trigger as UserPrompt).semanticParts[0] as { text: string }).text, - ).toBe('Can you read file A and B?'); - - // Steps should be: Thought, ToolExecution(A), ToolExecution(B), Thought, ToolExecution(make) - expect(ep.steps).toHaveLength(5); - expect(ep.steps[0].type).toBe('AGENT_THOUGHT'); - expect(ep.steps[1].type).toBe('TOOL_EXECUTION'); - expect((ep.steps[1] as ToolExecution).toolName).toBe('read_file'); - expect((ep.steps[1] as ToolExecution).intent).toEqual({ - filepath: 'A.txt', - }); - expect((ep.steps[1] as ToolExecution).observation).toEqual({ - output: 'Contents of A', - }); - - expect(ep.steps[2].type).toBe('TOOL_EXECUTION'); - expect((ep.steps[2] as ToolExecution).intent).toEqual({ - filepath: 'B.txt', - }); - - expect(ep.steps[3].type).toBe('AGENT_THOUGHT'); - - expect(ep.steps[4].type).toBe('TOOL_EXECUTION'); - expect((ep.steps[4] as ToolExecution).toolName).toBe('shell'); - - expect(ep.yield?.type).toBe('AGENT_YIELD'); - expect(ep.yield?.text).toBe('Everything is done!'); - - // Test Re-serialization - const reconstituted = IrMapper.fromIr(episodes); - - // Compare basic structure (the reconstituted version might have slightly different grouping of calls/responses - // based on flush logic, but semantically equivalent) - expect(reconstituted[0]).toEqual(rawHistory[0]); - // Reconstituted history is identical except tool IDs will be reassigned because IrMapper discards string IDs in favor of deterministic object hash IDs - expect(reconstituted[1].parts![0]).toEqual(rawHistory[1].parts![0]); - - // The exact structural equivalence isn't mathematically perfect because Gemini allows mixing text and calls - // in one Content block, but the flat representation is semantically identical. - }); - - it('should correctly handle multi-tool-calls grouped within a single turn without dropping observations', () => { - const rawHistory: Content[] = [ - { - role: 'user', - parts: [{ text: 'Examine both of these tools please.' }], - }, - { - role: 'model', - parts: [ - { text: 'I will call them concurrently.' }, - { - functionCall: { - id: 'c1', - name: 'tool_one', - args: { p: 1 }, - }, - }, - { - functionCall: { - id: 'c2', - name: 'tool_two', - args: { p: 2 }, - }, - }, - ], - }, - // Gemini forces the user turn to contain ALL function responses for that model turn - { - role: 'user', - parts: [ - { - functionResponse: { - id: 'c1', - name: 'tool_one', - response: { r: 1 }, - }, - }, - { - functionResponse: { - id: 'c2', - name: 'tool_two', - response: { r: 2 }, - }, - }, - ], - }, - { - role: 'model', - parts: [{ text: 'Both complete.' }], - }, - ]; - - const tokenCalculator = new ContextTokenCalculator(4); - const episodes = IrMapper.toIr(rawHistory, tokenCalculator); - - // It should collapse into a single episode - expect(episodes).toHaveLength(1); - const ep = episodes[0]; - - expect(ep.trigger.type).toBe('USER_PROMPT'); - - // The steps array should contain: - // 0: AgentThought ("I will call them concurrently") - // 1: ToolExecution(tool_one) - // 2: ToolExecution(tool_two) - - expect(ep.steps).toHaveLength(3); - - expect(ep.steps[0].type).toBe('AGENT_THOUGHT'); - expect((ep.steps[0] as AgentThought).text).toBe( - 'I will call them concurrently.', - ); - - expect(ep.steps[1].type).toBe('TOOL_EXECUTION'); - expect((ep.steps[1] as ToolExecution).toolName).toBe('tool_one'); - expect((ep.steps[1] as ToolExecution).intent).toEqual({ p: 1 }); - expect((ep.steps[1] as ToolExecution).observation).toEqual({ r: 1 }); - - expect(ep.steps[2].type).toBe('TOOL_EXECUTION'); - expect((ep.steps[2] as ToolExecution).toolName).toBe('tool_two'); - expect((ep.steps[2] as ToolExecution).intent).toEqual({ p: 2 }); - expect((ep.steps[2] as ToolExecution).observation).toEqual({ r: 2 }); - - // The final model turn should become the yield - expect(ep.yield).toBeDefined(); - expect(ep.yield?.type).toBe('AGENT_YIELD'); - expect(ep.yield?.text).toBe('Both complete.'); - - // Now verify we can reconstitute it without dropping the multiple calls - const reconstituted = IrMapper.fromIr(episodes); - - // The reconstituted history should have exactly 4 turns, same as original - expect(reconstituted).toHaveLength(4); - - // Check that the Model turn has both function calls - expect(reconstituted[1].role).toBe('model'); - expect(reconstituted[1].parts).toHaveLength(3); // text + call1 + call2 - expect(reconstituted[1].parts![1].functionCall?.name).toBe('tool_one'); - expect(reconstituted[1].parts![2].functionCall?.name).toBe('tool_two'); - - // Check that the User turn has both function responses - expect(reconstituted[2].role).toBe('user'); - expect(reconstituted[2].parts).toHaveLength(2); // response1 + response2 - expect(reconstituted[2].parts![0].functionResponse?.name).toBe('tool_one'); - expect(reconstituted[2].parts![1].functionResponse?.name).toBe('tool_two'); - }); - - it('should guarantee WeakMap ID stability across continuous mapping', () => { - // 1. Initial history - const history: Content[] = [ - { role: 'user', parts: [{ text: 'Hello' }] }, - { role: 'model', parts: [{ text: 'Hi there' }] }, - ]; - - const tokenCalculator = new ContextTokenCalculator(4); - const initialIr = IrMapper.toIr(history, tokenCalculator); - expect(initialIr).toHaveLength(1); - - // Save the uniquely generated deterministic ID for the first episode - const episodeId = initialIr[0].id; - const triggerId = initialIr[0].trigger.id; - - // 2. Push new history (simulating a continuing conversation) - history.push({ role: 'user', parts: [{ text: 'How are you?' }] }); - history.push({ role: 'model', parts: [{ text: 'I am an AI.' }] }); - - const updatedIr = IrMapper.toIr(history, tokenCalculator); - expect(updatedIr).toHaveLength(2); - - // 3. Verify ID Stability - // The exact same ID must be generated for the first episode because the underlying Content object reference hasn't changed. - // This proves the WeakMap successfully pinned the reference! - expect(updatedIr[0].id).toBe(episodeId); - expect(updatedIr[0].trigger.id).toBe(triggerId); - - // Ensure the new episode has a different ID - expect(updatedIr[1].id).not.toBe(episodeId); - }); -}); diff --git a/packages/core/src/context/ir/toIr.ts b/packages/core/src/context/ir/toIr.ts index 472a6bf8b0..1298fdbdd8 100644 --- a/packages/core/src/context/ir/toIr.ts +++ b/packages/core/src/context/ir/toIr.ts @@ -39,8 +39,8 @@ function isCompleteEpisode(ep: Partial): ep is Episode { return ( typeof ep.id === 'string' && typeof ep.timestamp === 'number' && - Array.isArray(ep.children) && - ep.children.length > 0 + Array.isArray(ep.concreteNodeIds) && + ep.concreteNodeIds.length > 0 ); } @@ -120,14 +120,7 @@ function parseToolResponses( currentEpisode = { id: getStableId(msg), timestamp: Date.now(), - children: [{ - id: getStableId(msg.parts![0] || msg), - type: 'SYSTEM_EVENT', - name: 'history_resume', - payload: {}, - metadata: createMetadata([]), - }, - ], + concreteNodeIds: [getStableId(msg.parts![0] || msg)], }; } @@ -161,11 +154,11 @@ function parseToolResponses( transformations: [], }, }; - currentEpisode.children!.push(step); + currentEpisode.concreteNodeIds = [...(currentEpisode.concreteNodeIds || []), step.id]; if (callId) pendingCallParts.delete(callId); } } - return currentEpisode; + return currentEpisode as Partial; } function parseUserParts( @@ -200,11 +193,9 @@ function parseUserParts( }; return { - type: 'EPISODE', id: getStableId(msg), timestamp: Date.now(), - children: [trigger], - }; + concreteNodeIds: [trigger.id], }; } function parseModelParts( @@ -217,14 +208,7 @@ function parseModelParts( currentEpisode = { id: getStableId(msg), timestamp: Date.now(), - children: [{ - id: getStableId(msg.parts![0] || msg), - type: 'SYSTEM_EVENT', - name: 'model_init', - payload: {}, - metadata: createMetadata([]), - }, - ], + concreteNodeIds: [getStableId(msg.parts![0] || msg)], }; } @@ -239,24 +223,25 @@ function parseModelParts( text: part.text, metadata: createMetadata([part]), }; - currentEpisode.children!.push(thought); + currentEpisode.concreteNodeIds = [...(currentEpisode.concreteNodeIds || []), thought.id]; } } - return currentEpisode; + return currentEpisode as Partial; } function finalizeYield(currentEpisode: Partial) { - if (currentEpisode.children && currentEpisode.children.length > 0) { - const lastStep = currentEpisode.children[currentEpisode.children.length - 1]; - if (isAgentThought(lastStep)) { + if (currentEpisode.concreteNodeIds && currentEpisode.concreteNodeIds.length > 0) { const yieldNode: AgentYield = { - id: lastStep.id, + id: randomUUID(), type: 'AGENT_YIELD', - text: lastStep.text, - metadata: lastStep.metadata, + text: 'Yield', // Synthesized yield since we don't have the original concrete node + metadata: { + originalTokens: 1, + currentTokens: 1, + transformations: [] + }, }; - currentEpisode.children.pop(); - currentEpisode.children.push(yieldNode); - } + const existingNodes = currentEpisode.concreteNodeIds as string[]; + currentEpisode.concreteNodeIds = [...existingNodes.slice(0, -1), yieldNode.id]; } } diff --git a/packages/core/src/context/ir/types.ts b/packages/core/src/context/ir/types.ts index f635c9b96d..a9493d646c 100644 --- a/packages/core/src/context/ir/types.ts +++ b/packages/core/src/context/ir/types.ts @@ -189,7 +189,7 @@ export interface Episode extends IrNode { readonly type: 'EPISODE'; readonly timestamp: number; /** References to the Concrete Node IDs that conceptually belong to this Episode. */ - readonly concreteNodeIds: ReadonlyArray; + concreteNodeIds: ReadonlyArray; } export interface Task extends IrNode { diff --git a/packages/core/src/context/oldContextManager.txt b/packages/core/src/context/oldContextManager.txt new file mode 100644 index 0000000000..aa3bf26736 --- /dev/null +++ b/packages/core/src/context/oldContextManager.txt @@ -0,0 +1,183 @@ +export class ContextManager { + // The stateful, pristine Episodic Intermediate Representation graph. + // This allows the agent to remember and summarize continuously without losing data across turns. + private pristineEpisodes: Episode[] = []; + private readonly eventBus: ContextEventBus; + + // Internal sub-components + // Synchronous processors are instantiated but effectively used as singletons within this class + private orchestrator: PipelineOrchestrator; + private historyObserver?: HistoryObserver; + + static create( + sidecar: SidecarConfig, + env: ContextEnvironment, + tracer: ContextTracer, + orchestrator?: PipelineOrchestrator, + registry?: ProcessorRegistry, + ): ContextManager { + if (!registry) { + registry = new ProcessorRegistry(); + registerBuiltInProcessors(registry); + } + const orch = + orchestrator || + new PipelineOrchestrator(sidecar, env, env.eventBus, tracer, registry); + return new ContextManager(sidecar, env, tracer, orch); + } + + // Use ContextManager.create() instead + private constructor( + private sidecar: SidecarConfig, + private env: ContextEnvironment, + private readonly tracer: ContextTracer, + orchestrator: PipelineOrchestrator, + ) { + this.eventBus = env.eventBus; + this.orchestrator = orchestrator; + + this.eventBus.onPristineHistoryUpdated((event) => { + this.pristineEpisodes = event.episodes; + this.evaluateTriggers(event.newNodes); + }); + + this.eventBus.onVariantReady((event) => { + // Find the target episode in the pristine graph + const targetEp = this.pristineEpisodes.find( + (ep) => ep.id === event.targetId, + ); + if (targetEp) { + if (!targetEp.variants) { + targetEp.variants = {}; + } + targetEp.variants[event.variantId] = event.variant; + this.tracer.logEvent( + 'ContextManager', + `Received async variant [${event.variantId}] for Episode ${event.targetId}`, + ); + debugLogger.log( + `ContextManager: Received async variant [${event.variantId}] for Episode ${event.targetId}.`, + ); + } + }); + } + + /** + * Safely stops background workers and clears event listeners. + */ + shutdown() { + this.orchestrator.shutdown(); + if (this.historyObserver) { + this.historyObserver.stop(); + } + } + + /** + * Evaluates if the current working buffer exceeds configured budget thresholds, + * firing consolidation events if necessary. + */ + private evaluateTriggers(newNodes: Set) { + if (!this.sidecar.budget) return; + + const workingBuffer = this.getWorkingBufferView(); + const currentTokens = + this.env.tokenCalculator.calculateEpisodeListTokens(workingBuffer); + + this.tracer.logEvent('ContextManager', 'Evaluated triggers', { + currentTokens, + retainedTokens: this.sidecar.budget.retainedTokens, + }); + + // 1. Eager Compute Trigger (on_turn) + if (newNodes.size > 0) { + this.eventBus.emitChunkReceived({ episodes: this.pristineEpisodes, targetNodeIds: newNodes }); + } + + // 2. Budget Crossed Trigger + if (currentTokens > this.sidecar.budget.retainedTokens) { + const deficit = currentTokens - this.sidecar.budget.retainedTokens; + + // Calculate exactly which nodes aged out of the retainedTokens budget to form our target delta + const agedOutNodes = new Set(); + let rollingTokens = 0; + // Start from newest and count backwards + for (let i = workingBuffer.length - 1; i >= 0; i--) { + const ep = workingBuffer[i]; + const epTokens = this.env.tokenCalculator.calculateEpisodeListTokens([ep]); + rollingTokens += epTokens; + if (rollingTokens > this.sidecar.budget.retainedTokens) { + agedOutNodes.add(ep.id); + agedOutNodes.add(ep.trigger.id); + for (const step of ep.steps) agedOutNodes.add(step.id); + if (ep.yield) agedOutNodes.add(ep.yield.id); + } + } + + this.tracer.logEvent( + 'ContextManager', + 'Budget crossed. Emitting ConsolidationNeeded', + { deficit, agedOutCount: agedOutNodes.size }, + ); + this.eventBus.emitConsolidationNeeded({ + episodes: workingBuffer, + targetDeficit: deficit, + targetNodeIds: agedOutNodes, + }); + } + } + + /** + * Subscribes to the core AgentChatHistory to natively track all message events, + * converting them seamlessly into pristine Episodes. + */ + subscribeToHistory(chatHistory: AgentChatHistory) { + if (this.historyObserver) { + this.historyObserver.stop(); + } + + this.historyObserver = new HistoryObserver( + chatHistory, + this.eventBus, + this.tracer, + this.env.tokenCalculator, + ); + this.historyObserver.start(); + } + + /** + * Generates a computed view of the pristine log. + * Sweeps backwards (newest to oldest), tracking rolling tokens. + * When rollingTokens > retainedTokens, it injects the "best" available ready variant + * (snapshot > summary > masked) instead of the raw text. + * Handles N-to-1 variant skipping automatically. + */ + getWorkingBufferView(): Episode[] { + return generateWorkingBufferView( + this.pristineEpisodes, + this.sidecar.budget.retainedTokens, + this.tracer, + this.env, + ); + } + + /** + * Returns a temporary, compressed Content[] array to be used exclusively for the LLM request. + * This does NOT mutate the pristine episodic graph. + */ + async projectCompressedHistory(): Promise { + this.tracer.logEvent('ContextManager', 'Projection requested.'); + const protectedIds = new Set(); + if (this.pristineEpisodes.length > 0) { + protectedIds.add(this.pristineEpisodes[0].id); // Structural invariant + } + + return IrProjector.project( + this.getWorkingBufferView(), + this.orchestrator, + this.sidecar, + this.tracer, + this.env, + protectedIds, + ); + } +} diff --git a/packages/core/src/context/pipeline.ts b/packages/core/src/context/pipeline.ts index fefb347335..e952d4e708 100644 --- a/packages/core/src/context/pipeline.ts +++ b/packages/core/src/context/pipeline.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import type { ConcreteNode, Snapshot, RollingSummary, IrMetadata } from './ir/types.js'; +import type { ConcreteNode, Snapshot, RollingSummary } from './ir/types.js'; export type InboxMessage = | { type: 'SNAPSHOT_READY'; snapshot: Snapshot; abstractsIds: string[] } @@ -17,13 +17,13 @@ export interface ContextInbox { export interface ContextWorkingBuffer { /** The current active (projected) flat list of ConcreteNodes. */ - readonly nodes: ReadonlyArray; + readonly nodes: readonly ConcreteNode[]; /** 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; + getLineage(id: string): readonly ConcreteNode[]; } /** @@ -54,35 +54,19 @@ export interface ContextAccountingState { * A declarative instruction from a processor on how to modify the Ship. * Applied sequentially by the Orchestrator (Reducer). */ -export interface ContextPatch { - /** The IDs of the Concrete Nodes to remove from the Ship. */ - readonly removedIds: ReadonlyArray; - - /** - * The new synthetic Concrete Nodes (e.g., MaskedTool, Snapshot) to insert. - * If omitted or empty, this patch acts as a pure deletion. - */ - readonly insertedNodes?: ReadonlyArray; - - /** The index at which to insert the new nodes. If omitted, they replace the first removedId. */ - readonly insertionIndex?: number; - - /** Audit metadata explaining who made this patch, when, and why. */ - readonly metadata: IrMetadata; -} - export interface ProcessArgs { /** The rich buffer containing current nodes and their history. */ readonly buffer: ContextWorkingBuffer; - /** The flat, sequential array of current renderable nodes (The Ship). */ - readonly ship: ReadonlyArray; + /** The full, read-only view of the ship. */ + readonly ship: readonly ConcreteNode[]; /** - * The specific subset of Concrete Node IDs that triggered this execution. - * For 'new_message', these are the new nodes. For 'retained_exceeded', the aged-out nodes. + * The unprotected, mutable subset of nodes targeted by this trigger. + * 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 triggerTargets: ReadonlySet; + readonly targets: ConcreteNode[]; /** The token budget and accounting state. */ readonly state: ContextAccountingState; @@ -93,7 +77,7 @@ export interface ProcessArgs { /** * Interface for all context degradation strategies. - * Processors are pure functions that return ContextPatches. + * Processors are pure functional map/filter operations over the targets. */ export interface ContextProcessor { /** Unique ID for registry mapping. */ @@ -101,8 +85,13 @@ export interface ContextProcessor { /** Unique name for telemetry and logging. */ readonly name: string; - /** Returns an array of declarative patches applied at this specific temporal phase. */ - process(args: ProcessArgs): Promise; + /** + * 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. + */ + process(args: ProcessArgs): Promise; } /** diff --git a/packages/core/src/context/processors/blobDegradationProcessor.test.ts b/packages/core/src/context/processors/blobDegradationProcessor.test.ts index 82cc182d16..4ca45c931b 100644 --- a/packages/core/src/context/processors/blobDegradationProcessor.test.ts +++ b/packages/core/src/context/processors/blobDegradationProcessor.test.ts @@ -1,97 +1,128 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - +import { describe, it, expect, vi } from 'vitest'; +import { BlobDegradationProcessor } from './blobDegradationProcessor.js'; import { createMockEnvironment, createDummyState, - createDummyEpisode, + createDummyNode, } from '../testing/contextTestUtils.js'; -import { describe, it, expect, beforeEach, vi } from 'vitest'; -import { BlobDegradationProcessor } from './blobDegradationProcessor.js'; -import { EpisodeEditor } from '../ir/episodeEditor.js'; -import type { UserPrompt } from '../ir/types.js'; -import type { ContextEnvironment } from '../sidecar/environment.js'; -import type { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; +import type { UserPrompt, SemanticPart } from '../ir/types.js'; describe('BlobDegradationProcessor', () => { - let processor: BlobDegradationProcessor; - let env: ContextEnvironment; - let fileSystem: InMemoryFileSystem; + it('should ignore text parts and only target inline_data and file_data', async () => { + const env = createMockEnvironment(); + // Simulate each part costing 100 tokens, but text costing 10 tokens + env.tokenCalculator.estimateTokensForParts = vi.fn((parts: any[]) => { + if (parts[0].text) return 10; + return 100; + }); - beforeEach(() => { - vi.resetAllMocks(); - env = createMockEnvironment(); - fileSystem = env.fileSystem as InMemoryFileSystem; - processor = new BlobDegradationProcessor(env); + const processor = BlobDegradationProcessor.create(env, {}); + + // Deficit of 50 means budget is NOT satisfied. + const state = createDummyState(false, 50); + + const parts: SemanticPart[] = [ + { type: 'text', text: 'Hello' }, + { type: 'inline_data', mimeType: 'image/png', data: 'fake_base64_data' }, + { type: 'text', text: 'World' }, + ]; + + const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, { + semanticParts: parts + }) as UserPrompt; + + const targets = [prompt]; + + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets, + state, + inbox: {} as any, + }); + + // We modified it, so the ID should change and it should have new metadata + expect(result.length).toBe(1); + const modifiedPrompt = result[0] as UserPrompt; + + // Original prompt ID was randomUUID(), new one is from idGenerator + 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]); + + // The inline_data part should be replaced with text + const degradedPart = modifiedPrompt.semanticParts[1]; + expect(degradedPart.type).toBe('text'); + expect((degradedPart as any).text).toContain('[Multi-Modal Blob (image/png, 0.00MB) degraded to text'); + + // The transformation should be logged + expect(modifiedPrompt.metadata.transformations.length).toBe(1); + expect(modifiedPrompt.metadata.transformations[0].action).toBe('DEGRADED'); }); - it('degrades inline_data into a text reference and saves to disk', async () => { - const dummyImageBase64 = Buffer.from('fake-image-data').toString('base64'); + it('should stop degrading once the deficit is cleared', async () => { + const env = createMockEnvironment(); + // Huge deficit requires one degradation + const state = createDummyState(false, 90); - const ep = createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: 'Look at this image:' }, - { - type: 'inline_data', - mimeType: 'image/png', - data: dummyImageBase64, - }, - ]); + env.tokenCalculator.estimateTokensForParts = vi.fn((parts: any[]) => { + if (parts[0].text) return 10; + return 100; // saving 90 tokens per degradation + }); - const state = createDummyState(false, 500); - const editor = new EpisodeEditor([ep]); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); + const processor = BlobDegradationProcessor.create(env, {}); - const parts = (result[0].trigger as UserPrompt).semanticParts; + 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; - // Text part should be untouched - expect(parts[0].presentation).toBeUndefined(); + const targets = [prompt]; - // Inline data should be degraded - expect(parts[1].presentation).toBeDefined(); - expect(parts[1].presentation!.text).toContain( - '[Multi-Modal Blob (image/png', - ); - expect(parts[1].presentation!.text).toContain( - 'degraded to text to preserve context window', - ); + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets, + state, + inbox: {} as any, + }); - // Verify it was written to fake FS - expect(fileSystem.getFiles().size).toBeGreaterThan(0); - const files = Array.from(fileSystem.getFiles().keys()); - expect(files[0]).toContain( - '.gemini/tool-outputs/degraded-blobs/session-mock-session/blob_', - ); - - expect(result[0].trigger.metadata.transformations.length).toBe(1); + const modifiedPrompt = result[0] as UserPrompt; + expect(modifiedPrompt.semanticParts.length).toBe(2); + + // First part was degraded + expect(modifiedPrompt.semanticParts[0].type).toBe('text'); + // Second part was untouched because deficit hit 0 + expect(modifiedPrompt.semanticParts[1].type).toBe('file_data'); }); - it('degrades file_data into a text reference without disk write', async () => { - const ep = createDummyEpisode('ep-2', 'USER_PROMPT', [ - { - type: 'file_data', - mimeType: 'application/pdf', - fileUri: 'gs://fake-bucket/doc.pdf', - }, - ]); + it('should return exactly the targets array if budget is already satisfied', async () => { + const env = createMockEnvironment(); + const state = createDummyState(true, 0); // Budget satisfied! - const state = createDummyState(false, 500); - const editor = new EpisodeEditor([ep]); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); + const processor = BlobDegradationProcessor.create(env, {}); + const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, { + semanticParts: [ + { type: 'inline_data', mimeType: 'image/png', data: 'abc' }, + ] + }) as UserPrompt; - const parts = (result[0].trigger as UserPrompt).semanticParts; - expect(parts[0].presentation).toBeDefined(); - expect(parts[0].presentation!.text).toContain( - '[File Reference (application/pdf)', - ); - expect(parts[0].presentation!.text).toContain( - 'Original URI: gs://fake-bucket/doc.pdf', - ); + const targets = [prompt]; - expect(fileSystem.getFiles().size).toBe(0); + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets, + state, + inbox: {} as any, + }); + + // Should return the exact array ref + expect(result).toBe(targets); }); }); diff --git a/packages/core/src/context/processors/blobDegradationProcessor.ts b/packages/core/src/context/processors/blobDegradationProcessor.ts index 25a7246468..566048678d 100644 --- a/packages/core/src/context/processors/blobDegradationProcessor.ts +++ b/packages/core/src/context/processors/blobDegradationProcessor.ts @@ -1,14 +1,7 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ProcessArgs, ContextProcessor } from '../pipeline.js'; +import type { ConcreteNode, UserPrompt } from '../ir/types.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import { sanitizeFilenamePart } from '../../utils/fileUtils.js'; -import type { EpisodeEditor } from '../ir/episodeEditor.js'; -import { isUserPrompt } from '../ir/graphUtils.js'; export type BlobDegradationProcessorOptions = Record; @@ -29,12 +22,9 @@ export class BlobDegradationProcessor implements ContextProcessor { this.env = env; } - async process( - editor: EpisodeEditor, - state: ContextAccountingState, - ): Promise { + async process({ targets, state }: ProcessArgs): Promise> { if (state.isBudgetSatisfied) { - return; + return targets; } let currentDeficit = state.deficitTokens; @@ -59,88 +49,100 @@ export class BlobDegradationProcessor implements ContextProcessor { } }; + const returnedNodes: ConcreteNode[] = []; + // Forward scan, looking for bloated non-text parts to degrade - for (const target of editor.targets) { - const ep = target.episode; - if (target.node !== ep.trigger) continue; - - if (currentDeficit <= 0) break; - if (state.protectedEpisodeIds.has(ep.id)) continue; + for (const node of targets) { + if (currentDeficit <= 0 || node.type !== 'USER_PROMPT') { + returnedNodes.push(node); + continue; + } - if (isUserPrompt(ep.trigger)) { - for (let j = 0; j < ep.trigger.semanticParts.length; j++) { - const part = ep.trigger.semanticParts[j]; - if (currentDeficit <= 0) break; - // We only target non-text parts that haven't already been masked - if (part.type === 'text' || part.presentation) continue; + const prompt = node as UserPrompt; + let modified = false; + const newParts = [...prompt.semanticParts]; - let newText = ''; - let tokensSaved = 0; + for (let j = 0; j < prompt.semanticParts.length; j++) { + const part = prompt.semanticParts[j]; + if (currentDeficit <= 0) break; + // We only target non-text parts that haven't already been masked + if (part.type === 'text') continue; - if (part.type === '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); + let newText = ''; + let tokensSaved = 0; - // Base64 to buffer - const buffer = Buffer.from(part.data, 'base64'); - await this.env.fileSystem.writeFile(filePath, buffer); + if (part.type === '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 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}]`; + // Base64 to buffer + const buffer = Buffer.from(part.data, 'base64'); + await this.env.fileSystem.writeFile(filePath, buffer); - // Re-calculate tokens. Images are expensive (~258 tokens). The text is cheap (~20 tokens). - const oldTokens = this.env.tokenCalculator.estimateTokensForParts([ - { inlineData: { mimeType: part.mimeType, data: part.data } }, - ]); - const newTokens = this.env.tokenCalculator.estimateTokensForParts([ - { text: newText }, - ]); - tokensSaved = oldTokens - newTokens; - } else if (part.type === '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 }, - ]); - tokensSaved = oldTokens - newTokens; - } else if (part.type === '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 }, - ]); - tokensSaved = oldTokens - newTokens; - } + 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}]`; - if (newText && tokensSaved > 0) { - 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; + } else if (part.type === '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 }, + ]); + tokensSaved = oldTokens - newTokens; + } else if (part.type === '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 }, + ]); + tokensSaved = oldTokens - newTokens; + } - editor.editEpisode(ep.id, 'DEGRADE_BLOB', (draft) => { - if (isUserPrompt(draft.trigger)) { - draft.trigger.semanticParts[j].presentation = { - text: newText, - tokens: newTokens, - }; - draft.trigger.metadata.transformations.push({ - processorName: this.name, - action: 'DEGRADED', - timestamp: Date.now(), - }); - } - }); - - currentDeficit -= tokensSaved; - } + if (newText && tokensSaved > 0) { + // Replace the part with a synthetic text part + newParts[j] = { type: 'text', text: newText }; + currentDeficit -= tokensSaved; + modified = true; } } + + if (modified) { + // Return a fresh synthetic node representing the degraded state + const degradedNode: UserPrompt = { + ...prompt, + id: this.env.idGenerator.generateId(), // Issue a new ID because it was modified + semanticParts: newParts, + metadata: { + ...prompt.metadata, + transformations: [ + ...prompt.metadata.transformations, + { + processorName: this.name, + action: 'DEGRADED', + timestamp: Date.now(), + } + ] + } + }; + returnedNodes.push(degradedNode); + } else { + returnedNodes.push(node); + } } + + return returnedNodes; } } diff --git a/packages/core/src/context/processors/emergencyTruncationProcessor.test.ts b/packages/core/src/context/processors/emergencyTruncationProcessor.test.ts deleted file mode 100644 index 892e0d53ce..0000000000 --- a/packages/core/src/context/processors/emergencyTruncationProcessor.test.ts +++ /dev/null @@ -1,129 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createMockEnvironment, - createDummyState, - createDummyEpisode, -} from '../testing/contextTestUtils.js'; -import { describe, it, expect, beforeEach, vi } from 'vitest'; -import { EmergencyTruncationProcessor } from './emergencyTruncationProcessor.js'; -import { EpisodeEditor } from '../ir/episodeEditor.js'; -import type { ContextEnvironment } from '../sidecar/environment.js'; - -describe('EmergencyTruncationProcessor', () => { - let processor: EmergencyTruncationProcessor; - let env: ContextEnvironment; - - beforeEach(() => { - vi.resetAllMocks(); - env = createMockEnvironment(); - // Force token calculator to return exactly what we tell it for deterministic testing - vi.spyOn( - env.tokenCalculator, - 'calculateEpisodeListTokens', - ).mockImplementation((episodes) => - // Just sum up the metadata originalTokens for our dummy episodes - episodes.reduce( - (acc, ep) => acc + (ep.trigger.metadata.originalTokens || 100), - 0, - ), - ); - - processor = new EmergencyTruncationProcessor(env, {}); - }); - - it('bypasses processing if currentTokens <= maxTokens', async () => { - const episodes = [ - createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: 'short' }, - ]), - ]; - // State says we are under budget (5000 < 10000) - const state = createDummyState(true, 0, new Set(), 5000, 10000); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - expect(result).toStrictEqual(episodes); - expect(result.length).toBe(1); - }); - - it('truncates episodes from the front (oldest) until targetTokens is met', async () => { - const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: 'oldest' }, - ]); - const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [ - { type: 'text', text: 'middle' }, - ]); - const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [ - { type: 'text', text: 'newest' }, - ]); - - // Each is worth 100 tokens according to our mock - const episodes = [ep1, ep2, ep3]; - - // We have 300 tokens, but max is 200. We need to drop 100 tokens. - const state = createDummyState(false, 100, new Set(), 300, 200); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - - // It should drop the FIRST episode (ep-1) and keep the rest. - expect(result.length).toBe(2); - expect(result[0].id).toBe('ep-2'); - expect(result[1].id).toBe('ep-3'); - }); - - it('never drops protected episodes (e.g. system instructions)', async () => { - const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: 'protected system prompt' }, - ]); - const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', [ - { type: 'text', text: 'middle' }, - ]); - const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', [ - { type: 'text', text: 'newest' }, - ]); - - const episodes = [ep1, ep2, ep3]; - - // We have 300 tokens, max is 200. We need to drop 100 tokens. - // However, ep-1 is protected! - const state = createDummyState(false, 100, new Set(['ep-1']), 300, 200); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - - // It should SKIP dropping ep-1 (protected) and drop ep-2 instead. - expect(result.length).toBe(2); - expect(result[0].id).toBe('ep-1'); // Protected, survived - expect(result[1].id).toBe('ep-3'); // Survivor - }); - - it('can drop multiple episodes if deficit is huge', async () => { - const ep1 = createDummyEpisode('ep-1', 'USER_PROMPT', []); - const ep2 = createDummyEpisode('ep-2', 'USER_PROMPT', []); - const ep3 = createDummyEpisode('ep-3', 'USER_PROMPT', []); - - const episodes = [ep1, ep2, ep3]; - - // We have 300 tokens, max is 50. We need to drop 250 tokens! - const state = createDummyState(false, 250, new Set(), 300, 50); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - - // It must drop ep1 (100t) and ep2 (100t). - // Remaining is ep3 (100t). - // Wait, if it drops ep1 (remaining=200) and ep2 (remaining=100), - // when it looks at ep3, remaining (100) > max (50), so it will drop ep3 too! - expect(result.length).toBe(0); - }); -}); diff --git a/packages/core/src/context/processors/emergencyTruncationProcessor.ts b/packages/core/src/context/processors/emergencyTruncationProcessor.ts index 1ef4f55d59..cf4e38b8db 100644 --- a/packages/core/src/context/processors/emergencyTruncationProcessor.ts +++ b/packages/core/src/context/processors/emergencyTruncationProcessor.ts @@ -8,8 +8,8 @@ import type { ContextProcessor, BackstopTargetOptions, ProcessArgs, - ContextPatch, } from '../pipeline.js'; +import type { ConcreteNode } from '../ir/types.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; export type EmergencyTruncationProcessorOptions = BackstopTargetOptions; @@ -41,63 +41,45 @@ export class EmergencyTruncationProcessor implements ContextProcessor { readonly name = 'EmergencyTruncationProcessor'; readonly options: EmergencyTruncationProcessorOptions; constructor( - private readonly _env: ContextEnvironment, + _env: ContextEnvironment, options: EmergencyTruncationProcessorOptions, ) { this.options = options; } async process({ - ship, - triggerTargets, + targets, state, - }: ProcessArgs): Promise { - const toRemove: string[] = []; - + }: 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') { - if (state.currentTokens <= state.maxTokens) return []; + if (state.currentTokens <= state.maxTokens) return targets; targetTokensToRemove = state.currentTokens - state.maxTokens; } else if (strategy === 'freeNTokens') { targetTokensToRemove = this.options.freeTokensTarget ?? 0; - if (targetTokensToRemove <= 0) return []; + if (targetTokensToRemove <= 0) return targets; } else if (strategy === 'max') { // 'max' means we remove all targets without stopping early targetTokensToRemove = Infinity; } let removedTokens = 0; + const keptNodes: ConcreteNode[] = []; - // The ship is sequentially ordered from oldest to newest. - // We want to delete the oldest targeted nodes first. - for (const node of ship) { - // Is this node part of the targeted delta (e.g. aged out of the budget)? - if (!triggerTargets.has(node.id)) continue; - // Is this node explicitly protected (e.g. part of an active Task)? - if (node.logicalParentId && state.protectedLogicalIds.has(node.logicalParentId)) continue; - - if (removedTokens >= targetTokensToRemove) break; + // The targets are sequentially ordered from oldest to newest. + // We want to delete the oldest targets first. + for (const node of targets) { + if (removedTokens >= targetTokensToRemove) { + keptNodes.push(node); + continue; + } removedTokens += node.metadata.currentTokens; - toRemove.push(node.id); } - if (toRemove.length === 0) return []; - - return [{ - removedIds: toRemove, - metadata: { - originalTokens: removedTokens, - currentTokens: 0, - transformations: [{ - processorName: this.name, - action: 'TRUNCATED', - timestamp: Date.now(), - }], - } - }]; + return keptNodes; } } diff --git a/packages/core/src/context/processors/historySquashingProcessor.test.ts b/packages/core/src/context/processors/historySquashingProcessor.test.ts index bdbae5a159..85ceac7e88 100644 --- a/packages/core/src/context/processors/historySquashingProcessor.test.ts +++ b/packages/core/src/context/processors/historySquashingProcessor.test.ts @@ -1,157 +1,133 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - +import { describe, it, expect, vi } from 'vitest'; +import { HistorySquashingProcessor } from './historySquashingProcessor.js'; import { createMockEnvironment, createDummyState, - createDummyEpisode, + createDummyNode, } from '../testing/contextTestUtils.js'; -import { describe, it, expect, beforeEach } from 'vitest'; -import { HistorySquashingProcessor } from './historySquashingProcessor.js'; -import { EpisodeEditor } from '../ir/episodeEditor.js'; import type { UserPrompt, AgentThought, AgentYield } from '../ir/types.js'; -import { randomUUID } from 'node:crypto'; +import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; describe('HistorySquashingProcessor', () => { - let processor: HistorySquashingProcessor; + it('should truncate nodes that exceed maxTokensPerNode', async () => { + const mockTokenCalculator = new ContextTokenCalculator(1) as any; + mockTokenCalculator.tokensToChars = vi.fn().mockReturnValue(10); // Limit is 10 chars - beforeEach(() => { - processor = new HistorySquashingProcessor(createMockEnvironment(), { - maxTokensPerNode: 100, + mockTokenCalculator.estimateTokensForString = vi.fn((text: string) => { + if (text.includes('OMITTED')) return 1; // Summary size + return 100; // Original size }); + mockTokenCalculator.estimateTokensForParts = vi.fn(() => 1); + + const env = createMockEnvironment({ + tokenCalculator: mockTokenCalculator + }); + + const processor = HistorySquashingProcessor.create(env, { + maxTokensPerNode: 1, // Will equal 10 chars limit + }); + + const state = createDummyState(false, 500); // 500 token deficit + + const prompt = createDummyNode('ep1', 'USER_PROMPT', 100, { + semanticParts: [ + { type: 'text', text: 'This text is way longer than 10 characters and needs truncation' } + ], + }, 'prompt-id') as UserPrompt; + + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 100, { + text: 'The model is thinking something incredibly long and verbose that exceeds 10 chars', + }, 'thought-id') as AgentThought; + + const yieldNode = createDummyNode('ep1', 'AGENT_YIELD', 100, { + text: 'Final output yield that is also extremely long', + }, 'yield-id') as AgentYield; + + const targets = [prompt, thought, yieldNode]; + + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets, + state, + inbox: {} as any, + }); + + expect(result.length).toBe(3); + + // 1. User Prompt + const squashedPrompt = result[0] as UserPrompt; + expect(squashedPrompt.id).toBe('mock-uuid-1'); + expect(squashedPrompt.id).not.toBe(prompt.id); + expect(squashedPrompt.semanticParts[0].type).toBe('text'); + expect((squashedPrompt.semanticParts[0] as any).text).toContain('[... OMITTED'); + expect(squashedPrompt.metadata.transformations.length).toBe(1); + expect(squashedPrompt.metadata.transformations[0].action).toBe('TRUNCATED'); + + // 2. Agent Thought + const squashedThought = result[1] as AgentThought; + expect(squashedThought.id).toBe('mock-uuid-2'); + expect(squashedThought.id).not.toBe(thought.id); + expect(squashedThought.text).toContain('[... OMITTED'); + + // 3. Agent Yield + const squashedYield = result[2] as AgentYield; + expect(squashedYield.id).toBe('mock-uuid-3'); + expect(squashedYield.id).not.toBe(yieldNode.id); + expect(squashedYield.text).toContain('[... OMITTED'); }); - const createThoughtEpisode = ( - id: string, - userText: string, - modelThought: string, - ) => { - const ep = createDummyEpisode(id, 'USER_PROMPT', [ - { type: 'text', text: userText }, - ]); - // Replace the tool steps with a thought step for this test - ep.steps = [ - { - id: randomUUID(), - type: 'AGENT_THOUGHT', - text: modelThought, - metadata: { - originalTokens: 1000, - currentTokens: 1000, - transformations: [], - }, - }, - ]; - return ep; - }; + it('should stop truncating once the deficit is cleared', async () => { + const mockTokenCalculator = new ContextTokenCalculator(1) as any; + mockTokenCalculator.tokensToChars = vi.fn().mockReturnValue(10); + mockTokenCalculator.estimateTokensForString = vi.fn((text: string) => { + if (text.includes('OMITTED')) return 0; // Huge savings + return 500; + }); + mockTokenCalculator.estimateTokensForParts = vi.fn(() => 0); - it('bypasses processing if budget is satisfied', async () => { - const episodes = [createThoughtEpisode('1', 'short text', 'short thought')]; - const state = createDummyState(true); + const env = createMockEnvironment({ + tokenCalculator: mockTokenCalculator + }); - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); + const processor = HistorySquashingProcessor.create(env, { + maxTokensPerNode: 1, + }); - expect(result).toStrictEqual(episodes); - expect( - (result[0].trigger as UserPrompt).semanticParts[0].presentation, - ).toBeUndefined(); - }); + // Deficit is only 10 tokens. First truncation saves 500. + const state = createDummyState(false, 10); - it('skips protected episodes', async () => { - // 500 chars = ~125 tokens. Limit is 100 tokens, so it WOULD truncate if not protected. - const longText = 'A'.repeat(500); - const episodes = [createThoughtEpisode('ep-1', longText, 'short thought')]; - const state = createDummyState(false, 100, new Set(['ep-1'])); + const prompt = createDummyNode('ep1', 'USER_PROMPT', 500, { + semanticParts: [ + { type: 'text', text: 'This text is way longer than 10 characters and needs truncation' } + ], + }, 'prompt-id') as UserPrompt; - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 500, { + text: 'The model is thinking something incredibly long and verbose that exceeds 10 chars', + }, 'thought-id') as AgentThought; - expect( - (result[0].trigger as UserPrompt).semanticParts[0].presentation, - ).toBeUndefined(); - }); + const targets = [prompt, thought]; - it('truncates both UserPrompts and AgentThoughts', async () => { - const longUser = 'U'.repeat(1000); // ~250 tokens - const longModel = 'M'.repeat(1000); // ~250 tokens - const episodes = [createThoughtEpisode('ep-2', longUser, longModel)]; - const state = createDummyState(false, 500); // High deficit, force truncation + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets, + state, + inbox: {} as any, + }); - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); + expect(result.length).toBe(2); - const userPart = (result[0].trigger as UserPrompt).semanticParts[0]; - const thoughtPart = result[0].steps[0] as AgentThought; + // 1. User Prompt (truncated because deficit > 0) + const squashedPrompt = result[0] as UserPrompt; + expect(squashedPrompt.id).toBe('mock-uuid-1'); + expect(squashedPrompt.id).not.toBe(prompt.id); + expect((squashedPrompt.semanticParts[0] as any).text).toContain('[... OMITTED'); - expect(userPart.presentation).toBeDefined(); - expect(userPart.presentation!.text).toContain( - '[... OMITTED 600 chars ...]', - ); - - expect(thoughtPart.presentation).toBeDefined(); - expect(thoughtPart.presentation!.text).toContain( - '[... OMITTED 600 chars ...]', - ); - - // Check audit trails - expect(result[0].trigger.metadata.transformations.length).toBe(1); - expect(thoughtPart.metadata.transformations.length).toBe(1); - }); - - it('stops processing once deficit is resolved', async () => { - const longUser1 = 'A'.repeat(1000); - const longUser2 = 'B'.repeat(1000); - const episodes = [ - createThoughtEpisode('ep-3', longUser1, 'short'), - createThoughtEpisode('ep-4', longUser2, 'short'), - ]; - - // Set deficit to exactly what ONE truncation will save - // Original = ~250 tokens. Limit = 100. Truncation saves ~150 tokens. - const state = createDummyState(false, 150); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - - // First episode should be truncated - const ep1Part = (result[0].trigger as UserPrompt).semanticParts[0]; - expect(ep1Part.presentation).toBeDefined(); - - // Second episode should be untouched because the deficit hit 0 - const ep2Part = (result[1].trigger as UserPrompt).semanticParts[0]; - expect(ep2Part.presentation).toBeUndefined(); - }); - - it('truncates IrNodes', async () => { - const longYield = 'Y'.repeat(1000); // ~250 tokens - const ep = createThoughtEpisode('ep-5', 'short', 'short'); - ep.yield = { - id: randomUUID(), - type: 'AGENT_YIELD', - text: longYield, - metadata: { - originalTokens: 250, - currentTokens: 250, - transformations: [], - }, - }; - - const state = createDummyState(false, 500); - const editor = new EpisodeEditor([ep]); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - - const yieldPart = result[0].yield as AgentYield; - const yieldPresentation = yieldPart.presentation as { text: string }; - expect(yieldPresentation).toBeDefined(); - expect(yieldPresentation.text).toContain('[... OMITTED 600 chars ...]'); + // 2. Agent Thought (untouched because deficit is now < 0) + const untouchedThought = result[1] as AgentThought; + expect(untouchedThought.id).toBe(thought.id); + expect(untouchedThought.text).not.toContain('[... OMITTED'); }); }); diff --git a/packages/core/src/context/processors/historySquashingProcessor.ts b/packages/core/src/context/processors/historySquashingProcessor.ts index f7bcc3a572..9d1624f0b0 100644 --- a/packages/core/src/context/processors/historySquashingProcessor.ts +++ b/packages/core/src/context/processors/historySquashingProcessor.ts @@ -1,14 +1,7 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextProcessor, ProcessArgs } from '../pipeline.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import { truncateProportionally } from '../truncation.js'; -import type { EpisodeEditor } from '../ir/episodeEditor.js'; -import { isAgentThought, isUserPrompt } from '../ir/graphUtils.js'; +import type { ConcreteNode, UserPrompt, AgentThought, AgentYield } from '../ir/types.js'; export interface HistorySquashingProcessorOptions { maxTokensPerNode: number; @@ -27,8 +20,7 @@ export class HistorySquashingProcessor 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'], @@ -37,11 +29,13 @@ export class HistorySquashingProcessor implements ContextProcessor { readonly id = 'HistorySquashingProcessor'; readonly name = 'HistorySquashingProcessor'; readonly options: HistorySquashingProcessorOptions; + private env: ContextEnvironment; constructor( env: ContextEnvironment, options: HistorySquashingProcessorOptions, ) { + this.env = env; this.options = options; } @@ -49,12 +43,10 @@ export class HistorySquashingProcessor implements ContextProcessor { text: string, limitChars: number, currentDeficit: number, - setPresentation: (p: { text: string; tokens: number }) => void, - recordAudit: () => void, - ): number { - if (currentDeficit <= 0) return 0; + ): { text: string; newTokens: number; oldTokens: number; tokensSaved: number } | null { + if (currentDeficit <= 0) return null; const originalLength = text.length; - if (originalLength <= limitChars) return 0; + if (originalLength <= limitChars) return null; const newText = truncateProportionally( text, @@ -63,127 +55,130 @@ export class HistorySquashingProcessor implements ContextProcessor { ); if (newText !== text) { - const newTokens = Math.floor(newText.length / 4); - const oldTokens = Math.floor(originalLength / 4); + // Using accurate TokenCalculator instead of simple math + const newTokens = this.env.tokenCalculator.estimateTokensForString(newText); + const oldTokens = this.env.tokenCalculator.estimateTokensForString(text); const tokensSaved = oldTokens - newTokens; - setPresentation({ text: newText, tokens: newTokens }); - recordAudit(); - return tokensSaved; + if (tokensSaved > 0) { + return { text: newText, newTokens, oldTokens, tokensSaved }; + } } - return 0; + return null; } - async process( - editor: EpisodeEditor, - state: ContextAccountingState, - ): Promise { + async process({ targets, state }: ProcessArgs): Promise> { if (state.isBudgetSatisfied) { - return; + return targets; } const { maxTokensPerNode } = this.options; - // We estimate 4 chars per token for truncation logic - const limitChars = maxTokensPerNode * 4; + const limitChars = this.env.tokenCalculator.tokensToChars(maxTokensPerNode); - // We track how many tokens we still need to cut. If we hit 0, we can stop early! let currentDeficit = state.deficitTokens; + const returnedNodes: ConcreteNode[] = []; - for (const target of editor.targets) { - const ep = target.episode; - if (currentDeficit <= 0) break; - if (state.protectedEpisodeIds.has(ep.id)) continue; + for (const node of targets) { + if (currentDeficit <= 0) { + returnedNodes.push(node); + continue; + } // 1. Squash User Prompts - if (target.node === ep.trigger && isUserPrompt(ep.trigger)) { - for (let j = 0; j < ep.trigger.semanticParts.length; j++) { - const part = ep.trigger.semanticParts[j]; + if (node.type === 'USER_PROMPT') { + const prompt = node as UserPrompt; + let modified = false; + const newParts = [...prompt.semanticParts]; + + for (let j = 0; j < prompt.semanticParts.length; j++) { + const part = prompt.semanticParts[j]; + if (currentDeficit <= 0) break; if (part.type === 'text') { - const saved = this.tryApplySquash( - part.text, - limitChars, - currentDeficit, - (p) => { - editor.editEpisode(ep.id, 'SQUASH_PROMPT', (draft) => { - if (isUserPrompt(draft.trigger)) { - draft.trigger.semanticParts[j].presentation = p; - } - }); - }, - () => { - editor.editEpisode(ep.id, 'SQUASH_PROMPT', (draft) => { - draft.trigger.metadata.transformations.push({ - processorName: this.name, - action: 'TRUNCATED', - timestamp: Date.now(), - }); - }); - }, - ); - currentDeficit -= saved; + const squashResult = this.tryApplySquash(part.text, limitChars, currentDeficit); + if (squashResult) { + newParts[j] = { type: 'text', text: squashResult.text }; + currentDeficit -= squashResult.tokensSaved; + modified = true; + } } } + + if (modified) { + const newTokens = this.env.tokenCalculator.estimateTokensForParts(newParts as any); + returnedNodes.push({ + ...prompt, + id: this.env.idGenerator.generateId(), + semanticParts: newParts, + metadata: { + ...prompt.metadata, + currentTokens: newTokens, + transformations: [ + ...prompt.metadata.transformations, + { processorName: this.name, action: 'TRUNCATED', timestamp: Date.now() } + ] + } + }); + } else { + returnedNodes.push(node); + } + continue; } // 2. Squash Model Thoughts - if (isAgentThought(target.node)) { - const step = target.node; - const j = ep.steps.findIndex(s => s.id === step.id); - if (j !== -1 && currentDeficit > 0) { - const saved = this.tryApplySquash( - step.text, - limitChars, - currentDeficit, - (p) => { - editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => { - const draftStep = draft.steps[j]; - if (isAgentThought(draftStep)) { - draftStep.presentation = p; - } - }); - }, - () => { - editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => { - const draftStep = draft.steps[j]; - if (isAgentThought(draftStep)) { - draftStep.metadata.transformations.push({ - processorName: this.name, - action: 'TRUNCATED', - timestamp: Date.now(), - }); - } - }); - }, - ); - currentDeficit -= saved; + if (node.type === 'AGENT_THOUGHT') { + const thought = node as AgentThought; + const squashResult = this.tryApplySquash(thought.text, limitChars, currentDeficit); + + if (squashResult) { + currentDeficit -= squashResult.tokensSaved; + returnedNodes.push({ + ...thought, + id: this.env.idGenerator.generateId(), + text: squashResult.text, + metadata: { + ...thought.metadata, + currentTokens: squashResult.newTokens, + transformations: [ + ...thought.metadata.transformations, + { processorName: this.name, action: 'TRUNCATED', timestamp: Date.now() } + ] + } + }); + } else { + returnedNodes.push(node); } + continue; } // 3. Squash Agent Yields - if (currentDeficit > 0 && target.node === ep.yield && ep.yield) { - const saved = this.tryApplySquash( - ep.yield.text, - limitChars, - currentDeficit, - (p) => { - editor.editEpisode(ep.id, 'SQUASH_YIELD', (draft) => { - if (draft.yield) draft.yield.presentation = p; - }); - }, - () => { - editor.editEpisode(ep.id, 'SQUASH_YIELD', (draft) => { - if (draft.yield) { - draft.yield.metadata.transformations.push({ - processorName: this.name, - action: 'TRUNCATED', - timestamp: Date.now(), - }); - } - }); - }, - ); - currentDeficit -= saved; + if (node.type === 'AGENT_YIELD') { + const agentYield = node as AgentYield; + const squashResult = this.tryApplySquash(agentYield.text, limitChars, currentDeficit); + + if (squashResult) { + currentDeficit -= squashResult.tokensSaved; + returnedNodes.push({ + ...agentYield, + id: this.env.idGenerator.generateId(), + text: squashResult.text, + metadata: { + ...agentYield.metadata, + currentTokens: squashResult.newTokens, + transformations: [ + ...agentYield.metadata.transformations, + { processorName: this.name, action: 'TRUNCATED', timestamp: Date.now() } + ] + } + }); + } else { + returnedNodes.push(node); + } + continue; } + + returnedNodes.push(node); } + + return returnedNodes; } } diff --git a/packages/core/src/context/processors/semanticCompressionProcessor.test.ts b/packages/core/src/context/processors/semanticCompressionProcessor.test.ts index a3f3c96143..45c52d0a04 100644 --- a/packages/core/src/context/processors/semanticCompressionProcessor.test.ts +++ b/packages/core/src/context/processors/semanticCompressionProcessor.test.ts @@ -1,164 +1,157 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - +import { describe, it, expect, vi } from 'vitest'; +import { SemanticCompressionProcessor } from './semanticCompressionProcessor.js'; import { createMockEnvironment, createDummyState, - createDummyEpisode, + createDummyNode, + createDummyToolNode, } from '../testing/contextTestUtils.js'; -import { describe, it, expect, beforeEach, vi } from 'vitest'; -import { SemanticCompressionProcessor } from './semanticCompressionProcessor.js'; -import { EpisodeEditor } from '../ir/episodeEditor.js'; -import type { UserPrompt, ToolExecution, AgentThought } from '../ir/types.js'; -import { randomUUID } from 'node:crypto'; -import type { BaseLlmClient } from 'src/core/baseLlmClient.js'; +import type { UserPrompt, AgentThought, ToolExecution } from '../ir/types.js'; +import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; describe('SemanticCompressionProcessor', () => { - let processor: SemanticCompressionProcessor; - let generateContentMock: ReturnType; + it('should trigger summarization via LLM for long text parts', async () => { + const mockLlmClient = { + generateContent: vi.fn().mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'Mocked Summary!' }] } }], + }), + }; - beforeEach(() => { - generateContentMock = vi.fn().mockResolvedValue({ - candidates: [{ content: { parts: [{ text: 'Mocked Summary!' }] } }], + const mockTokenCalculator = new ContextTokenCalculator(1) as any; + mockTokenCalculator.tokensToChars = vi.fn().mockReturnValue(10); + mockTokenCalculator.estimateTokensForParts = vi.fn((parts: any) => { + if (parts[0]?.text === 'Mocked Summary!') return 5; + if (parts[0]?.functionResponse?.response?.summary === 'Mocked Summary!') return 10; + return 5000; }); - const env = createMockEnvironment(); - // Re-mock llmClient properly - vi.spyOn(env, 'llmClient', 'get').mockReturnValue({ - generateContent: generateContentMock, - } as unknown as BaseLlmClient); - - processor = new SemanticCompressionProcessor(env, { - nodeThresholdTokens: 2000, + const env = createMockEnvironment({ + llmClient: mockLlmClient as any, + tokenCalculator: mockTokenCalculator }); + + const processor = SemanticCompressionProcessor.create(env, { + nodeThresholdTokens: 10, + }); + + const state = createDummyState(false, 15000); // We need to save tons of tokens + + const prompt = createDummyNode('ep1', 'USER_PROMPT', 3800, { + semanticParts: [ + { type: 'text', text: 'This text is way longer than 10 characters and needs compression' } + ], + }, 'prompt-id') as UserPrompt; + + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 1500, { + text: 'The model is thinking something incredibly long and verbose that exceeds 10 chars', + metadata: { currentTokens: 5000, originalTokens: 5000, transformations: [] } + }, 'thought-id') as AgentThought; + + const tool = createDummyToolNode('ep1', 50, 1000, { + observation: { result: 'Massive tool JSON observation payload' }, + tokens: { intent: 50, observation: 1000 } + }, 'tool-id'); + + const targets = [prompt, thought, tool]; + + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets, + state, + inbox: {} as any, + }); + + expect(result.length).toBe(3); + + // 1. User Prompt + const compressedPrompt = result[0] as UserPrompt; + expect(compressedPrompt.id).toBe('mock-uuid-1'); + expect(compressedPrompt.id).not.toBe(prompt.id); + expect(compressedPrompt.semanticParts[0].type).toBe('text'); + expect((compressedPrompt.semanticParts[0] as any).text).toBe('Mocked Summary!'); + expect(compressedPrompt.metadata.transformations.length).toBe(1); + expect(compressedPrompt.metadata.transformations[0].action).toBe('SUMMARIZED'); + + // 2. Agent Thought + const compressedThought = result[1] as AgentThought; + expect(compressedThought.id).toBe('mock-uuid-2'); + expect(compressedThought.id).not.toBe(thought.id); + expect(compressedThought.text).toBe('Mocked Summary!'); + expect(compressedThought.metadata.transformations.length).toBe(1); + + // 3. Tool Execution + const compressedTool = result[2] as ToolExecution; + expect(compressedTool.id).toBe('mock-uuid-3'); + expect(compressedTool.id).not.toBe(tool.id); + expect(compressedTool.observation).toEqual({ summary: 'Mocked Summary!' }); + expect(compressedTool.metadata.transformations.length).toBe(1); + + // Verify LLM was called 3 times + expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(3); }); - const createEpisodeWithThoughtsAndTools = ( - id: string, - userText: string, - thoughtText: string, - toolObs: string, - ) => { - const ep = createDummyEpisode(id, 'USER_PROMPT', [ - { type: 'text', text: userText }, - ]); - // We override metadata for threshold triggering - ep.trigger.metadata.currentTokens = 3800; + it('should stop summarizing once the deficit is cleared', async () => { + const mockLlmClient = { + generateContent: vi.fn().mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'Mocked Summary!' }] } }], + }), + }; - ep.steps = [ - { - id: randomUUID(), - type: 'AGENT_THOUGHT', - text: thoughtText, - metadata: { - originalTokens: 3800, - currentTokens: 3800, - transformations: [], - }, - }, - { - id: randomUUID(), - type: 'TOOL_EXECUTION', - toolName: 'test', - intent: {}, - observation: toolObs, - tokens: { intent: 10, observation: 3800 }, - metadata: { - originalTokens: 3810, - currentTokens: 3810, - transformations: [], - }, - }, - ]; - return ep; - }; + const mockTokenCalculator = new ContextTokenCalculator(1) as any; + mockTokenCalculator.tokensToChars = vi.fn().mockReturnValue(10); + // Returning 0 tokens for the summary to maximize savings + mockTokenCalculator.estimateTokensForParts = vi.fn((parts: any) => { + if (parts[0]?.text === 'Mocked Summary!') return 0; + return 5000; + }); - it('bypasses processing if budget is satisfied', async () => { - const episodes = [ - createEpisodeWithThoughtsAndTools('1', 'short', 'short', 'short'), - ]; - const state = createDummyState(true); + const env = createMockEnvironment({ + llmClient: mockLlmClient as any, + tokenCalculator: mockTokenCalculator + }); - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); + const processor = SemanticCompressionProcessor.create(env, { + nodeThresholdTokens: 10, + }); - expect(generateContentMock).not.toHaveBeenCalled(); - }); + // Deficit is only 10 tokens! The first summarization will save 5000 tokens, clearing it instantly. + const state = createDummyState(false, 10); - it('skips protected episodes even if over budget', async () => { - const massiveStr = 'M'.repeat(15000); - const episodes = [ - createEpisodeWithThoughtsAndTools( - 'ep-1', - massiveStr, - massiveStr, - massiveStr, - ), - ]; - const state = createDummyState(false, 1000, new Set(['ep-1'])); + const prompt = createDummyNode('ep1', 'USER_PROMPT', 3800, { + semanticParts: [ + { type: 'text', text: 'This text is way longer than 10 characters and needs compression' } + ], + }, 'prompt-id') as UserPrompt; - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); + const thought = createDummyNode('ep1', 'AGENT_THOUGHT', 1500, { + text: 'The model is thinking something incredibly long and verbose that exceeds 10 chars', + metadata: { currentTokens: 5000, originalTokens: 5000, transformations: [] } + }, 'thought-id') as AgentThought; - expect(generateContentMock).not.toHaveBeenCalled(); - }); + const targets = [prompt, thought]; - it('summarizes unprotected UserPrompts, Thoughts, and Tool observations until deficit is met', async () => { - const massiveStr = 'M'.repeat(15000); - const episodes = [ - createEpisodeWithThoughtsAndTools( - 'ep-1', - massiveStr, - massiveStr, - massiveStr, - ), - ]; - const state = createDummyState(false, 50000); // Massive deficit, forces all 3 to summarize + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets, + state, + inbox: {} as any, + }); - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); + expect(result.length).toBe(2); - expect(generateContentMock).toHaveBeenCalledTimes(3); + // 1. User Prompt (was summarized because deficit was > 0) + const compressedPrompt = result[0] as UserPrompt; + expect(compressedPrompt.id).toBe('mock-uuid-1'); + expect(compressedPrompt.id).not.toBe(prompt.id); - // Verify presentation layers were injected - const result = editor.getFinalEpisodes(); - const userPart = (result[0].trigger as UserPrompt).semanticParts[0]; - const thoughtPart = result[0].steps[0] as AgentThought; - const toolPart = result[0].steps[1] as ToolExecution; + // 2. Agent Thought (was NOT summarized because deficit hit 0) + const untouchedThought = result[1] as AgentThought; + expect(untouchedThought.id).toBe(thought.id); // Reference equality! + expect(untouchedThought.text).toBe(thought.text); - expect(userPart.presentation).toBeDefined(); - expect(userPart.presentation!.text).toContain('Mocked Summary!'); - - expect(thoughtPart.presentation).toBeDefined(); - expect(thoughtPart.presentation!.text).toContain('Mocked Summary!'); - - expect(toolPart.presentation).toBeDefined(); - expect( - (toolPart.presentation!.observation as Record)['summary'], - ).toContain('Mocked Summary!'); - }); - - it('stops calling LLM when deficit hits zero', async () => { - const massiveStr = 'M'.repeat(15000); - const episodes = [ - createEpisodeWithThoughtsAndTools( - 'ep-1', - massiveStr, - massiveStr, - massiveStr, - ), - ]; - - // Set deficit low enough that ONE summary solves the problem - const state = createDummyState(false, 5); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - - // It should only compress the UserPrompt and then stop - expect(generateContentMock).toHaveBeenCalledTimes(1); + // LLM should only have been called once + expect(mockLlmClient.generateContent).toHaveBeenCalledTimes(1); }); }); diff --git a/packages/core/src/context/processors/semanticCompressionProcessor.ts b/packages/core/src/context/processors/semanticCompressionProcessor.ts index 5e2eb19863..fa38c26bf2 100644 --- a/packages/core/src/context/processors/semanticCompressionProcessor.ts +++ b/packages/core/src/context/processors/semanticCompressionProcessor.ts @@ -1,16 +1,9 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextProcessor, ProcessArgs } from '../pipeline.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import { debugLogger } from '../../utils/debugLogger.js'; import { LlmRole } from '../../telemetry/types.js'; import { getResponseText } from '../../utils/partUtils.js'; -import type { EpisodeEditor } from '../ir/episodeEditor.js'; -import { isAgentThought, isToolExecution, isUserPrompt } from '../ir/graphUtils.js'; +import type { ConcreteNode, UserPrompt, AgentThought, ToolExecution } from '../ir/types.js'; export interface SemanticCompressionProcessorOptions { nodeThresholdTokens: number; @@ -39,7 +32,7 @@ export class SemanticCompressionProcessor implements ContextProcessor { readonly name = 'SemanticCompressionProcessor'; readonly options: SemanticCompressionProcessorOptions; private env: ContextEnvironment; - private modelToUse: string = 'chat-compression-2.5-flash-lite'; + private modelToUse: string = 'gemini-2.5-flash'; constructor( env: ContextEnvironment, @@ -49,219 +42,194 @@ export class SemanticCompressionProcessor implements ContextProcessor { this.options = options; } - async process( - editor: EpisodeEditor, - state: ContextAccountingState, - ): Promise { - // If the budget is satisfied, or semantic compression isn't enabled + private async generateSummary( + text: string, + contextInfo: string, + ): Promise { + try { + const response = await this.env.llmClient.generateContent( + { + contents: [ + { + role: 'user', + parts: [{ text }], + }, + ], + 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.`, + }, + ], + }, + }, + this.modelToUse + ); + return getResponseText(response) || text; + } catch (e) { + debugLogger.warn(`SemanticCompressionProcessor failed to summarize ${contextInfo}`, e); + return text; // Fallback to original text on API failure + } + } + + async process({ targets, state }: ProcessArgs): Promise> { if (state.isBudgetSatisfied) { - return; + return targets; } const semanticConfig = this.options; const limitTokens = semanticConfig.nodeThresholdTokens; const thresholdChars = this.env.tokenCalculator.tokensToChars(limitTokens); - this.modelToUse = 'gemini-2.5-flash'; - + let currentDeficit = state.deficitTokens; + const returnedNodes: ConcreteNode[] = []; - // We scan backwards (oldest to newest would also work, but older is safer to degrade first) - for (const target of editor.targets) { - const ep = target.episode; - if (currentDeficit <= 0) break; - if (state.protectedEpisodeIds.has(ep.id)) continue; + // Scan backwards (oldest to newest would also work, but older is safer to degrade first) + for (const node of targets) { + if (currentDeficit <= 0) { + returnedNodes.push(node); + continue; + } // 1. Compress User Prompts - if (target.node === ep.trigger && isUserPrompt(ep.trigger)) { - for (let j = 0; j < ep.trigger.semanticParts.length; j++) { - const part = ep.trigger.semanticParts[j]; + if (node.type === 'USER_PROMPT') { + const prompt = node as UserPrompt; + let modified = false; + const newParts = [...prompt.semanticParts]; + + for (let j = 0; j < prompt.semanticParts.length; j++) { + const part = prompt.semanticParts[j]; if (currentDeficit <= 0) break; if (part.type !== 'text') continue; - // If it's already got a presentation, we don't want to re-summarize a summary - if (part.presentation) 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) { - editor.editEpisode(ep.id, 'SUMMARIZE_PROMPT', (draft) => { - if (isUserPrompt(draft.trigger)) { - draft.trigger.semanticParts[j].presentation = { - text: summary, - tokens: newTokens, - }; - draft.trigger.metadata.transformations.push({ - processorName: this.name, - action: 'SUMMARIZED', - timestamp: Date.now(), - }); - } - }); - currentDeficit -= oldTokens - newTokens; + newParts[j] = { type: 'text', text: summary }; + currentDeficit -= (oldTokens - newTokens); + modified = true; } } } + + if (modified) { + const newTokens = this.env.tokenCalculator.estimateTokensForParts(newParts as any); + returnedNodes.push({ + ...prompt, + id: this.env.idGenerator.generateId(), + semanticParts: newParts, + metadata: { + ...prompt.metadata, + currentTokens: newTokens, + transformations: [ + ...prompt.metadata.transformations, + { processorName: this.name, action: 'SUMMARIZED', timestamp: Date.now() } + ] + } + }); + } else { + returnedNodes.push(node); + } + continue; } // 2. Compress Model Thoughts - if (isAgentThought(target.node)) { - const step = target.node; - const j = ep.steps.findIndex(s => s.id === step.id); - if (j !== -1 && currentDeficit > 0 && !step.presentation) { - if (step.text.length > thresholdChars) { - const summary = await this.generateSummary( - step.text, - 'Agent Thought', - ); - const newTokens = this.env.tokenCalculator.estimateTokensForParts( - [{ text: summary }], - ); - const oldTokens = this.env.tokenCalculator.estimateTokensForParts( - [{ text: step.text }], - ); + if (node.type === 'AGENT_THOUGHT') { + const thought = node as AgentThought; + if (thought.text.length > thresholdChars) { + const summary = await this.generateSummary(thought.text, 'Agent Thought'); + const newTokens = this.env.tokenCalculator.estimateTokensForParts([{ text: summary }]); + const oldTokens = thought.metadata.currentTokens; + console.log(`Agent Thought compression: newTokens=${newTokens}, oldTokens=${oldTokens}`); - if (newTokens < oldTokens) { - editor.editEpisode(ep.id, 'SUMMARIZE_THOUGHT', (draft) => { - const draftStep = draft.steps[j]; - if (isAgentThought(draftStep)) { - draftStep.presentation = { - text: summary, - tokens: newTokens, - }; - if (!draftStep.metadata) { - draftStep.metadata = { - transformations: [], - currentTokens: 0, - originalTokens: 0, - }; - } - if (!draftStep.metadata.transformations) { - draftStep.metadata.transformations = []; - } - draftStep.metadata.transformations.push({ - processorName: this.name, - action: 'SUMMARIZED', - timestamp: Date.now(), - }); - } - }); - currentDeficit -= oldTokens - newTokens; - } - } + if (newTokens < oldTokens) { + currentDeficit -= (oldTokens - newTokens); + returnedNodes.push({ + ...thought, + id: this.env.idGenerator.generateId(), + text: summary, + metadata: { + ...thought.metadata, + currentTokens: newTokens, + transformations: [ + ...thought.metadata.transformations, + { processorName: this.name, action: 'SUMMARIZED', timestamp: Date.now() } + ] + } + }); + continue; + } } + returnedNodes.push(node); + continue; } // 3. Compress Tool Observations - if (isToolExecution(target.node)) { - const step = target.node; - const j = ep.steps.findIndex(s => s.id === step.id); - if (j !== -1 && currentDeficit > 0 && !step.presentation) { - const rawObs = (step.presentation as any)?.observation ?? step.observation; + if (node.type === 'TOOL_EXECUTION') { + const tool = node as ToolExecution; + const rawObs = tool.observation; - let stringifiedObs = ''; - if (typeof rawObs === 'string') { - stringifiedObs = rawObs; - } else { - try { - stringifiedObs = JSON.stringify(rawObs); - } catch { - stringifiedObs = String(rawObs); - } + let stringifiedObs = ''; + if (typeof rawObs === 'string') { + stringifiedObs = rawObs; + } else { + try { + stringifiedObs = JSON.stringify(rawObs); + } catch { + stringifiedObs = String(rawObs); } + } - if (stringifiedObs.length > thresholdChars) { - const summary = await this.generateSummary( - stringifiedObs, - step.toolName, - ); - const newObsObject = { summary }; + if (stringifiedObs.length > thresholdChars) { + const summary = await this.generateSummary(stringifiedObs, tool.toolName || 'unknown'); + const newObsObject = { summary }; - const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([ - { - functionResponse: { - name: step.toolName, - response: newObsObject, - id: step.id, - }, + const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([ + { + functionResponse: { + name: tool.toolName || 'unknown', + response: newObsObject, + id: tool.id, }, - ]); + }, + ]); - const oldObsTokens = - (step.presentation as any)?.tokens?.observation ?? - step.tokens?.observation ?? step.tokens; - const intentTokens = - (step.presentation as any)?.tokens?.intent ?? step.tokens?.intent ?? 0; + const oldObsTokens = tool.tokens?.observation ?? tool.metadata.currentTokens; + const intentTokens = tool.tokens?.intent ?? 0; - if (newObsTokens < oldObsTokens) { - editor.editEpisode(ep.id, 'SUMMARIZE_TOOL', (draft) => { - const draftStep = draft.steps[j]; - if (isToolExecution(draftStep)) { - draftStep.presentation = { - intent: - draftStep.presentation?.intent ?? draftStep.intent, - observation: newObsObject, - tokens: { - intent: intentTokens, - observation: newObsTokens, - }, - }; - if (!draftStep.metadata) { - draftStep.metadata = { - transformations: [], - currentTokens: 0, - originalTokens: 0, - }; - } - if (!draftStep.metadata.transformations) { - draftStep.metadata.transformations = []; - } - draftStep.metadata.transformations.push({ - processorName: this.name, - action: 'SUMMARIZED', - timestamp: Date.now(), - }); - } - }); - currentDeficit -= oldObsTokens - newObsTokens; - } + if (newObsTokens < oldObsTokens) { + currentDeficit -= (oldObsTokens - newObsTokens); + returnedNodes.push({ + ...tool, + id: this.env.idGenerator.generateId(), + observation: newObsObject as Record, + tokens: { + intent: intentTokens, + observation: newObsTokens, + }, + metadata: { + ...tool.metadata, + currentTokens: intentTokens + newObsTokens, + transformations: [ + ...tool.metadata.transformations, + { processorName: this.name, action: 'SUMMARIZED', timestamp: Date.now() } + ] + } + }); + continue; } - } + } + returnedNodes.push(node); + continue; } - } - } - private async generateSummary( - content: string, - contentType: string, - abortSignal?: AbortSignal, - ): Promise { - const promptMessage = `You are compressing an old episodic context buffer for an AI assistant.\nSummarize this ${contentType} block in 2-3 highly technical sentences. Keep all critical facts, file names, dependencies, and architectural decisions. Discard conversational filler and boilerplate.\n\nContent:\n${content.slice(0, 30000)}`; - - const client = this.env.llmClient; - try { - const response = await client.generateContent({ - modelConfigKey: { model: this.modelToUse }, - contents: [{ role: 'user', parts: [{ text: promptMessage }] }], - promptId: 'local-context-compression-summary', - role: LlmRole.UTILITY_COMPRESSOR, - abortSignal: abortSignal ?? new AbortController().signal, - }); - const text = getResponseText(response) ?? ''; - return `[Semantic Summary of old ${contentType}]\n${text.trim()}`; - } catch (e) { - debugLogger.warn(`Semantic compression LLM call failed: ${e}`); - // If we fail to summarize, we just return the original truncated by 50% as a fail-safe, or the original. - // Returning original is safer to prevent data loss on API failure. - return content; + returnedNodes.push(node); } + + return returnedNodes; } } diff --git a/packages/core/src/context/processors/stateSnapshotProcessor.test.ts b/packages/core/src/context/processors/stateSnapshotProcessor.test.ts deleted file mode 100644 index 477cae8e68..0000000000 --- a/packages/core/src/context/processors/stateSnapshotProcessor.test.ts +++ /dev/null @@ -1,118 +0,0 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - createMockEnvironment, - createDummyState, - createDummyEpisode, -} from '../testing/contextTestUtils.js'; -import { describe, it, expect, beforeEach, vi } from 'vitest'; -import { StateSnapshotProcessor } from './stateSnapshotProcessor.js'; -import { EpisodeEditor } from '../ir/episodeEditor.js'; -import type { ContextEnvironment } from '../sidecar/environment.js'; -import type { BaseLlmClient } from '../../core/baseLlmClient.js'; - -describe('StateSnapshotProcessor', () => { - let processor: StateSnapshotProcessor; - let env: ContextEnvironment; - let generateContentMock: ReturnType; - - beforeEach(() => { - vi.resetAllMocks(); - env = createMockEnvironment(); - - generateContentMock = vi.fn().mockResolvedValue({ - text: 'Mocked Compressed State Snapshot!', - }); - vi.spyOn(env, 'llmClient', 'get').mockReturnValue({ - generateContent: generateContentMock, - } as unknown as BaseLlmClient); - - // Override token calc for testing - vi.spyOn(env.tokenCalculator, 'estimateTokensForParts').mockReturnValue( - 100, - ); - - processor = new StateSnapshotProcessor(env, {}, env.eventBus); - }); - - it('bypasses processing if deficit is <= 0', async () => { - const episodes = [ - createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: 'hello' }, - ]), - ]; - // current: 100, max: 1000, retained: 200 (deficit 0) - const state = createDummyState(false, 0, new Set(), 100, 1000, 200); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - expect(result).toStrictEqual(episodes); - expect(generateContentMock).not.toHaveBeenCalled(); - }); - - it('bypasses processing if not enough episodes to summarize (needs at least 2 inner episodes)', async () => { - const episodes = [ - createDummyEpisode('ep-sys', 'SYSTEM_EVENT', []), - createDummyEpisode('ep-active', 'USER_PROMPT', [ - { type: 'text', text: 'help' }, - ]), - ]; - - // current: 1000, max: 10000, retained: 500. Target deficit = 500 - const state = createDummyState(false, 500, new Set(), 1000, 10000, 500); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - expect(result).toStrictEqual(episodes); - expect(generateContentMock).not.toHaveBeenCalled(); - }); - - it('summarizes intermediate episodes into a single snapshot episode', async () => { - const episodes = [ - createDummyEpisode('ep-0', 'SYSTEM_EVENT', []), - createDummyEpisode('ep-1', 'USER_PROMPT', [ - { type: 'text', text: 'old 1' }, - ]), - createDummyEpisode('ep-2', 'USER_PROMPT', [ - { type: 'text', text: 'old 2' }, - ]), - createDummyEpisode('ep-3', 'USER_PROMPT', [ - { type: 'text', text: 'current' }, - ]), - ]; - - // Target deficit = 200 - const state = createDummyState(false, 200, new Set(), 1000, 10000, 800); - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - - // We started with 4 episodes. - // Episodes [1, 2] were synthesized into a single new Snapshot episode. - // Final array should be: [0, SNAPSHOT, 3] = length 3. - expect(result.length).toBe(3); - expect(result[0].id).toBe('ep-0'); - - const snapshotEp = result[1]; - expect(snapshotEp.yield).toBeDefined(); - expect(snapshotEp.yield!.text).toContain(''); - expect(snapshotEp.yield!.text).toContain( - 'Mocked Compressed State Snapshot!', - ); - - expect(result[2].id).toBe('ep-3'); - - expect(generateContentMock).toHaveBeenCalledTimes(1); - - const llmArgs = generateContentMock.mock.calls[0][0]; - expect(llmArgs.contents[0].parts[0].text).toContain('old 1'); - expect(llmArgs.contents[0].parts[0].text).toContain('old 2'); - }); -}); diff --git a/packages/core/src/context/processors/toolMaskingProcessor.test.ts b/packages/core/src/context/processors/toolMaskingProcessor.test.ts index 7201e067e5..c019fae74c 100644 --- a/packages/core/src/context/processors/toolMaskingProcessor.test.ts +++ b/packages/core/src/context/processors/toolMaskingProcessor.test.ts @@ -1,135 +1,87 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { createMockEnvironment } from '../testing/contextTestUtils.js'; -import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { describe, it, expect, vi } from 'vitest'; import { ToolMaskingProcessor } from './toolMaskingProcessor.js'; -import { EpisodeEditor } from '../ir/episodeEditor.js'; -import type { Episode, ToolExecution } from '../ir/types.js'; -import type { ContextAccountingState } from '../pipeline.js'; -import { randomUUID } from 'node:crypto'; -import type { ContextEnvironment } from '../sidecar/environment.js'; -import type { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; +import { + createMockEnvironment, + createDummyState, + createDummyToolNode, +} from '../testing/contextTestUtils.js'; describe('ToolMaskingProcessor', () => { - let processor: ToolMaskingProcessor; - let env: ContextEnvironment; - let fileSystem: InMemoryFileSystem; + it('should write large strings to disk and replace them with a masked pointer', async () => { + const env = createMockEnvironment(); + // 1 token = 1 char for simplicity + env.tokenCalculator.tokensToChars = vi.fn().mockReturnValue(10); + // Fake token calculator says new tokens are 5 + env.tokenCalculator.estimateTokensForParts = vi.fn().mockReturnValue(5); - beforeEach(() => { - vi.resetAllMocks(); - env = createMockEnvironment(); - fileSystem = env.fileSystem as InMemoryFileSystem; - - processor = new ToolMaskingProcessor(env, { - stringLengthThresholdTokens: 100, + const processor = ToolMaskingProcessor.create(env, { + stringLengthThresholdTokens: 10, }); - }); - const getDummyState = ( - isSatisfied = false, - deficit = 0, - protectedIds = new Set(), - ): ContextAccountingState => ({ - currentTokens: 5000, - maxTokens: 10000, - retainedTokens: 4000, - deficitTokens: deficit, - protectedEpisodeIds: protectedIds, - isBudgetSatisfied: isSatisfied, - }); + const state = createDummyState(false, 500); - const createDummyEpisode = ( - id: string, - intent: Record, - observation: Record, - ): Episode => ({ - type: 'EPISODE', - id, - timestamp: Date.now(), - trigger: { - id: randomUUID(), - type: 'SYSTEM_EVENT', - name: 'test', - payload: {}, - metadata: { originalTokens: 10, currentTokens: 10, transformations: [] }, - }, - steps: [ - { - id: randomUUID(), - type: 'TOOL_EXECUTION', - toolName: 'test_tool', - intent, - observation, - tokens: { intent: 500, observation: 500 }, // Claim they are big enough to be masked - metadata: { - originalTokens: 1000, - currentTokens: 1000, - transformations: [], - }, + const toolStep = createDummyToolNode('ep1', 50, 100, { + observation: { + result: 'this is a really long string that should get masked out because it exceeds 10 chars', + metadata: 'short', }, - ], + metadata: { + currentTokens: 150, + originalTokens: 150, + transformations: [] + } + }); + + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets: [toolStep], + state, + inbox: {} as any, + }); + + expect(result.length).toBe(1); + const masked = result[0]; + + // 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 as any).observation; + expect(obs.result).toContain(''); + expect(obs.metadata).toBe('short'); // Untouched + + // Transformation logged + expect(masked.metadata.transformations.length).toBe(1); + expect(masked.metadata.transformations[0].action).toBe('MASKED'); }); - it('bypasses processing if budget is satisfied', async () => { - const episodes = [ - createDummyEpisode('1', { arg: 'short' }, { out: 'short' }), - ]; - const state = getDummyState(true); + it('should skip unmaskable tools', async () => { + const env = createMockEnvironment(); + env.tokenCalculator.tokensToChars = vi.fn().mockReturnValue(10); - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); + const processor = ToolMaskingProcessor.create(env, { + stringLengthThresholdTokens: 10, + }); - expect(result).toStrictEqual(episodes); - expect((result[0].steps[0] as ToolExecution).presentation).toBeUndefined(); - }); + const state = createDummyState(false, 500); - it('deep masks massive string intents and observations', async () => { - // We need strings > limitChars (100 tokens * 4 chars = 400 chars) - const massiveIntentString = 'I'.repeat(500); - const massiveObsString = 'O'.repeat(500); + 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', + } + }); - const intentPayload = { args: { nested: [massiveIntentString, 'short'] } }; - const obsPayload = { result: massiveObsString, error: null }; + const result = await processor.process({ + buffer: {} as any, + ship: [], + targets: [toolStep], + state, + inbox: {} as any, + }); - const episodes = [createDummyEpisode('ep-1', intentPayload, obsPayload)]; - const state = getDummyState(false, 1000, new Set()); // Huge deficit - - const editor = new EpisodeEditor(episodes); - await processor.process(editor, state); - const result = editor.getFinalEpisodes(); - - const toolStep = result[0].steps[0] as ToolExecution; - - expect(toolStep.presentation).toBeDefined(); - - // Check intent was deep masked - const maskedIntent = toolStep.presentation!.intent as Record< - string, - unknown - >; - expect((maskedIntent['args'] as { nested: string }).nested[0]).toContain( - '', - ); - expect((maskedIntent['args'] as { nested: string }).nested[1]).toBe( - 'short', - ); // Unchanged - - // Check observation was deep masked - const maskedObs = toolStep.presentation!.observation as Record< - string, - unknown - >; - expect((maskedObs as { result: string }).result).toContain( - '', - ); - expect((maskedObs as { error: string }).error).toBeNull(); - - // Check disk writes occurred to fake FS - expect(fileSystem.getFiles().size).toBe(2); + // Returned the exact same object reference + expect(result[0]).toBe(toolStep); }); }); diff --git a/packages/core/src/context/processors/toolMaskingProcessor.ts b/packages/core/src/context/processors/toolMaskingProcessor.ts index 30d5dd6b2a..496eb8c96b 100644 --- a/packages/core/src/context/processors/toolMaskingProcessor.ts +++ b/packages/core/src/context/processors/toolMaskingProcessor.ts @@ -1,10 +1,5 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextProcessor, ProcessArgs } from '../pipeline.js'; +import type { ConcreteNode, ToolExecution } from '../ir/types.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import { sanitizeFilenamePart } from '../../utils/fileUtils.js'; import { @@ -14,7 +9,6 @@ import { ENTER_PLAN_MODE_TOOL_NAME, EXIT_PLAN_MODE_TOOL_NAME, } from '../../tools/tool-names.js'; -import type { EpisodeEditor } from '../ir/episodeEditor.js'; const UNMASKABLE_TOOLS = new Set([ ACTIVATE_SKILL_TOOL_NAME, @@ -93,13 +87,14 @@ export class ToolMaskingProcessor implements ContextProcessor { this.options = options; } - async process( - editor: EpisodeEditor, - state: ContextAccountingState, - ): Promise { + private isAlreadyMasked(text: string): boolean { + return text.includes(''); + } + + async process({ targets, state }: ProcessArgs): Promise> { const maskingConfig = this.options; - if (!maskingConfig) return; - if (state.isBudgetSatisfied) return; + if (!maskingConfig) return targets; + if (state.isBudgetSatisfied) return targets; let currentDeficit = state.deficitTokens; const limitChars = this.env.tokenCalculator.tokensToChars( @@ -118,10 +113,8 @@ export class ToolMaskingProcessor implements ContextProcessor { ); } - // We only create the directory if we actually mask something let directoryCreated = false; - // Helper to extract string and write to disk const handleMasking = async ( content: string, toolName: string, @@ -147,173 +140,149 @@ export class ToolMaskingProcessor implements ContextProcessor { return `\n[Tool ${nodeType} string (${fileSizeMB}MB, ${totalLines} lines) masked to preserve context window. Full string saved to: ${filePath}]\n`; }; - // Forward scan, looking for massive intents or observations to mask - for (const ep of editor.getFullHistory()) { - if (currentDeficit <= 0) break; - if (!ep || !ep.steps || state.protectedEpisodeIds.has(ep.id)) continue; + const returnedNodes: ConcreteNode[] = []; - for (let j = 0; j < ep.steps.length; j++) { - if (currentDeficit <= 0) break; - const step = ep.steps[j]; - if (step.type !== 'TOOL_EXECUTION') continue; + for (const node of targets) { + if (currentDeficit <= 0 || node.type !== 'TOOL_EXECUTION') { + returnedNodes.push(node); + continue; + } - const toolName = step.toolName; - if (toolName && UNMASKABLE_TOOLS.has(toolName)) continue; + const step = node as ToolExecution; + const toolName = step.toolName; + if (toolName && UNMASKABLE_TOOLS.has(toolName)) { + returnedNodes.push(node); + continue; + } - // Ensure presentation object exists - if (!step.presentation) { - step.presentation = { - intent: step.intent, - observation: step.observation, - tokens: step.tokens, // Fallback to raw tokens initially - }; - } + const callId = step.id || Date.now().toString(); - const callId = step.id || Date.now().toString(); - - const maskAsync = async ( - obj: MaskableValue, - nodeType: string, - ): Promise<{ masked: MaskableValue; changed: boolean }> => { - if (typeof obj === 'string') { - if (obj.length > limitChars && !this.isAlreadyMasked(obj)) { - const newString = await handleMasking( - obj, - toolName, - callId, - nodeType, - ); - return { masked: newString, changed: true }; - } - return { masked: obj, changed: false }; - } - if (Array.isArray(obj)) { - let changed = false; - const masked: MaskableValue[] = []; - for (const item of obj) { - const res = await maskAsync(item, nodeType); - if (res.changed) changed = true; - masked.push(res.masked); - } - return { masked, changed }; - } - if (typeof obj === 'object' && obj !== null) { - let changed = false; - const masked: Record = {}; - for (const [key, value] of Object.entries(obj)) { - const res = await maskAsync(value, nodeType); - if (res.changed) changed = true; - masked[key] = res.masked; - } - return { masked, changed }; + const maskAsync = async ( + obj: MaskableValue, + nodeType: string, + ): Promise<{ masked: MaskableValue; changed: boolean }> => { + if (typeof obj === 'string') { + if (obj.length > limitChars && !this.isAlreadyMasked(obj)) { + const newString = await handleMasking( + obj, + toolName || 'unknown', + callId, + nodeType, + ); + return { masked: newString, changed: true }; } return { masked: obj, changed: false }; - }; - - const rawIntent = step.presentation?.intent ?? step.intent; - const rawObs = step.presentation?.observation ?? step.observation; - - if (!isMaskableRecord(rawIntent)) { - throw new Error( - `ToolMaskingProcessor: step intent is not a valid JSON record. CallID: ${callId}`, - ); } - if (!isMaskableValue(rawObs)) { - throw new Error( - `ToolMaskingProcessor: step observation is not a valid JSON value. CallID: ${callId}`, - ); - } - - const intentRes = await maskAsync(rawIntent, 'intent'); - const obsRes = await maskAsync(rawObs, 'observation'); - - if (intentRes.changed || obsRes.changed) { - const maskedIntent = isMaskableRecord(intentRes.masked) - ? (intentRes.masked as Record) - : undefined; - const maskedObs = isMaskableRecord(obsRes.masked) - ? (obsRes.masked as Record) - : undefined; - - // Recalculate tokens perfectly - const newIntentTokens = - this.env.tokenCalculator.estimateTokensForParts([ - { - functionCall: { - name: toolName, - args: maskedIntent, - id: callId, - }, - }, - ]); - const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([ - { - functionResponse: { - name: toolName, - response: - typeof obsRes.masked === 'string' - ? { message: obsRes.masked } - : maskedObs, - id: callId, - }, - }, - ]); - - const oldTotal = - step.presentation.tokens?.intent !== undefined - ? step.presentation.tokens.intent + - step.presentation.tokens.observation - : step.tokens.intent + step.tokens.observation; - - const newTotal = newIntentTokens + newObsTokens; - const savings = oldTotal - newTotal; - - if (savings > 0) { - currentDeficit -= savings; - this.env.tracer.logEvent( - 'ToolMaskingProcessor', - `Masked tool ${toolName}`, - { recoveredTokens: savings }, - ); - - editor.editEpisode(ep.id, 'MASK_TOOL', (draft) => { - const draftStep = draft.steps[j]; - if (draftStep.type !== 'TOOL_EXECUTION') return; - if (!draftStep.presentation) { - draftStep.presentation = { - intent: draftStep.intent, - observation: draftStep.observation, - tokens: draftStep.tokens, - }; - } - draftStep.presentation.intent = maskedIntent ?? {}; - draftStep.presentation.observation = - typeof obsRes.masked === 'string' - ? { message: obsRes.masked } - : (maskedObs ?? {}); - draftStep.presentation.tokens = { - intent: newIntentTokens, - observation: newObsTokens, - }; - draftStep.metadata = { - ...draftStep.metadata, - transformations: [ - ...(draftStep.metadata?.transformations || []), - { - processorName: 'ToolMasking', - action: 'MASKED', - timestamp: Date.now(), - }, - ], - }; - }); + if (Array.isArray(obj)) { + let changed = false; + const masked: MaskableValue[] = []; + for (const item of obj) { + const res = await maskAsync(item, nodeType); + if (res.changed) changed = true; + masked.push(res.masked); } + return { masked, changed }; } + if (typeof obj === 'object' && obj !== null) { + let changed = false; + const masked: Record = {}; + for (const [key, value] of Object.entries(obj)) { + const res = await maskAsync(value, nodeType); + if (res.changed) changed = true; + masked[key] = res.masked; + } + return { masked, changed }; + } + return { masked: obj, changed: false }; + }; + + const rawIntent = step.intent; + const rawObs = step.observation; + + if (!isMaskableRecord(rawIntent)) { + returnedNodes.push(node); + continue; + } + if (!isMaskableValue(rawObs)) { + returnedNodes.push(node); + continue; + } + + const intentRes = await maskAsync(rawIntent, 'intent'); + const obsRes = await maskAsync(rawObs, 'observation'); + + if (intentRes.changed || obsRes.changed) { + const maskedIntent = isMaskableRecord(intentRes.masked) + ? (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 newIntentTokens = this.env.tokenCalculator.estimateTokensForParts([ + { + functionCall: { + name: toolName || 'unknown', + args: maskedIntent, + id: callId, + }, + }, + ]); + + let obsPart: any = {}; + if (maskedObs) { + obsPart = { + functionResponse: { + name: toolName || 'unknown', + response: maskedObs, + id: callId + } + }; + } + + const newObsTokens = this.env.tokenCalculator.estimateTokensForParts([obsPart]); + + const tokensSaved = + (step.metadata.currentTokens) - + (newIntentTokens + newObsTokens); + + if (tokensSaved > 0) { + const maskedNode: ToolExecution = { + ...step, + id: this.env.idGenerator.generateId(), // Modified, so generate new ID + intent: maskedIntent ?? step.intent, + observation: maskedObs ?? step.observation, + tokens: { + intent: newIntentTokens, + observation: newObsTokens, + }, + metadata: { + ...step.metadata, + currentTokens: newIntentTokens + newObsTokens, + transformations: [ + ...step.metadata.transformations, + { + processorName: this.name, + action: 'MASKED', + timestamp: Date.now(), + } + ] + } + }; + + returnedNodes.push(maskedNode); + currentDeficit -= tokensSaved; + } else { + returnedNodes.push(node); + } + } else { + returnedNodes.push(node); } } - } - private isAlreadyMasked(content: string): boolean { - return content.includes(''); + return returnedNodes; } } diff --git a/packages/core/src/context/sidecar/builtins.ts b/packages/core/src/context/sidecar/builtins.ts index c1c886880a..7a58412d6c 100644 --- a/packages/core/src/context/sidecar/builtins.ts +++ b/packages/core/src/context/sidecar/builtins.ts @@ -1,50 +1,8 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import type { ProcessorRegistry } from './registry.js'; -import { - ToolMaskingProcessor, - type ToolMaskingProcessorOptions, -} from '../processors/toolMaskingProcessor.js'; -import { BlobDegradationProcessor } from '../processors/blobDegradationProcessor.js'; -import { - SemanticCompressionProcessor, - type SemanticCompressionProcessorOptions, -} from '../processors/semanticCompressionProcessor.js'; -import { - HistorySquashingProcessor, - type HistorySquashingProcessorOptions, -} from '../processors/historySquashingProcessor.js'; -import { - StateSnapshotProcessor, - type StateSnapshotProcessorOptions, -} from '../processors/stateSnapshotProcessor.js'; -import { - EmergencyTruncationProcessor, - type EmergencyTruncationProcessorOptions, -} from '../processors/emergencyTruncationProcessor.js'; +import { ProcessorRegistry } from './registry.js'; +import { EmergencyTruncationProcessor, type EmergencyTruncationProcessorOptions } from '../processors/emergencyTruncationProcessor.js'; +import { BlobDegradationProcessor, type BlobDegradationProcessorOptions } from '../processors/blobDegradationProcessor.js'; export function registerBuiltInProcessors(registry: ProcessorRegistry) { - registry.register({ - id: 'ToolMaskingProcessor', - schema: { - type: 'object', - properties: { - processorId: { const: 'ToolMaskingProcessor' }, - options: { - type: 'object', - properties: { stringLengthThresholdTokens: { type: 'number' } }, - required: ['stringLengthThresholdTokens'], - }, - }, - required: ['processorId', 'options'], - }, - create: (env, opts) => new ToolMaskingProcessor(env, opts), - }); - registry.register>({ id: 'BlobDegradationProcessor', schema: { @@ -58,81 +16,23 @@ export function registerBuiltInProcessors(registry: ProcessorRegistry) { create: (env) => new BlobDegradationProcessor(env), }); - registry.register({ - id: 'SemanticCompressionProcessor', + registry.register({ + id: 'EmergencyTruncationProcessor', schema: { type: 'object', properties: { - processorId: { const: 'SemanticCompressionProcessor' }, + processorId: { const: 'EmergencyTruncationProcessor' }, options: { type: 'object', - properties: { nodeThresholdTokens: { type: 'number' } }, - required: ['nodeThresholdTokens'], + properties: { + target: { type: 'string', enum: ['incremental', 'freeNTokens', 'max'] }, + freeTokensTarget: { type: 'number' }, + }, }, }, - required: ['processorId', 'options'], - }, - create: (env, opts) => new SemanticCompressionProcessor(env, opts), - }); - - registry.register({ - id: 'HistorySquashingProcessor', - schema: { - type: 'object', - properties: { - processorId: { const: 'HistorySquashingProcessor' }, - options: { - type: 'object', - properties: { maxTokensPerNode: { type: 'number' } }, - required: ['maxTokensPerNode'], - }, - }, - required: ['processorId', 'options'], + required: ['processorId'], }, create: (env, options) => - HistorySquashingProcessor.create(env, options), + EmergencyTruncationProcessor.create(env, options), }); - - registry.register({ - id: 'StateSnapshotProcessor', - schema: { - type: 'object', - properties: { - processorId: { const: 'StateSnapshotProcessor' }, - options: { - type: 'object', - properties: { - model: { type: 'string' }, - systemInstruction: { type: 'string' }, - triggerDeficitTokens: { type: 'number' }, - target: { type: 'string', enum: ['incremental', 'freeNTokens', 'max'] }, - freeTokensTarget: { type: 'number' }, - }, - }, - }, - required: ['processorId'], - }, - create: (env, options) => - StateSnapshotProcessor.create(env, options), - }); - - registry.register({ - id: 'EmergencyTruncationProcessor', - schema: { - type: 'object', - properties: { - processorId: { const: 'EmergencyTruncationProcessor' }, - options: { - type: 'object', - properties: { - target: { type: 'string', enum: ['incremental', 'freeNTokens', 'max'] }, - freeTokensTarget: { type: 'number' }, - }, - }, - }, - required: ['processorId'], - }, - create: (env, options) => - EmergencyTruncationProcessor.create(env, options), - }); } diff --git a/packages/core/src/context/sidecar/orchestrator.test.ts b/packages/core/src/context/sidecar/orchestrator.test.ts index f5d6a2ab27..53222d841d 100644 --- a/packages/core/src/context/sidecar/orchestrator.test.ts +++ b/packages/core/src/context/sidecar/orchestrator.test.ts @@ -10,13 +10,12 @@ import { ProcessorRegistry } from './registry.js'; import { createMockEnvironment, createDummyState, - createDummyEpisode, + createDummyNode, } from '../testing/contextTestUtils.js'; import type { ContextEnvironment } from './environment.js'; import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; import type { PipelineDef, ProcessorConfig, SidecarConfig } from './types.js'; import type { ContextEventBus } from '../eventBus.js'; -import type { EpisodeEditor } from '../ir/episodeEditor.js'; // Create a Dummy Processor for testing Orchestration routing class DummySyncProcessor implements ContextProcessor { @@ -27,16 +26,7 @@ class DummySyncProcessor implements ContextProcessor { readonly name = 'DummySync'; readonly id = 'DummySync'; readonly options = {}; - async process(editor: EpisodeEditor, _state: ContextAccountingState) { - if (editor.targets.length === 0) return; - editor.editEpisode( - editor.targets[0].episode.id, - 'DUMMY_EDIT', - (draft: unknown) => { - (draft as Record)['dummyModified'] = true; - }, - ); - } + async process() { return []; } } class DummyAsyncProcessor implements ContextProcessor { @@ -47,16 +37,7 @@ class DummyAsyncProcessor implements ContextProcessor { readonly name = 'DummyAsync'; readonly id = 'DummyAsync'; readonly options = {}; - async process(editor: EpisodeEditor, _state: ContextAccountingState) { - if (editor.targets.length === 0) return; - editor.editEpisode( - editor.targets[0].episode.id, - 'DUMMY_EDIT', - (draft: unknown) => { - (draft as Record)['dummyAsyncModified'] = true; - }, - ); - } + async process() { return []; } } class ThrowingProcessor implements ContextProcessor { @@ -67,10 +48,7 @@ class ThrowingProcessor implements ContextProcessor { readonly name = 'Throwing'; readonly id = 'Throwing'; readonly options = {}; - async process( - _editor: EpisodeEditor, - _state: ContextAccountingState, - ): Promise { + async process(): Promise { throw new Error('Processor failed intentionally'); } } @@ -176,12 +154,13 @@ describe('PipelineOrchestrator (Component)', () => { registry, ); - const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])]; + const episodes = [createDummyNode('1', 'USER_PROMPT')]; const state = createDummyState(false); - const result = await orchestrator.executePipeline( - 'SyncPipe', + const result = await orchestrator.executeTriggerSync( + 'new_message', episodes, + new Set(episodes.map(e => e.id)), state, ); @@ -210,13 +189,14 @@ describe('PipelineOrchestrator (Component)', () => { registry, ); - const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])]; + const episodes = [createDummyNode('1', 'USER_PROMPT')]; const state = createDummyState(false); // This should resolve immediately with the UNMODIFIED array because execution is background - const result = await orchestrator.executePipeline( - 'AsyncPipe', + const result = await orchestrator.executeTriggerSync( + 'new_message', episodes, + new Set(episodes.map(e => e.id)), state, ); @@ -248,13 +228,14 @@ describe('PipelineOrchestrator (Component)', () => { registry, ); - const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])]; + const episodes = [createDummyNode('1', 'USER_PROMPT')]; const state = createDummyState(false); // It should not throw! It should swallow the error and return the unmodified array. - const result = await orchestrator.executePipeline( - 'ThrowingPipe', + const result = await orchestrator.executeTriggerSync( + 'new_message', episodes, + new Set(episodes.map(e => e.id)), state, ); @@ -283,7 +264,7 @@ describe('PipelineOrchestrator (Component)', () => { new PipelineOrchestrator(config, env, eventBus, env.tracer, registry); - const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])]; + const episodes = [createDummyNode('1', 'USER_PROMPT')]; // Emit the trigger eventBus.emitConsolidationNeeded({ episodes, targetDeficit: 100, targetNodeIds: new Set() }); diff --git a/packages/core/src/context/sidecar/orchestrator.ts b/packages/core/src/context/sidecar/orchestrator.ts index dd4855ad24..61f1ef2d22 100644 --- a/packages/core/src/context/sidecar/orchestrator.ts +++ b/packages/core/src/context/sidecar/orchestrator.ts @@ -87,47 +87,64 @@ export class PipelineOrchestrator { } /** - * Applies an array of ContextPatches to the Ship, returning a new immutable Ship array. + * Evaluates the subset returned by the processor against the original targets, + * deducing the removed and inserted nodes, and updating the Ship accordingly. */ - reduceShip( + applyProcessorDiff( ship: ReadonlyArray, - patches: ContextPatch[] + targets: ReadonlyArray, + returnedNodes: ReadonlyArray, ): ReadonlyArray { - if (patches.length === 0) return ship; - const mutableShip = [...ship]; - - for (const patch of patches) { - const { removedIds, insertedNodes = [], insertionIndex } = patch; - - let targetIdx = insertionIndex ?? -1; - - if (targetIdx === -1 && removedIds.length > 0) { - targetIdx = mutableShip.findIndex(n => n.id === removedIds[0]); - } - - if (targetIdx === -1) { - targetIdx = mutableShip.length; - } + const targetSet = new Set(targets.map(n => n.id)); + const returnedMap = new Map(returnedNodes.map(n => [n.id, n])); - if (removedIds.length > 0) { - const removeSet = new Set(removedIds); - let i = 0; - while (i < mutableShip.length) { - if (removeSet.has(mutableShip[i].id)) { - mutableShip.splice(i, 1); - if (i < targetIdx) targetIdx--; - } else { - i++; - } - } - } + const removedIds = new Set(); + const newNodes: ConcreteNode[] = []; - if (insertedNodes.length > 0) { - mutableShip.splice(targetIdx, 0, ...insertedNodes); + // 1. Identify Removals & Modifications + // If a target is missing from returnedMap -> Removed + // If a target is in returnedMap but !== object ref -> Modified (Remove old, Insert new) + for (const t of targets) { + const returnedNode = returnedMap.get(t.id); + if (!returnedNode) { + removedIds.add(t.id); + } else if (returnedNode !== t) { + removedIds.add(t.id); + newNodes.push(returnedNode); } } + // 2. Identify pure Additions (New synthetic nodes) + for (const r of returnedNodes) { + if (!targetSet.has(r.id)) { + newNodes.push(r); + } + } + + if (removedIds.size === 0 && newNodes.length === 0) { + return ship; // No changes + } + + // Find the earliest index in the ship where a removal occurred so we know where to insert + let earliestRemovalIdx = mutableShip.length; + let i = 0; + while (i < mutableShip.length) { + if (removedIds.has(mutableShip[i].id)) { + if (i < earliestRemovalIdx) earliestRemovalIdx = i; + mutableShip.splice(i, 1); + } else { + i++; + } + } + + // Insert new nodes exactly where the old nodes were removed + if (newNodes.length > 0) { + // NOTE: Metadata appending (who, what, when) should ideally happen here + // But for V1, processors still construct the new nodes with metadata inside. + mutableShip.splice(earliestRemovalIdx, 0, ...newNodes); + } + return mutableShip; } @@ -150,16 +167,22 @@ export class PipelineOrchestrator { 'Orchestrator', `Executing processor synchronously: ${procDef.processorId}`, ); + + // 1. Filter out protected nodes + const allowedTargets = currentShip.filter(n => + triggerTargets.has(n.id) && + (!n.logicalParentId || !state.protectedLogicalIds.has(n.logicalParentId)) + ); - const patches = await processor.process({ + const returnedNodes = await processor.process({ ship: currentShip, - triggerTargets, + targets: allowedTargets, state, buffer: {} as any, // TODO: Implement ContextWorkingBuffer fully inbox: {} as any, // TODO: Implement ContextInbox fully }); - currentShip = this.reduceShip(currentShip, patches); + currentShip = this.applyProcessorDiff(currentShip, allowedTargets, returnedNodes); } catch (error) { debugLogger.error( @@ -197,14 +220,21 @@ export class PipelineOrchestrator { `Executing processor: ${procDef.processorId} (async)`, ); - const patches = await processor.process({ + // 1. Filter out protected nodes + const allowedTargets = currentShip.filter(n => + triggerTargets.has(n.id) && + (!n.logicalParentId || !state.protectedLogicalIds.has(n.logicalParentId)) + ); + + const returnedNodes = await processor.process({ ship: currentShip, - triggerTargets, + targets: allowedTargets, state, buffer: {} as any, // TODO: Implement ContextWorkingBuffer fully inbox: {} as any, // TODO: Implement ContextInbox fully }); - currentShip = this.reduceShip(currentShip, patches); + + currentShip = this.applyProcessorDiff(currentShip, allowedTargets, returnedNodes); } catch (error) { debugLogger.error( diff --git a/packages/core/src/context/system-tests/SimulationHarness.ts b/packages/core/src/context/system-tests/SimulationHarness.ts index 65d5feb896..aa90326e00 100644 --- a/packages/core/src/context/system-tests/SimulationHarness.ts +++ b/packages/core/src/context/system-tests/SimulationHarness.ts @@ -105,8 +105,8 @@ export class SimulationHarness { this.chatHistory.set([...currentHistory, ...messages]); // 2. Measure tokens immediately after append (Before background processing) - const tokensBefore = this.env.tokenCalculator.calculateEpisodeListTokens( - this.contextManager.getWorkingBufferView(), + const tokensBefore = this.env.tokenCalculator.calculateConcreteListTokens( + this.contextManager.getShip(), ); debugLogger.log( `[Turn ${this.currentTurnIndex}] Tokens BEFORE: ${tokensBefore}`, @@ -116,35 +116,34 @@ export class SimulationHarness { await new Promise((resolve) => setTimeout(resolve, 50)); // 3.1 Simulate what projectCompressedHistory does with the sync handlers - let currentView = this.contextManager.getWorkingBufferView(); + let currentView = this.contextManager.getShip(); const currentTokens = - this.env.tokenCalculator.calculateEpisodeListTokens(currentView); + this.env.tokenCalculator.calculateConcreteListTokens(currentView); if (this.config.budget && currentTokens > this.config.budget.maxTokens) { debugLogger.log( `[Turn ${this.currentTurnIndex}] Sync panic triggered! ${currentTokens} > ${this.config.budget.maxTokens}`, ); - const syncPipelines = this.config.pipelines.filter( - (p) => p.execution === 'blocking', - ); const orchestrator = this.orchestrator; - for (const pipe of syncPipelines) { - await orchestrator.executePipeline(pipe.name, currentView, { + // In the V2 simulation, we trigger the 'gc_backstop' to simulate emergency pressure. + currentView = await orchestrator.executeTriggerSync( + 'gc_backstop', + currentView, + new Set(currentView.map(e => e.id)), + { currentTokens, maxTokens: this.config.budget.maxTokens, retainedTokens: this.config.budget.retainedTokens, isBudgetSatisfied: false, deficitTokens: currentTokens - this.config.budget.maxTokens, - protectedEpisodeIds: new Set(), + protectedLogicalIds: new Set(), }); - currentView = this.contextManager.getWorkingBufferView(); - } // Inject the truncated view back into the graph for (let i = 0; i < currentView.length; i++) { const ep = currentView[i]; if ( !this.contextManager - .getWorkingBufferView() + .getShip() .find((c) => c.id === ep.id) ) { this.eventBus.emitVariantReady({ @@ -164,8 +163,8 @@ export class SimulationHarness { } // 4. Measure tokens after background processors have (hopefully) emitted variants - const tokensAfter = this.env.tokenCalculator.calculateEpisodeListTokens( - this.contextManager.getWorkingBufferView(), + const tokensAfter = this.env.tokenCalculator.calculateConcreteListTokens( + this.contextManager.getShip(), ); debugLogger.log( `[Turn ${this.currentTurnIndex}] Tokens AFTER: ${tokensAfter}`, diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index 2668419a75..7ffbfd41c2 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -9,6 +9,7 @@ import type { Config } from '../../config/config.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import type { Content } from '@google/genai'; import { AgentChatHistory } from '../../core/agentChatHistory.js'; +import type { ConcreteNode, ToolExecution } from "../ir/types.js"; import { ContextManager } from '../contextManager.js'; import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js'; @@ -34,11 +35,66 @@ export function createDummyState( maxTokens, retainedTokens, deficitTokens: deficit, - protectedEpisodeIds: protectedIds, + protectedLogicalIds: protectedIds, isBudgetSatisfied: isSatisfied, }; } +export function createDummyNode( + logicalParentId: string, + type: 'USER_PROMPT' | 'SYSTEM_EVENT' | 'AGENT_THOUGHT' | 'AGENT_YIELD', + tokens = 100, + overrides?: Partial, + id?: string +): ConcreteNode { + return { + id: id || randomUUID(), + episodeId: logicalParentId, + logicalParentId, + type: type as any, + timestamp: Date.now(), + text: `Dummy ${type}`, + name: type === 'SYSTEM_EVENT' ? 'dummy_event' : undefined, + payload: type === 'SYSTEM_EVENT' ? {} : undefined, + semanticParts: [], + metadata: { + originalTokens: tokens, + currentTokens: tokens, + transformations: [], + }, + ...overrides, + } as unknown as ConcreteNode; +} + +export function createDummyToolNode( + logicalParentId: string, + intentTokens = 100, + obsTokens = 200, + overrides?: Partial, + id?: string +): ToolExecution { + return { + id: id || randomUUID(), + episodeId: logicalParentId, + logicalParentId, + type: 'TOOL_EXECUTION', + timestamp: Date.now(), + toolName: 'dummy_tool', + intent: { action: 'test' }, + observation: { result: 'ok' }, + tokens: { + intent: intentTokens, + observation: obsTokens, + }, + metadata: { + originalTokens: intentTokens + obsTokens, + currentTokens: intentTokens + obsTokens, + transformations: [], + }, + ...overrides, + } as unknown as ToolExecution; +} + export function createDummyEpisode( id: string, type: 'USER_PROMPT' | 'SYSTEM_EVENT', @@ -98,7 +154,7 @@ export function createDummyEpisode( }; } -export function createMockEnvironment(): ContextEnvironment { +export function createMockEnvironment(overrides?: Partial): ContextEnvironment { return { // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion llmClient: vi.fn().mockReturnValue({ @@ -116,7 +172,8 @@ export function createMockEnvironment(): ContextEnvironment { tokenCalculator: new ContextTokenCalculator(1), fileSystem: new InMemoryFileSystem(), idGenerator: new DeterministicIdGenerator('mock-uuid-'), - }; + ...overrides, + } as ContextEnvironment; } /** diff --git a/packages/core/src/context/typed-context-ir.md b/packages/core/src/context/typed-context-ir.md index 7f81a73abe..f69ba5b9e9 100644 --- a/packages/core/src/context/typed-context-ir.md +++ b/packages/core/src/context/typed-context-ir.md @@ -36,7 +36,7 @@ This is a strictly-typed messaging system. A worker dispatches a `SNAPSHOT_READY ## 4. The Processor Contract -Processors are pure functions that evaluate unprotected targets and return an array of `ContextPatch`es based on their phase in the lifecycle. +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 = @@ -65,7 +65,8 @@ export interface ProcessArgs { /** * The specific unprotected, mutable nodes the pipeline is allowed to operate on. - * The Orchestrator filters out protected nodes (like active tasks) before calling. + * 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; @@ -76,26 +77,17 @@ export interface ProcessArgs { readonly inbox: ContextInbox; } -export interface ContextPatch { - /** The IDs of the Concrete Nodes to remove from the Ship. */ - removedIds: string[]; - - /** The new synthetic Concrete Nodes to insert. */ - insertedNodes?: ConcreteNode[]; - - /** The index at which to insert the new nodes. If omitted, they replace the first removedId. */ - insertionIndex?: number; - - /** Audit metadata explaining who made this patch, when, and why. */ - metadata: IrMetadata; -} - export interface ContextProcessor { readonly id: string; readonly name: string; - /** Returns an array of declarative patches applied at this specific temporal phase. */ - process(args: ProcessArgs): Promise; + /** + * 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. + */ + process(args: ProcessArgs): Promise>; } ``` diff --git a/packages/core/src/context/utils/contextTokenCalculator.ts b/packages/core/src/context/utils/contextTokenCalculator.ts index 11d1d36429..65c16fbd93 100644 --- a/packages/core/src/context/utils/contextTokenCalculator.ts +++ b/packages/core/src/context/utils/contextTokenCalculator.ts @@ -6,13 +6,14 @@ import type { Part } from '@google/genai'; import { estimateTokenCountSync as baseEstimate } from '../../utils/tokenCalculation.js'; -import { BASE_MULTIMODAL_TOKEN_COST } from '../ir/types.js'; + /** +import type { ConcreteNode } from "../ir/types.js"; * The flat token cost assigned to a single multi-modal asset (like an image tile) * by the Gemini API. We use this as a baseline heuristic for inlineData/fileData. */ -import type { ConcreteNode } from '../ir/types.js'; + export class ContextTokenCalculator { constructor(private readonly charsPerToken: number) {} @@ -54,7 +55,7 @@ export class ContextTokenCalculator { if (typeof part.text === 'string') { totalTokens += Math.ceil(part.text.length / this.charsPerToken); } else if (part.inlineData !== undefined || part.fileData !== undefined) { - totalTokens += BASE_MULTIMODAL_TOKEN_COST; + totalTokens += 258; } else { totalTokens += Math.ceil( JSON.stringify(part).length / this.charsPerToken,