diff --git a/packages/core/src/context/contextManager.golden.test.ts b/packages/core/src/context/contextManager.golden.test.ts index 473c2d99cb..907add6a5a 100644 --- a/packages/core/src/context/contextManager.golden.test.ts +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -18,6 +18,8 @@ import { ContextEnvironmentImpl } from './sidecar/environmentImpl.js'; import { SidecarLoader } from './sidecar/SidecarLoader.js'; import { ContextTracer } from './tracer.js'; import { ContextEventBus } from './eventBus.js'; +import { PipelineOrchestrator } from './sidecar/orchestrator.js'; +import { AgentChatHistory } from '../core/agentChatHistory.js'; import type { Content } from '@google/genai'; import type { BaseLlmClient } from '../core/baseLlmClient.js'; import type { Episode } from './ir/types.js'; @@ -96,12 +98,21 @@ describe('ContextManager Golden Tests', () => { 4, eventBus, ); - contextManager = ContextManager.create( + const chatHistory = new AgentChatHistory(); + const orchestrator = new PipelineOrchestrator( + sidecar, + env, + eventBus, + tracer, + registry + ); + + contextManager = new ContextManager( sidecar, env, tracer, - undefined, - registry, + orchestrator, + chatHistory ); }); diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index c632bba116..5ae2616bb2 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -15,8 +15,6 @@ import type { SidecarConfig } from './sidecar/types.js'; import { PipelineOrchestrator } from './sidecar/orchestrator.js'; import { HistoryObserver } from './historyObserver.js'; import { IrProjector } from './ir/projector.js'; -import { registerBuiltInProcessors } from './sidecar/builtins.js'; -import { SidecarRegistry } from './sidecar/registry.js'; export class ContextManager { // The stateful, pristine flat graph. @@ -25,36 +23,28 @@ export class ContextManager { private readonly eventBus: ContextEventBus; // Internal sub-components - private orchestrator: PipelineOrchestrator; - private historyObserver?: HistoryObserver; + private readonly orchestrator: PipelineOrchestrator; + private readonly historyObserver: HistoryObserver; - static create( - sidecar: SidecarConfig, - env: ContextEnvironment, - tracer: ContextTracer, - orchestrator?: PipelineOrchestrator, - registry?: SidecarRegistry, - ): ContextManager { - if (!registry) { - registry = new SidecarRegistry(); - registerBuiltInProcessors(registry); - } - const orch = - orchestrator || - new PipelineOrchestrator(sidecar, env, env.eventBus, tracer, registry); - return new ContextManager(sidecar, env, tracer, orch); - } - - // Use ContextManager.create() instead - private constructor( - private sidecar: SidecarConfig, - private env: ContextEnvironment, + constructor( + private readonly sidecar: SidecarConfig, + private readonly env: ContextEnvironment, private readonly tracer: ContextTracer, orchestrator: PipelineOrchestrator, + chatHistory: AgentChatHistory, ) { this.eventBus = env.eventBus; this.orchestrator = orchestrator; + this.historyObserver = new HistoryObserver( + chatHistory, + this.env.eventBus, + this.tracer, + this.env.tokenCalculator, + this.env.irMapper, + ); + this.historyObserver.start(); + this.eventBus.onPristineHistoryUpdated((event) => { this.pristineNodes = event.nodes; // In V2, we assume currentNodes updates sequentially via Orchestrator patches. @@ -89,9 +79,7 @@ export class ContextManager { */ shutdown() { this.orchestrator.shutdown(); - if (this.historyObserver) { - this.historyObserver.stop(); - } + this.historyObserver.stop(); } /** @@ -132,24 +120,6 @@ export class ContextManager { } } - /** - * Starts tracking the raw agent history and translating it to Episodic IR. - */ - subscribeToHistory(chatHistory: AgentChatHistory) { - if (this.historyObserver) { - this.historyObserver.stop(); - } - - this.historyObserver = new HistoryObserver( - chatHistory, - this.env.eventBus, - this.tracer, - this.env.tokenCalculator, - this.env.irMapper, - ); - this.historyObserver.start(); - } - /** * Retrieves the raw, uncompressed Episodic IR graph. * Useful for internal tool rendering (like the trace viewer). diff --git a/packages/core/src/context/historyObserver.ts b/packages/core/src/context/historyObserver.ts index 3ac9eb6eee..7fef1f26fc 100644 --- a/packages/core/src/context/historyObserver.ts +++ b/packages/core/src/context/historyObserver.ts @@ -13,6 +13,8 @@ import type { ContextTokenCalculator } from './utils/contextTokenCalculator.js'; import type { ContextEventBus } from './eventBus.js'; import type { ContextTracer } from './tracer.js'; +import type { ConcreteNode } from './ir/types.js'; + /** * Connects the raw AgentChatHistory to the ContextManager. * It maps raw messages into Episodic Intermediate Representation (IR) @@ -46,7 +48,7 @@ export class HistoryObserver { this.tokenCalculator, ); - const nodes: Array = []; + const nodes: Array = []; for (const ep of pristineEpisodes) { if (ep.concreteNodes) { for (const child of ep.concreteNodes) { diff --git a/packages/core/src/context/ir/behaviorRegistry.ts b/packages/core/src/context/ir/behaviorRegistry.ts index 934af53bb8..ea39644fcc 100644 --- a/packages/core/src/context/ir/behaviorRegistry.ts +++ b/packages/core/src/context/ir/behaviorRegistry.ts @@ -22,13 +22,13 @@ export interface IrNodeBehavior { } export class IrNodeBehaviorRegistry { - private readonly behaviors = new Map>(); + private readonly behaviors = new Map>(); register(behavior: IrNodeBehavior) { - this.behaviors.set(behavior.type, behavior); + this.behaviors.set(behavior.type, behavior as IrNodeBehavior); } - get(type: string): IrNodeBehavior { + get(type: string): IrNodeBehavior { const behavior = this.behaviors.get(type); if (!behavior) { throw new Error(`Unregistered IrNode type: ${type}`); diff --git a/packages/core/src/context/ir/builtinBehaviors.ts b/packages/core/src/context/ir/builtinBehaviors.ts index e84d347335..8840b7d00a 100644 --- a/packages/core/src/context/ir/builtinBehaviors.ts +++ b/packages/core/src/context/ir/builtinBehaviors.ts @@ -8,6 +8,7 @@ import type { AgentYield, Snapshot, RollingSummary, + SystemEvent, } from './types.js'; export const UserPromptBehavior: IrNodeBehavior = { @@ -15,10 +16,20 @@ export const UserPromptBehavior: IrNodeBehavior = { getEstimatableParts(prompt) { const parts: Part[] = []; for (const sp of prompt.semanticParts) { - if (sp.type === 'text') parts.push({ text: sp.text }); - else if (sp.type === 'inline_data') parts.push({ inlineData: { mimeType: sp.mimeType, data: sp.data } }); - else if (sp.type === 'file_data') parts.push({ fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri } }); - else if (sp.type === 'raw_part') parts.push(sp.part); + switch (sp.type) { + case 'text': + parts.push({ text: sp.text }); + break; + case 'inline_data': + parts.push({ inlineData: { mimeType: sp.mimeType, data: sp.data } }); + break; + case 'file_data': + parts.push({ fileData: { mimeType: sp.mimeType, fileUri: sp.fileUri } }); + break; + case 'raw_part': + parts.push(sp.part); + break; + } } return parts; }, @@ -84,7 +95,7 @@ export const AgentYieldBehavior: IrNodeBehavior = { } }; -export const SystemEventBehavior: IrNodeBehavior = { +export const SystemEventBehavior: IrNodeBehavior = { type: 'SYSTEM_EVENT', getEstimatableParts() { return []; }, serialize(node, writer) { diff --git a/packages/core/src/context/ir/graphUtils.ts b/packages/core/src/context/ir/graphUtils.ts index 6c4b542ef7..de10a70783 100644 --- a/packages/core/src/context/ir/graphUtils.ts +++ b/packages/core/src/context/ir/graphUtils.ts @@ -41,5 +41,5 @@ export function isSnapshot(node: IrNode): node is Snapshot { return node.type === 'SNAPSHOT'; } export function isRollingSummary(node: IrNode): node is RollingSummary { - return node.type === 'USER_PROMPT'; + return node.type === 'ROLLING_SUMMARY'; } \ No newline at end of file diff --git a/packages/core/src/context/ir/toIr.ts b/packages/core/src/context/ir/toIr.ts index 4c330e8456..8129908244 100644 --- a/packages/core/src/context/ir/toIr.ts +++ b/packages/core/src/context/ir/toIr.ts @@ -113,7 +113,8 @@ function parseToolResponses( }; } - for (const part of msg.parts!) { + const parts = msg.parts || []; + for (const part of parts) { if (part.functionResponse) { const callId = part.functionResponse.id || ''; const matchingCall = pendingCallParts.get(callId); @@ -153,7 +154,8 @@ function parseUserParts( nodeIdentityMap: WeakMap ): Partial { const semanticParts: SemanticPart[] = []; - for (const p of msg.parts!) { + const parts = msg.parts || []; + for (const p of parts) { if (p.text !== undefined) semanticParts.push({ type: 'text', text: p.text }); else if (p.inlineData) @@ -172,8 +174,9 @@ function parseUserParts( semanticParts.push({ type: 'raw_part', part: p }); // Preserve unknowns } + const baseObj = parts.length > 0 ? parts[0] : msg; const trigger: UserPrompt = { - id: getStableId(msg.parts![0] || msg, nodeIdentityMap), + id: getStableId(baseObj, nodeIdentityMap), type: 'USER_PROMPT', semanticParts, }; @@ -198,7 +201,8 @@ function parseModelParts( }; } - for (const part of msg.parts!) { + const parts = msg.parts || []; + for (const part of parts) { if (part.functionCall) { const callId = part.functionCall.id || ''; if (callId) pendingCallParts.set(callId, part); diff --git a/packages/core/src/context/system-tests/SimulationHarness.ts b/packages/core/src/context/system-tests/SimulationHarness.ts index 90834f8e49..ba170ca967 100644 --- a/packages/core/src/context/system-tests/SimulationHarness.ts +++ b/packages/core/src/context/system-tests/SimulationHarness.ts @@ -85,14 +85,13 @@ export class SimulationHarness { this.tracer, registry, ); - this.contextManager = ContextManager.create( + this.contextManager = new ContextManager( config, this.env, this.tracer, this.orchestrator, - registry, + this.chatHistory ); - this.contextManager.subscribeToHistory(this.chatHistory); } /** diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index f9bcf84278..4b33e53c3a 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -20,6 +20,7 @@ import { registerBuiltInBehaviors } from '../ir/builtinBehaviors.js'; import { IrMapper } from '../ir/mapper.js'; import { SidecarRegistry } from '../sidecar/registry.js'; import { registerBuiltInProcessors } from '../sidecar/builtins.js'; +import { PipelineOrchestrator } from '../sidecar/orchestrator.js'; import type { ConcreteNode, ToolExecution } from '../ir/types.js'; import type { ContextEnvironment } from '../sidecar/environment.js'; import type { Config } from '../../config/config.js'; @@ -209,18 +210,24 @@ export function setupContextComponentTest( 1, eventBus, ); - const contextManager = ContextManager.create( + + const orchestrator = new PipelineOrchestrator( + sidecar, + env, + eventBus, + tracer, + registry + ); + + const contextManager = new ContextManager( sidecar, env, tracer, - undefined, - registry, + orchestrator, + chatHistory ); // The async worker is now internally managed by ContextManager - // Subscribe to history to enable the Eager/Opportunistic triggers - contextManager.subscribeToHistory(chatHistory); - return { chatHistory, contextManager }; }