From ac6dc1d477e040c65d2f410b7b0a1b4da816487c Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 6 Apr 2026 16:43:04 +0000 Subject: [PATCH] feat(core): introduce decoupled ContextManager and Sidecar architecture --- packages/core/src/config/config.ts | 8 + .../contextManager.golden.test.ts.snap | 52 ++ .../src/context/contextManager.async.test.ts | 113 ++++ .../context/contextManager.barrier.test.ts | 74 +++ .../src/context/contextManager.golden.test.ts | 149 +++++ packages/core/src/context/contextManager.ts | 586 ++++++++++++++++++ packages/core/src/context/eventBus.ts | 49 ++ packages/core/src/context/ir/mapper.test.ts | 161 +++++ packages/core/src/context/ir/mapper.ts | 307 +++++++++ packages/core/src/context/ir/types.ts | 202 ++++++ packages/core/src/context/pipeline.ts | 48 ++ .../blobDegradationProcessor.test.ts | 123 ++++ .../processors/blobDegradationProcessor.ts | 123 ++++ .../historySquashingProcessor.test.ts | 167 +++++ .../processors/historySquashingProcessor.ts | 132 ++++ .../semanticCompressionProcessor.test.ts | 154 +++++ .../semanticCompressionProcessor.ts | 202 ++++++ .../processors/toolMaskingProcessor.test.ts | 130 ++++ .../processors/toolMaskingProcessor.ts | 243 ++++++++ .../core/src/context/sidecar/SidecarLoader.ts | 41 ++ .../core/src/context/sidecar/environment.ts | 17 + .../src/context/sidecar/environmentImpl.ts | 44 ++ packages/core/src/context/sidecar/profiles.ts | 51 ++ packages/core/src/context/sidecar/registry.ts | 42 ++ packages/core/src/context/sidecar/types.ts | 55 ++ .../src/context/testing/contextTestUtils.ts | 113 ++++ packages/core/src/context/tracer.ts | 70 +++ .../context/utils/contextTokenCalculator.ts | 31 + .../src/context/workers/asyncContextWorker.ts | 18 + .../context/workers/stateSnapshotWorker.ts | 229 +++++++ packages/core/src/core/agentChatHistory.ts | 77 +++ 31 files changed, 3811 insertions(+) create mode 100644 packages/core/src/context/__snapshots__/contextManager.golden.test.ts.snap create mode 100644 packages/core/src/context/contextManager.async.test.ts create mode 100644 packages/core/src/context/contextManager.barrier.test.ts create mode 100644 packages/core/src/context/contextManager.golden.test.ts create mode 100644 packages/core/src/context/contextManager.ts create mode 100644 packages/core/src/context/eventBus.ts create mode 100644 packages/core/src/context/ir/mapper.test.ts create mode 100644 packages/core/src/context/ir/mapper.ts create mode 100644 packages/core/src/context/ir/types.ts create mode 100644 packages/core/src/context/pipeline.ts create mode 100644 packages/core/src/context/processors/blobDegradationProcessor.test.ts create mode 100644 packages/core/src/context/processors/blobDegradationProcessor.ts create mode 100644 packages/core/src/context/processors/historySquashingProcessor.test.ts create mode 100644 packages/core/src/context/processors/historySquashingProcessor.ts create mode 100644 packages/core/src/context/processors/semanticCompressionProcessor.test.ts create mode 100644 packages/core/src/context/processors/semanticCompressionProcessor.ts create mode 100644 packages/core/src/context/processors/toolMaskingProcessor.test.ts create mode 100644 packages/core/src/context/processors/toolMaskingProcessor.ts create mode 100644 packages/core/src/context/sidecar/SidecarLoader.ts create mode 100644 packages/core/src/context/sidecar/environment.ts create mode 100644 packages/core/src/context/sidecar/environmentImpl.ts create mode 100644 packages/core/src/context/sidecar/profiles.ts create mode 100644 packages/core/src/context/sidecar/registry.ts create mode 100644 packages/core/src/context/sidecar/types.ts create mode 100644 packages/core/src/context/testing/contextTestUtils.ts create mode 100644 packages/core/src/context/tracer.ts create mode 100644 packages/core/src/context/utils/contextTokenCalculator.ts create mode 100644 packages/core/src/context/workers/asyncContextWorker.ts create mode 100644 packages/core/src/context/workers/stateSnapshotWorker.ts create mode 100644 packages/core/src/core/agentChatHistory.ts diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index efb3e296df..b145281754 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -699,6 +699,7 @@ export interface ConfigParameters { experimentalJitContext?: boolean; autoDistillation?: boolean; experimentalMemoryManager?: boolean; + experimentalContextSidecarConfig?: string; experimentalAgentHistoryTruncation?: boolean; experimentalAgentHistoryTruncationThreshold?: number; experimentalAgentHistoryRetainedMessages?: number; @@ -940,6 +941,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; @@ -1151,6 +1153,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, @@ -2413,6 +2417,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..b83f6d751c --- /dev/null +++ b/packages/core/src/context/contextManager.async.test.ts @@ -0,0 +1,113 @@ +import { IrMapper } from './ir/mapper.js'; +/** + * @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. + IrMapper.setConfig({ charsPerToken: 1 }); + + 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 + await 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); + // console.dir(projection, {depth: null}); + // 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. + IrMapper.setConfig({ charsPerToken: 1 }); + 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!!' }] }, + ]; + await 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..ce1a6c7fb5 --- /dev/null +++ b/packages/core/src/context/contextManager.barrier.test.ts @@ -0,0 +1,74 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { IrMapper } from './ir/mapper.js'; +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; + IrMapper.setConfig({ charsPerToken: 1 }); + + // 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..f0de3635b8 --- /dev/null +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -0,0 +1,149 @@ +/** + * @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 type { Content } from '@google/genai'; + +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), + 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 sidecar = SidecarLoader.fromLegacyConfig(mockConfig as any); + const tracer = new ContextTracer('/tmp', 'test-session'); + const env = new ContextEnvironmentImpl( + {} as any, + 'test', + '/tmp', + '/tmp', + tracer, + 4, + ); + contextManager = new ContextManager(sidecar, env, tracer); + }); + + 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 any).pristineEpisodes = ( + await import('./ir/mapper.js') + ).IrMapper.toIr(history); + const result = await contextManager.projectCompressedHistory(); + expect(result).toMatchSnapshot(); + }); + + it('should not modify history when under budget', async () => { + const history = createLargeHistory(); + (contextManager as any).pristineEpisodes = ( + await import('./ir/mapper.js') + ).IrMapper.toIr(history); + // 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('/tmp', 'test2'); + contextManager = new ContextManager( + { + pipelines: { + eagerBackground: [], + normalProcessingGraph: [], + retainedProcessingGraph: [], + }, + } as any, + {} as any, + tracer2, + ); + + (contextManager as any).pristineEpisodes = ( + await import('./ir/mapper.js') + ).IrMapper.toIr(history); + 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..2823c6debe --- /dev/null +++ b/packages/core/src/context/contextManager.ts @@ -0,0 +1,586 @@ +/** + * @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 { IrMapper } from './ir/mapper.js'; +import type { Episode } from './ir/types.js'; + +import { ContextEventBus } from './eventBus.js'; +import { ContextTracer } from './tracer.js'; + +import { StateSnapshotWorker } from './workers/stateSnapshotWorker.js'; + +import type { ContextEnvironment } from './sidecar/environment.js'; + +import type { SidecarConfig } from './sidecar/types.js'; +import { ProcessorRegistry } from './sidecar/registry.js'; +import type { ContextProcessor } from './pipeline.js'; +import type { AsyncContextWorker } from './workers/asyncContextWorker.js'; + +import { ToolMaskingProcessor } from './processors/toolMaskingProcessor.js'; +import { BlobDegradationProcessor } from './processors/blobDegradationProcessor.js'; +import { SemanticCompressionProcessor } from './processors/semanticCompressionProcessor.js'; +import { HistorySquashingProcessor } from './processors/historySquashingProcessor.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 unsubscribeHistory?: () => void; + private readonly eventBus: ContextEventBus; + + // Internal sub-components + // Synchronous processors are instantiated but effectively used as singletons within this class + private workers: AsyncContextWorker[] = []; + + constructor( + private sidecar: SidecarConfig, + private env: ContextEnvironment, + private readonly tracer: ContextTracer, + ) { + this.eventBus = new ContextEventBus(); + + // Register built-ins + ProcessorRegistry.register({ + id: 'ToolMaskingProcessor', + create: (env, opts) => new ToolMaskingProcessor(env, opts as any), + }); + ProcessorRegistry.register({ + id: 'BlobDegradationProcessor', + create: (env, opts) => new BlobDegradationProcessor(env), + }); + ProcessorRegistry.register({ + id: 'SemanticCompressionProcessor', + create: (env, opts) => new SemanticCompressionProcessor(env, opts as any), + }); + ProcessorRegistry.register({ + id: 'HistorySquashingProcessor', + create: (env, opts) => new HistorySquashingProcessor(env, opts as any), + }); + ProcessorRegistry.register({ + id: 'StateSnapshotWorker', + create: (env, opts) => new StateSnapshotWorker(env), + }); + + 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}.`, + ); + } + }); + + // Initialize synchronous fallback processors + // Order matters: Fast, lossless masking -> Intelligent degradation -> Brutal truncation fallback + + // Initialize and start background subconscious workers + for (const bgDef of this.sidecar.pipelines.eagerBackground) { + const worker = ProcessorRegistry.get(bgDef.processorId).create( + this.env, + bgDef.options, + ) as AsyncContextWorker; + worker.start(this.eventBus); + this.workers.push(worker); + } + } + + /** + * Safely stops background workers and clears event listeners. + */ + shutdown() { + for (const worker of this.workers) { + worker.stop(); + } + if (this.unsubscribeHistory) { + this.unsubscribeHistory(); + } + } + + /** + * Subscribes to the core AgentChatHistory to natively track all message events, + * converting them seamlessly into pristine Episodes. + */ + subscribeToHistory(chatHistory: AgentChatHistory) { + if (this.unsubscribeHistory) { + this.unsubscribeHistory(); + } + + this.unsubscribeHistory = chatHistory.subscribe((event) => { + // Rebuild the pristine IR graph from the full source history on every change. + // We must map the FULL array at once because IrMapper groups adjacent + // function calls and responses into unified Episodes. Pushing messages + // individually would shatter these episodic boundaries. + this.pristineEpisodes = IrMapper.toIr(chatHistory.get()); + this.tracer.logEvent( + 'ContextManager', + 'Rebuilt pristine graph from chat history update', + { episodeCount: this.pristineEpisodes.length }, + ); + this.checkTriggers(); + }); + } + + private checkTriggers() { + if (!this.sidecar.budget) return; + + const mngConfig = this.sidecar; + + // Calculate tokens based on the *Working Buffer View*, not the raw pristine log. + // This solves Bug 2: The View shrinks when variants are applied, preventing infinite GC loops. + const workingBuffer = this.getWorkingBufferView(); + const currentTokens = this.calculateIrTokens(workingBuffer); + + this.tracer.logEvent('ContextManager', 'Evaluated triggers', { + currentTokens, + retainedTokens: mngConfig.budget.retainedTokens, + }); + + // 1. Eager Compute Trigger (Continuous Streaming) + // Broadcast the full pristine log to the async workers so they can proactively summarize partial massive files. + this.eventBus.emitChunkReceived({ episodes: this.pristineEpisodes }); + + // 2. The Ship of Theseus Trigger (retainedTokens crossed) + // If we exceed 65k, tell the background processors to opportunistically synthesize the oldest nodes. + if (currentTokens > mngConfig.budget.retainedTokens) { + const deficit = currentTokens - mngConfig.budget.retainedTokens; + this.tracer.logEvent( + 'ContextManager', + 'Budget crossed. Emitting ConsolidationNeeded', + { deficit }, + ); + console.log( + 'EMITTING CONSOLIDATION. Buffer:', + workingBuffer.length, + 'Deficit:', + deficit, + ); + this.eventBus.emitConsolidationNeeded({ + episodes: workingBuffer, // Pass the working buffer so they know what still needs compression + targetDeficit: deficit, + }); + } + } + + /** + * 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. + */ + /** + * Applies the data-driven Sidecar configuration graphs. + * Splits the episodes into the 'retained' and 'normal' ranges, + * runs their respective processor pipelines sequentially, and recombines them. + */ + private async applyProcessorGraphs(episodes: Episode[]): Promise { + const mngConfig = this.sidecar; + const retainedLimit = mngConfig.budget.retainedTokens; + + // If we're incredibly small, maybe we just run the retained graph on everything? + // Let's divide the episodes exactly at the retained boundary. + const retainedWindow: Episode[] = []; + const normalWindow: Episode[] = []; + let rollingTokens = 0; + + // Scan backwards to fill the retained window + for (let i = episodes.length - 1; i >= 0; i--) { + const ep = episodes[i]; + const epTokens = this.calculateIrTokens([ep]); + if ( + (rollingTokens + epTokens <= retainedLimit && + normalWindow.length === 0) || + retainedWindow.length === 0 + ) { + // We always put at least the latest episode in the retained window. + // We only add to retainedWindow if we haven't already started the normalWindow (contiguous block). + retainedWindow.unshift(ep); + rollingTokens += epTokens; + } else { + normalWindow.unshift(ep); + } + } + + const protectedIds = new Set(); + // We must protect the System Episode, which is always index 0 of pristineEpisodes. + if (this.pristineEpisodes.length > 0) { + protectedIds.add(this.pristineEpisodes[0].id); // Structural invariant + } + + const createAccountingState = (currentTotal: number) => ({ + currentTokens: currentTotal, + maxTokens: mngConfig.budget.maxTokens, + retainedTokens: mngConfig.budget.retainedTokens, + deficitTokens: Math.max(0, currentTotal - mngConfig.budget.maxTokens), + protectedEpisodeIds: protectedIds, + isBudgetSatisfied: currentTotal <= mngConfig.budget.maxTokens, // We use maxTokens here so processors don't prematurely short-circuit if they are trying to prevent a barrier hit + }); + + // Run Retained Graph + let processedRetained = [...retainedWindow]; + for (const def of mngConfig.pipelines.retainedProcessingGraph) { + const processor = ProcessorRegistry.get(def.processorId).create( + this.env, + def.options, + ) as ContextProcessor; + this.tracer.logEvent( + 'ContextManager', + `Running ${processor.name} on retained window.`, + ); + const state = createAccountingState( + this.calculateIrTokens([...normalWindow, ...processedRetained]), + ); + processedRetained = await processor.process(processedRetained, state); + } + + // Run Normal Graph + let processedNormal = [...normalWindow]; + for (const def of mngConfig.pipelines.normalProcessingGraph) { + const processor = ProcessorRegistry.get(def.processorId).create( + this.env, + def.options, + ) as ContextProcessor; + this.tracer.logEvent( + 'ContextManager', + `Running ${processor.name} on normal window.`, + ); + const state = createAccountingState( + this.calculateIrTokens([...processedNormal, ...processedRetained]), + ); + processedNormal = await processor.process(processedNormal, state); + } + + return [...processedNormal, ...processedRetained]; + } + + public getWorkingBufferView(): Episode[] { + const mngConfig = this.sidecar; + const retainedTokens = mngConfig.budget.retainedTokens; + + let currentEpisodes: Episode[] = []; + let rollingTokens = 0; + const skippedIds = new Set(); + this.tracer.logEvent('ViewGenerator', 'Generating Working Buffer View'); + + for (let i = this.pristineEpisodes.length - 1; i >= 0; i--) { + const ep = this.pristineEpisodes[i]; + + // If this episode was already replaced by an N-to-1 Snapshot injected earlier in the sweep, skip it entirely! + // This solves Bug 1 (Duplicate Projection). + if (skippedIds.has(ep.id)) { + this.tracer.logEvent( + 'ViewGenerator', + `Skipping episode [${ep.id}] due to N-to-1 replacement.`, + ); + continue; + } + + let projectedEp = { + ...ep, + trigger: { + ...ep.trigger, + metadata: { + ...ep.trigger.metadata, + transformations: [...ep.trigger.metadata.transformations], + }, + semanticParts: + ep.trigger.type === 'USER_PROMPT' + ? [...ep.trigger.semanticParts.map((sp) => ({ ...sp }))] + : undefined, + } as any, + steps: ep.steps.map( + (step) => + ({ + ...step, + metadata: { + ...step.metadata, + transformations: [...step.metadata.transformations], + }, + }) as any, + ), + yield: ep.yield + ? { + ...ep.yield, + metadata: { + ...ep.yield.metadata, + transformations: [...ep.yield.metadata.transformations], + }, + } + : undefined, + }; + + const epTokens = this.calculateIrTokens([projectedEp]); + + if (ep.variants) { + console.log( + 'Checking variants for', + ep.id, + 'rollingTokens:', + rollingTokens, + 'retained:', + retainedTokens, + ); + } + if (rollingTokens > retainedTokens && ep.variants) { + console.log('EVALUATING VARIANTS FOR', ep.id); + 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 as any; + // Mark all the episodes this snapshot covers to be skipped by the backwards sweep. + for (const id of snapshot.replacedEpisodeIds) { + skippedIds.add(id); + } + this.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 any; + projectedEp.yield = undefined; + this.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.length > 0 + ) { + projectedEp.trigger.semanticParts[0].presentation = { + text: masked.text, + tokens: masked.recoveredTokens || 10, + }; + } + this.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 += this.calculateIrTokens([projectedEp]); + } + + return currentEpisodes; + } + + /** + * 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 { + if (!this.sidecar.budget) { + return this._projectAndDump(IrMapper.fromIr(this.pristineEpisodes)); + } + + const mngConfig = this.sidecar; + const maxTokens = mngConfig.budget.maxTokens; + this.tracer.logEvent('ContextManager', 'Projection requested.'); + + // Get the dynamically computed Working Buffer View + let currentEpisodes = this.getWorkingBufferView(); + + currentEpisodes = await this.applyProcessorGraphs(currentEpisodes); + + let currentTokens = this.calculateIrTokens(currentEpisodes); + + if (currentTokens <= maxTokens) { + this.tracer.logEvent( + 'ContextManager', + `View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`, + ); + return this._projectAndDump(IrMapper.fromIr(currentEpisodes)); + } + + this.tracer.logEvent( + 'ContextManager', + `View exceeds maxTokens (${currentTokens} > ${maxTokens}). Hitting Synchronous Pressure Barrier. Strategy: ${mngConfig.gcBackstop.strategy}`, + ); + // --- The Synchronous Pressure Barrier --- + // The background eager workers couldn't keep up, or a massive file was pasted. + // The Working Buffer View is still over the absolute hard limit (maxTokens). + // We MUST reduce tokens before returning, or the API request will 400. + + debugLogger.log( + `Context Manager Synchronous Barrier triggered: View at ${currentTokens} tokens (limit: ${maxTokens}). Strategy: ${mngConfig.gcBackstop.strategy}`, + ); + + // Calculate target based on gcTarget + let targetTokens = maxTokens; + + if (mngConfig.gcBackstop.target === 'max') { + targetTokens = mngConfig.budget.retainedTokens; + } else if (mngConfig.gcBackstop.target === 'freeNTokens') { + targetTokens = + maxTokens - (mngConfig.gcBackstop.freeTokensTarget ?? 10000); + } + + // Structural invariant: We ALWAYS protect the architectural initialization turn (Turn 0) + // We do NOT arbitrarily protect recent episodes (like currentEpisodes.length - 1) + // because an episode can be unboundedly large, and protecting it would crash the LLM. + const protectedEpisodeId = + this.pristineEpisodes.length > 0 ? this.pristineEpisodes[0].id : null; + + let remainingTokens = currentTokens; + + const truncated: Episode[] = []; + + const strategy = mngConfig.gcBackstop.strategy; + + for (const ep of currentEpisodes) { + const epTokens = this.calculateIrTokens([ep]); + if (remainingTokens > targetTokens && ep.id !== protectedEpisodeId) { + console.log( + 'DROPPING EPISODE:', + ep.id, + 'rem:', + remainingTokens, + 'tgt:', + targetTokens, + ); + + remainingTokens -= epTokens; + if (strategy === 'truncate') { + this.tracer.logEvent('Barrier', `Truncating episode [${ep.id}].`); + + debugLogger.log(`Barrier (truncate): Dropped Episode ${ep.id}`); + } else if (strategy === 'compress') { + this.tracer.logEvent( + 'Barrier', + `Compress fallback to truncate for [${ep.id}].`, + ); + debugLogger.warn( + `Synchronous compress barrier not fully implemented, truncating Episode ${ep.id}.`, + ); + } else if (strategy === 'rollingSummarizer') { + this.tracer.logEvent( + 'Barrier', + `RollingSummarizer fallback to truncate for [${ep.id}].`, + ); + debugLogger.warn( + `Synchronous rollingSummarizer barrier not fully implemented, truncating Episode ${ep.id}.`, + ); + } + } else { + console.log( + 'KEEPING EPISODE:', + ep.id, + 'rem:', + remainingTokens, + 'tgt:', + targetTokens, + ); + truncated.push(ep); + } + } + currentEpisodes = truncated; + + const finalTokens = this.calculateIrTokens(currentEpisodes); + this.tracer.logEvent( + 'ContextManager', + `Finished projection. Final token count: ${finalTokens}.`, + ); + debugLogger.log( + `Context Manager finished. Final actual token count: ${finalTokens}.`, + ); + + return this._projectAndDump(IrMapper.fromIr(currentEpisodes)); + } + + private async _projectAndDump(contents: Content[]): Promise { + if (process.env['GEMINI_DUMP_CONTEXT'] === 'true') { + try { + const fs = await import('node:fs/promises'); + const path = await import('node:path'); + const dumpPath = path.join( + this.env.getTraceDir(), + '.gemini', + 'projected_context.json', + ); + await fs.mkdir(path.dirname(dumpPath), { recursive: true }); + await fs.writeFile( + dumpPath, + JSON.stringify(contents, null, 2), + 'utf-8', + ); + debugLogger.log( + `[Observability] Context successfully dumped to ${dumpPath}`, + ); + } catch (e) { + debugLogger.error(`Failed to dump context: ${e}`); + } + } + return contents; + } + + private calculateIrTokens(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; + } +} diff --git a/packages/core/src/context/eventBus.ts b/packages/core/src/context/eventBus.ts new file mode 100644 index 0000000000..dc9e8b5b74 --- /dev/null +++ b/packages/core/src/context/eventBus.ts @@ -0,0 +1,49 @@ +/** + * @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 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 { + 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/ir/mapper.test.ts b/packages/core/src/context/ir/mapper.test.ts new file mode 100644 index 0000000000..27b401dcd8 --- /dev/null +++ b/packages/core/src/context/ir/mapper.test.ts @@ -0,0 +1,161 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect } from 'vitest'; +import { IrMapper } from './mapper.js'; +import type { Content } from '@google/genai'; +import type { UserPrompt, ToolExecution } 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 episodes = IrMapper.toIr(rawHistory); + + 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 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 initialIr = IrMapper.toIr(history); + 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); + 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..9c18473c99 --- /dev/null +++ b/packages/core/src/context/ir/mapper.ts @@ -0,0 +1,307 @@ +/** + * @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, +} from './types.js'; +import { estimateContextTokenCountSync as estimateTokenCountSync } from '../utils/contextTokenCalculator.js'; + +// WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references +const nodeIdentityMap = new WeakMap(); + +function getStableId(obj: object): string { + let id = nodeIdentityMap.get(obj); + if (!id) { + id = randomUUID(); + nodeIdentityMap.set(obj, id); + } + return id; +} + +export class IrMapper { + static setConfig(cfg: { charsPerToken?: number }) { + this.config = cfg; + } + private static config: { charsPerToken?: number } | undefined; + + /** + * 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[]): Episode[] { + const episodes: Episode[] = []; + let currentEpisode: Partial | null = null; + const pendingCallParts: Map = new Map(); + + const createMetadata = (parts: Part[]): IrMetadata => { + const tokens = estimateTokenCountSync(parts, 0, IrMapper.config); + return { + originalTokens: tokens, + currentTokens: tokens, + transformations: [], + }; + }; + + const finalizeEpisode = () => { + if (currentEpisode && currentEpisode.trigger) { + episodes.push(currentEpisode as unknown as Episode); // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion + } + 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) { + 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([]), + }, + steps: [], + }; + } + + for (const part of msg.parts) { + if (part.functionResponse) { + const callId = part.functionResponse.id || ''; + const matchingCall = pendingCallParts.get(callId); + + const intentTokens = matchingCall + ? estimateTokenCountSync([matchingCall]) + : 0; + const obsTokens = estimateTokenCountSync([part]); + + const step: ToolExecution = { + id: getStableId(part), + type: 'TOOL_EXECUTION', + toolName: part.functionResponse.name || 'unknown', + intent: + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + (matchingCall?.functionCall?.args as unknown as Record< + string, + unknown + >) || {}, + observation: + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + (part.functionResponse.response as unknown as Record< + string, + unknown + >) || {}, + tokens: { + intent: intentTokens, + observation: obsTokens, + }, + metadata: { + originalTokens: intentTokens + obsTokens, + currentTokens: intentTokens + obsTokens, + transformations: [], + }, + }; + currentEpisode.steps!.push(step); + if (callId) pendingCallParts.delete(callId); + } + } + } + + if (hasUserParts) { + finalizeEpisode(); + + 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), + ), + }; + + currentEpisode = { + id: getStableId(msg), + timestamp: Date.now(), + trigger, + steps: [], + }; + } + } else if (msg.role === 'model') { + 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([]), + }, + 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); + } + } + } + } + + if (currentEpisode) { + 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; + } + } + finalizeEpisode(); + } + + return episodes; + } + + /** + * Re-serializes the Episodic IR back into a flat Gemini Content[] array. + */ + static fromIr(episodes: Episode[]): Content[] { + const history: Content[] = []; + + for (const ep of episodes) { + // 1. Serialize Trigger + if (ep.trigger.type === 'USER_PROMPT') { + const parts: Part[] = []; + for (const sp of ep.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') { + // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion, @typescript-eslint/no-unsafe-type-assertion + parts.push(sp.part as unknown as Part); + } + } + if (parts.length > 0) history.push({ role: 'user', parts }); + } + + // 2. Serialize Steps + 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 ep.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 as unknown as Record, // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion + id: step.id, + }, + }); + const observation = step.presentation + ? step.presentation.observation + : step.observation; + pendingUserParts.push({ + functionResponse: { + name: step.toolName, + response: observation as unknown as Record, // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion + id: step.id, + }, + }); + } + } + flushPending(); + + // 3. Serialize Yield + if (ep.yield) { + history.push({ + role: 'model', + parts: [{ text: ep.yield.presentation?.text ?? ep.yield.text }], + }); + } + } + + return history; + } +} diff --git a/packages/core/src/context/ir/types.ts b/packages/core/src/context/ir/types.ts new file mode 100644 index 0000000000..e60964304f --- /dev/null +++ b/packages/core/src/context/ir/types.ts @@ -0,0 +1,202 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * 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: unknown; + 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..77fa0c2cc4 --- /dev/null +++ b/packages/core/src/context/pipeline.ts @@ -0,0 +1,48 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Episode } from './ir/types.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 based on the current accounting state. + * Processors should return a new or mutated array of episodes. + */ + process( + episodes: Episode[], + 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..9c57095113 --- /dev/null +++ b/packages/core/src/context/processors/blobDegradationProcessor.test.ts @@ -0,0 +1,123 @@ +/** + * @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 { BlobDegradationProcessor } from './blobDegradationProcessor.js'; +import type { Episode, UserPrompt } from '../ir/types.js'; +import type { ContextAccountingState } from '../pipeline.js'; +import { randomUUID } from 'node:crypto'; +import * as fsPromises from 'node:fs/promises'; + +vi.mock('node:fs/promises'); + +describe('BlobDegradationProcessor', () => { + let processor: BlobDegradationProcessor; + + beforeEach(() => { + vi.resetAllMocks(); + + processor = new BlobDegradationProcessor(createMockEnvironment()); + }); + + const getDummyState = ( + isSatisfied = false, + deficit = 0, + protectedIds = new Set(), + ): ContextAccountingState => ({ + currentTokens: 5000, + maxTokens: 10000, + retainedTokens: 4000, + deficitTokens: deficit, + protectedEpisodeIds: protectedIds, + isBudgetSatisfied: isSatisfied, + }); + + it('degrades inline_data into a text reference and saves to disk', async () => { + const dummyImageBase64 = Buffer.from('fake-image-data').toString('base64'); + + const ep: Episode = { + id: 'ep-1', + timestamp: Date.now(), + trigger: { + id: randomUUID(), + type: 'USER_PROMPT', + semanticParts: [ + { type: 'text', text: 'Look at this image:' }, + { + type: 'inline_data', + mimeType: 'image/png', + data: dummyImageBase64, + }, + ], + metadata: { + originalTokens: 300, + currentTokens: 300, + transformations: [], + }, + }, + steps: [], + }; + + // Fake token calculator says inlineData costs 258 tokens, text costs 10 + const state = getDummyState(false, 500, new Set()); + const result = await processor.process([ep], state); + + 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', + ); + + expect(fsPromises.writeFile).toHaveBeenCalledTimes(1); + expect(result[0].trigger.metadata.transformations.length).toBe(1); + }); + + it('degrades file_data into a text reference without disk write', async () => { + const ep: Episode = { + id: 'ep-2', + timestamp: Date.now(), + trigger: { + id: randomUUID(), + type: 'USER_PROMPT', + semanticParts: [ + { + type: 'file_data', + mimeType: 'application/pdf', + fileUri: 'gs://fake-bucket/doc.pdf', + }, + ], + metadata: { + originalTokens: 300, + currentTokens: 300, + transformations: [], + }, + }, + steps: [], + }; + + const state = getDummyState(false, 500, new Set()); + const result = await processor.process([ep], state); + + 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(fsPromises.writeFile).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/core/src/context/processors/blobDegradationProcessor.ts b/packages/core/src/context/processors/blobDegradationProcessor.ts new file mode 100644 index 0000000000..7dc54848cf --- /dev/null +++ b/packages/core/src/context/processors/blobDegradationProcessor.ts @@ -0,0 +1,123 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ +import type { Episode } from '../ir/types.js'; +import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextEnvironment } from '../sidecar/environment.js'; +import { estimateContextTokenCountSync as estimateTokenCountSync } from '../utils/contextTokenCalculator.js'; +import { sanitizeFilenamePart } from '../../utils/fileUtils.js'; +import * as fsPromises from 'node:fs/promises'; +import path from 'node:path'; +import type { Part } from '@google/genai'; + +export class BlobDegradationProcessor implements ContextProcessor { + readonly name = 'BlobDegradation'; + private env: ContextEnvironment; + + constructor(env: ContextEnvironment, options: Record = {}) { + this.env = env; + } + + async process( + episodes: Episode[], + state: ContextAccountingState, + ): Promise { + if (state.isBudgetSatisfied) { + return episodes; + } + + let currentDeficit = state.deficitTokens; + const newEpisodes = [...episodes]; + let directoryCreated = false; + + let blobOutputsDir = path.join( + this.env.getProjectTempDir(), + 'degraded-blobs', + ); + const sessionId = this.env.getSessionId(); + if (sessionId) { + blobOutputsDir = path.join( + blobOutputsDir, + `session-${sanitizeFilenamePart(sessionId)}`, + ); + } + + const ensureDir = async () => { + if (!directoryCreated) { + await fsPromises.mkdir(blobOutputsDir, { recursive: true }); + directoryCreated = true; + } + }; + + // Forward scan, looking for bloated non-text parts to degrade + for (let i = 0; i < newEpisodes.length; i++) { + if (currentDeficit <= 0) break; + const ep = newEpisodes[i]; + if (state.protectedEpisodeIds.has(ep.id)) continue; + + if (ep.trigger.type === 'USER_PROMPT') { + for (const part of ep.trigger.semanticParts) { + 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()}_${Math.random().toString(36).substring(7)}.${ext}`; + const filePath = path.join(blobOutputsDir, fileName); + + // Base64 to buffer + const buffer = Buffer.from(part.data, 'base64'); + await fsPromises.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 = estimateTokenCountSync([ + { inlineData: { mimeType: part.mimeType, data: part.data } }, + ]); + const newTokens = estimateTokenCountSync([{ 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 = estimateTokenCountSync([ + { fileData: { mimeType: part.mimeType, fileUri: part.fileUri } }, + ]); + const newTokens = estimateTokenCountSync([{ text: newText }]); + tokensSaved = oldTokens - newTokens; + } else if (part.type === 'raw_part') { + newText = `[Unknown Part degraded to text to preserve context window.]`; + // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion + const oldTokens = estimateTokenCountSync([part.part as Part]); + const newTokens = estimateTokenCountSync([{ text: newText }]); + tokensSaved = oldTokens - newTokens; + } + + if (newText && tokensSaved > 0) { + const newTokens = estimateTokenCountSync([{ text: newText }], 0, { + charsPerToken: this.env.getCharsPerToken(), + }); + part.presentation = { text: newText, tokens: newTokens }; + + ep.trigger.metadata.transformations.push({ + processorName: this.name, + action: 'DEGRADED', + timestamp: Date.now(), + }); + + currentDeficit -= tokensSaved; + } + } + } + } + + return newEpisodes; + } +} 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..06cb313872 --- /dev/null +++ b/packages/core/src/context/processors/historySquashingProcessor.test.ts @@ -0,0 +1,167 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ +import { createMockEnvironment } from '../testing/contextTestUtils.js'; +import { describe, it, expect, beforeEach } from 'vitest'; +import { HistorySquashingProcessor } from './historySquashingProcessor.js'; +import type { + Episode, + UserPrompt, + AgentThought, + AgentYield, +} from '../ir/types.js'; +import type { ContextAccountingState } from '../pipeline.js'; +import { randomUUID } from 'node:crypto'; + +describe('HistorySquashingProcessor', () => { + let processor: HistorySquashingProcessor; + + beforeEach(() => { + processor = new HistorySquashingProcessor(createMockEnvironment(), { + maxTokensPerNode: 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, + userText: string, + modelThought: string, + ): Episode => ({ + id, + timestamp: Date.now(), + trigger: { + id: randomUUID(), + type: 'USER_PROMPT', + semanticParts: [{ type: 'text', text: userText }], + metadata: { + originalTokens: 1000, + currentTokens: 1000, + transformations: [], + }, + }, + steps: [ + { + id: randomUUID(), + type: 'AGENT_THOUGHT', + text: modelThought, + metadata: { + originalTokens: 1000, + currentTokens: 1000, + transformations: [], + }, + }, + ], + }); + + it('bypasses processing if budget is satisfied', async () => { + const episodes = [createDummyEpisode('1', 'short text', 'short thought')]; + const state = getDummyState(true); + + const result = await processor.process(episodes, state); + + 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 = [createDummyEpisode('ep-1', longText, 'short thought')]; + const state = getDummyState(false, 100, new Set(['ep-1'])); + + const result = await processor.process(episodes, state); + + 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 = [createDummyEpisode('ep-2', longUser, longModel)]; + const state = getDummyState(false, 500, new Set()); // High deficit, force truncation + + const result = await processor.process(episodes, state); + + 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 = [ + createDummyEpisode('ep-3', longUser1, 'short'), + createDummyEpisode('ep-4', longUser2, 'short'), + ]; + + // Set deficit to exactly what ONE truncation will save + // Original = ~250 tokens. Limit = 100. Truncation saves ~150 tokens. + const state = getDummyState(false, 150, new Set()); + + const result = await processor.process(episodes, state); + + // 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 = createDummyEpisode('ep-5', 'short', 'short'); + ep.yield = { + id: randomUUID(), + type: 'AGENT_YIELD', + text: longYield, + metadata: { + originalTokens: 250, + currentTokens: 250, + transformations: [], + }, + }; + + const state = getDummyState(false, 500, new Set()); + const result = await processor.process([ep], state); + + 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..beda502644 --- /dev/null +++ b/packages/core/src/context/processors/historySquashingProcessor.ts @@ -0,0 +1,132 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Episode } from '../ir/types.js'; +import type { ContextAccountingState, ContextProcessor } from '../pipeline.js'; +import type { ContextEnvironment } from '../sidecar/environment.js'; +import { truncateProportionally } from '../truncation.js'; + +export class HistorySquashingProcessor implements ContextProcessor { + readonly name = 'HistorySquashing'; + private options: { maxTokensPerNode: number }; + + constructor(env: ContextEnvironment, options: { maxTokensPerNode: number }) { + 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( + episodes: Episode[], + state: ContextAccountingState, + ): Promise { + if (state.isBudgetSatisfied) { + return episodes; + } + + 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; + const newEpisodes = [...episodes]; + + for (let i = 0; i < newEpisodes.length; i++) { + if (currentDeficit <= 0) break; + if (state.protectedEpisodeIds.has(newEpisodes[i].id)) continue; + + const ep = newEpisodes[i]; + + // 1. Squash User Prompts + if (ep.trigger.type === 'USER_PROMPT') { + for (const part of ep.trigger.semanticParts) { + if (part.type === 'text') { + const saved = this.tryApplySquash( + part.text, + limitChars, + currentDeficit, + (p) => (part.presentation = p), + () => + ep.trigger.metadata.transformations.push({ + processorName: this.name, + action: 'TRUNCATED', + timestamp: Date.now(), + }), + ); + currentDeficit -= saved; + } + } + } + + // 2. Squash Model Thoughts + for (const step of ep.steps) { + if (currentDeficit <= 0) break; + if (step.type === 'AGENT_THOUGHT') { + const saved = this.tryApplySquash( + step.text, + limitChars, + currentDeficit, + (p) => (step.presentation = p), + () => + step.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) => (ep.yield!.presentation = p), + () => + ep.yield!.metadata.transformations.push({ + processorName: this.name, + action: 'TRUNCATED', + timestamp: Date.now(), + }), + ); + currentDeficit -= saved; + } + } + + return newEpisodes; + } +} 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..5af28443ce --- /dev/null +++ b/packages/core/src/context/processors/semanticCompressionProcessor.test.ts @@ -0,0 +1,154 @@ +/** + * @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 { SemanticCompressionProcessor } from './semanticCompressionProcessor.js'; +import type { + Episode, + UserPrompt, + ToolExecution, + AgentThought, +} from '../ir/types.js'; +import type { ContextAccountingState } from '../pipeline.js'; +import { randomUUID } from 'node:crypto'; + +describe('SemanticCompressionProcessor', () => { + let processor: SemanticCompressionProcessor; + let generateContentMock: ReturnType; + + beforeEach(() => { + generateContentMock = vi.fn().mockResolvedValue({ + candidates: [{ content: { parts: [{ text: 'Mocked Summary!' }] } }], + }); + + const env = createMockEnvironment(); + env.getLlmClient = vi + .fn() + .mockReturnValue({ generateContent: generateContentMock }) as any; + processor = new SemanticCompressionProcessor(env, { + nodeThresholdTokens: 2000, + }); + }); + + 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, + userText: string, + thoughtText: string, + toolObs: string, + ): Episode => ({ + id, + timestamp: Date.now(), + trigger: { + id: randomUUID(), + type: 'USER_PROMPT', + semanticParts: [{ type: 'text', text: userText }], + metadata: { + originalTokens: 3800, + currentTokens: 3800, + transformations: [], + }, + }, + steps: [ + { + id: randomUUID(), + type: 'AGENT_THOUGHT', + text: thoughtText, + metadata: { + originalTokens: 100, + currentTokens: 100, + transformations: [], + }, + }, + { + id: randomUUID(), + type: 'TOOL_EXECUTION', + toolName: 'test', + intent: {}, + observation: toolObs, + tokens: { intent: 10, observation: 3800 }, + metadata: { + originalTokens: 3810, + currentTokens: 3810, + transformations: [], + }, + }, + ], + }); + + it('bypasses processing if budget is satisfied', async () => { + const episodes = [createDummyEpisode('1', 'short', 'short', 'short')]; + const state = getDummyState(true); + + await processor.process(episodes, state); + expect(generateContentMock).not.toHaveBeenCalled(); + }); + + it('skips protected episodes even if over budget', async () => { + const massiveStr = 'M'.repeat(15000); // Exceeds threshold (10 * 4 = 40) + const episodes = [ + createDummyEpisode('ep-1', massiveStr, massiveStr, massiveStr), + ]; + const state = getDummyState(false, 1000, new Set(['ep-1'])); + + await processor.process(episodes, 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 = [ + createDummyEpisode('ep-1', massiveStr, massiveStr, massiveStr), + ]; + const state = getDummyState(false, 50000, new Set()); // Massive deficit, forces all 3 to summarize + + const result = await processor.process(episodes, state); + expect(generateContentMock).toHaveBeenCalledTimes(3); + + // Verify presentation layers were injected + 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 = [ + createDummyEpisode('ep-1', massiveStr, massiveStr, massiveStr), + ]; + + // Set deficit low enough that ONE summary solves the problem + const state = getDummyState(false, 5, new Set()); + + await processor.process(episodes, 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..74509d2791 --- /dev/null +++ b/packages/core/src/context/processors/semanticCompressionProcessor.ts @@ -0,0 +1,202 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { Episode } from '../ir/types.js'; +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 { estimateTokenCountSync } from '../../utils/tokenCalculation.js'; + +export class SemanticCompressionProcessor implements ContextProcessor { + readonly name = 'SemanticCompression'; + private env: ContextEnvironment; + private options: { nodeThresholdTokens: number }; + private modelToUse: string = 'chat-compression-2.5-flash-lite'; + + constructor( + env: ContextEnvironment, + options: { nodeThresholdTokens: number }, + ) { + this.env = env; + this.options = options; + } + + async process( + episodes: Episode[], + state: ContextAccountingState, + ): Promise { + require('fs').appendFileSync( + '/tmp/debug2.json', + 'SEMANTIC PROCESS: First episode ID: ' + + episodes[0]?.id + + '\nProtected IDs: ' + + Array.from(state.protectedEpisodeIds).join(', ') + + '\n', + ); + // If the budget is satisfied, or semantic compression isn't enabled + if (state.isBudgetSatisfied) { + return episodes; + } + + const semanticConfig = this.options; + // We estimate 4 chars per token for truncation logic + const thresholdChars = semanticConfig.nodeThresholdTokens * 4; + this.modelToUse = 'gemini-2.5-flash'; + + let currentDeficit = state.deficitTokens; + const newEpisodes = [...episodes]; + + // We scan backwards (oldest to newest would also work, but older is safer to degrade first) + for (let i = 0; i < newEpisodes.length; i++) { + if (currentDeficit <= 0) break; + const ep = newEpisodes[i]; + if (state.protectedEpisodeIds.has(ep.id)) continue; + + // 1. Compress User Prompts + if (ep.trigger.type === 'USER_PROMPT') { + for (const part of ep.trigger.semanticParts) { + 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 = estimateTokenCountSync([{ text: summary }]); + const oldTokens = estimateTokenCountSync([{ text: part.text }]); + + if (newTokens < oldTokens) { + part.presentation = { text: summary, tokens: newTokens }; + ep.trigger.metadata.transformations.push({ + processorName: this.name, + action: 'SUMMARIZED', + timestamp: Date.now(), + }); + currentDeficit -= oldTokens - newTokens; + } + } + } + } + + // 2. Compress Model Thoughts + for (const step of ep.steps) { + 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 = estimateTokenCountSync([{ text: summary }]); + const oldTokens = estimateTokenCountSync([{ text: step.text }]); + + if (newTokens < oldTokens) { + step.presentation = { text: summary, tokens: newTokens }; + step.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 (_e) { + 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 = estimateTokenCountSync([ + { + functionResponse: { + name: step.toolName, + response: newObsObject as unknown as Record, // eslint-disable-line @typescript-eslint/no-unsafe-type-assertion + id: step.id, + }, + }, + ]); + + const oldObsTokens = + step.presentation?.tokens.observation ?? step.tokens.observation; + const intentTokens = + step.presentation?.tokens.intent ?? step.tokens.intent; + + if (newObsTokens < oldObsTokens) { + step.presentation = { + intent: step.presentation?.intent ?? step.intent, + observation: newObsObject, + tokens: { intent: intentTokens, observation: newObsTokens }, + }; + step.metadata.transformations.push({ + processorName: this.name, + action: 'SUMMARIZED', + timestamp: Date.now(), + }); + currentDeficit -= oldObsTokens - newObsTokens; + } + } + } + } + } + + return newEpisodes; + } + + 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.getLlmClient(); + 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: ' + String(_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/toolMaskingProcessor.test.ts b/packages/core/src/context/processors/toolMaskingProcessor.test.ts new file mode 100644 index 0000000000..6ba240710b --- /dev/null +++ b/packages/core/src/context/processors/toolMaskingProcessor.test.ts @@ -0,0 +1,130 @@ +/** + * @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 type { Episode, ToolExecution } from '../ir/types.js'; +import type { ContextAccountingState } from '../pipeline.js'; +import { randomUUID } from 'node:crypto'; +import * as fsPromises from 'node:fs/promises'; + +vi.mock('node:fs/promises'); + +describe('ToolMaskingProcessor', () => { + let processor: ToolMaskingProcessor; + + beforeEach(() => { + vi.resetAllMocks(); + + processor = new ToolMaskingProcessor(createMockEnvironment(), { + 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 result = await processor.process(episodes, state); + require('fs').appendFileSync( + '/tmp/debug.json', + '\n\n' + JSON.stringify({ res: result[0].steps[0] }, null, 2), + ); + + 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 result = await processor.process(episodes, state); + + 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 + expect(fsPromises.writeFile).toHaveBeenCalledTimes(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..f68095fced --- /dev/null +++ b/packages/core/src/context/processors/toolMaskingProcessor.ts @@ -0,0 +1,243 @@ +/** + * @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 { estimateTokenCountSync } from '../../utils/tokenCalculation.js'; +import { sanitizeFilenamePart } from '../../utils/fileUtils.js'; +import * as fsPromises from 'node:fs/promises'; +import path from 'node:path'; +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 { Episode } from '../ir/types.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 class ToolMaskingProcessor implements ContextProcessor { + readonly name = 'ToolMasking'; + private env: ContextEnvironment; + private options: { stringLengthThresholdTokens: number }; + + constructor( + env: ContextEnvironment, + options: { stringLengthThresholdTokens: number }, + ) { + this.env = env; + this.options = options; + } + + async process( + episodes: Episode[], + state: ContextAccountingState, + ): Promise { + const maskingConfig = this.options; + if (!maskingConfig) return episodes; + if (state.isBudgetSatisfied) return episodes; + + const newEpisodes = [...episodes]; + let currentDeficit = state.deficitTokens; + const limitChars = maskingConfig.stringLengthThresholdTokens * 4; + + let toolOutputsDir = path.join( + this.env.getProjectTempDir(), + 'tool-outputs', + ); + const sessionId = this.env.getSessionId(); + if (sessionId) { + toolOutputsDir = path.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 fsPromises.mkdir(toolOutputsDir, { recursive: true }); + directoryCreated = true; + } + + const fileName = `${sanitizeFilenamePart(toolName).toLowerCase()}_${sanitizeFilenamePart(callId).toLowerCase()}_${nodeType}_${Math.random().toString(36).substring(7)}.txt`; + const filePath = path.join(toolOutputsDir, fileName); + + await fsPromises.writeFile(filePath, content, 'utf-8'); + + 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 (let i = 0; i < newEpisodes.length; i++) { + if (currentDeficit <= 0) break; + const ep = newEpisodes[i]; + 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(); + + /* eslint-disable @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-assignment */ + + const maskAsync = async ( + obj: any, + nodeType: string, + ): Promise<{ masked: any; changed: boolean }> => { + if (typeof obj === 'string') { + require('fs').appendFileSync( + '/tmp/debug.json', + 'STRING FOUND. length: ' + + obj.length + + ' limitChars: ' + + limitChars + + '\n', + ); + if (obj.length > 1000) + console.log( + 'Found string of length:', + obj.length, + 'limitChars is:', + limitChars, + 'isAlreadyMasked:', + this.isAlreadyMasked(obj), + ); + 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 = []; + 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 intentRes = await maskAsync( + step.presentation.intent ?? step.intent, + 'intent', + ); + const obsRes = await maskAsync( + step.presentation.observation ?? step.observation, + 'observation', + ); + + if (intentRes.changed || obsRes.changed) { + step.presentation.intent = intentRes.masked; + step.presentation.observation = obsRes.masked; + + // Recalculate tokens perfectly + const newIntentTokens = estimateTokenCountSync([ + { + functionCall: { + name: toolName, + args: intentRes.masked, + id: callId, + }, + }, + ]); + const newObsTokens = estimateTokenCountSync([ + { + functionResponse: { + name: toolName, + response: obsRes.masked, + 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) { + step.presentation.tokens = { + intent: newIntentTokens, + observation: newObsTokens, + }; + step.metadata.transformations.push({ + processorName: 'ToolMasking', + action: 'MASKED', + timestamp: Date.now(), + }); + currentDeficit -= savings; + } + } + } + } + + return newEpisodes; + } + + private isAlreadyMasked(content: string): boolean { + return content.includes(''); + } +} diff --git a/packages/core/src/context/sidecar/SidecarLoader.ts b/packages/core/src/context/sidecar/SidecarLoader.ts new file mode 100644 index 0000000000..3040c3d7b2 --- /dev/null +++ b/packages/core/src/context/sidecar/SidecarLoader.ts @@ -0,0 +1,41 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs'; +import type { Config } from '../../config/config.js'; +import type { SidecarConfig } from './types.js'; +import { defaultSidecarProfile } from './profiles.js'; + +export class SidecarLoader { + /** + * Generates a Sidecar JSON graph from the experimental config file path or defaults. + */ + static fromConfig(config: Config): SidecarConfig { + const sidecarPath = + typeof (config as any).getExperimentalContextSidecarConfig === 'function' + ? (config as any).getExperimentalContextSidecarConfig() + : undefined; + + if (sidecarPath && fs.existsSync(sidecarPath)) { + try { + const fileContent = fs.readFileSync(sidecarPath, 'utf8'); + return JSON.parse(fileContent) as SidecarConfig; + } catch (error) { + console.error( + `Failed to parse Sidecar configuration file at ${sidecarPath}:`, + error, + ); + // Fallback to default + } + } + + return defaultSidecarProfile; + } + + static fromLegacyConfig(config: Config): SidecarConfig { + return SidecarLoader.fromConfig(config); + } +} diff --git a/packages/core/src/context/sidecar/environment.ts b/packages/core/src/context/sidecar/environment.ts new file mode 100644 index 0000000000..a313926538 --- /dev/null +++ b/packages/core/src/context/sidecar/environment.ts @@ -0,0 +1,17 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { BaseLlmClient } from '../../core/baseLlmClient.js'; +import type { ContextTracer } from '../tracer.js'; + +export interface ContextEnvironment { + getLlmClient(): BaseLlmClient; + getSessionId(): string; + getTraceDir(): string; + getProjectTempDir(): string; + getTracer(): ContextTracer; + getCharsPerToken(): number; +} diff --git a/packages/core/src/context/sidecar/environmentImpl.ts b/packages/core/src/context/sidecar/environmentImpl.ts new file mode 100644 index 0000000000..1ea179d919 --- /dev/null +++ b/packages/core/src/context/sidecar/environmentImpl.ts @@ -0,0 +1,44 @@ +/** + * @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'; + +export class ContextEnvironmentImpl implements ContextEnvironment { + constructor( + private llmClient: BaseLlmClient, + private sessionId: string, + private traceDir: string, + private tempDir: string, + private tracer: ContextTracer, + private charsPerToken: number, + ) {} + + getLlmClient(): BaseLlmClient { + return this.llmClient; + } + + getSessionId(): string { + return this.sessionId; + } + + getTraceDir(): string { + return this.traceDir; + } + + getProjectTempDir(): string { + return this.tempDir; + } + + getTracer(): ContextTracer { + return this.tracer; + } + + getCharsPerToken(): number { + return this.charsPerToken; + } +} diff --git a/packages/core/src/context/sidecar/profiles.ts b/packages/core/src/context/sidecar/profiles.ts new file mode 100644 index 0000000000..10b9e71cfa --- /dev/null +++ b/packages/core/src/context/sidecar/profiles.ts @@ -0,0 +1,51 @@ +/** + * @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: { + eagerBackground: [ + { + processorId: 'StateSnapshotWorker', + options: { pollingIntervalMs: 5000 }, + }, + ], + retainedProcessingGraph: [ + { + processorId: 'HistorySquashingProcessor', + options: { maxTokensPerNode: 3000 }, + }, + ], + normalProcessingGraph: [ + { + processorId: 'ToolMaskingProcessor', + options: { stringLengthThresholdTokens: 8000 }, + }, + { + processorId: 'BlobDegradationProcessor', + options: {}, + }, + { + processorId: 'SemanticCompressionProcessor', + options: { nodeThresholdTokens: 5000, contextWindowPercentage: 0.2 }, + }, + ], + }, +}; diff --git a/packages/core/src/context/sidecar/registry.ts b/packages/core/src/context/sidecar/registry.ts new file mode 100644 index 0000000000..1ad9115b42 --- /dev/null +++ b/packages/core/src/context/sidecar/registry.ts @@ -0,0 +1,42 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextProcessor } from '../pipeline.js'; +import type { AsyncContextWorker } from '../workers/asyncContextWorker.js'; +import type { ContextEnvironment } from './environment.js'; + +export interface ContextProcessorDef< + TOptions extends Record = any, +> { + readonly id: string; + create( + env: ContextEnvironment, + options: TOptions, + ): ContextProcessor | AsyncContextWorker; +} + +/** + * Registry for mapping declarative sidecar configs to running Processor instances. + */ +export class ProcessorRegistry { + private static processors = new Map(); + + static register(def: ContextProcessorDef) { + this.processors.set(def.id, def); + } + + static get(id: string): ContextProcessorDef { + const def = this.processors.get(id); + if (!def) { + throw new Error(`Context Processor [${id}] is not registered.`); + } + return def; + } + + static clear() { + this.processors.clear(); + } +} diff --git a/packages/core/src/context/sidecar/types.ts b/packages/core/src/context/sidecar/types.ts new file mode 100644 index 0000000000..fbb62625a5 --- /dev/null +++ b/packages/core/src/context/sidecar/types.ts @@ -0,0 +1,55 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Definition of a processor or worker to be instantiated in the graph. + */ +export interface ProcessorConfig { + /** The registered ID of the processor (e.g. 'SemanticCompressionProcessor') */ + processorId: string; + + /** Dynamic, processor-specific hyperparameters */ + options: Record; +} + +/** + * 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: { + /** + * Eagerly executes in the background when the 'retainedTokens' boundary is crossed. + * Contains AsyncContextWorkers (e.g. StateSnapshotWorker). + */ + eagerBackground: ProcessorConfig[]; + + /** + * Executes sequentially to protect the pristine outliers within the retained window. + * Contains ContextProcessors (e.g. HistorySquashingProcessor). + */ + retainedProcessingGraph: ProcessorConfig[]; + + /** + * Executes sequentially to opportunistically degrade messages older than the retained window. + * Contains ContextProcessors (e.g. ToolMaskingProcessor, SemanticCompressionProcessor). + */ + normalProcessingGraph: ProcessorConfig[]; + }; +} diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts new file mode 100644 index 0000000000..440e01b070 --- /dev/null +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -0,0 +1,113 @@ +/** + * @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'; + +export function createMockEnvironment(): ContextEnvironment { + return { + getLlmClient: vi.fn().mockReturnValue({ + generateContent: vi.fn().mockResolvedValue({ + text: 'Mock LLM summary response', + }), + }) as any, + getSessionId: vi.fn().mockReturnValue('mock-session'), + getTraceDir: vi.fn().mockReturnValue('/tmp/.gemini/trace'), + getProjectTempDir: vi.fn().mockReturnValue('/tmp'), + getTracer: vi.fn().mockReturnValue({ + logEvent: vi.fn(), + saveAsset: vi.fn().mockReturnValue('mock-asset-id'), + }) as any, + getCharsPerToken: vi.fn().mockReturnValue(1), + }; +} + +import type { Content } from '@google/genai'; +import { AgentChatHistory } from '../../core/agentChatHistory.js'; +import { ContextManager } from '../contextManager.js'; + +/** + * 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'), + }; + + 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'; + +export function setupContextComponentTest(config: Config) { + const chatHistory = new AgentChatHistory(); + const sidecar = SidecarLoader.fromLegacyConfig(config); + const tracer = new ContextTracer('/tmp', 'test-session'); + const env = new ContextEnvironmentImpl( + config.getBaseLlmClient() as any, + 'test-session', + '/tmp', + '/tmp/gemini-test', + tracer, + 1, + ); + const contextManager = new ContextManager(sidecar, env, tracer); + + // 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.ts b/packages/core/src/context/tracer.ts new file mode 100644 index 0000000000..20baced9fc --- /dev/null +++ b/packages/core/src/context/tracer.ts @@ -0,0 +1,70 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import { randomUUID } from 'node:crypto'; + +export class ContextTracer { + private traceDir: string; + private assetsDir: string; + private enabled: boolean; + + constructor(targetDir: string, sessionId: string) { + this.enabled = process.env['GEMINI_CONTEXT_TRACE'] === 'true'; + this.traceDir = path.join(targetDir, '.gemini', 'context_trace', sessionId); + this.assetsDir = path.join(this.traceDir, 'assets'); + + if (this.enabled) { + try { + fs.mkdirSync(this.assetsDir, { recursive: true }); + this.logEvent('SYSTEM', 'Context Tracer Initialized', { sessionId }); + } catch (e) { + console.error('Failed to initialize ContextTracer', e); + this.enabled = false; + } + } + } + + logEvent( + component: string, + action: string, + details?: Record, + ) { + if (!this.enabled) return; + try { + const timestamp = new Date().toISOString(); + const detailsStr = details + ? ` | Details: ${JSON.stringify(details)}` + : ''; + const logLine = `[${timestamp}] [${component}] ${action}${detailsStr}\n`; + fs.appendFileSync( + path.join(this.traceDir, 'trace.log'), + logLine, + 'utf-8', + ); + } catch (e) { + // fail silently in trace + } + } + + saveAsset(component: string, assetName: string, data: unknown): string { + if (!this.enabled) return 'asset-recording-disabled'; + try { + const assetId = `${Date.now()}-${randomUUID().slice(0, 6)}-${assetName}.json`; + const assetPath = path.join(this.assetsDir, assetId); + + fs.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..d120e8fc63 --- /dev/null +++ b/packages/core/src/context/utils/contextTokenCalculator.ts @@ -0,0 +1,31 @@ +/** + * @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'; + +export function estimateContextTokenCountSync( + parts: Part[], + depth: number = 0, + config?: { charsPerToken?: number }, +): number { + if (config?.charsPerToken !== undefined && config.charsPerToken !== 4) { + let totalTokens = 0; + for (const part of parts) { + if (typeof part.text === 'string') { + totalTokens += Math.ceil(part.text.length / config.charsPerToken); + } else { + totalTokens += Math.ceil( + JSON.stringify(part).length / config.charsPerToken, + ); + } + } + return totalTokens; + } + + // The baseEstimate no longer accepts config because we forked it! + return baseEstimate(parts, depth); +} diff --git a/packages/core/src/context/workers/asyncContextWorker.ts b/packages/core/src/context/workers/asyncContextWorker.ts new file mode 100644 index 0000000000..6727def2bf --- /dev/null +++ b/packages/core/src/context/workers/asyncContextWorker.ts @@ -0,0 +1,18 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { ContextEventBus } from '../eventBus.js'; + +export interface AsyncContextWorker { + /** The unique name of the worker (e.g., 'StateSnapshotWorker') */ + readonly name: string; + + /** Starts listening to the ContextEventBus for background tasks */ + start(bus: ContextEventBus): void; + + /** Stops listening and aborts any pending background tasks */ + stop(): void; +} diff --git a/packages/core/src/context/workers/stateSnapshotWorker.ts b/packages/core/src/context/workers/stateSnapshotWorker.ts new file mode 100644 index 0000000000..0a7e5e4aa0 --- /dev/null +++ b/packages/core/src/context/workers/stateSnapshotWorker.ts @@ -0,0 +1,229 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { randomUUID } from 'node:crypto'; +import type { ContextEnvironment } from '../sidecar/environment.js'; +import type { Episode, SnapshotVariant } from '../ir/types.js'; +import type { AsyncContextWorker } from './asyncContextWorker.js'; +import type { + ContextEventBus, + ContextConsolidationEvent, +} from '../eventBus.js'; +import { debugLogger } from '../../utils/debugLogger.js'; +import { estimateContextTokenCountSync as estimateTokenCountSync } from '../utils/contextTokenCalculator.js'; +import { IrMapper } from '../ir/mapper.js'; +import { LlmRole } from '../../telemetry/llmRole.js'; +import type { ContextTracer } from '../tracer.js'; + +export class StateSnapshotWorker implements AsyncContextWorker { + name = 'StateSnapshotWorker'; + private bus?: ContextEventBus; + private tracer?: ContextTracer; + private isSynthesizing = false; + + constructor(private readonly env: ContextEnvironment) {} + + start(bus: ContextEventBus, tracer?: ContextTracer): void { + console.log('Worker start() called with bus:', !!bus); + this.bus = bus; + this.tracer = tracer; + this.bus.onConsolidationNeeded(this.handleConsolidation.bind(this)); + } + + stop(): void { + if (this.bus) { + // In a real implementation we would `removeListener` here + this.bus = undefined; + } + } + + private async handleConsolidation( + event: ContextConsolidationEvent, + ): Promise { + console.log( + `Worker handling consolidation. targetDeficit: ${event.targetDeficit}, isSynthesizing: ${this.isSynthesizing}`, + ); + if (this.isSynthesizing || event.targetDeficit <= 0) return; + + // Identify the "dying" block of episodes that need to be collected. + // For now, we assume older episodes are at the front of the array. + // We only want episodes that don't already have a snapshot variant computing/ready. + const unprotectedOldest = event.episodes.filter( + (ep) => !ep.variants?.['snapshot'], + ); + + if (unprotectedOldest.length === 0) { + return; + } + + let targetDeficit = event.targetDeficit; + const episodesToSynthesize: Episode[] = []; + let tokensToSynthesize = 0; + + for (const ep of unprotectedOldest) { + console.log('Worker considering episode:', ep.id); + if (tokensToSynthesize >= targetDeficit) break; + episodesToSynthesize.push(ep); + // Rough estimate of tokens in this episode + const epTokens = ep.steps.reduce( + (sum, step) => sum + step.metadata.currentTokens, + ep.trigger.metadata.currentTokens + + (ep.yield?.metadata.currentTokens || 0), + ); + tokensToSynthesize += epTokens; + } + + if (episodesToSynthesize.length === 0) return; + + console.log( + `Worker synthesized logic loop complete. Selected ${episodesToSynthesize.length} episodes for ~${tokensToSynthesize} tokens.`, + ); + this.isSynthesizing = true; + + try { + debugLogger.log( + `StateSnapshotWorker: Asynchronously synthesizing ${episodesToSynthesize.length} episodes to recover ~${tokensToSynthesize} tokens.`, + ); + this.tracer?.logEvent( + 'StateSnapshotWorker', + `Consolidation requested. Synthesizing ${episodesToSynthesize.length} episodes for ~${tokensToSynthesize} tokens.`, + ); + + const client = this.env.getLlmClient(); + const rawContents = IrMapper.fromIr(episodesToSynthesize); + const rawAssetId = this.tracer?.saveAsset( + 'StateSnapshotWorker', + 'episodes_to_synthesize', + rawContents, + ); + this.tracer?.logEvent( + 'StateSnapshotWorker', + 'Dispatching LLM request for snapshot generation', + { rawAssetId }, + ); + + const promptText = ` +You are a background memory consolidation worker for an AI assistant. +Your task is to review the following block of the oldest conversation history and synthesize it into a highly dense, accurate "World State Snapshot". +This snapshot will completely replace these old memories. +Preserve all critical facts, technical decisions, file paths, and outstanding tasks. Discard all conversational filler. + +Conversation History to Synthesize: +${JSON.stringify(rawContents, null, 2).slice(0, 50000)} + +Output the snapshot as a dense, structured summary.`; + + const response = await client.generateContent({ + modelConfigKey: { model: 'gemini-2.5-flash' }, // Fast and cheap for background tasks + contents: [{ role: 'user', parts: [{ text: promptText }] }], + promptId: 'async-world-state-snapshot', + role: LlmRole.UTILITY_COMPRESSOR, + abortSignal: new AbortController().signal, // Run in background, could add cancellation logic later + }); + + // Extract text safely from the GenAI response + const snapshotText = response.text; + const responseAssetId = this.tracer?.saveAsset( + 'StateSnapshotWorker', + 'snapshot_response', + snapshotText || '', + ); + this.tracer?.logEvent('StateSnapshotWorker', 'Received LLM response', { + responseAssetId, + }); + if (!snapshotText) { + debugLogger.warn( + 'StateSnapshotWorker: LLM returned empty response for snapshot generation.', + ); + } + + const mockSnapshotText = ` + +${snapshotText || '[Failed to generate snapshot]'} +`; + + const snapshotTokens = estimateTokenCountSync( + [{ text: mockSnapshotText }], + 0, + { charsPerToken: this.env.getCharsPerToken() }, + ); + + const replacedEpisodeIds = episodesToSynthesize.map((e) => e.id); + + const snapshotEpisode: Episode = { + id: randomUUID(), + timestamp: Date.now(), + trigger: { + id: randomUUID(), + type: 'SYSTEM_EVENT', + name: 'world_state_snapshot', + payload: { + originalEpisodeCount: episodesToSynthesize.length, + recoveredTokens: tokensToSynthesize, + }, + metadata: { + originalTokens: snapshotTokens, + currentTokens: snapshotTokens, + transformations: [ + { + processorName: 'StateSnapshotWorker', + action: 'SYNTHESIZED', + timestamp: Date.now(), + }, + ], + }, + }, + steps: [ + { + id: randomUUID(), + type: 'AGENT_THOUGHT', + text: mockSnapshotText, + metadata: { + originalTokens: snapshotTokens, + currentTokens: snapshotTokens, + transformations: [], + }, + }, + ], + }; + + const variant: SnapshotVariant = { + type: 'snapshot', + status: 'ready', + recoveredTokens: tokensToSynthesize, + episode: snapshotEpisode, + replacedEpisodeIds, + }; + + // Emit the variant for the MOST RECENT episode in the batch, + // since the Opportunistic Swapper sweeps from newest to oldest. + const targetId = replacedEpisodeIds[replacedEpisodeIds.length - 1]; + + if (this.bus) { + this.tracer?.logEvent( + 'StateSnapshotWorker', + `Emitting VARIANT_READY for targetId [${targetId}]`, + ); + + this.bus.emitVariantReady({ + targetId, + variantId: 'snapshot', + variant, + }); + } else { + debugLogger.warn( + 'StateSnapshotWorker: Event bus disconnected before variant could be emitted.', + ); + } + } catch (error) { + debugLogger.error( + `StateSnapshotWorker: Critical failure during snapshot synthesis: ${error instanceof Error ? error.message : String(error)}`, + ); + } finally { + this.isSynthesizing = false; + } + } +} 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; + } +}