diff --git a/packages/core/src/context/contextManager.golden.test.ts b/packages/core/src/context/contextManager.golden.test.ts index e6063d2330..cb00a36b60 100644 --- a/packages/core/src/context/contextManager.golden.test.ts +++ b/packages/core/src/context/contextManager.golden.test.ts @@ -25,7 +25,6 @@ import type { Episode } from './ir/types.js'; import type { SidecarConfig } from './sidecar/types.js'; import { ProcessorRegistry } from './sidecar/registry.js'; import { registerBuiltInProcessors } from './sidecar/builtins.js'; -import { IrMapper } from './ir/mapper.js'; import { createMockContextConfig, setupContextComponentTest } from './testing/contextTestUtils.js'; expect.addSnapshotSerializer({ @@ -141,7 +140,7 @@ describe('ContextManager Golden Tests', () => { const history = createLargeHistory(); ( contextManager as unknown as { pristineEpisodes: Episode[] } - ).pristineEpisodes = IrMapper.toIr(history, new ContextTokenCalculator(4)); + ).pristineEpisodes = (contextManager as any).env.irMapper.toIr(history, new ContextTokenCalculator(4)); const result = await contextManager.projectCompressedHistory(); expect(result).toMatchSnapshot(); }); diff --git a/packages/core/src/context/contextManager.ts b/packages/core/src/context/contextManager.ts index 7c0e201a59..497f05a064 100644 --- a/packages/core/src/context/contextManager.ts +++ b/packages/core/src/context/contextManager.ts @@ -141,10 +141,11 @@ export class ContextManager { } this.historyObserver = new HistoryObserver( - chatHistory, - this.eventBus, + this.chatHistory, + this.env.eventBus, this.tracer, this.env.tokenCalculator, + this.env.irMapper, ); this.historyObserver.start(); } diff --git a/packages/core/src/context/historyObserver.ts b/packages/core/src/context/historyObserver.ts index c92b343e33..e8898b7e65 100644 --- a/packages/core/src/context/historyObserver.ts +++ b/packages/core/src/context/historyObserver.ts @@ -28,6 +28,7 @@ export class HistoryObserver { private readonly eventBus: ContextEventBus, private readonly tracer: ContextTracer, private readonly tokenCalculator: ContextTokenCalculator, + private readonly irMapper: IrMapper, ) {} start() { @@ -40,7 +41,7 @@ export class HistoryObserver { // Rebuild the pristine IR graph from the full source history on every change. // Wait, toIr still returns an Episode[]. // We actually need to map the Episode[] to a flat ConcreteNode[] here to form the 'ship'. - const pristineEpisodes = IrMapper.toIr( + const pristineEpisodes = this.irMapper.toIr( this.chatHistory.get(), this.tokenCalculator, ); diff --git a/packages/core/src/context/ir/behaviorRegistry.ts b/packages/core/src/context/ir/behaviorRegistry.ts new file mode 100644 index 0000000000..934af53bb8 --- /dev/null +++ b/packages/core/src/context/ir/behaviorRegistry.ts @@ -0,0 +1,38 @@ +import type { Content, Part } from '@google/genai'; +import type { ConcreteNode } from './types.js'; + +export interface IrSerializationWriter { + appendContent(content: Content): void; + appendModelPart(part: Part): void; + appendUserPart(part: Part): void; + flushModelParts(): void; +} + +export interface IrNodeBehavior { + readonly type: T['type']; + + /** Serializes the node into the Gemini Content structure. */ + serialize(node: T, writer: IrSerializationWriter): void; + + /** + * Generates a structural representation of the node for the purpose + * of estimating its token cost. + */ + getEstimatableParts(node: T): Part[]; +} + +export class IrNodeBehaviorRegistry { + private readonly behaviors = new Map>(); + + register(behavior: IrNodeBehavior) { + this.behaviors.set(behavior.type, behavior); + } + + get(type: string): IrNodeBehavior { + const behavior = this.behaviors.get(type); + if (!behavior) { + throw new Error(`Unregistered IrNode type: ${type}`); + } + return behavior; + } +} diff --git a/packages/core/src/context/ir/builtinBehaviors.ts b/packages/core/src/context/ir/builtinBehaviors.ts new file mode 100644 index 0000000000..cb3a462e3e --- /dev/null +++ b/packages/core/src/context/ir/builtinBehaviors.ts @@ -0,0 +1,122 @@ +import type { Part } from '@google/genai'; +import type { IrNodeBehavior } from './behaviorRegistry.js'; +import type { + UserPrompt, + AgentThought, + ToolExecution, + MaskedTool, + AgentYield, + Snapshot, + RollingSummary, +} from './types.js'; + +export const UserPromptBehavior: IrNodeBehavior = { + type: 'USER_PROMPT', + 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); + } + return parts; + }, + serialize(prompt, writer) { + const parts = this.getEstimatableParts(prompt); + if (parts.length > 0) { + writer.flushModelParts(); + writer.appendContent({ role: 'user', parts }); + } + } +}; + +export const AgentThoughtBehavior: IrNodeBehavior = { + type: 'AGENT_THOUGHT', + getEstimatableParts(thought) { + return [{ text: thought.text }]; + }, + serialize(thought, writer) { + writer.appendModelPart({ text: thought.text }); + } +}; + +export const ToolExecutionBehavior: IrNodeBehavior = { + type: 'TOOL_EXECUTION', + getEstimatableParts(tool) { + return [ + { functionCall: { id: tool.id, name: tool.toolName, args: tool.intent } }, + { functionResponse: { id: tool.id, name: tool.toolName, response: typeof tool.observation === 'string' ? { message: tool.observation } : tool.observation } } + ]; + }, + serialize(tool, writer) { + const parts = this.getEstimatableParts(tool); + writer.appendModelPart(parts[0]); + writer.flushModelParts(); + writer.appendUserPart(parts[1]); + } +}; + +export const MaskedToolBehavior: IrNodeBehavior = { + type: 'MASKED_TOOL', + getEstimatableParts(tool) { + return [ + { functionCall: { id: tool.id, name: tool.toolName, args: tool.intent ?? {} } }, + { functionResponse: { id: tool.id, name: tool.toolName, response: typeof tool.observation === 'string' ? { message: tool.observation } : (tool.observation ?? {}) } } + ]; + }, + serialize(tool, writer) { + const parts = this.getEstimatableParts(tool); + writer.appendModelPart(parts[0]); + writer.flushModelParts(); + writer.appendUserPart(parts[1]); + } +}; + +export const AgentYieldBehavior: IrNodeBehavior = { + type: 'AGENT_YIELD', + getEstimatableParts(yieldNode) { + return [{ text: yieldNode.text }]; + }, + serialize(yieldNode, writer) { + writer.appendModelPart({ text: yieldNode.text }); + writer.flushModelParts(); + } +}; + +export const SystemEventBehavior: IrNodeBehavior = { + type: 'SYSTEM_EVENT', + getEstimatableParts() { return []; }, + serialize(node, writer) { + writer.flushModelParts(); + } +}; + +export const SnapshotBehavior: IrNodeBehavior = { + type: 'SNAPSHOT', + getEstimatableParts(node) { return [{ text: node.text }]; }, + serialize(node, writer) { + writer.flushModelParts(); + writer.appendUserPart({ text: node.text }); + } +}; + +export const RollingSummaryBehavior: IrNodeBehavior = { + type: 'ROLLING_SUMMARY', + getEstimatableParts(node) { return [{ text: node.text }]; }, + serialize(node, writer) { + writer.flushModelParts(); + writer.appendUserPart({ text: node.text }); + } +}; + +export function registerBuiltInBehaviors(registry: import('./behaviorRegistry.js').IrNodeBehaviorRegistry) { + registry.register(UserPromptBehavior); + registry.register(AgentThoughtBehavior); + registry.register(ToolExecutionBehavior); + registry.register(MaskedToolBehavior); + registry.register(AgentYieldBehavior); + registry.register(SystemEventBehavior); + registry.register(SnapshotBehavior); + registry.register(RollingSummaryBehavior); +} diff --git a/packages/core/src/context/ir/mapper.ts b/packages/core/src/context/ir/mapper.ts index f35c1377ab..f32e2fe507 100644 --- a/packages/core/src/context/ir/mapper.ts +++ b/packages/core/src/context/ir/mapper.ts @@ -1,31 +1,23 @@ -/** - * @license - * Copyright 2026 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - import type { Content } from '@google/genai'; import type { Episode, ConcreteNode } from './types.js'; import { toIr } from './toIr.js'; import { fromIr } from './fromIr.js'; import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; +import type { IrNodeBehaviorRegistry } from './behaviorRegistry.js'; export class IrMapper { - /** - * Translates a flat Gemini Content[] array into our rich Episodic Intermediate Representation. - * Groups adjacent function calls and responses into unified ToolExecution nodes. - */ - static toIr( + private readonly nodeIdentityMap = new WeakMap(); + + constructor(private readonly registry: IrNodeBehaviorRegistry) {} + + toIr( history: readonly Content[], tokenCalculator: ContextTokenCalculator, ): Episode[] { - return toIr(history, tokenCalculator); + return toIr(history, tokenCalculator, this.nodeIdentityMap); } - /** - * Re-serializes a flat array of ConcreteNodes back into a flat Gemini Content[] array. - */ - static fromIr(ship: readonly ConcreteNode[]): Content[] { - return fromIr(ship); + fromIr(ship: readonly ConcreteNode[]): Content[] { + return fromIr(ship, this.registry); } } diff --git a/packages/core/src/context/ir/projector.ts b/packages/core/src/context/ir/projector.ts index 9238b165e1..a6f9c6c830 100644 --- a/packages/core/src/context/ir/projector.ts +++ b/packages/core/src/context/ir/projector.ts @@ -5,7 +5,6 @@ */ import type { Content } from '@google/genai'; -import { IrMapper } from './mapper.js'; import type { ConcreteNode } from './types.js'; import { debugLogger } from '../../utils/debugLogger.js'; import type { @@ -29,7 +28,7 @@ export class IrProjector { protectedIds: Set, ): Promise { if (!sidecar.budget) { - const contents = IrMapper.fromIr(ship); + const contents = env.irMapper.fromIr(ship); tracer.logEvent('IrProjector', 'Projected Context to LLM (No Budget)', { projectedContext: contents, }); @@ -54,7 +53,7 @@ export class IrProjector { 'IrProjector', `View is within maxTokens (${currentTokens} <= ${maxTokens}). Returning view.`, ); - const contents = IrMapper.fromIr(ship); + const contents = env.irMapper.fromIr(ship); tracer.logEvent('IrProjector', 'Projected Context to LLM', { projectedContext: contents, }); @@ -116,7 +115,7 @@ export class IrProjector { const visibleShip = processedShip.filter((n) => !skipList.has(n.id)); - const contents = IrMapper.fromIr(visibleShip); + const contents = env.irMapper.fromIr(visibleShip); tracer.logEvent('IrProjector', 'Projected Sanitized Context to LLM', { projectedContextSanitized: contents, }); diff --git a/packages/core/src/context/ir/toIr.ts b/packages/core/src/context/ir/toIr.ts index 7a627202f4..4c330e8456 100644 --- a/packages/core/src/context/ir/toIr.ts +++ b/packages/core/src/context/ir/toIr.ts @@ -16,10 +16,8 @@ import type { } from './types.js'; import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; -// WeakMap to provide stable, deterministic identity across parses for the exact same Content/Part references -const nodeIdentityMap = new WeakMap(); - -export function getStableId(obj: object): string { +// We remove the global nodeIdentityMap and instead rely on one passed from IrMapper +export function getStableId(obj: object, nodeIdentityMap: WeakMap): string { let id = nodeIdentityMap.get(obj); if (!id) { id = randomUUID(); @@ -44,6 +42,7 @@ function isCompleteEpisode(ep: Partial): ep is Episode { export function toIr( history: readonly Content[], tokenCalculator: ContextTokenCalculator, + nodeIdentityMap: WeakMap ): Episode[] { const episodes: Episode[] = []; let currentEpisode: Partial | null = null; @@ -73,20 +72,21 @@ export function toIr( currentEpisode, pendingCallParts, tokenCalculator, - + nodeIdentityMap ); } if (hasUserParts) { finalizeEpisode(); - currentEpisode = parseUserParts(msg); + currentEpisode = parseUserParts(msg, nodeIdentityMap); } } else if (msg.role === 'model') { currentEpisode = parseModelParts( msg, currentEpisode, pendingCallParts, - ); + nodeIdentityMap + ); } } @@ -103,10 +103,11 @@ function parseToolResponses( currentEpisode: Partial | null, pendingCallParts: Map, tokenCalculator: ContextTokenCalculator, + nodeIdentityMap: WeakMap ): Partial { if (!currentEpisode) { currentEpisode = { - id: getStableId(msg), + id: getStableId(msg, nodeIdentityMap), timestamp: Date.now(), concreteNodes: [], }; @@ -123,7 +124,7 @@ function parseToolResponses( const obsTokens = tokenCalculator.estimateTokensForParts([part]); const step: ToolExecution = { - id: getStableId(part), + id: getStableId(part, nodeIdentityMap), type: 'TOOL_EXECUTION', toolName: part.functionResponse.name || 'unknown', intent: isRecord(matchingCall?.functionCall?.args) @@ -149,6 +150,7 @@ function parseToolResponses( function parseUserParts( msg: Content, + nodeIdentityMap: WeakMap ): Partial { const semanticParts: SemanticPart[] = []; for (const p of msg.parts!) { @@ -171,12 +173,12 @@ function parseUserParts( } const trigger: UserPrompt = { - id: getStableId(msg.parts![0] || msg), + id: getStableId(msg.parts![0] || msg, nodeIdentityMap), type: 'USER_PROMPT', semanticParts, }; return { - id: getStableId(msg), + id: getStableId(msg, nodeIdentityMap), timestamp: Date.now(), concreteNodes: [trigger], }; @@ -186,10 +188,11 @@ function parseModelParts( msg: Content, currentEpisode: Partial | null, pendingCallParts: Map, + nodeIdentityMap: WeakMap ): Partial { if (!currentEpisode) { currentEpisode = { - id: getStableId(msg), + id: getStableId(msg, nodeIdentityMap), timestamp: Date.now(), concreteNodes: [], }; @@ -201,7 +204,7 @@ function parseModelParts( if (callId) pendingCallParts.set(callId, part); } else if (part.text) { const thought: AgentThought = { - id: getStableId(part), + id: getStableId(part, nodeIdentityMap), type: 'AGENT_THOUGHT', text: part.text, }; diff --git a/packages/core/src/context/sidecar/environment.ts b/packages/core/src/context/sidecar/environment.ts index 11bf6f64fc..44805eb50c 100644 --- a/packages/core/src/context/sidecar/environment.ts +++ b/packages/core/src/context/sidecar/environment.ts @@ -26,4 +26,6 @@ export interface ContextEnvironment { readonly idGenerator: IIdGenerator; readonly eventBus: ContextEventBus; readonly inbox: LiveInbox; + readonly behaviorRegistry: import('../ir/behaviorRegistry.js').IrNodeBehaviorRegistry; + readonly irMapper: import('../ir/mapper.js').IrMapper; } diff --git a/packages/core/src/context/sidecar/environmentImpl.ts b/packages/core/src/context/sidecar/environmentImpl.ts index d14aeb1ca0..64783ae86e 100644 --- a/packages/core/src/context/sidecar/environmentImpl.ts +++ b/packages/core/src/context/sidecar/environmentImpl.ts @@ -16,11 +16,17 @@ import { NodeIdGenerator } from '../system/NodeIdGenerator.js'; import { LiveInbox } from './inbox.js'; +import { IrNodeBehaviorRegistry } from '../ir/behaviorRegistry.js'; +import { registerBuiltInBehaviors } from '../ir/builtinBehaviors.js'; +import { IrMapper } from '../ir/mapper.js'; + export class ContextEnvironmentImpl implements ContextEnvironment { readonly tokenCalculator: ContextTokenCalculator; readonly fileSystem: IFileSystem; readonly idGenerator: IIdGenerator; readonly inbox: LiveInbox; + readonly behaviorRegistry: import('../ir/behaviorRegistry.js').IrNodeBehaviorRegistry; + readonly irMapper: import('../ir/mapper.js').IrMapper; constructor( readonly llmClient: BaseLlmClient, @@ -38,5 +44,9 @@ export class ContextEnvironmentImpl implements ContextEnvironment { this.fileSystem = fileSystem || new NodeFileSystem(); this.idGenerator = idGenerator || new NodeIdGenerator(); this.inbox = new LiveInbox(); + + this.behaviorRegistry = new IrNodeBehaviorRegistry(); + registerBuiltInBehaviors(this.behaviorRegistry); + this.irMapper = new IrMapper(this.behaviorRegistry); } } diff --git a/packages/core/src/context/testing/contextTestUtils.ts b/packages/core/src/context/testing/contextTestUtils.ts index 292da2a7c6..3254f958b3 100644 --- a/packages/core/src/context/testing/contextTestUtils.ts +++ b/packages/core/src/context/testing/contextTestUtils.ts @@ -15,6 +15,9 @@ import { ContextEnvironmentImpl } from '../sidecar/environmentImpl.js'; import { SidecarLoader } from '../sidecar/SidecarLoader.js'; import { ContextEventBus } from '../eventBus.js'; import { ContextTokenCalculator } from '../utils/contextTokenCalculator.js'; +import { IrNodeBehaviorRegistry } from '../ir/behaviorRegistry.js'; +import { registerBuiltInBehaviors } from '../ir/builtinBehaviors.js'; +import { IrMapper } from '../ir/mapper.js'; import { ProcessorRegistry } from '../sidecar/registry.js'; import { registerBuiltInProcessors } from '../sidecar/builtins.js'; import type { ContextAccountingState } from '../pipeline.js'; @@ -102,6 +105,10 @@ export function createDummyToolNode( export function createMockEnvironment( overrides?: Partial, ): ContextEnvironment { + const registry = new IrNodeBehaviorRegistry(); + registerBuiltInBehaviors(registry); + const irMapper = new IrMapper(registry); + return { // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion llmClient: vi.fn().mockReturnValue({ @@ -119,6 +126,8 @@ export function createMockEnvironment( tokenCalculator: new ContextTokenCalculator(1), fileSystem: new InMemoryFileSystem(), idGenerator: new DeterministicIdGenerator('mock-uuid-'), + behaviorRegistry: registry, + irMapper, ...overrides, } as ContextEnvironment; }