test(context): inject IdGenerator into irMapper to restore snapshot determinism

packages/core/src/context/ir/mapper.ts

@@ -9,17 +9,26 @@ import { toIr } from './toIr.js';
 import { fromIr } from './fromIr.js';
 import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
 import type { IrNodeBehaviorRegistry } from './behaviorRegistry.js';
+import type { IIdGenerator } from '../system/IIdGenerator.js';
 
 export class IrMapper {
   private readonly nodeIdentityMap = new WeakMap<object, string>();
 
-  constructor(private readonly registry: IrNodeBehaviorRegistry) {}
+  constructor(
+    private readonly registry: IrNodeBehaviorRegistry,
+    private readonly idGenerator: IIdGenerator,
+  ) {}
 
   toIr(
     history: readonly Content[],
     tokenCalculator: ContextTokenCalculator,
   ): Episode[] {
-    return toIr(history, tokenCalculator, this.nodeIdentityMap);
+    return toIr(
+      history,
+      tokenCalculator,
+      this.nodeIdentityMap,
+      this.idGenerator,
+    );
   }
 
   fromIr(nodes: readonly ConcreteNode[]): Content[] {

packages/core/src/context/ir/toIr.ts

@@ -5,7 +5,6 @@
  */
 
 import type { Content, Part } from '@google/genai';
-import { randomUUID } from 'node:crypto';
 import type {
   Episode,
   SemanticPart,
@@ -15,15 +14,17 @@ import type {
   UserPrompt,
 } from './types.js';
 import type { ContextTokenCalculator } from '../utils/contextTokenCalculator.js';
+import type { IIdGenerator } from '../system/IIdGenerator.js';
 
 // We remove the global nodeIdentityMap and instead rely on one passed from IrMapper
 export function getStableId(
   obj: object,
   nodeIdentityMap: WeakMap<object, string>,
+  idGenerator: IIdGenerator,
 ): string {
   let id = nodeIdentityMap.get(obj);
   if (!id) {
-    id = randomUUID();
+    id = idGenerator.generateId();
     nodeIdentityMap.set(obj, id);
   }
   return id;
@@ -46,6 +47,7 @@ export function toIr(
   history: readonly Content[],
   tokenCalculator: ContextTokenCalculator,
   nodeIdentityMap: WeakMap<object, string>,
+  idGenerator: IIdGenerator,
 ): Episode[] {
   const episodes: Episode[] = [];
   let currentEpisode: Partial<Episode> | null = null;
@@ -74,12 +76,13 @@ export function toIr(
           pendingCallParts,
           tokenCalculator,
           nodeIdentityMap,
+          idGenerator,
         );
       }
 
       if (hasUserParts) {
         finalizeEpisode();
-        currentEpisode = parseUserParts(msg, nodeIdentityMap);
+        currentEpisode = parseUserParts(msg, nodeIdentityMap, idGenerator);
       }
     } else if (msg.role === 'model') {
       currentEpisode = parseModelParts(
@@ -87,12 +90,13 @@ export function toIr(
         currentEpisode,
         pendingCallParts,
         nodeIdentityMap,
+        idGenerator,
       );
     }
   }
 
   if (currentEpisode) {
-    finalizeYield(currentEpisode);
+    finalizeYield(currentEpisode, idGenerator);
     finalizeEpisode();
   }
 
@@ -105,10 +109,11 @@ function parseToolResponses(
   pendingCallParts: Map<string, Part>,
   tokenCalculator: ContextTokenCalculator,
   nodeIdentityMap: WeakMap<object, string>,
+  idGenerator: IIdGenerator,
 ): Partial<Episode> {
   if (!currentEpisode) {
     currentEpisode = {
-      id: getStableId(msg, nodeIdentityMap),
+      id: getStableId(msg, nodeIdentityMap, idGenerator),
       timestamp: Date.now(),
       concreteNodes: [],
     };
@@ -126,7 +131,7 @@ function parseToolResponses(
       const obsTokens = tokenCalculator.estimateTokensForParts([part]);
 
       const step: ToolExecution = {
-        id: getStableId(part, nodeIdentityMap),
+        id: getStableId(part, nodeIdentityMap, idGenerator),
         type: 'TOOL_EXECUTION',
         toolName: part.functionResponse.name || 'unknown',
         intent: isRecord(matchingCall?.functionCall?.args)
@@ -153,6 +158,7 @@ function parseToolResponses(
 function parseUserParts(
   msg: Content,
   nodeIdentityMap: WeakMap<object, string>,
+  idGenerator: IIdGenerator,
 ): Partial<Episode> {
   const semanticParts: SemanticPart[] = [];
   const parts = msg.parts || [];
@@ -177,12 +183,12 @@ function parseUserParts(
 
   const baseObj = parts.length > 0 ? parts[0] : msg;
   const trigger: UserPrompt = {
-    id: getStableId(baseObj, nodeIdentityMap),
+    id: getStableId(baseObj, nodeIdentityMap, idGenerator),
     type: 'USER_PROMPT',
     semanticParts,
   };
   return {
-    id: getStableId(msg, nodeIdentityMap),
+    id: getStableId(msg, nodeIdentityMap, idGenerator),
     timestamp: Date.now(),
     concreteNodes: [trigger],
   };
@@ -193,10 +199,11 @@ function parseModelParts(
   currentEpisode: Partial<Episode> | null,
   pendingCallParts: Map<string, Part>,
   nodeIdentityMap: WeakMap<object, string>,
+  idGenerator: IIdGenerator,
 ): Partial<Episode> {
   if (!currentEpisode) {
     currentEpisode = {
-      id: getStableId(msg, nodeIdentityMap),
+      id: getStableId(msg, nodeIdentityMap, idGenerator),
       timestamp: Date.now(),
       concreteNodes: [],
     };
@@ -209,7 +216,7 @@ function parseModelParts(
       if (callId) pendingCallParts.set(callId, part);
     } else if (part.text) {
       const thought: AgentThought = {
-        id: getStableId(part, nodeIdentityMap),
+        id: getStableId(part, nodeIdentityMap, idGenerator),
         type: 'AGENT_THOUGHT',
         text: part.text,
       };
@@ -223,10 +230,13 @@ function parseModelParts(
   return currentEpisode;
 }
 
-function finalizeYield(currentEpisode: Partial<Episode>) {
+function finalizeYield(
+  currentEpisode: Partial<Episode>,
+  idGenerator: IIdGenerator,
+) {
   if (currentEpisode.concreteNodes && currentEpisode.concreteNodes.length > 0) {
     const yieldNode: AgentYield = {
-      id: randomUUID(),
+      id: idGenerator.generateId(),
       type: 'AGENT_YIELD',
       text: 'Yield', // Synthesized yield since we don't have the original concrete node
     };

packages/core/src/context/pipeline/environmentImpl.ts

@@ -47,6 +47,6 @@ export class ContextEnvironmentImpl implements ContextEnvironment {
     this.fileSystem = fileSystem || new NodeFileSystem();
     this.idGenerator = idGenerator || new NodeIdGenerator();
     this.inbox = new LiveInbox();
-    this.irMapper = new IrMapper(this.behaviorRegistry);
+    this.irMapper = new IrMapper(this.behaviorRegistry, this.idGenerator);
   }
 }