diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 0edd4af7b0..286bbdb37c 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -700,6 +700,7 @@ export interface ConfigParameters { experimentalJitContext?: boolean; autoDistillation?: boolean; experimentalMemoryManager?: boolean; + experimentalContextSidecarConfig?: string; experimentalAgentHistoryTruncation?: boolean; experimentalAgentHistoryTruncationThreshold?: number; experimentalAgentHistoryRetainedMessages?: number; @@ -942,6 +943,7 @@ export class Config implements McpContext, AgentLoopContext { private readonly adminSkillsEnabled: boolean; private readonly experimentalJitContext: boolean; private readonly experimentalMemoryManager: boolean; + private readonly experimentalContextSidecarConfig?: string; private readonly memoryBoundaryMarkers: readonly string[]; private readonly topicUpdateNarration: boolean; private readonly disableLLMCorrection: boolean; @@ -1153,6 +1155,8 @@ export class Config implements McpContext, AgentLoopContext { this.experimentalJitContext = params.experimentalJitContext ?? false; this.experimentalMemoryManager = params.experimentalMemoryManager ?? false; + this.experimentalContextSidecarConfig = + params.experimentalContextSidecarConfig; this.memoryBoundaryMarkers = params.memoryBoundaryMarkers ?? ['.git']; this.contextManagement = { enabled: params.contextManagement?.enabled ?? false, @@ -2427,6 +2431,10 @@ export class Config implements McpContext, AgentLoopContext { return this.experimentalMemoryManager; } + getExperimentalContextSidecarConfig(): string | undefined { + return this.experimentalContextSidecarConfig; + } + getContextManagementConfig(): ContextManagementConfig { return this.contextManagement; } diff --git a/packages/core/src/context/__snapshots__/contextManager.golden.test.ts.snap b/packages/core/src/context/__snapshots__/contextManager.golden.test.ts.snap new file mode 100644 index 0000000000..a03382c70e --- /dev/null +++ b/packages/core/src/context/__snapshots__/contextManager.golden.test.ts.snap @@ -0,0 +1,52 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[`ContextManager Golden Tests > should process history and match golden snapshot 1`] = ` +[ + { + "parts": [ + { + "text": "A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, A long long time ago, ", + }, + ], + "role": "user", + }, + { + "parts": [ + { + "text": "in a galaxy far far away...", + }, + { + "functionCall": { + "args": {}, + "id": "", + "name": "some_tool", + }, + }, + ], + "role": "model", + }, + { + "parts": [ + { + "functionResponse": { + "id": "", + "name": "some_tool", + "response": { + "output": "TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA TOOL OUTPUT DATA ", + }, + }, + }, + ], + "role": "user", + }, + { + "parts": [ + { + "text": "--- test_file.txt --- +FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA FILE DATA ", + }, + ], + "role": "user", + }, +] +`; diff --git a/packages/core/src/context/contextManager.async.test.ts b/packages/core/src/context/contextManager.async.test.ts new file mode 100644 index 0000000000..b07512b914 --- /dev/null +++ b/packages/core/src/context/contextManager.async.test.ts @@ -0,0 +1,110 @@ +/** + * @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: { + 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.barrier.test.ts b/packages/core/src/context/contextManager.barrier.test.ts new file mode 100644 index 0000000000..41cc54015e --- /dev/null +++ b/packages/core/src/context/contextManager.barrier.test.ts @@ -0,0 +1,73 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; + +import { + createSyntheticHistory, + createMockContextConfig, + setupContextComponentTest, +} from './testing/contextTestUtils.js'; + +describe('ContextManager Sync Pressure Barrier Tests', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + it('should instantly truncate history when maxTokens is exceeded using truncate strategy', async () => { + // 1. Setup + const config = createMockContextConfig(); + const { chatHistory, contextManager } = setupContextComponentTest(config); + + // 2. Add System Prompt (Episode 0 - Protected) + chatHistory.set([ + { role: 'user', parts: [{ text: 'System prompt' }] }, + { role: 'model', parts: [{ text: 'Understood.' }] }, + ]); + + // 3. Add massive history that blows past the 150k maxTokens limit + // 20 turns * 10,000 tokens/turn = ~200,000 tokens + const massiveHistory = createSyntheticHistory(20, 35000); + chatHistory.set([...chatHistory.get(), ...massiveHistory]); + + // 4. Add the Latest Turn (Protected) + chatHistory.set([ + ...chatHistory.get(), + { role: 'user', parts: [{ text: 'Final question.' }] }, + { role: 'model', parts: [{ text: 'Final answer.' }] }, + ]); + + const rawHistoryLength = chatHistory.get().length; + + // 5. Project History (Triggers Sync Barrier) + const projection = await contextManager.projectCompressedHistory(); + + // 6. Assertions + // The barrier should have dropped several older episodes to get under 150k. + + expect(projection.length).toBeLessThan(rawHistoryLength); + + // Verify Episode 0 (System) is perfectly preserved at the front + + expect(projection[0].role).toBe('user'); + expect(projection[0].parts![0].text).toBe('System prompt'); + + // Verify the latest turn is perfectly preserved at the back + const lastUser = projection[projection.length - 2]; + const lastModel = projection[projection.length - 1]; + + expect(lastUser.role).toBe('user'); + expect(lastUser.parts![0].text).toBe('Final question.'); + + expect(lastModel.role).toBe('model'); + expect(lastModel.parts![0].text).toBe('Final answer.'); + }); +}); diff --git a/packages/core/src/context/contextManager.golden.test.ts b/packages/core/src/context/contextManager.golden.test.ts new file mode 100644 index 0000000000..b7791d4f9b --- /dev/null +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -0,0 +1,189 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + describe, + it, + expect, + vi, + beforeEach, + beforeAll, + afterAll, +} from 'vitest'; +import { ContextManager } from './contextManager.js'; +import { ContextEnvironmentImpl } from './sidecar/environmentImpl.js'; +import { SidecarLoader } from './sidecar/SidecarLoader.js'; +import { ContextTracer } from './tracer.js'; +import { ContextEventBus } from './eventBus.js'; +import { ContextTokenCalculator } from './utils/contextTokenCalculator.js'; +import type { Content } from '@google/genai'; +import type { BaseLlmClient } from '../core/baseLlmClient.js'; +import type { Episode } from './ir/types.js'; +import type { SidecarConfig } from './sidecar/types.js'; +import { ProcessorRegistry } from './sidecar/registry.js'; +import { registerBuiltInProcessors } from './sidecar/builtins.js'; +import { IrMapper } from './ir/mapper.js'; + +expect.addSnapshotSerializer({ + test: (val) => + typeof val === 'string' && + /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test(val), + print: () => '""', +}); + +describe('ContextManager Golden Tests', () => { + beforeAll(() => { + vi.useFakeTimers(); + vi.setSystemTime(new Date(2026, 3, 2).getTime()); + vi.spyOn(Math, 'random').mockReturnValue(0.5); + }); + + afterAll(() => { + vi.useRealTimers(); + vi.restoreAllMocks(); + }); + + let mockConfig: any; // eslint-disable-line @typescript-eslint/no-explicit-any + let contextManager: ContextManager; + + beforeEach(() => { + mockConfig = { + isContextManagementEnabled: vi.fn().mockReturnValue(true), + getExperimentalContextSidecarConfig: vi.fn().mockReturnValue(undefined), + getTargetDir: vi.fn().mockReturnValue('/tmp'), + getSessionId: vi.fn().mockReturnValue('test-session'), + getToolOutputMaskingConfig: vi.fn().mockResolvedValue({ + enabled: true, + minPrunableThresholdTokens: 50, + protectLatestTurn: false, + protectionThresholdTokens: 100, + }), + storage: { getProjectTempDir: vi.fn().mockReturnValue('/tmp') }, + getUsageStatisticsEnabled: vi.fn().mockReturnValue(false), + getBaseLlmClient: vi.fn().mockReturnValue({ + generateJson: vi.fn().mockResolvedValue({ + 'test_file.txt': { level: 'SUMMARY' }, + }), + generateContent: vi.fn().mockResolvedValue({ + candidates: [ + { content: { parts: [{ text: 'This is a summary.' }] } }, + ], + }), + }), + }; + + const registry = new ProcessorRegistry(); + registerBuiltInProcessors(registry); + + const sidecar = SidecarLoader.fromConfig(mockConfig, registry); + const tracer = new ContextTracer({ + targetDir: '/tmp', + sessionId: 'test-session', + }); + const eventBus = new ContextEventBus(); + const env = new ContextEnvironmentImpl( + { + generateContent: async () => ({}), + generateJson: async () => ({}), + } as unknown as BaseLlmClient, + 'test-prompt-id', + 'test', + '/tmp', + '/tmp', + tracer, + 4, + eventBus, + ); + contextManager = ContextManager.create( + sidecar, + env, + tracer, + undefined, + registry, + ); + }); + + const createLargeHistory = (): Content[] => [ + { + role: 'user', + parts: [ + { text: 'A long long time ago, '.repeat(500) }, // Squashing target + ], + }, + { + role: 'model', + parts: [{ text: 'in a galaxy far far away...' }], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + name: 'some_tool', + response: { output: 'TOOL OUTPUT DATA '.repeat(500) }, // Masking target + }, + }, + ], + }, + { + role: 'user', + parts: [ + { text: '--- test_file.txt ---\n' + 'FILE DATA '.repeat(1000) }, // Semantic target + ], + }, + ]; + + it('should process history and match golden snapshot', async () => { + const history = createLargeHistory(); + ( + contextManager as unknown as { pristineEpisodes: Episode[] } + ).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4)); + const result = await contextManager.projectCompressedHistory(); + expect(result).toMatchSnapshot(); + }); + + it('should not modify history when under budget', async () => { + const history = createLargeHistory(); + ( + contextManager as unknown as { pristineEpisodes: Episode[] } + ).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4)); + // In Golden Tests, we just want to ensure the logic doesn't throw or alter unprotected history in weird ways. + // Since we're skipping processors due to being under budget, it should equal history. + const tracer2 = new ContextTracer({ + targetDir: '/tmp', + sessionId: 'test2', + }); + const eventBus2 = new ContextEventBus(); + const env2 = new ContextEnvironmentImpl( + { + generateContent: async () => ({}), + generateJson: async () => ({}), + } as unknown as BaseLlmClient, + 'test-prompt-id', + 'test', + '/tmp', + '/tmp', + tracer2, + 4, + eventBus2, + ); + contextManager = ContextManager.create( + { + budget: { retainedTokens: 100000, maxTokens: 150000 }, + pipelines: [], + } as unknown as SidecarConfig, + env2, + tracer2, + ); + + ( + contextManager as unknown as { pristineEpisodes: Episode[] } + ).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4)); + const result = await contextManager.projectCompressedHistory(); + + expect(result.length).toEqual(history.length); + }); +}); diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts new file mode 100644 index 0000000000..7a1078a15b --- /dev/null +++ b/packages/core/src/context/contextManager.ts @@ -0,0 +1,184 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +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 { 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[] = []; + 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(); + }); + + 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() { + 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 + this.eventBus.emitChunkReceived({ episodes: this.pristineEpisodes }); + + // 2. Budget Crossed Trigger + if (currentTokens > this.sidecar.budget.retainedTokens) { + const deficit = currentTokens - this.sidecar.budget.retainedTokens; + this.tracer.logEvent( + 'ContextManager', + 'Budget crossed. Emitting ConsolidationNeeded', + { deficit }, + ); + this.eventBus.emitConsolidationNeeded({ + episodes: workingBuffer, + targetDeficit: deficit, + }); + } + } + + /** + * 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/eventBus.ts b/packages/core/src/context/eventBus.ts new file mode 100644 index 0000000000..d7e8a9e0c5 --- /dev/null +++ b/packages/core/src/context/eventBus.ts @@ -0,0 +1,63 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { EventEmitter } from 'node:events'; +import type { Episode, Variant } from './ir/types.js'; + +export interface PristineHistoryUpdatedEvent { + episodes: Episode[]; +} + +export interface ContextConsolidationEvent { + episodes: Episode[]; + targetDeficit: number; +} + +export interface IrChunkReceivedEvent { + episodes: Episode[]; +} + +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; +} + +export class ContextEventBus extends EventEmitter { + emitPristineHistoryUpdated(event: PristineHistoryUpdatedEvent) { + this.emit('PRISTINE_HISTORY_UPDATED', event); + } + + onPristineHistoryUpdated( + listener: (event: PristineHistoryUpdatedEvent) => void, + ) { + this.on('PRISTINE_HISTORY_UPDATED', listener); + } + + emitChunkReceived(event: IrChunkReceivedEvent) { + this.emit('IR_CHUNK_RECEIVED', event); + } + + onChunkReceived(listener: (event: IrChunkReceivedEvent) => void) { + this.on('IR_CHUNK_RECEIVED', listener); + } + + emitConsolidationNeeded(event: ContextConsolidationEvent) { + this.emit('BUDGET_RETAINED_CROSSED', event); + } + + onConsolidationNeeded(listener: (event: ContextConsolidationEvent) => void) { + this.on('BUDGET_RETAINED_CROSSED', listener); + } + + emitVariantReady(event: VariantReadyEvent) { + this.emit('VARIANT_READY', event); + } + + onVariantReady(listener: (event: VariantReadyEvent) => void) { + this.on('VARIANT_READY', listener); + } +} diff --git a/packages/core/src/context/historyObserver.ts b/packages/core/src/context/historyObserver.ts new file mode 100644 index 0000000000..012a1f2e27 --- /dev/null +++ b/packages/core/src/context/historyObserver.ts @@ -0,0 +1,62 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { + AgentChatHistory, + HistoryEvent, +} from '../core/agentChatHistory.js'; +import { IrMapper } from './ir/mapper.js'; +import type { ContextTokenCalculator } from './utils/contextTokenCalculator.js'; +import type { ContextEventBus } from './eventBus.js'; +import type { ContextTracer } from './tracer.js'; + +/** + * Connects the raw AgentChatHistory to the ContextManager. + * It maps raw messages into Episodic Intermediate Representation (IR) + * and evaluates background triggers whenever history changes. + */ +export class HistoryObserver { + private unsubscribeHistory?: () => void; + + constructor( + private readonly chatHistory: AgentChatHistory, + private readonly eventBus: ContextEventBus, + private readonly tracer: ContextTracer, + private readonly tokenCalculator: ContextTokenCalculator, + ) {} + + start() { + if (this.unsubscribeHistory) { + this.unsubscribeHistory(); + } + + this.unsubscribeHistory = this.chatHistory.subscribe( + (_event: HistoryEvent) => { + // Rebuild the pristine IR graph from the full source history on every change. + const pristineEpisodes = IrMapper.toIr( + this.chatHistory.get(), + this.tokenCalculator, + ); + this.tracer.logEvent( + 'HistoryObserver', + 'Rebuilt pristine graph from chat history update', + { episodeCount: pristineEpisodes.length }, + ); + + this.eventBus.emitPristineHistoryUpdated({ + episodes: pristineEpisodes, + }); + }, + ); + } + + stop() { + if (this.unsubscribeHistory) { + this.unsubscribeHistory(); + this.unsubscribeHistory = undefined; + } + } +} diff --git a/packages/core/src/context/ir/episodeEditor.ts b/packages/core/src/context/ir/episodeEditor.ts new file mode 100644 index 0000000000..8e71d55df4 --- /dev/null +++ b/packages/core/src/context/ir/episodeEditor.ts @@ -0,0 +1,137 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Episode } from './types.js'; + +export interface MutationRecord { + episodeId: string; + type: 'modified' | 'inserted' | 'replaced' | 'deleted'; + action: string; + originalIds?: string[]; // If replaced + episode?: Episode; // For new or modified +} + +export class EpisodeEditor { + private originalMap: Map; + private workingOrder: string[]; + private workingMap: Map; + private mutations: MutationRecord[] = []; + + constructor(episodes: Episode[]) { + this.originalMap = new Map(episodes.map((e) => [e.id, e])); + this.workingOrder = episodes.map((e) => e.id); + this.workingMap = new Map(episodes.map((e) => [e.id, e])); + } + + /** + * Provides a readonly view of the current working state of the episodes. + * Processors should iterate over this to decide what to mutate. + */ + get episodes(): readonly Episode[] { + return this.workingOrder.map((id) => this.workingMap.get(id)!); + } + + /** + * Safely edits an existing episode. + * The framework will handle deeply cloning the episode before passing it to the mutator, + * guaranteeing that original references are never modified. + */ + editEpisode(id: string, action: string, mutator: (draft: Episode) => void) { + const ep = this.workingMap.get(id); + if (!ep) return; + + // Lazy deep clone only if it's the original reference + if (ep === this.originalMap.get(id)) { + const clone = structuredClone(ep); + this.workingMap.set(id, clone); + } + + const draft = this.workingMap.get(id)!; + mutator(draft); + + // Log mutation if not already tracked as modified/inserted/replaced + if (!this.mutations.find((m) => m.episodeId === id)) { + this.mutations.push({ + episodeId: id, + type: 'modified', + action, + episode: draft, + }); + } + } + + /** + * Inserts a brand new episode into the graph at the specified index. + */ + insertEpisode(index: number, newEpisode: Episode, action: string) { + this.workingMap.set(newEpisode.id, newEpisode); + this.workingOrder.splice(index, 0, newEpisode.id); + this.mutations.push({ + episodeId: newEpisode.id, + type: 'inserted', + action, + episode: newEpisode, + }); + } + + /** + * Replaces a set of older episodes with a single new episode (e.g., a Summary or Snapshot). + * It inserts the new episode at the lowest index of the removed episodes. + */ + replaceEpisodes(oldIds: string[], newEpisode: Episode, action: string) { + const indices = oldIds + .map((id) => this.workingOrder.indexOf(id)) + .filter((i) => i !== -1); + if (indices.length === 0) return; + + const insertIndex = Math.min(...indices); + + // Remove old + this.workingOrder = this.workingOrder.filter((id) => !oldIds.includes(id)); + for (const id of oldIds) { + this.workingMap.delete(id); + } + + // Insert new + this.workingOrder.splice(insertIndex, 0, newEpisode.id); + this.workingMap.set(newEpisode.id, newEpisode); + + this.mutations.push({ + episodeId: newEpisode.id, + type: 'replaced', + action, + originalIds: oldIds, + episode: newEpisode, + }); + } + + /** + * Removes episodes from the graph completely (e.g., emergency truncation). + */ + removeEpisodes(oldIds: string[], action: string) { + this.workingOrder = this.workingOrder.filter((id) => !oldIds.includes(id)); + for (const id of oldIds) { + this.workingMap.delete(id); + this.mutations.push({ episodeId: id, type: 'deleted', action }); + } + } + + /** + * Retrieves the final, finalized array of episodes. + * Called by the Orchestrator. + */ + getFinalEpisodes(): Episode[] { + return this.workingOrder.map((id) => this.workingMap.get(id)!); + } + + /** + * Retrieves a log of all structural and property mutations performed by this editor. + * Called by the Orchestrator to emit VariantReady events. + */ + getMutations(): MutationRecord[] { + return this.mutations; + } +} diff --git a/packages/core/src/context/ir/fromIr.ts b/packages/core/src/context/ir/fromIr.ts new file mode 100644 index 0000000000..b1d2be18b5 --- /dev/null +++ b/packages/core/src/context/ir/fromIr.ts @@ -0,0 +1,107 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Content, Part } from '@google/genai'; +import type { Episode, EpisodeStep, UserPrompt, AgentYield } from './types.js'; + +export function fromIr(episodes: Episode[]): Content[] { + const history: Content[] = []; + + for (const ep of episodes) { + if (ep.trigger.type === 'USER_PROMPT') { + const triggerContent = serializeTrigger(ep.trigger); + if (triggerContent) history.push(triggerContent); + } + + const stepContents = serializeSteps(ep.steps); + history.push(...stepContents); + + if (ep.yield) { + history.push(serializeYield(ep.yield)); + } + } + + return history; +} + +function serializeTrigger(trigger: UserPrompt): Content | null { + const parts: Part[] = []; + for (const sp of trigger.semanticParts) { + if (sp.presentation) { + parts.push({ text: sp.presentation.text }); + } else if (sp.type === 'text') { + parts.push({ text: sp.text }); + } else if (sp.type === 'inline_data') { + parts.push({ + inlineData: { mimeType: sp.mimeType, data: sp.data }, + }); + } else if (sp.type === 'file_data') { + parts.push({ + fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri }, + }); + } else if (sp.type === 'raw_part') { + parts.push(sp.part); + } + } + return parts.length > 0 ? { role: 'user', parts } : null; +} + +function serializeSteps(steps: EpisodeStep[]): Content[] { + const history: Content[] = []; + let pendingModelParts: Part[] = []; + let pendingUserParts: Part[] = []; + + const flushPending = () => { + if (pendingModelParts.length > 0) { + history.push({ role: 'model', parts: [...pendingModelParts] }); + pendingModelParts = []; + } + if (pendingUserParts.length > 0) { + history.push({ role: 'user', parts: [...pendingUserParts] }); + pendingUserParts = []; + } + }; + + for (const step of steps) { + if (step.type === 'AGENT_THOUGHT') { + if (pendingUserParts.length > 0) flushPending(); + pendingModelParts.push({ + text: step.presentation?.text ?? step.text, + }); + } else if (step.type === 'TOOL_EXECUTION') { + pendingModelParts.push({ + functionCall: { + name: step.toolName, + args: step.intent, + id: step.id, + }, + }); + const observation = step.presentation + ? step.presentation.observation + : step.observation; + pendingUserParts.push({ + functionResponse: { + name: step.toolName, + response: + typeof observation === 'string' + ? { message: observation } + : observation, + id: step.id, + }, + }); + } + } + flushPending(); + + return history; +} + +function serializeYield(yieldNode: AgentYield): Content { + return { + role: 'model', + parts: [{ text: yieldNode.presentation?.text ?? yieldNode.text }], + }; +} diff --git a/packages/core/src/context/ir/graphUtils.test.ts b/packages/core/src/context/ir/graphUtils.test.ts new file mode 100644 index 0000000000..9ba63db80f --- /dev/null +++ b/packages/core/src/context/ir/graphUtils.test.ts @@ -0,0 +1,175 @@ +/** + * @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 new file mode 100644 index 0000000000..62e9d0a4c2 --- /dev/null +++ b/packages/core/src/context/ir/graphUtils.ts @@ -0,0 +1,174 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Episode } from './types.js'; +import type { ContextTracer } from '../tracer.js'; +import { debugLogger } from '../../utils/debugLogger.js'; +import type { ContextEnvironment } from '../sidecar/environment.js'; + +/** + * 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( + 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 (ep.trigger.type === 'USER_PROMPT') { + 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, + 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 ( + projectedEp.trigger.type === 'USER_PROMPT' && + 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 new file mode 100644 index 0000000000..d74befd62b --- /dev/null +++ b/packages/core/src/context/ir/mapper.test.ts @@ -0,0 +1,271 @@ +/** + * @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/mapper.ts b/packages/core/src/context/ir/mapper.ts new file mode 100644 index 0000000000..bf2c09100b --- /dev/null +++ b/packages/core/src/context/ir/mapper.ts @@ -0,0 +1,31 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Content } from '@google/genai'; +import type { Episode } from './types.js'; +import { toIr } from './toIr.js'; +import { fromIr } from './fromIr.js'; +import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; + +export class IrMapper { + /** + * Translates a flat Gemini Content[] array into our rich Episodic Intermediate Representation. + * Groups adjacent function calls and responses into unified ToolExecution nodes. + */ + static toIr( + history: readonly Content[], + tokenCalculator: ContextTokenCalculator, + ): Episode[] { + return toIr(history, tokenCalculator); + } + + /** + * Re-serializes the Episodic IR back into a flat Gemini Content[] array. + */ + static fromIr(episodes: Episode[]): Content[] { + return fromIr(episodes); + } +} diff --git a/packages/core/src/context/ir/projector.ts b/packages/core/src/context/ir/projector.ts new file mode 100644 index 0000000000..b98d494a90 --- /dev/null +++ b/packages/core/src/context/ir/projector.ts @@ -0,0 +1,92 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Content } from '@google/genai'; +import { IrMapper } from './mapper.js'; +import type { Episode } from './types.js'; +import { debugLogger } from '../../utils/debugLogger.js'; +import type { + ContextEnvironment, + ContextTracer, +} from '../sidecar/environment.js'; +import type { PipelineOrchestrator } from '../sidecar/orchestrator.js'; +import type { SidecarConfig } from '../sidecar/types.js'; + +export class IrProjector { + /** + * Orchestrates the final projection: takes a working buffer view, + * applies the Immediate Sanitization pipeline, and enforces token boundaries. + */ + static async project( + workingBuffer: Episode[], + orchestrator: PipelineOrchestrator, + sidecar: SidecarConfig, + tracer: ContextTracer, + env: ContextEnvironment, + protectedIds: Set, + ): Promise { + if (!sidecar.budget) { + const contents = IrMapper.fromIr(workingBuffer); + tracer.logEvent('IrProjector', 'Projected Context to LLM (No Budget)', { + projectedContext: contents, + }); + return contents; + } + + const maxTokens = sidecar.budget.maxTokens; + const currentTokens = + env.tokenCalculator.calculateEpisodeListTokens(workingBuffer); + + if (currentTokens <= maxTokens) { + tracer.logEvent( + 'IrProjector', + `View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`, + ); + const contents = IrMapper.fromIr(workingBuffer); + tracer.logEvent('IrProjector', 'Projected Context to LLM', { + projectedContext: contents, + }); + return contents; + } + + tracer.logEvent( + 'IrProjector', + `View exceeds maxTokens (${currentTokens} > ${maxTokens}). Hitting Synchronous Pressure Barrier.`, + ); + debugLogger.log( + `Context Manager Synchronous Barrier triggered: View at ${currentTokens} tokens (limit: ${maxTokens}).`, + ); + + const processedEpisodes = await orchestrator.executePipeline( + 'Immediate Sanitization', + workingBuffer, + { + currentTokens, + maxTokens: sidecar.budget.maxTokens, + retainedTokens: sidecar.budget.retainedTokens, + deficitTokens: Math.max(0, currentTokens - sidecar.budget.maxTokens), + protectedEpisodeIds: protectedIds, + isBudgetSatisfied: currentTokens <= sidecar.budget.maxTokens, + }, + ); + + const finalTokens = + env.tokenCalculator.calculateEpisodeListTokens(processedEpisodes); + tracer.logEvent( + 'IrProjector', + `Finished projection. Final token count: ${finalTokens}.`, + ); + debugLogger.log( + `Context Manager finished. Final actual token count: ${finalTokens}.`, + ); + + const contents = IrMapper.fromIr(processedEpisodes); + tracer.logEvent('IrProjector', 'Projected Sanitized Context to LLM', { + projectedContextSanitized: contents, + }); + return contents; + } +} diff --git a/packages/core/src/context/ir/toIr.ts b/packages/core/src/context/ir/toIr.ts new file mode 100644 index 0000000000..7081d0817a --- /dev/null +++ b/packages/core/src/context/ir/toIr.ts @@ -0,0 +1,261 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Content, Part } from '@google/genai'; +import { randomUUID } from 'node:crypto'; +import type { + Episode, + IrMetadata, + SemanticPart, + ToolExecution, + AgentThought, + AgentYield, + UserPrompt, + SystemEvent, +} from './types.js'; +import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; + +// WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references +const nodeIdentityMap = new WeakMap(); + +export function getStableId(obj: object): string { + let id = nodeIdentityMap.get(obj); + if (!id) { + id = randomUUID(); + nodeIdentityMap.set(obj, id); + } + return id; +} + +function isRecord(v: unknown): v is Record { + return typeof v === 'object' && v !== null && !Array.isArray(v); +} + +function isCompleteEpisode(ep: Partial): ep is Episode { + return ( + typeof ep.id === 'string' && + typeof ep.timestamp === 'number' && + !!ep.trigger && + Array.isArray(ep.steps) + ); +} + +export function toIr( + history: readonly Content[], + tokenCalculator: ContextTokenCalculator, +): Episode[] { + const episodes: Episode[] = []; + let currentEpisode: Partial | null = null; + const pendingCallParts: Map = new Map(); + + const createMetadata = (parts: Part[]): IrMetadata => { + const tokens = tokenCalculator.estimateTokensForParts(parts, 0); + return { + originalTokens: tokens, + currentTokens: tokens, + transformations: [], + }; + }; + + const finalizeEpisode = () => { + if (currentEpisode && isCompleteEpisode(currentEpisode)) { + episodes.push(currentEpisode); + } + currentEpisode = null; + }; + + for (const msg of history) { + if (!msg.parts) continue; + + if (msg.role === 'user') { + const hasToolResponses = msg.parts.some((p) => !!p.functionResponse); + const hasUserParts = msg.parts.some( + (p) => !!p.text || !!p.inlineData || !!p.fileData, + ); + + if (hasToolResponses) { + currentEpisode = parseToolResponses( + msg, + currentEpisode, + pendingCallParts, + tokenCalculator, + createMetadata, + ); + } + + if (hasUserParts) { + finalizeEpisode(); + currentEpisode = parseUserParts(msg, createMetadata); + } + } else if (msg.role === 'model') { + currentEpisode = parseModelParts( + msg, + currentEpisode, + pendingCallParts, + createMetadata, + ); + } + } + + if (currentEpisode) { + finalizeYield(currentEpisode); + finalizeEpisode(); + } + + return episodes; +} + +function parseToolResponses( + msg: Content, + currentEpisode: Partial | null, + pendingCallParts: Map, + tokenCalculator: ContextTokenCalculator, + createMetadata: (parts: Part[]) => IrMetadata, +): Partial { + if (!currentEpisode) { + currentEpisode = { + id: getStableId(msg), + timestamp: Date.now(), + trigger: { + id: getStableId(msg.parts![0] || msg), + type: 'SYSTEM_EVENT', + name: 'history_resume', + payload: {}, + metadata: createMetadata([]), + } as SystemEvent, + steps: [], + }; + } + + for (const part of msg.parts!) { + if (part.functionResponse) { + const callId = part.functionResponse.id || ''; + const matchingCall = pendingCallParts.get(callId); + + const intentTokens = matchingCall + ? tokenCalculator.estimateTokensForParts([matchingCall]) + : 0; + const obsTokens = tokenCalculator.estimateTokensForParts([part]); + + const step: ToolExecution = { + id: getStableId(part), + type: 'TOOL_EXECUTION', + toolName: part.functionResponse.name || 'unknown', + intent: isRecord(matchingCall?.functionCall?.args) + ? matchingCall.functionCall.args + : {}, + observation: isRecord(part.functionResponse.response) + ? part.functionResponse.response + : {}, + tokens: { + intent: intentTokens, + observation: obsTokens, + }, + metadata: { + originalTokens: intentTokens + obsTokens, + currentTokens: intentTokens + obsTokens, + transformations: [], + }, + }; + currentEpisode.steps!.push(step); + if (callId) pendingCallParts.delete(callId); + } + } + return currentEpisode; +} + +function parseUserParts( + msg: Content, + createMetadata: (parts: Part[]) => IrMetadata, +): Partial { + const semanticParts: SemanticPart[] = []; + for (const p of msg.parts!) { + if (p.text !== undefined) + semanticParts.push({ type: 'text', text: p.text }); + else if (p.inlineData) + semanticParts.push({ + type: 'inline_data', + mimeType: p.inlineData.mimeType || '', + data: p.inlineData.data || '', + }); + else if (p.fileData) + semanticParts.push({ + type: 'file_data', + mimeType: p.fileData.mimeType || '', + fileUri: p.fileData.fileUri || '', + }); + else if (!p.functionResponse) + semanticParts.push({ type: 'raw_part', part: p }); // Preserve unknowns + } + + const trigger: UserPrompt = { + id: getStableId(msg.parts![0] || msg), + type: 'USER_PROMPT', + semanticParts, + metadata: createMetadata(msg.parts!.filter((p) => !p.functionResponse)), + }; + + return { + id: getStableId(msg), + timestamp: Date.now(), + trigger, + steps: [], + }; +} + +function parseModelParts( + msg: Content, + currentEpisode: Partial | null, + pendingCallParts: Map, + createMetadata: (parts: Part[]) => IrMetadata, +): Partial { + if (!currentEpisode) { + currentEpisode = { + id: getStableId(msg), + timestamp: Date.now(), + trigger: { + id: getStableId(msg.parts![0] || msg), + type: 'SYSTEM_EVENT', + name: 'model_init', + payload: {}, + metadata: createMetadata([]), + } as SystemEvent, + steps: [], + }; + } + + for (const part of msg.parts!) { + if (part.functionCall) { + const callId = part.functionCall.id || ''; + if (callId) pendingCallParts.set(callId, part); + } else if (part.text) { + const thought: AgentThought = { + id: getStableId(part), + type: 'AGENT_THOUGHT', + text: part.text, + metadata: createMetadata([part]), + }; + currentEpisode.steps!.push(thought); + } + } + return currentEpisode; +} + +function finalizeYield(currentEpisode: Partial) { + if (currentEpisode.steps && currentEpisode.steps.length > 0) { + const lastStep = currentEpisode.steps[currentEpisode.steps.length - 1]; + if (lastStep.type === 'AGENT_THOUGHT') { + const yieldNode: AgentYield = { + id: lastStep.id, + type: 'AGENT_YIELD', + text: lastStep.text, + metadata: lastStep.metadata, + }; + currentEpisode.steps.pop(); + currentEpisode.yield = yieldNode; + } + } +} diff --git a/packages/core/src/context/ir/types.ts b/packages/core/src/context/ir/types.ts new file mode 100644 index 0000000000..fddf55197b --- /dev/null +++ b/packages/core/src/context/ir/types.ts @@ -0,0 +1,204 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Part } from '@google/genai'; + +/** + * Universal Audit Metadata + * Tracks the lifecycle and transformations of a node or part within the IR. + * This guarantees perfect reversibility and enables long-term memory offloading. + */ +export interface IrMetadata { + /** The estimated number of tokens this entity originally consumed. */ + originalTokens: number; + /** The current estimated number of tokens this entity consumes in its degraded state. */ + currentTokens: number; + /** An audit trail of all transformations applied by ContextProcessors. */ + transformations: Array<{ + processorName: string; + action: + | 'MASKED' + | 'TRUNCATED' + | 'DEGRADED' + | 'SUMMARIZED' + | 'EVICTED' + | 'SYNTHESIZED'; + timestamp: number; + /** Pointer to where the original uncompressed payload was saved (if applicable) */ + diskPointer?: string; + }>; +} + +export type IrNodeType = + | 'USER_PROMPT' + | 'SYSTEM_EVENT' + | 'AGENT_THOUGHT' + | 'TOOL_EXECUTION' + | 'AGENT_YIELD'; + +/** Base interface for all nodes in the Episodic IR */ +export type VariantStatus = 'computing' | 'ready' | 'failed'; + +export interface BaseVariant { + status: VariantStatus; + recoveredTokens?: number; + error?: string; +} + +export interface SummaryVariant extends BaseVariant { + type: 'summary'; + text: string; +} + +export interface MaskedVariant extends BaseVariant { + type: 'masked'; + text: string; +} + +export interface SnapshotVariant extends BaseVariant { + type: 'snapshot'; + episode: Episode; + replacedEpisodeIds: string[]; +} + +export type Variant = SummaryVariant | MaskedVariant | SnapshotVariant; + +/** Base interface for all nodes in the Episodic IR */ +export interface IrNode { + readonly id: string; + readonly type: IrNodeType; + metadata: IrMetadata; + variants?: Record; +} + +/** + * Semantic Parts for User Prompts + * Ensures we can safely truncate text without deleting multi-modal parts (like images). + */ +export type SemanticPart = + | { + type: 'text'; + text: string; + presentation?: { text: string; tokens: number }; + } + | { + type: 'inline_data'; + mimeType: string; + data: string; + presentation?: { text: string; tokens: number }; + } + | { + type: 'file_data'; + mimeType: string; + fileUri: string; + presentation?: { text: string; tokens: number }; + } + | { + type: 'raw_part'; + part: Part; + presentation?: { text: string; tokens: number }; + }; + +/** + * Trigger Nodes + * Events that wake the agent up and initiate an Episode. + */ +export interface UserPrompt extends IrNode { + readonly type: 'USER_PROMPT'; + /** The semantic breakdown of the user's multi-modal input */ + semanticParts: SemanticPart[]; +} + +export interface SystemEvent extends IrNode { + readonly type: 'SYSTEM_EVENT'; + name: string; + payload: Record; +} + +export type EpisodeTrigger = UserPrompt | SystemEvent; + +/** + * Step Nodes + * The internal autonomous actions taken by the agent during its loop. + */ +export interface AgentThought extends IrNode { + readonly type: 'AGENT_THOUGHT'; + text: string; + /** Overrides the rendered output for this thought */ + presentation?: { + text: string; + tokens: number; + }; +} + +export interface ToolExecution extends IrNode { + readonly type: 'TOOL_EXECUTION'; + /** The name of the tool invoked */ + toolName: string; + + /** The arguments passed to the tool (The 'FunctionCall') */ + intent: Record; + + /** The result returned by the tool (The 'FunctionResponse') */ + observation: string | Record; + + /** Granular token tracking for the different lifecycle phases of the tool */ + tokens: { + intent: number; + observation: number; + }; + + /** + * The presentation layer. If defined, the IrMapper uses this instead of the + * raw observation to build the functionResponse. + * This preserves the immutable raw data for semantic queries while modifying the rendered output. + */ + presentation?: { + intent?: Record; + observation?: string | Record; + tokens: { + intent: number; + observation: number; + }; + }; +} + +export type EpisodeStep = AgentThought | ToolExecution; + +/** + * Resolution Node + * The final message where the agent yields control back to the user. + */ +export interface AgentYield extends IrNode { + readonly type: 'AGENT_YIELD'; + text: string; + presentation?: { + text: string; + tokens: number; + }; +} + +/** + * The Episode + * A discrete, continuous run of the agent. Represents the full cycle from + * taking control (Trigger) to returning control (Yield), encompassing all + * internal reasoning and observations (Steps). + */ +export interface Episode { + readonly id: string; + /** When the episode began */ + readonly timestamp: number; + variants?: Record; + + /** The event that initiated this run */ + trigger: EpisodeTrigger; + + /** The sequence of autonomous actions and observations */ + steps: EpisodeStep[]; + + /** The final handover back to the user (can be undefined if the episode was aborted/errored) */ + yield?: AgentYield; +} diff --git a/packages/core/src/context/pipeline.ts b/packages/core/src/context/pipeline.ts new file mode 100644 index 0000000000..b114098a74 --- /dev/null +++ b/packages/core/src/context/pipeline.ts @@ -0,0 +1,45 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { EpisodeEditor } from './ir/episodeEditor.js'; + +/** + * State object passed through the processing pipeline. + * Contains global accounting logic and semantic protection rules. + */ +export interface ContextAccountingState { + readonly currentTokens: number; + readonly maxTokens: number; + readonly retainedTokens: number; + + /** The exact number of tokens that need to be trimmed to reach the retainedTokens goal */ + readonly deficitTokens: number; + + /** + * Set of Episode IDs that the orchestrator has deemed highly protected. + * Processors should generally skip mutating these episodes unless doing proactive/required transforms. + */ + readonly protectedEpisodeIds: Set; + + /** + * True if currentTokens <= retainedTokens. + */ + readonly isBudgetSatisfied: boolean; +} + +/** + * Interface for all context degradation strategies. + */ +export interface ContextProcessor { + /** Unique name for telemetry and logging. */ + readonly name: string; + + /** + * Processes the episodic history payload via the provided EpisodeEditor, based on the current accounting state. + * Processors should safely mutate or replace episodes using the editor's API. + */ + process(editor: EpisodeEditor, state: ContextAccountingState): Promise; +} diff --git a/packages/core/src/context/processors/blobDegradationProcessor.test.ts b/packages/core/src/context/processors/blobDegradationProcessor.test.ts new file mode 100644 index 0000000000..82cc182d16 --- /dev/null +++ b/packages/core/src/context/processors/blobDegradationProcessor.test.ts @@ -0,0 +1,97 @@ +/** + * @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 { 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'; + +describe('BlobDegradationProcessor', () => { + let processor: BlobDegradationProcessor; + let env: ContextEnvironment; + let fileSystem: InMemoryFileSystem; + + beforeEach(() => { + vi.resetAllMocks(); + env = createMockEnvironment(); + fileSystem = env.fileSystem as InMemoryFileSystem; + processor = new BlobDegradationProcessor(env); + }); + + it('degrades inline_data into a text reference and saves to disk', async () => { + const dummyImageBase64 = Buffer.from('fake-image-data').toString('base64'); + + const ep = createDummyEpisode('ep-1', 'USER_PROMPT', [ + { type: 'text', text: 'Look at this image:' }, + { + type: 'inline_data', + mimeType: 'image/png', + data: dummyImageBase64, + }, + ]); + + const state = createDummyState(false, 500); + const editor = new EpisodeEditor([ep]); + await processor.process(editor, state); + const result = editor.getFinalEpisodes(); + + const parts = (result[0].trigger as UserPrompt).semanticParts; + + // Text part should be untouched + expect(parts[0].presentation).toBeUndefined(); + + // 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', + ); + + // 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); + }); + + 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', + }, + ]); + + const state = createDummyState(false, 500); + const editor = new EpisodeEditor([ep]); + await processor.process(editor, state); + const result = editor.getFinalEpisodes(); + + 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', + ); + + expect(fileSystem.getFiles().size).toBe(0); + }); +}); diff --git a/packages/core/src/context/processors/blobDegradationProcessor.ts b/packages/core/src/context/processors/blobDegradationProcessor.ts new file mode 100644 index 0000000000..2981a5d7e8 --- /dev/null +++ b/packages/core/src/context/processors/blobDegradationProcessor.ts @@ -0,0 +1,142 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextEnvironment } from '../sidecar/environment.js'; +import { sanitizeFilenamePart } from '../../utils/fileUtils.js'; +import type { EpisodeEditor } from '../ir/episodeEditor.js'; + +export type BlobDegradationProcessorOptions = Record; + +export class BlobDegradationProcessor implements ContextProcessor { + static create( + env: ContextEnvironment, + _options: BlobDegradationProcessorOptions, + ): BlobDegradationProcessor { + return new BlobDegradationProcessor(env); + } + + readonly id = 'BlobDegradationProcessor'; + readonly name = 'BlobDegradationProcessor'; + readonly options = {}; + private env: ContextEnvironment; + + constructor(env: ContextEnvironment) { + this.env = env; + } + + async process( + editor: EpisodeEditor, + state: ContextAccountingState, + ): Promise { + if (state.isBudgetSatisfied) { + return; + } + + let currentDeficit = state.deficitTokens; + let directoryCreated = false; + + let blobOutputsDir = this.env.fileSystem.join( + this.env.projectTempDir, + 'degraded-blobs', + ); + const sessionId = this.env.sessionId; + if (sessionId) { + blobOutputsDir = this.env.fileSystem.join( + blobOutputsDir, + `session-${sanitizeFilenamePart(sessionId)}`, + ); + } + + const ensureDir = async () => { + if (!directoryCreated) { + await this.env.fileSystem.mkdir(blobOutputsDir, { recursive: true }); + directoryCreated = true; + } + }; + + // Forward scan, looking for bloated non-text parts to degrade + for (const ep of editor.episodes) { + if (currentDeficit <= 0) break; + if (state.protectedEpisodeIds.has(ep.id)) continue; + + if (ep.trigger.type === 'USER_PROMPT') { + 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; + + let newText = ''; + let tokensSaved = 0; + + 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); + + // Base64 to buffer + const buffer = Buffer.from(part.data, 'base64'); + await this.env.fileSystem.writeFile(filePath, buffer); + + 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}]`; + + // 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; + } + + if (newText && tokensSaved > 0) { + const newTokens = this.env.tokenCalculator.estimateTokensForParts([ + { text: newText }, + ]); + + editor.editEpisode(ep.id, 'DEGRADE_BLOB', (draft) => { + if (draft.trigger.type === 'USER_PROMPT') { + draft.trigger.semanticParts[j].presentation = { + text: newText, + tokens: newTokens, + }; + draft.trigger.metadata.transformations.push({ + processorName: this.name, + action: 'DEGRADED', + timestamp: Date.now(), + }); + } + }); + + currentDeficit -= tokensSaved; + } + } + } + } + } +} diff --git a/packages/core/src/context/processors/emergencyTruncationProcessor.test.ts b/packages/core/src/context/processors/emergencyTruncationProcessor.test.ts new file mode 100644 index 0000000000..892e0d53ce --- /dev/null +++ b/packages/core/src/context/processors/emergencyTruncationProcessor.test.ts @@ -0,0 +1,129 @@ +/** + * @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 new file mode 100644 index 0000000000..ed1e120dd3 --- /dev/null +++ b/packages/core/src/context/processors/emergencyTruncationProcessor.ts @@ -0,0 +1,60 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextProcessor, ContextAccountingState } from '../pipeline.js'; +import type { ContextEnvironment } from '../sidecar/environment.js'; +import type { EpisodeEditor } from '../ir/episodeEditor.js'; + +export type EmergencyTruncationProcessorOptions = Record; + +export class EmergencyTruncationProcessor implements ContextProcessor { + static create( + env: ContextEnvironment, + options: EmergencyTruncationProcessorOptions, + ): EmergencyTruncationProcessor { + return new EmergencyTruncationProcessor(env, options); + } + + readonly id = 'EmergencyTruncationProcessor'; + readonly name = 'EmergencyTruncationProcessor'; + readonly options: EmergencyTruncationProcessorOptions; + constructor( + private readonly _env: ContextEnvironment, + options: EmergencyTruncationProcessorOptions, + ) { + this.options = options; + } + + async process( + editor: EpisodeEditor, + state: ContextAccountingState, + ): Promise { + if (state.currentTokens <= state.maxTokens) return; + + let remainingTokens = state.currentTokens; + const targetTokens = state.maxTokens; + const toRemove: string[] = []; + + // We respect the global protected Episode IDs (like the system prompt at index 0) + for (const ep of editor.episodes) { + const epTokens = this._env.tokenCalculator.calculateEpisodeListTokens([ + ep, + ]); + + if ( + remainingTokens > targetTokens && + !state.protectedEpisodeIds.has(ep.id) + ) { + remainingTokens -= epTokens; + toRemove.push(ep.id); + } + } + + if (toRemove.length > 0) { + editor.removeEpisodes(toRemove, 'TRUNCATED'); + } + } +} diff --git a/packages/core/src/context/processors/historySquashingProcessor.test.ts b/packages/core/src/context/processors/historySquashingProcessor.test.ts new file mode 100644 index 0000000000..bdbae5a159 --- /dev/null +++ b/packages/core/src/context/processors/historySquashingProcessor.test.ts @@ -0,0 +1,157 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + createMockEnvironment, + createDummyState, + createDummyEpisode, +} 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'; + +describe('HistorySquashingProcessor', () => { + let processor: HistorySquashingProcessor; + + beforeEach(() => { + processor = new HistorySquashingProcessor(createMockEnvironment(), { + maxTokensPerNode: 100, + }); + }); + + 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('bypasses processing if budget is satisfied', async () => { + const episodes = [createThoughtEpisode('1', 'short text', 'short thought')]; + const state = createDummyState(true); + + const editor = new EpisodeEditor(episodes); + await processor.process(editor, state); + const result = editor.getFinalEpisodes(); + + expect(result).toStrictEqual(episodes); + expect( + (result[0].trigger as UserPrompt).semanticParts[0].presentation, + ).toBeUndefined(); + }); + + 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 editor = new EpisodeEditor(episodes); + await processor.process(editor, state); + const result = editor.getFinalEpisodes(); + + expect( + (result[0].trigger as UserPrompt).semanticParts[0].presentation, + ).toBeUndefined(); + }); + + 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 editor = new EpisodeEditor(episodes); + await processor.process(editor, state); + const result = editor.getFinalEpisodes(); + + const userPart = (result[0].trigger as UserPrompt).semanticParts[0]; + const thoughtPart = result[0].steps[0] as AgentThought; + + 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 ...]'); + }); +}); diff --git a/packages/core/src/context/processors/historySquashingProcessor.ts b/packages/core/src/context/processors/historySquashingProcessor.ts new file mode 100644 index 0000000000..add0c813fb --- /dev/null +++ b/packages/core/src/context/processors/historySquashingProcessor.ts @@ -0,0 +1,189 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextEnvironment } from '../sidecar/environment.js'; +import { truncateProportionally } from '../truncation.js'; +import type { EpisodeEditor } from '../ir/episodeEditor.js'; + +export interface HistorySquashingProcessorOptions { + maxTokensPerNode: number; +} + +export class HistorySquashingProcessor implements ContextProcessor { + static create( + env: ContextEnvironment, + options: HistorySquashingProcessorOptions, + ): HistorySquashingProcessor { + return new HistorySquashingProcessor(env, options); + } + + static readonly schema = { + type: 'object', + properties: { + maxTokensPerNode: { + type: 'number', + description: + 'The maximum tokens a node can have before being truncated.', + }, + }, + required: ['maxTokensPerNode'], + }; + + readonly id = 'HistorySquashingProcessor'; + readonly name = 'HistorySquashingProcessor'; + readonly options: HistorySquashingProcessorOptions; + + constructor( + env: ContextEnvironment, + options: HistorySquashingProcessorOptions, + ) { + this.options = options; + } + + private tryApplySquash( + text: string, + limitChars: number, + currentDeficit: number, + setPresentation: (p: { text: string; tokens: number }) => void, + recordAudit: () => void, + ): number { + if (currentDeficit <= 0) return 0; + const originalLength = text.length; + if (originalLength <= limitChars) return 0; + + const newText = truncateProportionally( + text, + limitChars, + `\n\n[... OMITTED ${originalLength - limitChars} chars ...]\n\n`, + ); + + if (newText !== text) { + const newTokens = Math.floor(newText.length / 4); + const oldTokens = Math.floor(originalLength / 4); + const tokensSaved = oldTokens - newTokens; + + setPresentation({ text: newText, tokens: newTokens }); + recordAudit(); + return tokensSaved; + } + return 0; + } + + async process( + editor: EpisodeEditor, + state: ContextAccountingState, + ): Promise { + if (state.isBudgetSatisfied) { + return; + } + + const { maxTokensPerNode } = this.options; + // We estimate 4 chars per token for truncation logic + const limitChars = maxTokensPerNode * 4; + + // We track how many tokens we still need to cut. If we hit 0, we can stop early! + let currentDeficit = state.deficitTokens; + + for (const ep of editor.episodes) { + if (currentDeficit <= 0) break; + if (state.protectedEpisodeIds.has(ep.id)) continue; + + // 1. Squash User Prompts + if (ep.trigger.type === 'USER_PROMPT') { + for (let j = 0; j < ep.trigger.semanticParts.length; j++) { + const part = ep.trigger.semanticParts[j]; + if (part.type === 'text') { + const saved = this.tryApplySquash( + part.text, + limitChars, + currentDeficit, + (p) => { + editor.editEpisode(ep.id, 'SQUASH_PROMPT', (draft) => { + if (draft.trigger.type === 'USER_PROMPT') { + 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; + } + } + } + + // 2. Squash Model Thoughts + if (ep.steps) { + for (let j = 0; j < ep.steps.length; j++) { + const step = ep.steps[j]; + if (currentDeficit <= 0) break; + if (step.type === 'AGENT_THOUGHT') { + const saved = this.tryApplySquash( + step.text, + limitChars, + currentDeficit, + (p) => { + editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => { + const draftStep = draft.steps[j]; + if (draftStep.type === 'AGENT_THOUGHT') { + draftStep.presentation = p; + } + }); + }, + () => { + editor.editEpisode(ep.id, 'SQUASH_THOUGHT', (draft) => { + const draftStep = draft.steps[j]; + if (draftStep.type === 'AGENT_THOUGHT') { + draftStep.metadata.transformations.push({ + processorName: this.name, + action: 'TRUNCATED', + timestamp: Date.now(), + }); + } + }); + }, + ); + currentDeficit -= saved; + } + } + } + + // 3. Squash Agent Yields + if (currentDeficit > 0 && 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; + } + } + } +} diff --git a/packages/core/src/context/processors/semanticCompressionProcessor.test.ts b/packages/core/src/context/processors/semanticCompressionProcessor.test.ts new file mode 100644 index 0000000000..a3f3c96143 --- /dev/null +++ b/packages/core/src/context/processors/semanticCompressionProcessor.test.ts @@ -0,0 +1,164 @@ +/** + * @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 { 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'; + +describe('SemanticCompressionProcessor', () => { + let processor: SemanticCompressionProcessor; + let generateContentMock: ReturnType; + + beforeEach(() => { + generateContentMock = vi.fn().mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'Mocked Summary!' }] } }], + }); + + 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 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; + + 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; + }; + + it('bypasses processing if budget is satisfied', async () => { + const episodes = [ + createEpisodeWithThoughtsAndTools('1', 'short', 'short', 'short'), + ]; + const state = createDummyState(true); + + const editor = new EpisodeEditor(episodes); + await processor.process(editor, state); + + expect(generateContentMock).not.toHaveBeenCalled(); + }); + + 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 editor = new EpisodeEditor(episodes); + await processor.process(editor, state); + + expect(generateContentMock).not.toHaveBeenCalled(); + }); + + 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 editor = new EpisodeEditor(episodes); + await processor.process(editor, state); + + expect(generateContentMock).toHaveBeenCalledTimes(3); + + // 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; + + 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); + }); +}); diff --git a/packages/core/src/context/processors/semanticCompressionProcessor.ts b/packages/core/src/context/processors/semanticCompressionProcessor.ts new file mode 100644 index 0000000000..9ba737124f --- /dev/null +++ b/packages/core/src/context/processors/semanticCompressionProcessor.ts @@ -0,0 +1,261 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextAccountingState, ContextProcessor } 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'; + +export interface SemanticCompressionProcessorOptions { + nodeThresholdTokens: number; +} + +export class SemanticCompressionProcessor implements ContextProcessor { + static create( + env: ContextEnvironment, + options: SemanticCompressionProcessorOptions, + ): SemanticCompressionProcessor { + return new SemanticCompressionProcessor(env, options); + } + + static readonly schema = { + type: 'object', + properties: { + nodeThresholdTokens: { + type: 'number', + description: 'The token threshold above which nodes are summarized.', + }, + }, + required: ['nodeThresholdTokens'], + }; + + readonly id = 'SemanticCompressionProcessor'; + readonly name = 'SemanticCompressionProcessor'; + readonly options: SemanticCompressionProcessorOptions; + private env: ContextEnvironment; + private modelToUse: string = 'chat-compression-2.5-flash-lite'; + + constructor( + env: ContextEnvironment, + options: SemanticCompressionProcessorOptions, + ) { + this.env = env; + this.options = options; + } + + async process( + editor: EpisodeEditor, + state: ContextAccountingState, + ): Promise { + // If the budget is satisfied, or semantic compression isn't enabled + if (state.isBudgetSatisfied) { + return; + } + + 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; + + // We scan backwards (oldest to newest would also work, but older is safer to degrade first) + for (const ep of editor.episodes) { + if (currentDeficit <= 0) break; + if (state.protectedEpisodeIds.has(ep.id)) continue; + + // 1. Compress User Prompts + if (ep.trigger.type === 'USER_PROMPT') { + for (let j = 0; j < ep.trigger.semanticParts.length; j++) { + const part = ep.trigger.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 }, + ]); + + if (newTokens < oldTokens) { + editor.editEpisode(ep.id, 'SUMMARIZE_PROMPT', (draft) => { + if (draft.trigger.type === 'USER_PROMPT') { + 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; + } + } + } + } + + // 2. Compress Model Thoughts + if (ep.steps) { + for (let j = 0; j < ep.steps.length; j++) { + const step = ep.steps[j]; + if (currentDeficit <= 0) break; + if (step.type === 'AGENT_THOUGHT') { + if (step.presentation) continue; + 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 (newTokens < oldTokens) { + editor.editEpisode(ep.id, 'SUMMARIZE_THOUGHT', (draft) => { + const draftStep = draft.steps[j]; + if (draftStep.type === 'AGENT_THOUGHT') { + draftStep.presentation = { + text: summary, + tokens: newTokens, + }; + draftStep.metadata.transformations.push({ + processorName: this.name, + action: 'SUMMARIZED', + timestamp: Date.now(), + }); + } + }); + currentDeficit -= oldTokens - newTokens; + } + } + } + + // 3. Compress Tool Observations + if (step.type === 'TOOL_EXECUTION') { + const rawObs = step.presentation?.observation ?? step.observation; + + let stringifiedObs = ''; + if (typeof rawObs === 'string') { + stringifiedObs = rawObs; + } else { + try { + stringifiedObs = JSON.stringify(rawObs); + } catch { + stringifiedObs = String(rawObs); + } + } + + if ( + stringifiedObs.length > thresholdChars && + !stringifiedObs.includes('') + ) { + const summary = await this.generateSummary( + stringifiedObs, + `Tool Output (${step.toolName})`, + ); + + // Wrap the summary in an object so the Gemini API accepts it as a valid functionResponse.response + const newObsObject = { summary }; + + const newObsTokens = + this.env.tokenCalculator.estimateTokensForParts([ + { + functionResponse: { + name: step.toolName, + response: newObsObject, + id: step.id, + }, + }, + ]); + + const oldObsTokens = + step.presentation?.tokens?.observation ?? + step.tokens?.observation ?? + step.tokens; + const intentTokens = + step.presentation?.tokens?.intent ?? step.tokens?.intent ?? 0; + + if (newObsTokens < oldObsTokens) { + editor.editEpisode(ep.id, 'SUMMARIZE_TOOL', (draft) => { + const draftStep = draft.steps[j]; + if (draftStep.type === 'TOOL_EXECUTION') { + 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; + } + } + } + } + } + } + } + + 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; + } + } +} diff --git a/packages/core/src/context/processors/stateSnapshotProcessor.test.ts b/packages/core/src/context/processors/stateSnapshotProcessor.test.ts new file mode 100644 index 0000000000..477cae8e68 --- /dev/null +++ b/packages/core/src/context/processors/stateSnapshotProcessor.test.ts @@ -0,0 +1,118 @@ +/** + * @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/stateSnapshotProcessor.ts b/packages/core/src/context/processors/stateSnapshotProcessor.ts new file mode 100644 index 0000000000..6872288130 --- /dev/null +++ b/packages/core/src/context/processors/stateSnapshotProcessor.ts @@ -0,0 +1,182 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextProcessor, ContextAccountingState } from '../pipeline.js'; +import type { Episode } from '../ir/types.js'; +import type { + ContextEnvironment, + ContextEventBus, +} from '../sidecar/environment.js'; +import { v4 as uuidv4 } from 'uuid'; +import { LlmRole } from '../../telemetry/llmRole.js'; +import { debugLogger } from 'src/utils/debugLogger.js'; +import type { EpisodeEditor } from '../ir/episodeEditor.js'; + +export interface StateSnapshotProcessorOptions { + model?: string; + systemInstruction?: string; + triggerDeficitTokens?: number; +} + +export class StateSnapshotProcessor implements ContextProcessor { + static create( + env: ContextEnvironment, + options: StateSnapshotProcessorOptions, + ): StateSnapshotProcessor { + return new StateSnapshotProcessor(env, options, env.eventBus); + } + readonly id = 'StateSnapshotProcessor'; + readonly name = 'StateSnapshotProcessor'; + readonly options: StateSnapshotProcessorOptions; + private readonly env: ContextEnvironment; + private isSynthesizing = false; + + constructor( + env: ContextEnvironment, + options: StateSnapshotProcessorOptions, + _eventBus: ContextEventBus, + ) { + this.env = env; + this.options = options; + } + + async process( + editor: EpisodeEditor, + state: ContextAccountingState, + ): Promise { + const targetDeficit = Math.max( + 0, + state.currentTokens - state.retainedTokens, + ); + if (this.isSynthesizing || targetDeficit <= 0) return; + + this.isSynthesizing = true; + try { + let deficitAccumulator = 0; + const selectedEpisodes: Episode[] = []; + + for (let i = 1; i < editor.episodes.length - 1; i++) { + const ep = editor.episodes[i]; + selectedEpisodes.push(ep); + let triggerText = ''; + if (ep.trigger?.type === 'USER_PROMPT') { + const firstPart = ep.trigger.semanticParts?.[0]; + if (firstPart) { + triggerText = + firstPart.type === 'text' + ? firstPart.text + : (firstPart.presentation?.text ?? ''); + } + } + deficitAccumulator += this.env.tokenCalculator.estimateTokensForParts([ + { text: triggerText }, + { text: ep.yield?.text ?? '' }, + ]); + if (deficitAccumulator >= targetDeficit) break; + } + + if (selectedEpisodes.length < 2) return; // Not enough context to summarize + + // Optimization: Do NOT emit VariantComputing, let the Orchestrator handle caching the final result. + const snapshotEp: Episode = + await this.synthesizeSnapshot(selectedEpisodes); + + const oldIds = selectedEpisodes.map((ep) => ep.id); + editor.replaceEpisodes(oldIds, snapshotEp, 'STATE_SNAPSHOT'); + } finally { + this.isSynthesizing = false; + } + } + + private async synthesizeSnapshot(episodes: Episode[]): Promise { + const client = this.env.llmClient; + const systemPrompt = + this.options.systemInstruction ?? + `You are an expert Context Memory Manager. You will be provided with a raw transcript of older conversation turns between a user and an AI assistant. +Your task is to synthesize these turns into a single, dense, factual snapshot that preserves all critical context, preferences, active tasks, and factual knowledge, but discards conversational filler, pleasantries, and redundant back-and-forth iterations. + +Output ONLY the raw factual snapshot, formatted compactly. Do not include markdown wrappers, prefixes like "Here is the snapshot", or conversational elements.`; + + let userPromptText = 'TRANSCRIPT TO SNAPSHOT:\n\n'; + for (const ep of episodes) { + if (ep.trigger?.type === 'USER_PROMPT') { + const partsText = ep.trigger.semanticParts + .map((p) => { + if (p.type === 'text') return p.text; + if (p.presentation) return p.presentation.text; + return ''; + }) + .join(''); + userPromptText += `USER: ${partsText}\n`; + } else if (ep.trigger?.type === 'SYSTEM_EVENT') { + userPromptText += `[SYSTEM EVENT: ${ep.trigger.name}]\n`; + } + for (const step of ep.steps) { + if (step.type === 'TOOL_EXECUTION') { + userPromptText += `[Tool Called: ${step.toolName}]\n`; + } + } + if (ep.yield) { + userPromptText += `ASSISTANT: ${ep.yield.text}\n`; + } + userPromptText += '\n'; + } + + try { + const response = await client.generateContent({ + modelConfigKey: { model: 'state-snapshot-processor' }, + contents: [{ role: 'user', parts: [{ text: userPromptText }] }], + systemInstruction: { role: 'system', parts: [{ text: systemPrompt }] }, + promptId: this.env.promptId, + role: LlmRole.UTILITY_STATE_SNAPSHOT_PROCESSOR, + abortSignal: new AbortController().signal, + }); + + const snapshotText = response.text; + + // Synthesize a new "Episode" representing this compressed block + const newId = uuidv4(); + const contentTokens = this.env.tokenCalculator.estimateTokensForParts([ + { text: snapshotText }, + ]); + + return { + id: newId, + timestamp: Date.now(), + trigger: { + id: `${newId}-t`, + type: 'USER_PROMPT', + semanticParts: [], + metadata: { + originalTokens: 0, + currentTokens: 0, + transformations: [], + }, + }, + steps: [], + yield: { + id: `${newId}-y`, + type: 'AGENT_YIELD', + text: `\n${snapshotText}\n`, + metadata: { + originalTokens: contentTokens, + currentTokens: contentTokens, + transformations: [ + { + processorName: 'StateSnapshotProcessor', + action: 'SYNTHESIZED', + timestamp: Date.now(), + }, + ], + }, + }, + }; + } catch (error) { + debugLogger.error('Failed to synthesize snapshot:', error); + throw error; + } + } +} diff --git a/packages/core/src/context/processors/toolMaskingProcessor.test.ts b/packages/core/src/context/processors/toolMaskingProcessor.test.ts new file mode 100644 index 0000000000..f74b02eb56 --- /dev/null +++ b/packages/core/src/context/processors/toolMaskingProcessor.test.ts @@ -0,0 +1,134 @@ +/** + * @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 { 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'; + +describe('ToolMaskingProcessor', () => { + let processor: ToolMaskingProcessor; + let env: ContextEnvironment; + let fileSystem: InMemoryFileSystem; + + beforeEach(() => { + vi.resetAllMocks(); + env = createMockEnvironment(); + fileSystem = env.fileSystem as InMemoryFileSystem; + + processor = new ToolMaskingProcessor(env, { + stringLengthThresholdTokens: 100, + }); + }); + + const getDummyState = ( + isSatisfied = false, + deficit = 0, + protectedIds = new Set(), + ): ContextAccountingState => ({ + currentTokens: 5000, + maxTokens: 10000, + retainedTokens: 4000, + deficitTokens: deficit, + protectedEpisodeIds: protectedIds, + isBudgetSatisfied: isSatisfied, + }); + + const createDummyEpisode = ( + id: string, + intent: Record, + observation: Record, + ): 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: [], + }, + }, + ], + }); + + it('bypasses processing if budget is satisfied', async () => { + const episodes = [ + createDummyEpisode('1', { arg: 'short' }, { out: 'short' }), + ]; + const state = getDummyState(true); + + const editor = new EpisodeEditor(episodes); + await processor.process(editor, state); + const result = editor.getFinalEpisodes(); + + expect(result).toStrictEqual(episodes); + expect((result[0].steps[0] as ToolExecution).presentation).toBeUndefined(); + }); + + 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 intentPayload = { args: { nested: [massiveIntentString, 'short'] } }; + const obsPayload = { result: massiveObsString, error: null }; + + 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); + }); +}); diff --git a/packages/core/src/context/processors/toolMaskingProcessor.ts b/packages/core/src/context/processors/toolMaskingProcessor.ts new file mode 100644 index 0000000000..15812d1629 --- /dev/null +++ b/packages/core/src/context/processors/toolMaskingProcessor.ts @@ -0,0 +1,319 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextEnvironment } from '../sidecar/environment.js'; +import { sanitizeFilenamePart } from '../../utils/fileUtils.js'; +import { + ACTIVATE_SKILL_TOOL_NAME, + MEMORY_TOOL_NAME, + ASK_USER_TOOL_NAME, + 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, + MEMORY_TOOL_NAME, + ASK_USER_TOOL_NAME, + ENTER_PLAN_MODE_TOOL_NAME, + EXIT_PLAN_MODE_TOOL_NAME, +]); + +export interface ToolMaskingProcessorOptions { + stringLengthThresholdTokens: number; +} + +type MaskableValue = + | string + | number + | boolean + | null + | MaskableValue[] + | { [key: string]: MaskableValue }; + +function isMaskableValue(val: unknown): val is MaskableValue { + if ( + val === null || + typeof val === 'string' || + typeof val === 'number' || + typeof val === 'boolean' + ) { + return true; + } + if (Array.isArray(val)) { + return val.every(isMaskableValue); + } + if (typeof val === 'object') { + return Object.values(val).every(isMaskableValue); + } + return false; +} + +function isMaskableRecord(val: unknown): val is Record { + return ( + typeof val === 'object' && + val !== null && + !Array.isArray(val) && + isMaskableValue(val) + ); +} + +export class ToolMaskingProcessor implements ContextProcessor { + static create( + env: ContextEnvironment, + options: ToolMaskingProcessorOptions, + ): ToolMaskingProcessor { + return new ToolMaskingProcessor(env, options); + } + + static readonly schema = { + type: 'object', + properties: { + stringLengthThresholdTokens: { + type: 'number', + description: + 'The token threshold above which tool intents/observations are masked.', + }, + }, + required: ['stringLengthThresholdTokens'], + }; + + readonly id = 'ToolMaskingProcessor'; + readonly name = 'ToolMaskingProcessor'; + readonly options: ToolMaskingProcessorOptions; + private env: ContextEnvironment; + + constructor(env: ContextEnvironment, options: ToolMaskingProcessorOptions) { + this.env = env; + this.options = options; + } + + async process( + editor: EpisodeEditor, + state: ContextAccountingState, + ): Promise { + const maskingConfig = this.options; + if (!maskingConfig) return; + if (state.isBudgetSatisfied) return; + + let currentDeficit = state.deficitTokens; + const limitChars = this.env.tokenCalculator.tokensToChars( + maskingConfig.stringLengthThresholdTokens, + ); + + let toolOutputsDir = this.env.fileSystem.join( + this.env.projectTempDir, + 'tool-outputs', + ); + const sessionId = this.env.sessionId; + if (sessionId) { + toolOutputsDir = this.env.fileSystem.join( + toolOutputsDir, + `session-${sanitizeFilenamePart(sessionId)}`, + ); + } + + // 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, + callId: string, + nodeType: string, + ): Promise => { + if (!directoryCreated) { + await this.env.fileSystem.mkdir(toolOutputsDir, { recursive: true }); + directoryCreated = true; + } + + const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${this.env.idGenerator.generateId()}.txt`; + const filePath = this.env.fileSystem.join(toolOutputsDir, fileName); + + await this.env.fileSystem.writeFile(filePath, content); + + const fileSizeMB = ( + Buffer.byteLength(content, 'utf8') / + 1024 / + 1024 + ).toFixed(2); + const totalLines = content.split('\n').length; + 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.episodes) { + if (currentDeficit <= 0) break; + if (!ep || !ep.steps || state.protectedEpisodeIds.has(ep.id)) continue; + + for (let j = 0; j < ep.steps.length; j++) { + if (currentDeficit <= 0) break; + const step = ep.steps[j]; + if (step.type !== 'TOOL_EXECUTION') continue; + + const toolName = step.toolName; + if (toolName && UNMASKABLE_TOOLS.has(toolName)) 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 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 }; + } + 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(), + }, + ], + }; + }); + } + } + } + } + } + + private isAlreadyMasked(content: string): boolean { + return content.includes(''); + } +} diff --git a/packages/core/src/context/sidecar/SidecarLoader.test.ts b/packages/core/src/context/sidecar/SidecarLoader.test.ts new file mode 100644 index 0000000000..88add76c20 --- /dev/null +++ b/packages/core/src/context/sidecar/SidecarLoader.test.ts @@ -0,0 +1,67 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ProcessorRegistry } from './registry.js'; +import { registerBuiltInProcessors } from './builtins.js'; +import { describe, it, expect, beforeEach } from 'vitest'; +import { SidecarLoader } from './SidecarLoader.js'; +import { defaultSidecarProfile } from './profiles.js'; +import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; +import type { Config } from 'src/config/config.js'; + +describe('SidecarLoader (Fake FS)', () => { + let fileSystem: InMemoryFileSystem; + let registry: ProcessorRegistry; + + beforeEach(() => { + fileSystem = new InMemoryFileSystem(); + registry = new ProcessorRegistry(); + registerBuiltInProcessors(registry); + }); + + const mockConfig = { + getExperimentalContextSidecarConfig: () => '/path/to/sidecar.json', + } as unknown as Config; + + it('returns default profile if file does not exist', () => { + const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem); + expect(result).toBe(defaultSidecarProfile); + }); + + it('returns default profile if file exists but is 0 bytes', () => { + fileSystem.setFile('/path/to/sidecar.json', ''); + const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem); + expect(result).toBe(defaultSidecarProfile); + }); + + it('throws an error if file is empty whitespace', () => { + fileSystem.setFile('/path/to/sidecar.json', ' \n '); + expect(() => + SidecarLoader.fromConfig(mockConfig, registry, fileSystem), + ).toThrow('is empty'); + }); + + it('returns parsed config if file is valid', () => { + const validConfig = { + budget: { retainedTokens: 1000, maxTokens: 2000 }, + gcBackstop: { strategy: 'truncate', target: 'max' }, + pipelines: [], + }; + fileSystem.setFile('/path/to/sidecar.json', JSON.stringify(validConfig)); + const result = SidecarLoader.fromConfig(mockConfig, registry, fileSystem); + expect(result.budget.maxTokens).toBe(2000); + }); + + it('throws validation error if file is invalid', () => { + const invalidConfig = { + budget: { retainedTokens: 1000 }, // missing maxTokens + }; + fileSystem.setFile('/path/to/sidecar.json', JSON.stringify(invalidConfig)); + expect(() => + SidecarLoader.fromConfig(mockConfig, registry, fileSystem), + ).toThrow('Validation error:'); + }); +}); diff --git a/packages/core/src/context/sidecar/SidecarLoader.ts b/packages/core/src/context/sidecar/SidecarLoader.ts new file mode 100644 index 0000000000..87f4dce8b6 --- /dev/null +++ b/packages/core/src/context/sidecar/SidecarLoader.ts @@ -0,0 +1,87 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Config } from '../../config/config.js'; +import type { SidecarConfig } from './types.js'; +import { defaultSidecarProfile } from './profiles.js'; +import { SchemaValidator } from '../../utils/schemaValidator.js'; +import { getSidecarConfigSchema } from './schema.js'; +import type { IFileSystem } from '../system/IFileSystem.js'; +import { NodeFileSystem } from '../system/NodeFileSystem.js'; +import type { ProcessorRegistry } from './registry.js'; + +export class SidecarLoader { + /** + * Loads and validates a sidecar config from a specific file path. + * Throws an error if the file cannot be read, parsed, or fails schema validation. + */ + static loadFromFile( + sidecarPath: string, + registry: ProcessorRegistry, + fileSystem: IFileSystem = new NodeFileSystem(), + ): SidecarConfig { + const fileContent = fileSystem.readFileSync(sidecarPath, 'utf8'); + + if (!fileContent.trim()) { + throw new Error(`Sidecar configuration file at ${sidecarPath} is empty.`); + } + + let parsed: unknown; + try { + parsed = JSON.parse(fileContent); + } catch (error) { + throw new Error( + `Failed to parse Sidecar configuration file at ${sidecarPath}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } + + const validationError = SchemaValidator.validate( + getSidecarConfigSchema(registry), + parsed, + ); + if (validationError) { + throw new Error( + `Invalid sidecar configuration in ${sidecarPath}. Validation error: ${validationError}`, + ); + } + + // Schema has been validated. + const isSidecarConfig = (val: unknown): val is SidecarConfig => true; + if (isSidecarConfig(parsed)) { + return parsed; + } + throw new Error( + 'Unreachable: schema validation passed but type predicate failed.', + ); + } + + /** + * Generates a Sidecar JSON graph from the experimental config file path or defaults. + * If a config file is present but invalid, this will THROW to prevent silent misconfiguration. + */ + static fromConfig( + config: Config, + registry: ProcessorRegistry, + fileSystem: IFileSystem = new NodeFileSystem(), + ): SidecarConfig { + const sidecarPath = config.getExperimentalContextSidecarConfig(); + + if (sidecarPath && fileSystem.existsSync(sidecarPath)) { + const size = fileSystem.statSyncSize(sidecarPath); + // If the file exists but is completely empty (0 bytes), it's safe to fallback. + if (size === 0) { + return defaultSidecarProfile; + } + + // If the file has content, enforce strict validation and throw on failure. + return this.loadFromFile(sidecarPath, registry, fileSystem); + } + + return defaultSidecarProfile; + } +} diff --git a/packages/core/src/context/sidecar/builtins.ts b/packages/core/src/context/sidecar/builtins.ts new file mode 100644 index 0000000000..7609b46567 --- /dev/null +++ b/packages/core/src/context/sidecar/builtins.ts @@ -0,0 +1,127 @@ +/** + * @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'; + +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: { + type: 'object', + properties: { + processorId: { const: 'BlobDegradationProcessor' }, + options: { type: 'object' }, + }, + required: ['processorId'], + }, + create: (env) => new BlobDegradationProcessor(env), + }); + + registry.register({ + id: 'SemanticCompressionProcessor', + schema: { + type: 'object', + properties: { + processorId: { const: 'SemanticCompressionProcessor' }, + options: { + type: 'object', + properties: { nodeThresholdTokens: { type: 'number' } }, + required: ['nodeThresholdTokens'], + }, + }, + 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'], + }, + create: (env, opts) => new HistorySquashingProcessor(env, opts), + }); + + registry.register({ + id: 'StateSnapshotProcessor', + schema: { + type: 'object', + properties: { + processorId: { const: 'StateSnapshotProcessor' }, + options: { + type: 'object', + properties: { + model: { type: 'string' }, + systemInstruction: { type: 'string' }, + triggerDeficitTokens: { type: 'number' }, + }, + }, + }, + required: ['processorId'], + }, + create: (env, opts) => StateSnapshotProcessor.create(env, opts), + }); + + registry.register({ + id: 'EmergencyTruncationProcessor', + schema: { + type: 'object', + properties: { + processorId: { const: 'EmergencyTruncationProcessor' }, + options: { type: 'object' }, + }, + required: ['processorId'], + }, + create: (env, opts) => EmergencyTruncationProcessor.create(env, opts), + }); +} diff --git a/packages/core/src/context/sidecar/environment.ts b/packages/core/src/context/sidecar/environment.ts new file mode 100644 index 0000000000..ee66ec13d0 --- /dev/null +++ b/packages/core/src/context/sidecar/environment.ts @@ -0,0 +1,27 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ +import type { BaseLlmClient } from '../../core/baseLlmClient.js'; +import type { ContextEventBus } from '../eventBus.js'; +import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; +import type { ContextTracer } from '../tracer.js'; +import type { IFileSystem } from '../system/IFileSystem.js'; +import type { IIdGenerator } from '../system/IIdGenerator.js'; + +export type { ContextTracer, ContextEventBus }; + +export interface ContextEnvironment { + readonly llmClient: BaseLlmClient; + readonly promptId: string; + readonly sessionId: string; + readonly traceDir: string; + readonly projectTempDir: string; + readonly tracer: ContextTracer; + readonly charsPerToken: number; + readonly tokenCalculator: ContextTokenCalculator; + readonly fileSystem: IFileSystem; + readonly idGenerator: IIdGenerator; + readonly eventBus: ContextEventBus; +} diff --git a/packages/core/src/context/sidecar/environmentImpl.ts b/packages/core/src/context/sidecar/environmentImpl.ts new file mode 100644 index 0000000000..0987e317de --- /dev/null +++ b/packages/core/src/context/sidecar/environmentImpl.ts @@ -0,0 +1,38 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { BaseLlmClient } from '../../core/baseLlmClient.js'; +import type { ContextTracer } from '../tracer.js'; +import type { ContextEnvironment } from './environment.js'; +import type { ContextEventBus } from '../eventBus.js'; +import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; +import type { IFileSystem } from '../system/IFileSystem.js'; +import { NodeFileSystem } from '../system/NodeFileSystem.js'; +import type { IIdGenerator } from '../system/IIdGenerator.js'; +import { NodeIdGenerator } from '../system/NodeIdGenerator.js'; + +export class ContextEnvironmentImpl implements ContextEnvironment { + readonly tokenCalculator: ContextTokenCalculator; + readonly fileSystem: IFileSystem; + readonly idGenerator: IIdGenerator; + + constructor( + readonly llmClient: BaseLlmClient, + readonly sessionId: string, + readonly promptId: string, + readonly traceDir: string, + readonly projectTempDir: string, + readonly tracer: ContextTracer, + readonly charsPerToken: number, + readonly eventBus: ContextEventBus, + fileSystem?: IFileSystem, + idGenerator?: IIdGenerator, + ) { + this.tokenCalculator = new ContextTokenCalculator(this.charsPerToken); + this.fileSystem = fileSystem || new NodeFileSystem(); + this.idGenerator = idGenerator || new NodeIdGenerator(); + } +} diff --git a/packages/core/src/context/sidecar/orchestrator.test.ts b/packages/core/src/context/sidecar/orchestrator.test.ts new file mode 100644 index 0000000000..3ecd342263 --- /dev/null +++ b/packages/core/src/context/sidecar/orchestrator.test.ts @@ -0,0 +1,292 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest'; +import { PipelineOrchestrator } from './orchestrator.js'; +import { ProcessorRegistry } from './registry.js'; +import { + createMockEnvironment, + createDummyState, + createDummyEpisode, +} 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 { + static create() { + return new DummySyncProcessor(); + } + constructor() {} + readonly name = 'DummySync'; + readonly id = 'DummySync'; + readonly options = {}; + async process(editor: EpisodeEditor, _state: ContextAccountingState) { + editor.editEpisode( + editor.episodes[0].id, + 'DUMMY_EDIT', + (draft: unknown) => { + (draft as Record)['dummyModified'] = true; + }, + ); + } +} + +class DummyAsyncProcessor implements ContextProcessor { + static create() { + return new DummyAsyncProcessor(); + } + constructor() {} + readonly name = 'DummyAsync'; + readonly id = 'DummyAsync'; + readonly options = {}; + async process(editor: EpisodeEditor, _state: ContextAccountingState) { + editor.editEpisode( + editor.episodes[0].id, + 'DUMMY_EDIT', + (draft: unknown) => { + (draft as Record)['dummyAsyncModified'] = true; + }, + ); + } +} + +class ThrowingProcessor implements ContextProcessor { + static create() { + return new ThrowingProcessor(); + } + constructor() {} + readonly name = 'Throwing'; + readonly id = 'Throwing'; + readonly options = {}; + async process( + _editor: EpisodeEditor, + _state: ContextAccountingState, + ): Promise { + throw new Error('Processor failed intentionally'); + } +} + +describe('PipelineOrchestrator (Component)', () => { + let env: ContextEnvironment; + let eventBus: ContextEventBus; + let registry: ProcessorRegistry; + + beforeEach(() => { + vi.resetAllMocks(); + env = createMockEnvironment(); + eventBus = env.eventBus; + registry = new ProcessorRegistry(); + + // Register our test processors + registry.register({ + id: 'DummySyncProcessor', + schema: {}, + create: () => new DummySyncProcessor(), + }); + registry.register({ + id: 'DummyAsyncProcessor', + schema: {}, + create: () => new DummyAsyncProcessor(), + }); + registry.register({ + id: 'ThrowingProcessor', + schema: {}, + create: () => new ThrowingProcessor(), + }); + }); + + afterEach(() => { + // Cleanup registry to not pollute other tests + registry.clear(); + }); + + const createConfig = (pipelines: PipelineDef[]): SidecarConfig => ({ + budget: { maxTokens: 100, retainedTokens: 50 }, + gcBackstop: { strategy: 'truncate', target: 'max' }, + pipelines, + }); + + it('instantiates processors from the registry on initialization', () => { + const config = createConfig([ + { + name: 'Sync', + execution: 'blocking', + triggers: [], + processors: [ + { processorId: 'DummySyncProcessor' } as unknown as ProcessorConfig, + ], + }, + ]); + + const orchestrator = new PipelineOrchestrator( + config, + env, + eventBus, + env.tracer, + registry, + ); + expect( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (orchestrator as any).instantiatedProcessors.has('DummySyncProcessor'), + ).toBe(true); + }); + + it('throws an error if a config requests an unknown processor', () => { + const config = createConfig([ + { + name: 'Bad', + execution: 'blocking', + triggers: [], + processors: [ + { processorId: 'DoesNotExist' } as unknown as ProcessorConfig, + ], + }, + ]); + + expect( + () => + new PipelineOrchestrator(config, env, eventBus, env.tracer, registry), + ).toThrow('Context Processor [DoesNotExist] is not registered.'); + }); + + it('executes blocking pipelines synchronously and returns the modified array', async () => { + const config = createConfig([ + { + name: 'SyncPipe', + execution: 'blocking', + triggers: [], + processors: [ + { processorId: 'DummySyncProcessor' } as unknown as ProcessorConfig, + ], + }, + ]); + const orchestrator = new PipelineOrchestrator( + config, + env, + eventBus, + env.tracer, + registry, + ); + + const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])]; + const state = createDummyState(false); + + const result = await orchestrator.executePipeline( + 'SyncPipe', + episodes, + state, + ); + + expect(result).toHaveLength(1); + expect( + (result[0] as unknown as { dummyModified: boolean }).dummyModified, + ).toBe(true); + }); + + it('executes background pipelines asynchronously without blocking the return', async () => { + const config = createConfig([ + { + name: 'AsyncPipe', + execution: 'background', + triggers: [], + processors: [ + { processorId: 'DummyAsyncProcessor' } as unknown as ProcessorConfig, + ], + }, + ]); + const orchestrator = new PipelineOrchestrator( + config, + env, + eventBus, + env.tracer, + registry, + ); + + const episodes = [createDummyEpisode('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', + episodes, + state, + ); + + expect(result).toHaveLength(1); + expect( + (result[0] as unknown as { asyncModified: unknown }).asyncModified, + ).toBeUndefined(); // Not modified yet! + + // Wait for the background task to complete (50ms delay in DummyAsyncProcessor) + await new Promise((resolve) => setTimeout(resolve, 60)); + }); + + it('gracefully handles and swallows processor errors in synchronous pipelines', async () => { + const config = createConfig([ + { + name: 'ThrowingPipe', + execution: 'blocking', + triggers: [], + processors: [ + { processorId: 'ThrowingProcessor' } as unknown as ProcessorConfig, + ], + }, + ]); + const orchestrator = new PipelineOrchestrator( + config, + env, + eventBus, + env.tracer, + registry, + ); + + const episodes = [createDummyEpisode('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', + episodes, + state, + ); + + expect(result).toHaveLength(1); + expect(result).toStrictEqual(episodes); + }); + + it('automatically binds to budget_exceeded trigger via EventBus', () => { + const config = createConfig([ + { + name: 'PressureRelief', + execution: 'background', + triggers: ['budget_exceeded'], + processors: [ + { processorId: 'DummyAsyncProcessor' } as unknown as ProcessorConfig, + ], + }, + ]); + + // Spy on the private method to see if the trigger fires it + const executeSpy = vi.spyOn( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + PipelineOrchestrator.prototype as any, + 'executePipelineAsync', + ); + + new PipelineOrchestrator(config, env, eventBus, env.tracer, registry); + + const episodes = [createDummyEpisode('1', 'USER_PROMPT', [])]; + + // Emit the trigger + eventBus.emitConsolidationNeeded({ episodes, targetDeficit: 100 }); + + expect(executeSpy).toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/context/sidecar/orchestrator.ts b/packages/core/src/context/sidecar/orchestrator.ts new file mode 100644 index 0000000000..10905f0015 --- /dev/null +++ b/packages/core/src/context/sidecar/orchestrator.ts @@ -0,0 +1,234 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Episode } from '../ir/types.js'; +import type { ContextProcessor, ContextAccountingState } from '../pipeline.js'; +import type { SidecarConfig, PipelineDef } from './types.js'; +import type { + ContextEnvironment, + ContextEventBus, + ContextTracer, +} from './environment.js'; +import type { ProcessorRegistry } from './registry.js'; +import { debugLogger } from '../../utils/debugLogger.js'; +import { EpisodeEditor } from '../ir/episodeEditor.js'; + +export class PipelineOrchestrator { + private activeTimers: NodeJS.Timeout[] = []; + private readonly instantiatedProcessors = new Map(); + + constructor( + private readonly config: SidecarConfig, + private readonly env: ContextEnvironment, + private readonly eventBus: ContextEventBus, + private readonly tracer: ContextTracer, + private readonly registry: ProcessorRegistry, + ) { + this.instantiateProcessors(); + this.registerTriggers(); + } + + /** + * Pre-loads and configures all processors defined in the sidecar config. + */ + private instantiateProcessors() { + for (const pipeline of this.config.pipelines) { + for (const procDef of pipeline.processors) { + if (!this.instantiatedProcessors.has(procDef.processorId)) { + const processorClass = this.registry.get(procDef.processorId); + if (!processorClass) { + throw new Error( + `Context Processor [${procDef.processorId}] is not registered.`, + ); + } + // The Orchestrator injects standard dependencies required by processors + // If a processor needs the eventBus (like Snapshot), it expects it via constructor. + const instance = processorClass.create( + this.env, + procDef.options ?? {}, + ); + this.instantiatedProcessors.set(procDef.processorId, instance); + } + } + } + } + + /** + * Sets up listeners for the triggers defined in the SidecarConfig. + */ + private registerTriggers() { + for (const pipeline of this.config.pipelines) { + for (const trigger of pipeline.triggers) { + if (typeof trigger === 'object' && trigger.type === 'timer') { + const timer = setInterval(() => { + // For background timers, we need a way to get the latest state + // But timers are generally disabled right now via the triggers config. + // If needed, we will pass it via event bus. + }, trigger.intervalMs); + this.activeTimers.push(timer); + } else if (trigger === 'budget_exceeded') { + this.eventBus.onConsolidationNeeded((event) => { + const state: ContextAccountingState = { + currentTokens: 0, + retainedTokens: this.config.budget.retainedTokens, + maxTokens: this.config.budget.maxTokens, + isBudgetSatisfied: false, + deficitTokens: event.targetDeficit, + protectedEpisodeIds: new Set(), + }; + void this.executePipelineAsync(pipeline, event.episodes, state); + }); + } + } + } + } + + shutdown() { + for (const timer of this.activeTimers) { + clearInterval(timer); + } + } + + /** + * Executes a pipeline asynchronously in the background. This is the "Eventual Consistency" path. + * When the pipeline resolves, it emits a VariantReady event to cache the new graph. + */ + /** + * Executes a pipeline based on its configured execution strategy ('blocking' or 'background'). + */ + async executePipeline( + pipelineName: string, + episodes: Episode[], + state: ContextAccountingState, + ): Promise { + const pipeline = this.config.pipelines.find((p) => p.name === pipelineName); + if (!pipeline) return episodes; + + if (pipeline.execution === 'background') { + this.executePipelineAsync(pipeline, episodes, state).catch((e) => { + debugLogger.error(`Background pipeline ${pipeline.name} failed:`, e); + }); + return episodes; // Return immediately + } + + // Blocking execution + this.tracer.logEvent( + 'Orchestrator', + `Triggering synchronous pipeline: ${pipeline.name}`, + ); + let currentEpisodes = [...episodes]; + for (let i = 0; i < pipeline.processors.length; i++) { + const procDef = pipeline.processors[i]; + const processor = this.instantiatedProcessors.get(procDef.processorId); + if (!processor) continue; + + try { + this.tracer.logEvent( + 'Orchestrator', + `Executing processor: ${procDef.processorId}`, + ); + const editor = new EpisodeEditor(currentEpisodes); + await processor.process(editor, state); + currentEpisodes = editor.getFinalEpisodes(); + } catch (error) { + debugLogger.error( + `Pipeline ${pipeline.name} failed synchronously at ${procDef.processorId}:`, + error, + ); + return currentEpisodes; // Return what we have so far + } + } + + return currentEpisodes; + } + + /** + * Internal method for running a pipeline entirely in the background. + */ + private async executePipelineAsync( + pipeline: PipelineDef, + currentState: Episode[], + state: ContextAccountingState, + ) { + this.tracer.logEvent( + 'Orchestrator', + `Triggering async pipeline: ${pipeline.name}`, + ); + if (!currentState || currentState.length === 0) return; + + let currentEpisodes = [...currentState]; + + for (const procDef of pipeline.processors) { + const processor = this.instantiatedProcessors.get(procDef.processorId); + if (!processor) continue; + + try { + this.tracer.logEvent( + 'Orchestrator', + `Executing processor: ${procDef.processorId} (async)`, + ); + + const editor = new EpisodeEditor(currentEpisodes); + await processor.process(editor, state); + currentEpisodes = editor.getFinalEpisodes(); + + // Synthesize VariantReady events for anything that changed or was newly created + for (const mutation of editor.getMutations()) { + // We only broadcast modifications or replacements + // (Insertions without replacement and deletions are not tracked as variants on an existing node) + if (mutation.type === 'modified' || mutation.type === 'replaced') { + const variantId = `v-${procDef.processorId.toLowerCase()}`; + + let vType: 'snapshot' | 'summary' | 'masked' = 'masked'; + if (procDef.processorId.includes('Snapshot')) vType = 'snapshot'; + else if (procDef.processorId.includes('Semantic')) + vType = 'summary'; + + const ep = mutation.episode!; + let fallbackText = ''; + if (ep.yield?.text) fallbackText = ep.yield.text; + else if (ep.trigger?.type === 'USER_PROMPT') { + const firstPart = ep.trigger.semanticParts?.[0]; + if (firstPart) { + fallbackText = + firstPart.type === 'text' + ? firstPart.presentation?.text || firstPart.text + : ''; + } + } + + this.eventBus.emitVariantReady({ + targetId: + mutation.type === 'replaced' ? mutation.originalIds![0] : ep.id, + variantId, + variant: + vType === 'snapshot' + ? { + status: 'ready', + type: 'snapshot', + episode: ep, + recoveredTokens: ep.yield?.metadata?.currentTokens || 10, + replacedEpisodeIds: mutation.originalIds || [], + } + : { + status: 'ready', + type: vType, + text: fallbackText, + recoveredTokens: ep.yield?.metadata?.currentTokens || 10, + }, + }); + } + } + } catch (error) { + debugLogger.error( + `Pipeline ${pipeline.name} failed at ${procDef.processorId}:`, + error, + ); + return; // Halt pipeline + } + } + } +} diff --git a/packages/core/src/context/sidecar/profiles.ts b/packages/core/src/context/sidecar/profiles.ts new file mode 100644 index 0000000000..f5012319f7 --- /dev/null +++ b/packages/core/src/context/sidecar/profiles.ts @@ -0,0 +1,54 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { SidecarConfig } from './types.js'; + +/** + * The standard default context management profile. + * Optimized for safety, precision, and reliable summarization. + */ +export const defaultSidecarProfile: SidecarConfig = { + budget: { + retainedTokens: 65000, + maxTokens: 150000, + }, + gcBackstop: { + strategy: 'truncate', + target: 'incremental', + freeTokensTarget: 10000, + }, + pipelines: [ + { + name: 'Immediate Sanitization', + triggers: ['on_turn'], + execution: 'blocking', + processors: [ + { + processorId: 'ToolMaskingProcessor', + options: { stringLengthThresholdTokens: 8000 }, + }, + { processorId: 'BlobDegradationProcessor', options: {} }, + { + processorId: 'SemanticCompressionProcessor', + options: { nodeThresholdTokens: 5000 }, + }, + { processorId: 'EmergencyTruncationProcessor', options: {} }, + ], + }, + { + name: 'Deep Background Compression', + triggers: [{ type: 'timer', intervalMs: 5000 }, 'budget_exceeded'], + execution: 'background', + processors: [ + { + processorId: 'HistorySquashingProcessor', + options: { maxTokensPerNode: 3000 }, + }, + { processorId: 'StateSnapshotProcessor', options: {} }, + ], + }, + ], +}; diff --git a/packages/core/src/context/sidecar/registry.ts b/packages/core/src/context/sidecar/registry.ts new file mode 100644 index 0000000000..6010ded765 --- /dev/null +++ b/packages/core/src/context/sidecar/registry.ts @@ -0,0 +1,47 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextProcessor } from '../pipeline.js'; +import type { ContextEnvironment } from './environment.js'; + +export interface ContextProcessorDef { + readonly id: string; + readonly schema: object; + create(env: ContextEnvironment, options: TOptions): ContextProcessor; +} + +/** + * Registry for mapping declarative sidecar configs to running Processor instances. + */ +export class ProcessorRegistry { + private processors = new Map>(); + + register(def: ContextProcessorDef) { + this.processors.set(def.id, def); + } + + get(id: string): ContextProcessorDef { + const def = this.processors.get(id); + if (!def) { + throw new Error(`Context Processor [${id}] is not registered.`); + } + return def; + } + + getSchemas(): object[] { + const schemas: object[] = []; + for (const def of this.processors.values()) { + if (def.schema) { + schemas.push(def.schema); + } + } + return schemas; + } + + clear() { + this.processors.clear(); + } +} diff --git a/packages/core/src/context/sidecar/schema.ts b/packages/core/src/context/sidecar/schema.ts new file mode 100644 index 0000000000..6f1efd57e7 --- /dev/null +++ b/packages/core/src/context/sidecar/schema.ts @@ -0,0 +1,103 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ProcessorRegistry } from './registry.js'; +import './builtins.js'; + +export function getSidecarConfigSchema(registry: ProcessorRegistry) { + return { + $schema: 'http://json-schema.org/draft-07/schema#', + title: 'SidecarConfig', + description: 'The Data-Driven Schema for the Context Manager.', + type: 'object', + required: ['budget', 'gcBackstop', 'pipelines'], + properties: { + budget: { + type: 'object', + description: 'Defines the token ceilings and limits for the pipeline.', + required: ['retainedTokens', 'maxTokens'], + properties: { + retainedTokens: { + type: 'number', + description: + 'The ideal token count the pipeline tries to shrink down to.', + }, + maxTokens: { + type: 'number', + description: + 'The absolute maximum token count allowed before synchronous truncation kicks in.', + }, + }, + }, + gcBackstop: { + type: 'object', + description: + "Defines what happens when the pipeline fails to compress under 'maxTokens'", + required: ['strategy', 'target'], + properties: { + strategy: { + type: 'string', + enum: ['truncate', 'compress', 'rollingSummarizer'], + }, + target: { + type: 'string', + enum: ['incremental', 'freeNTokens', 'max'], + }, + freeTokensTarget: { + type: 'number', + }, + }, + }, + pipelines: { + type: 'array', + description: 'The execution graphs for context manipulation.', + items: { + type: 'object', + required: ['name', 'triggers', 'execution', 'processors'], + properties: { + name: { + type: 'string', + }, + triggers: { + type: 'array', + items: { + anyOf: [ + { + type: 'string', + enum: ['on_turn', 'post_turn', 'budget_exceeded'], + }, + { + type: 'object', + required: ['type', 'intervalMs'], + properties: { + type: { + type: 'string', + const: 'timer', + }, + intervalMs: { + type: 'number', + }, + }, + }, + ], + }, + }, + execution: { + type: 'string', + enum: ['blocking', 'background'], + }, + processors: { + type: 'array', + items: { + oneOf: registry.getSchemas(), + }, + }, + }, + }, + }, + }, + }; +} diff --git a/packages/core/src/context/sidecar/types.ts b/packages/core/src/context/sidecar/types.ts new file mode 100644 index 0000000000..19e7a4f74a --- /dev/null +++ b/packages/core/src/context/sidecar/types.ts @@ -0,0 +1,67 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { StateSnapshotProcessorOptions } from '../processors/stateSnapshotProcessor.js'; + +/** + * Definition of a processor or worker to be instantiated in the graph. + */ +export type ProcessorConfig = + | { + processorId: 'ToolMaskingProcessor'; + options: { stringLengthThresholdTokens: number }; + } + | { processorId: 'BlobDegradationProcessor'; options?: object } + | { + processorId: 'SemanticCompressionProcessor'; + options: { nodeThresholdTokens: number }; + } + | { + processorId: 'HistorySquashingProcessor'; + options: { maxTokensPerNode: number }; + } + | { + processorId: 'StateSnapshotProcessor'; + options: StateSnapshotProcessorOptions; + } + | { + processorId: 'EmergencyTruncationProcessor'; + options?: Record; + }; + +export type PipelineTrigger = + | 'on_turn' + | 'post_turn' + | 'budget_exceeded' + | { type: 'timer'; intervalMs: number }; + +export interface PipelineDef { + name: string; + triggers: PipelineTrigger[]; + execution: 'blocking' | 'background'; + processors: ProcessorConfig[]; +} + +/** + * The Data-Driven Schema for the Context Manager. + */ +export interface SidecarConfig { + /** Defines the token ceilings and limits for the pipeline. */ + budget: { + retainedTokens: number; + maxTokens: number; + }; + + /** Defines what happens when the pipeline fails to compress under 'maxTokens' */ + gcBackstop: { + strategy: 'truncate' | 'compress' | 'rollingSummarizer'; + target: 'incremental' | 'freeNTokens' | 'max'; + freeTokensTarget?: number; + }; + + /** The execution graphs for context manipulation */ + pipelines: PipelineDef[]; +} diff --git a/packages/core/src/context/system-tests/SimulationHarness.ts b/packages/core/src/context/system-tests/SimulationHarness.ts new file mode 100644 index 0000000000..65d5feb896 --- /dev/null +++ b/packages/core/src/context/system-tests/SimulationHarness.ts @@ -0,0 +1,189 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ContextManager } from '../contextManager.js'; +import { AgentChatHistory } from '../../core/agentChatHistory.js'; +import type { Content } from '@google/genai'; +import type { SidecarConfig } from '../sidecar/types.js'; +import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js'; +import { ContextTracer } from '../tracer.js'; +import { ContextEventBus } from '../eventBus.js'; +import { PipelineOrchestrator } from '../sidecar/orchestrator.js'; +import { registerBuiltInProcessors } from '../sidecar/builtins.js'; +import { debugLogger } from '../../utils/debugLogger.js'; +import { ProcessorRegistry } from '../sidecar/registry.js'; +import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js'; +import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; +import type { BaseLlmClient } from '../../core/baseLlmClient.js'; + +export interface TurnSummary { + turnIndex: number; + tokensBeforeBackground: number; + tokensAfterBackground: number; +} + +export class SimulationHarness { + readonly chatHistory: AgentChatHistory; + contextManager!: ContextManager; + env!: ContextEnvironmentImpl; + orchestrator!: PipelineOrchestrator; + readonly eventBus: ContextEventBus; + config!: SidecarConfig; + private tracer!: ContextTracer; + private currentTurnIndex = 0; + private tokenTrajectory: TurnSummary[] = []; + + static async create( + config: SidecarConfig, + mockLlmClient: BaseLlmClient, + mockTempDir = '/tmp/sim', + ): Promise { + const harness = new SimulationHarness(); + await harness.init(config, mockLlmClient, mockTempDir); + return harness; + } + + private constructor() { + this.chatHistory = new AgentChatHistory(); + this.eventBus = new ContextEventBus(); + } + + private async init( + config: SidecarConfig, + mockLlmClient: BaseLlmClient, + mockTempDir: string, + ) { + this.config = config; + const registry = new ProcessorRegistry(); + // Register all standard processors + registerBuiltInProcessors(registry); + + this.tracer = new ContextTracer({ + targetDir: mockTempDir, + sessionId: 'sim-session', + }); + this.env = new ContextEnvironmentImpl( + mockLlmClient, + 'sim-prompt', + 'sim-session', + mockTempDir, + mockTempDir, + this.tracer, + 4, // 4 chars per token average + this.eventBus, + new InMemoryFileSystem(), + new DeterministicIdGenerator(), + ); + + this.orchestrator = new PipelineOrchestrator( + config, + this.env, + this.eventBus, + this.tracer, + registry, + ); + this.contextManager = ContextManager.create( + config, + this.env, + this.tracer, + this.orchestrator, + registry, + ); + this.contextManager.subscribeToHistory(this.chatHistory); + } + + /** + * Simulates a single "Turn" (User input + Model/Tool outputs) + * A turn might consist of multiple Content messages (e.g. user prompt -> model call -> user response -> model answer) + */ + async simulateTurn(messages: Content[]) { + // 1. Append the new messages + const currentHistory = this.chatHistory.get(); + this.chatHistory.set([...currentHistory, ...messages]); + + // 2. Measure tokens immediately after append (Before background processing) + const tokensBefore = this.env.tokenCalculator.calculateEpisodeListTokens( + this.contextManager.getWorkingBufferView(), + ); + debugLogger.log( + `[Turn ${this.currentTurnIndex}] Tokens BEFORE: ${tokensBefore}`, + ); + + // 3. Yield to event loop to allow internal async subscribers and orchestrator to finish + await new Promise((resolve) => setTimeout(resolve, 50)); + + // 3.1 Simulate what projectCompressedHistory does with the sync handlers + let currentView = this.contextManager.getWorkingBufferView(); + const currentTokens = + this.env.tokenCalculator.calculateEpisodeListTokens(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, { + currentTokens, + maxTokens: this.config.budget.maxTokens, + retainedTokens: this.config.budget.retainedTokens, + isBudgetSatisfied: false, + deficitTokens: currentTokens - this.config.budget.maxTokens, + protectedEpisodeIds: 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() + .find((c) => c.id === ep.id) + ) { + this.eventBus.emitVariantReady({ + targetId: ep.id, + variantId: 'v-emergency', + variant: { + status: 'ready', + type: 'masked', // Truncation is technically a mask + text: ep.yield?.text || '', + recoveredTokens: 0, + }, + }); + } + } + // Wait for variant propagation + await new Promise((resolve) => setTimeout(resolve, 50)); + } + + // 4. Measure tokens after background processors have (hopefully) emitted variants + const tokensAfter = this.env.tokenCalculator.calculateEpisodeListTokens( + this.contextManager.getWorkingBufferView(), + ); + debugLogger.log( + `[Turn ${this.currentTurnIndex}] Tokens AFTER: ${tokensAfter}`, + ); + + this.tokenTrajectory.push({ + turnIndex: this.currentTurnIndex++, + tokensBeforeBackground: tokensBefore, + tokensAfterBackground: tokensAfter, + }); + } + + async getGoldenState() { + const finalProjection = + await this.contextManager.projectCompressedHistory(); + return { + tokenTrajectory: this.tokenTrajectory, + finalProjection, + }; + } +} diff --git a/packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap b/packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap new file mode 100644 index 0000000000..cab629a597 --- /dev/null +++ b/packages/core/src/context/system-tests/__snapshots__/lifecycle.golden.test.ts.snap @@ -0,0 +1,89 @@ +// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html + +exports[`System Lifecycle Golden Tests > Scenario 1: Organic Growth with Huge Tool Output & Images 1`] = ` +{ + "finalProjection": [ + { + "parts": [ + { + "text": "System Instructions", + }, + ], + "role": "user", + }, + { + "parts": [ + { + "text": "Ack.", + }, + ], + "role": "model", + }, + { + "parts": [ + { + "text": "Look at this architecture diagram:", + }, + { + "inlineData": { + "data": "fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_fake_base64_data_", + "mimeType": "image/png", + }, + }, + ], + "role": "user", + }, + { + "parts": [ + { + "text": "Nice diagram.", + }, + ], + "role": "model", + }, + { + "parts": [ + { + "text": "Can we refactor?", + }, + ], + "role": "user", + }, + { + "parts": [ + { + "text": "Yes we can.", + }, + ], + "role": "model", + }, + ], + "tokenTrajectory": [ + { + "tokensAfterBackground": 11, + "tokensBeforeBackground": 11, + "turnIndex": 0, + }, + { + "tokensAfterBackground": 23, + "tokensBeforeBackground": 23, + "turnIndex": 1, + }, + { + "tokensAfterBackground": 10067, + "tokensBeforeBackground": 10067, + "turnIndex": 2, + }, + { + "tokensAfterBackground": 13349, + "tokensBeforeBackground": 13349, + "turnIndex": 3, + }, + { + "tokensAfterBackground": 13362, + "tokensBeforeBackground": 13362, + "turnIndex": 4, + }, + ], +} +`; diff --git a/packages/core/src/context/system-tests/lifecycle.golden.test.ts b/packages/core/src/context/system-tests/lifecycle.golden.test.ts new file mode 100644 index 0000000000..fb8dc5b28a --- /dev/null +++ b/packages/core/src/context/system-tests/lifecycle.golden.test.ts @@ -0,0 +1,146 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest'; +import { SimulationHarness } from './SimulationHarness.js'; +import type { SidecarConfig } from '../sidecar/types.js'; +import type { BaseLlmClient } from '../../core/baseLlmClient.js'; + +expect.addSnapshotSerializer({ + test: (val) => + typeof val === 'string' && + (/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i.test( + val, + ) || + /^\/tmp\/sim/.test(val)), // Mask temp directories and UUIDs + print: (val) => + typeof val === 'string' && /^\/tmp\/sim/.test(val) + ? '""' + : '""', +}); + +describe('System Lifecycle Golden Tests', () => { + beforeAll(() => { + vi.spyOn(Math, 'random').mockReturnValue(0.5); + }); + + afterAll(() => { + vi.restoreAllMocks(); + }); + + const getAggressiveConfig = (): SidecarConfig => ({ + budget: { maxTokens: 4000, retainedTokens: 2000 }, // Extremely tight limits + gcBackstop: { strategy: 'truncate', target: 'max' }, + pipelines: [ + { + name: 'Pressure Relief', // Emits from eventBus 'budget_exceeded' + execution: 'background', + triggers: ['budget_exceeded'], + processors: [ + { processorId: 'BlobDegradationProcessor' }, + { + processorId: 'ToolMaskingProcessor', + options: { stringLengthThresholdTokens: 50 }, + }, // Mask any tool string > 200 chars + { processorId: 'StateSnapshotProcessor', options: {} }, // Squash old history + ], + }, + { + name: 'Immediate Sanitization', // The magic string the projector is hardcoded to use + execution: 'blocking', + triggers: ['budget_exceeded'], + processors: [ + { processorId: 'EmergencyTruncationProcessor', options: {} }, + ], + }, + ], + }); + + const mockLlmClient = { + generateContent: vi.fn().mockResolvedValue({ + text: '', + }), + } as unknown as BaseLlmClient; + + it('Scenario 1: Organic Growth with Huge Tool Output & Images', async () => { + const harness = await SimulationHarness.create( + getAggressiveConfig(), + mockLlmClient, + ); + + // Turn 0: System Prompt + await harness.simulateTurn([ + { role: 'user', parts: [{ text: 'System Instructions' }] }, + { role: 'model', parts: [{ text: 'Ack.' }] }, + ]); + + // Turn 1: Normal conversation + await harness.simulateTurn([ + { role: 'user', parts: [{ text: 'Hello!' }] }, + { role: 'model', parts: [{ text: 'Hi, how can I help?' }] }, + ]); + + // Turn 2: Massive Tool Output (Should trigger ToolMaskingProcessor in background) + await harness.simulateTurn([ + { role: 'user', parts: [{ text: 'Read the logs.' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + name: 'run_shell_command', + args: { cmd: 'cat server.log' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + name: 'run_shell_command', + response: { output: 'LOG '.repeat(5000) }, + }, + }, + ], + }, + { role: 'model', parts: [{ text: 'The logs are very long.' }] }, + ]); + + // Turn 3: Multi-modal blob (Should trigger BlobDegradationProcessor) + await harness.simulateTurn([ + { + role: 'user', + parts: [ + { text: 'Look at this architecture diagram:' }, + { + inlineData: { + mimeType: 'image/png', + data: 'fake_base64_data_'.repeat(1000), + }, + }, + ], + }, + { role: 'model', parts: [{ text: 'Nice diagram.' }] }, + ]); + + // Turn 4: More conversation to trigger StateSnapshot + await harness.simulateTurn([ + { role: 'user', parts: [{ text: 'Can we refactor?' }] }, + { role: 'model', parts: [{ text: 'Yes we can.' }] }, + ]); + + // Get final state + const goldenState = await harness.getGoldenState(); + + // In a perfectly functioning opportunistic system, the token trajectory should show + // the massive spikes in Turn 2 and 3 being immediately resolved by the background tasks. + // The final projection should fit neatly under the Max Tokens limit. + + expect(goldenState).toMatchSnapshot(); + }); +}); diff --git a/packages/core/src/context/system/DeterministicIdGenerator.ts b/packages/core/src/context/system/DeterministicIdGenerator.ts new file mode 100644 index 0000000000..ae3dad6d33 --- /dev/null +++ b/packages/core/src/context/system/DeterministicIdGenerator.ts @@ -0,0 +1,18 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { IIdGenerator } from './IIdGenerator.js'; + +export class DeterministicIdGenerator implements IIdGenerator { + private counter = 0; + + constructor(private prefix: string = 'id-') {} + + generateId(): string { + this.counter++; + return `${this.prefix}${this.counter}`; + } +} diff --git a/packages/core/src/context/system/IFileSystem.ts b/packages/core/src/context/system/IFileSystem.ts new file mode 100644 index 0000000000..cfeca1f3ff --- /dev/null +++ b/packages/core/src/context/system/IFileSystem.ts @@ -0,0 +1,20 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export interface IFileSystem { + existsSync(path: string): boolean; + statSyncSize(path: string): number; + readFileSync(path: string, encoding: 'utf8'): string; + writeFileSync(path: string, data: string | Buffer, encoding?: 'utf-8'): void; + appendFileSync(path: string, data: string, encoding: 'utf-8'): void; + mkdirSync(path: string, options?: { recursive?: boolean }): void; + + writeFile(path: string, data: string | Buffer): Promise; + mkdir(path: string, options?: { recursive?: boolean }): Promise; + + join(...paths: string[]): string; + dirname(path: string): string; +} diff --git a/packages/core/src/context/system/IIdGenerator.ts b/packages/core/src/context/system/IIdGenerator.ts new file mode 100644 index 0000000000..2f5cb38449 --- /dev/null +++ b/packages/core/src/context/system/IIdGenerator.ts @@ -0,0 +1,9 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +export interface IIdGenerator { + generateId(): string; +} diff --git a/packages/core/src/context/system/InMemoryFileSystem.ts b/packages/core/src/context/system/InMemoryFileSystem.ts new file mode 100644 index 0000000000..b407ae31f5 --- /dev/null +++ b/packages/core/src/context/system/InMemoryFileSystem.ts @@ -0,0 +1,79 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { IFileSystem } from './IFileSystem.js'; + +export class InMemoryFileSystem implements IFileSystem { + private files = new Map(); + + getFiles(): ReadonlyMap { + return this.files; + } + + setFile(path: string, content: string | Buffer) { + this.files.set(this.normalize(path), content); + } + + private normalize(p: string): string { + return p.replace(/\/+/g, '/'); + } + + existsSync(p: string): boolean { + return this.files.has(this.normalize(p)); + } + + statSyncSize(p: string): number { + const content = this.files.get(this.normalize(p)); + if (content === undefined) { + throw new Error(`ENOENT: no such file or directory, stat '${p}'`); + } + return Buffer.isBuffer(content) + ? content.byteLength + : Buffer.byteLength(content, 'utf8'); + } + + readFileSync(p: string, encoding: 'utf8'): string { + const content = this.files.get(this.normalize(p)); + if (content === undefined) { + throw new Error(`ENOENT: no such file or directory, open '${p}'`); + } + if (Buffer.isBuffer(content)) { + return content.toString(encoding); + } + return content; + } + + writeFileSync(p: string, data: string | Buffer, _encoding?: 'utf-8'): void { + this.files.set(this.normalize(p), data); + } + + appendFileSync(p: string, data: string, _encoding: 'utf-8'): void { + const norm = this.normalize(p); + const existing = this.files.get(norm) || ''; + const existingStr = Buffer.isBuffer(existing) + ? existing.toString('utf8') + : existing; + this.files.set(norm, existingStr + data); + } + + mkdirSync(_p: string, _options?: { recursive?: boolean }): void {} + + async writeFile(p: string, data: string | Buffer): Promise { + this.writeFileSync(p, data); + } + + async mkdir(_p: string, _options?: { recursive?: boolean }): Promise {} + + join(...paths: string[]): string { + return this.normalize(paths.join('/')); + } + + dirname(p: string): string { + const parts = this.normalize(p).split('/'); + parts.pop(); + return parts.length === 0 ? '.' : parts.join('/') || '/'; + } +} diff --git a/packages/core/src/context/system/NodeFileSystem.ts b/packages/core/src/context/system/NodeFileSystem.ts new file mode 100644 index 0000000000..a2d71c468c --- /dev/null +++ b/packages/core/src/context/system/NodeFileSystem.ts @@ -0,0 +1,56 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs'; +import * as fsPromises from 'node:fs/promises'; +import * as path from 'node:path'; +import type { IFileSystem } from './IFileSystem.js'; + +export class NodeFileSystem implements IFileSystem { + existsSync(p: string): boolean { + return fs.existsSync(p); + } + + statSyncSize(p: string): number { + return fs.statSync(p).size; + } + + readFileSync(p: string, encoding: 'utf8'): string { + return fs.readFileSync(p, encoding); + } + + writeFileSync(p: string, data: string | Buffer, encoding?: 'utf-8'): void { + if (Buffer.isBuffer(data)) { + fs.writeFileSync(p, data); + } else { + fs.writeFileSync(p, data, encoding); + } + } + + appendFileSync(p: string, data: string, encoding: 'utf-8'): void { + fs.appendFileSync(p, data, encoding); + } + + mkdirSync(p: string, options?: { recursive?: boolean }): void { + fs.mkdirSync(p, options); + } + + async writeFile(p: string, data: string | Buffer): Promise { + await fsPromises.writeFile(p, data); + } + + async mkdir(p: string, options?: { recursive?: boolean }): Promise { + await fsPromises.mkdir(p, options); + } + + join(...paths: string[]): string { + return path.join(...paths); + } + + dirname(p: string): string { + return path.dirname(p); + } +} diff --git a/packages/core/src/context/system/NodeIdGenerator.ts b/packages/core/src/context/system/NodeIdGenerator.ts new file mode 100644 index 0000000000..540ec673ba --- /dev/null +++ b/packages/core/src/context/system/NodeIdGenerator.ts @@ -0,0 +1,14 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { randomUUID } from 'node:crypto'; +import type { IIdGenerator } from './IIdGenerator.js'; + +export class NodeIdGenerator implements IIdGenerator { + generateId(): string { + return randomUUID(); + } +} diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts new file mode 100644 index 0000000000..fa15f61a05 --- /dev/null +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -0,0 +1,221 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { vi } from 'vitest'; +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 { ContextManager } from '../contextManager.js'; +import { InMemoryFileSystem } from '../system/InMemoryFileSystem.js'; +import { DeterministicIdGenerator } from '../system/DeterministicIdGenerator.js'; +import type { + Episode, + UserPrompt, + SystemEvent, + SemanticPart, +} from '../ir/types.js'; +import type { ContextAccountingState } from '../pipeline.js'; +import { randomUUID } from 'node:crypto'; + +export function createDummyState( + isSatisfied = false, + deficit = 0, + protectedIds = new Set(), + currentTokens = 5000, + maxTokens = 10000, + retainedTokens = 4000, +): ContextAccountingState { + return { + currentTokens, + maxTokens, + retainedTokens, + deficitTokens: deficit, + protectedEpisodeIds: protectedIds, + isBudgetSatisfied: isSatisfied, + }; +} + +export function createDummyEpisode( + id: string, + type: 'USER_PROMPT' | 'SYSTEM_EVENT', + parts: SemanticPart[] = [], + toolSteps: Array<{ + intent: Record; + observation: Record; + toolName?: string; + tokens?: { intent: number; observation: number }; + }> = [], +): Episode { + let trigger: UserPrompt | SystemEvent; + + if (type === 'USER_PROMPT') { + trigger = { + id: randomUUID(), + type: 'USER_PROMPT', + semanticParts: parts, + metadata: { + originalTokens: 100, + currentTokens: 100, + transformations: [], + }, + }; + } else { + trigger = { + id: randomUUID(), + type: 'SYSTEM_EVENT', + name: 'dummy_event', + payload: {}, + metadata: { + originalTokens: 100, + currentTokens: 100, + transformations: [], + }, + }; + } + + return { + id, + timestamp: Date.now(), + trigger, + steps: toolSteps.map((step) => ({ + id: randomUUID(), + type: 'TOOL_EXECUTION', + toolName: step.toolName || 'test_tool', + intent: step.intent, + observation: step.observation, + tokens: step.tokens || { intent: 50, observation: 50 }, + metadata: { + originalTokens: 100, + currentTokens: 100, + transformations: [], + }, + })), + }; +} + +export function createMockEnvironment(): ContextEnvironment { + return { + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + llmClient: vi.fn().mockReturnValue({ + generateContent: vi.fn().mockResolvedValue({ + text: 'Mock LLM summary response', + }), + })() as unknown as BaseLlmClient, + promptId: 'mock-prompt-id', + sessionId: 'mock-session', + traceDir: '/tmp/.gemini/trace', + projectTempDir: '/tmp/.gemini/tool-outputs', + eventBus: new ContextEventBus(), + tracer: new ContextTracer({ targetDir: '/tmp', sessionId: 'mock-session' }), + charsPerToken: 1, + tokenCalculator: new ContextTokenCalculator(1), + fileSystem: new InMemoryFileSystem(), + idGenerator: new DeterministicIdGenerator('mock-uuid-'), + }; +} + +/** + * Creates a block of synthetic conversation history designed to consume a specific number of tokens. + * Assumes roughly 4 characters per token for standard English text. + */ +export function createSyntheticHistory( + numTurns: number, + tokensPerTurn: number, +): Content[] { + const history: Content[] = []; + const charsPerTurn = tokensPerTurn * 1; + + for (let i = 0; i < numTurns; i++) { + history.push({ + role: 'user', + parts: [{ text: `User turn ${i}. ` + 'A'.repeat(charsPerTurn) }], + }); + history.push({ + role: 'model', + parts: [{ text: `Model response ${i}. ` + 'B'.repeat(charsPerTurn) }], + }); + } + + return history; +} + +/** + * Creates a fully mocked Config object tailored for Context Component testing. + */ +export function createMockContextConfig( + overrides?: Record, + llmClientOverride?: unknown, +): Config { + const defaultConfig = { + isContextManagementEnabled: vi.fn().mockReturnValue(true), + storage: { + getProjectTempDir: vi.fn().mockReturnValue('/tmp/gemini-test'), + }, + getBaseLlmClient: vi.fn().mockReturnValue( + llmClientOverride || { + generateContent: vi.fn().mockResolvedValue({ + text: 'Synthesized state', + }), + }, + ), + getUsageStatisticsEnabled: vi.fn().mockReturnValue(false), + getTargetDir: vi.fn().mockReturnValue('/tmp'), + getSessionId: vi.fn().mockReturnValue('test-session'), + getExperimentalContextSidecarConfig: vi.fn().mockReturnValue(undefined), + }; + + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + return { ...defaultConfig, ...overrides } as unknown as Config; +} + +/** + * Wires up a full ContextManager component with an AgentChatHistory and active background workers. + */ +import { ContextTracer } from '../tracer.js'; +import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js'; +import { SidecarLoader } from '../sidecar/SidecarLoader.js'; +import { ContextEventBus } from '../eventBus.js'; +import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; +import type { BaseLlmClient } from 'src/core/baseLlmClient.js'; +import { ProcessorRegistry } from '../sidecar/registry.js'; +import { registerBuiltInProcessors } from '../sidecar/builtins.js'; + +export function setupContextComponentTest(config: Config) { + const chatHistory = new AgentChatHistory(); + const registry = new ProcessorRegistry(); + registerBuiltInProcessors(registry); + const sidecar = SidecarLoader.fromConfig(config, registry); + const tracer = new ContextTracer({ + targetDir: '/tmp', + sessionId: 'test-session', + }); + const eventBus = new ContextEventBus(); + const env = new ContextEnvironmentImpl( + config.getBaseLlmClient(), + 'test prompt-id', + 'test-session', + '/tmp', + '/tmp/gemini-test', + tracer, + 1, + eventBus, + ); + const contextManager = ContextManager.create( + sidecar, + env, + tracer, + undefined, + registry, + ); + + // The async worker is now internally managed by ContextManager + + // Subscribe to history to enable the Eager/Opportunistic triggers + contextManager.subscribeToHistory(chatHistory); + + return { chatHistory, contextManager }; +} diff --git a/packages/core/src/context/tracer.test.ts b/packages/core/src/context/tracer.test.ts new file mode 100644 index 0000000000..11d602a963 --- /dev/null +++ b/packages/core/src/context/tracer.test.ts @@ -0,0 +1,85 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { ContextTracer } from './tracer.js'; +import { InMemoryFileSystem } from './system/InMemoryFileSystem.js'; +import { DeterministicIdGenerator } from './system/DeterministicIdGenerator.js'; + +describe('ContextTracer (Fake FS & ID Gen)', () => { + let fileSystem: InMemoryFileSystem; + let idGenerator: DeterministicIdGenerator; + + beforeEach(() => { + fileSystem = new InMemoryFileSystem(); + idGenerator = new DeterministicIdGenerator('mock-uuid-'); + + // We must mock Date.now() to ensure asset file names are perfectly deterministic + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T12:00:00Z')); + }); + + it('initializes, logs events, and auto-saves large assets deterministically', () => { + const tracer = new ContextTracer( + { enabled: true, targetDir: '/fake/target', sessionId: 'test-session' }, + fileSystem, + idGenerator, + ); + + // Verify Initialization + const initTraceLog = fileSystem.readFileSync( + '/fake/target/.gemini/context_trace/test-session/trace.log', + 'utf8', + ); + expect(initTraceLog).toContain('[SYSTEM] Context Tracer Initialized'); + + // Small logging: shouldn't trigger saveAsset + tracer.logEvent('TestComponent', 'TestAction', { key: 'value' }); + + const smallTraceLog = fileSystem.readFileSync( + '/fake/target/.gemini/context_trace/test-session/trace.log', + 'utf8', + ); + expect(smallTraceLog).toContain('[TestComponent] TestAction'); + expect(smallTraceLog).toContain('{"key":"value"}'); + + // Large logging: should trigger auto-asset save + const hugeString = 'a'.repeat(2000); + tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString }); + + // 1767268800000 is 2026-01-01T12:00:00Z + const expectedAssetPath = + '/fake/target/.gemini/context_trace/test-session/assets/1767268800000-mock-uuid-1-largeKey.json'; + + // Assert asset was written to FS + expect(fileSystem.existsSync(expectedAssetPath)).toBe(true); + + const largeTraceLog = fileSystem.readFileSync( + '/fake/target/.gemini/context_trace/test-session/trace.log', + 'utf8', + ); + expect(largeTraceLog).toContain('[TestComponent] LargeAction'); + expect(largeTraceLog).toContain( + `{"largeKey":{"$asset":"1767268800000-mock-uuid-1-largeKey.json"}}`, + ); + }); + + it('silently ignores logging when disabled', () => { + const tracer = new ContextTracer( + { enabled: false, targetDir: '/fake/target', sessionId: 'test-session' }, + fileSystem, + idGenerator, + ); + + tracer.logEvent('TestComponent', 'TestAction'); + + const hugeString = 'a'.repeat(2000); + tracer.logEvent('TestComponent', 'LargeAction', { largeKey: hugeString }); + + // FS should be completely empty + expect(fileSystem.getFiles().size).toBe(0); + }); +}); diff --git a/packages/core/src/context/tracer.ts b/packages/core/src/context/tracer.ts new file mode 100644 index 0000000000..cf7299f0fa --- /dev/null +++ b/packages/core/src/context/tracer.ts @@ -0,0 +1,120 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { debugLogger } from '../utils/debugLogger.js'; +import type { IFileSystem } from './system/IFileSystem.js'; +import { NodeFileSystem } from './system/NodeFileSystem.js'; +import type { IIdGenerator } from './system/IIdGenerator.js'; +import { NodeIdGenerator } from './system/NodeIdGenerator.js'; + +export interface ContextTracerOptions { + enabled?: boolean; + targetDir: string; + sessionId: string; +} + +export class ContextTracer { + private traceDir: string; + private assetsDir: string; + private enabled: boolean; + private fileSystem: IFileSystem; + private idGenerator: IIdGenerator; + + private readonly MAX_INLINE_SIZE = 1000; + + constructor( + options: ContextTracerOptions, + fileSystem: IFileSystem = new NodeFileSystem(), + idGenerator: IIdGenerator = new NodeIdGenerator(), + ) { + this.enabled = options.enabled ?? false; + this.fileSystem = fileSystem; + this.idGenerator = idGenerator; + + this.traceDir = this.fileSystem.join( + options.targetDir, + '.gemini', + 'context_trace', + options.sessionId, + ); + this.assetsDir = this.fileSystem.join(this.traceDir, 'assets'); + + if (this.enabled) { + try { + this.fileSystem.mkdirSync(this.assetsDir, { recursive: true }); + this.logEvent('SYSTEM', 'Context Tracer Initialized', { + sessionId: options.sessionId, + }); + } catch (e) { + debugLogger.error('Failed to initialize ContextTracer', e); + this.enabled = false; + } + } + } + + logEvent( + component: string, + action: string, + details?: Record, + ) { + if (!this.enabled) return; + try { + let processedDetails: Record | undefined; + + if (details) { + processedDetails = {}; + for (const [key, value] of Object.entries(details)) { + const strValue = + typeof value === 'string' ? value : JSON.stringify(value); + if (strValue && strValue.length > this.MAX_INLINE_SIZE) { + const assetId = this.saveAsset(component, key, value); + processedDetails[key] = { $asset: assetId }; + } else { + processedDetails[key] = value; + } + } + } + + const timestamp = new Date().toISOString(); + const detailsStr = processedDetails + ? ` | Details: ${JSON.stringify(processedDetails)}` + : ''; + const logLine = `[${timestamp}] [${component}] ${action}${detailsStr}\n`; + this.fileSystem.appendFileSync( + this.fileSystem.join(this.traceDir, 'trace.log'), + logLine, + 'utf-8', + ); + } catch (e) { + debugLogger.warn(`Tracing failed: ${e}`); + } + } + + private saveAsset( + component: string, + assetName: string, + data: unknown, + ): string { + if (!this.enabled) return 'asset-recording-disabled'; + try { + const assetId = `${Date.now()}-${this.idGenerator.generateId()}-${assetName}.json`; + const assetPath = this.fileSystem.join(this.assetsDir, assetId); + + this.fileSystem.writeFileSync( + assetPath, + JSON.stringify(data, null, 2), + 'utf-8', + ); + this.logEvent(component, `Saved asset: ${assetName}`, { assetId }); + return assetId; + } catch (e) { + this.logEvent(component, `Failed to save asset: ${assetName}`, { + error: String(e), + }); + return 'asset-save-failed'; + } + } +} diff --git a/packages/core/src/context/utils/contextTokenCalculator.ts b/packages/core/src/context/utils/contextTokenCalculator.ts new file mode 100644 index 0000000000..cc73fc5de9 --- /dev/null +++ b/packages/core/src/context/utils/contextTokenCalculator.ts @@ -0,0 +1,71 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Part } from '@google/genai'; +import { estimateTokenCountSync as baseEstimate } from '../../utils/tokenCalculation.js'; +import type { Episode } 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. + */ +const BASE_MULTIMODAL_TOKEN_COST = 258; + +export class ContextTokenCalculator { + constructor(private readonly charsPerToken: number) {} + + /** + * Fast, simple heuristic estimation for a raw string. + */ + estimateTokensForString(text: string): number { + return Math.ceil(text.length / this.charsPerToken); + } + + /** + * Fast, simple heuristic conversion from tokens to expected character length. + * Useful for calculating truncation thresholds. + */ + tokensToChars(tokens: number): number { + return tokens * this.charsPerToken; + } + + /** + * Calculates the total token count for a complete Episodic IR graph. + * This is fast because it relies on pre-computed metadata where available. + */ + calculateEpisodeListTokens(episodes: Episode[]): number { + let tokens = 0; + for (const ep of episodes) { + if (ep.trigger) tokens += ep.trigger.metadata.currentTokens; + for (const step of ep.steps) { + tokens += step.metadata.currentTokens; + } + if (ep.yield) tokens += ep.yield.metadata.currentTokens; + } + return tokens; + } + + /** + * Slower, precise estimation for a Gemini Content/Part graph. + * Deeply inspects the nested structure and uses the base tokenization math. + */ + estimateTokensForParts(parts: Part[], depth: number = 0): number { + let totalTokens = 0; + for (const part of parts) { + 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; + } else { + totalTokens += Math.ceil( + JSON.stringify(part).length / this.charsPerToken, + ); + } + } + // Also include structural overhead + return totalTokens + baseEstimate(parts, depth); + } +} diff --git a/packages/core/src/core/agentChatHistory.ts b/packages/core/src/core/agentChatHistory.ts new file mode 100644 index 0000000000..ffff5a67a2 --- /dev/null +++ b/packages/core/src/core/agentChatHistory.ts @@ -0,0 +1,77 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Content } from '@google/genai'; + +export type HistoryEventType = 'PUSH' | 'SYNC_FULL' | 'CLEAR'; + +export interface HistoryEvent { + type: HistoryEventType; + payload: readonly Content[]; +} + +export type HistoryListener = (event: HistoryEvent) => void; + +export class AgentChatHistory { + private history: Content[]; + private listeners: Set = new Set(); + + constructor(initialHistory: Content[] = []) { + this.history = [...initialHistory]; + } + + subscribe(listener: HistoryListener): () => void { + this.listeners.add(listener); + // Emit initial state to new subscriber + listener({ type: 'SYNC_FULL', payload: this.history }); + return () => this.listeners.delete(listener); + } + + private notify(type: HistoryEventType, payload: readonly Content[]) { + const event: HistoryEvent = { type, payload }; + for (const listener of this.listeners) { + listener(event); + } + } + + push(content: Content) { + this.history.push(content); + this.notify('PUSH', [content]); + } + + set(history: readonly Content[]) { + this.history = [...history]; + this.notify('SYNC_FULL', this.history); + } + + clear() { + this.history = []; + this.notify('CLEAR', []); + } + + get(): readonly Content[] { + return this.history; + } + + map(callback: (value: Content, index: number, array: Content[]) => Content) { + this.history = this.history.map(callback); + this.notify('SYNC_FULL', this.history); + } + + flatMap( + callback: ( + value: Content, + index: number, + array: Content[], + ) => U | readonly U[], + ): U[] { + return this.history.flatMap(callback); + } + + get length(): number { + return this.history.length; + } +} diff --git a/packages/core/src/telemetry/llmRole.ts b/packages/core/src/telemetry/llmRole.ts index 843ac4123c..7d8f5d8df6 100644 --- a/packages/core/src/telemetry/llmRole.ts +++ b/packages/core/src/telemetry/llmRole.ts @@ -16,4 +16,5 @@ export enum LlmRole { UTILITY_EDIT_CORRECTOR = 'utility_edit_corrector', UTILITY_AUTOCOMPLETE = 'utility_autocomplete', UTILITY_FAST_ACK_HELPER = 'utility_fast_ack_helper', + UTILITY_STATE_SNAPSHOT_PROCESSOR = 'utility_state_snapshot_processr', }