decentralized nodes

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();
   });

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();
   }

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,
         );

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<T extends ConcreteNode = ConcreteNode> {
+  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<string, IrNodeBehavior<any>>();
+
+  register<T extends ConcreteNode>(behavior: IrNodeBehavior<T>) {
+    this.behaviors.set(behavior.type, behavior);
+  }
+
+  get(type: string): IrNodeBehavior<any> {
+    const behavior = this.behaviors.get(type);
+    if (!behavior) {
+      throw new Error(`Unregistered IrNode type: ${type}`);
+    }
+    return behavior;
+  }
+}

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<UserPrompt> = {
+  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<AgentThought> = {
+  type: 'AGENT_THOUGHT',
+  getEstimatableParts(thought) {
+    return [{ text: thought.text }];
+  },
+  serialize(thought, writer) {
+    writer.appendModelPart({ text: thought.text });
+  }
+};
+
+export const ToolExecutionBehavior: IrNodeBehavior<ToolExecution> = {
+  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<MaskedTool> = {
+  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<AgentYield> = {
+  type: 'AGENT_YIELD',
+  getEstimatableParts(yieldNode) {
+    return [{ text: yieldNode.text }];
+  },
+  serialize(yieldNode, writer) {
+    writer.appendModelPart({ text: yieldNode.text });
+    writer.flushModelParts();
+  }
+};
+
+export const SystemEventBehavior: IrNodeBehavior<any> = {
+  type: 'SYSTEM_EVENT',
+  getEstimatableParts() { return []; },
+  serialize(node, writer) {
+    writer.flushModelParts();
+  }
+};
+
+export const SnapshotBehavior: IrNodeBehavior<Snapshot> = {
+  type: 'SNAPSHOT',
+  getEstimatableParts(node) { return [{ text: node.text }]; },
+  serialize(node, writer) {
+    writer.flushModelParts();
+    writer.appendUserPart({ text: node.text });
+  }
+};
+
+export const RollingSummaryBehavior: IrNodeBehavior<RollingSummary> = {
+  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);
+}

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<object, string>();
+
+  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);
   }
 }

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<string>,
   ): Promise<Content[]> {
     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,
     });

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<object, string>();
-
-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<object, string>): string {
   let id = nodeIdentityMap.get(obj);
   if (!id) {
     id = randomUUID();
@@ -44,6 +42,7 @@ function isCompleteEpisode(ep: Partial<Episode>): ep is Episode {
 export function toIr(
   history: readonly Content[],
   tokenCalculator: ContextTokenCalculator,
+  nodeIdentityMap: WeakMap<object, string>
 ): Episode[] {
   const episodes: Episode[] = [];
   let currentEpisode: Partial<Episode> | 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<Episode> | null,
   pendingCallParts: Map<string, Part>,
   tokenCalculator: ContextTokenCalculator,
+  nodeIdentityMap: WeakMap<object, string>
 ): Partial<Episode> {
   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<object, string>
 ): Partial<Episode> {
   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<Episode> | null,
   pendingCallParts: Map<string, Part>,
+  nodeIdentityMap: WeakMap<object, string>
 ): Partial<Episode> {
   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,
       };

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;
 }

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);
   }
 }

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>,
 ): 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;
 }