feat(core): wire up the new ContextManager and AgentChatHistory (#25409)

2026-05-14 13:53:02 -07:00 · 2026-04-27 11:50:00 -07:00
parent 98aca28985
commit 71f313b51a
35 changed files with 707 additions and 377 deletions
@@ -0,0 +1,117 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { Config } from '../config/config.js';
+import type { GeminiChat } from '../core/geminiChat.js';
+import { ContextProcessorRegistry } from './config/registry.js';
+import { loadContextManagementConfig } from './config/configLoader.js';
+import { ContextTracer } from './tracer.js';
+import { ContextEventBus } from './eventBus.js';
+import { ContextEnvironmentImpl } from './pipeline/environmentImpl.js';
+import { PipelineOrchestrator } from './pipeline/orchestrator.js';
+import { ContextManager } from './contextManager.js';
+import { debugLogger } from '../utils/debugLogger.js';
+import { NodeTruncationProcessorOptionsSchema } from './processors/nodeTruncationProcessor.js';
+import { ToolMaskingProcessorOptionsSchema } from './processors/toolMaskingProcessor.js';
+import { HistoryTruncationProcessorOptionsSchema } from './processors/historyTruncationProcessor.js';
+import { BlobDegradationProcessorOptionsSchema } from './processors/blobDegradationProcessor.js';
+import { NodeDistillationProcessorOptionsSchema } from './processors/nodeDistillationProcessor.js';
+import { StateSnapshotProcessorOptionsSchema } from './processors/stateSnapshotProcessor.js';
+import { StateSnapshotAsyncProcessorOptionsSchema } from './processors/stateSnapshotAsyncProcessor.js';
+import { RollingSummaryProcessorOptionsSchema } from './processors/rollingSummaryProcessor.js';
+
+export async function initializeContextManager(
+  config: Config,
+  chat: GeminiChat,
+  lastPromptId: string,
+): Promise<ContextManager | undefined> {
+  const isV1Enabled = config.getContextManagementConfig().enabled;
+  debugLogger.log(
+    `[initializer] called with enabled=${isV1Enabled}, GEMINI_CONTEXT_TRACE_DIR=${process.env['GEMINI_CONTEXT_TRACE_DIR']}`,
+  );
+
+  if (!isV1Enabled) {
+    return undefined;
+  }
+
+  const registry = new ContextProcessorRegistry();
+  registry.registerProcessor({
+    id: 'NodeTruncationProcessor',
+    schema: NodeTruncationProcessorOptionsSchema,
+  });
+  registry.registerProcessor({
+    id: 'ToolMaskingProcessor',
+    schema: ToolMaskingProcessorOptionsSchema,
+  });
+  registry.registerProcessor({
+    id: 'HistoryTruncationProcessor',
+    schema: HistoryTruncationProcessorOptionsSchema,
+  });
+  registry.registerProcessor({
+    id: 'BlobDegradationProcessor',
+    schema: BlobDegradationProcessorOptionsSchema,
+  });
+  registry.registerProcessor({
+    id: 'NodeDistillationProcessor',
+    schema: NodeDistillationProcessorOptionsSchema,
+  });
+  registry.registerProcessor({
+    id: 'StateSnapshotProcessor',
+    schema: StateSnapshotProcessorOptionsSchema,
+  });
+  registry.registerProcessor({
+    id: 'StateSnapshotAsyncProcessor',
+    schema: StateSnapshotAsyncProcessorOptionsSchema,
+  });
+  registry.registerProcessor({
+    id: 'RollingSummaryProcessor',
+    schema: RollingSummaryProcessorOptionsSchema,
+  });
+
+  const sidecarProfile = await loadContextManagementConfig(
+    config.getExperimentalContextManagementConfig(),
+    registry,
+  );
+
+  const storage = config.storage;
+  const logDir = storage.getProjectTempLogsDir();
+  const projectTempDir = storage.getProjectTempDir();
+
+  const tracer = new ContextTracer({
+    enabled: !!process.env['GEMINI_CONTEXT_TRACE_DIR'],
+    targetDir: projectTempDir,
+    sessionId: lastPromptId,
+  });
+
+  const eventBus = new ContextEventBus();
+
+  const env = new ContextEnvironmentImpl(
+    () => config.getBaseLlmClient(),
+    config.getSessionId(),
+    lastPromptId,
+    logDir,
+    projectTempDir,
+    tracer,
+    4,
+    eventBus,
+  );
+
+  const orchestrator = new PipelineOrchestrator(
+    sidecarProfile.buildPipelines(env),
+    sidecarProfile.buildAsyncPipelines(env),
+    env,
+    eventBus,
+    tracer,
+  );
+
+  return new ContextManager(
+    sidecarProfile,
+    env,
+    tracer,
+    orchestrator,
+    chat.agentHistory,
+  );
+}