@librechat/agents 3.8.5 → 3.8.7

package/src/tools/ToolNode.ts

@@ -49,10 +49,21 @@ import type {
   PostToolBatchEntry,
   ToolApprovalReplayKey,
 } from '@/hooks';
-import type { PreparedSubagents } from '@/tools/preparedSubagents';
 import type { SettledToolBatchResult } from '@/tools/toolBatchReplay';
+import type { PreparedSubagents } from '@/tools/preparedSubagents';
 import type { RunBreakerScope } from '@/llm/streamLimits';
 import type * as t from '@/types';
+import {
+  TOOL_BATCH_REPLAY_KEY,
+  getToolBatchReplayOwner,
+  getToolBatchReplayScope,
+  attachToolBatchReplayState,
+  restoreToolBatchReplayState,
+  getToolBatchReplayState,
+  clearStaleToolApprovalConfig,
+  isChildToolReplayOwner,
+  getToolReplayResumeStatus,
+} from '@/tools/toolBatchReplay';
 import {
   type CallerCapabilityProjection,
   createCallerCapabilityProjectionSnapshot,
@@ -121,25 +132,14 @@ import { stripCodeSessionFileSummary } from '@/tools/CodeSessionFileSummary';
 import { Constants, GraphEvents, CODE_EXECUTION_TOOLS } from '@/common';
 import { PreparedSubagentError } from '@/tools/preparedSubagents';
 import { attachRunStepResumeState } from '@/tools/runStepResume';
-import {
-  TOOL_BATCH_REPLAY_KEY,
-  getToolBatchReplayOwner,
-  getToolBatchReplayScope,
-  attachToolBatchReplayState,
-  restoreToolBatchReplayState,
-  getToolBatchReplayState,
-  clearStaleToolApprovalConfig,
-  isChildToolReplayOwner,
-  getToolReplayResumeStatus,
-} from '@/tools/toolBatchReplay';
 
 function stripToolApprovalReviewConfig(
   configurable: Record<string, unknown> | undefined
 ): Record<string, unknown> | undefined {
   if (
     configurable == null ||
-    !(TOOL_APPROVAL_REVIEW_CONFIG_KEY in configurable) &&
-    !(TOOL_BATCH_REPLAY_KEY in configurable)
+    (!(TOOL_APPROVAL_REVIEW_CONFIG_KEY in configurable) &&
+      !(TOOL_BATCH_REPLAY_KEY in configurable))
   ) {
     return configurable;
   }
@@ -221,6 +221,8 @@ type RunToolBatchContext<T = unknown> = {
   batchScopeId?: string;
   /** Batch-local sink for post-substitution args. */
   resolvedArgsByCallId?: ResolvedArgsByCallId;
+  /** File identities actually injected into each direct code call. */
+  codeSessionBaselineByCallId?: Map<string, ReadonlyMap<string, string>>;
   /**
    * Frozen pre-batch view of the tool-output registry. When supplied,
    * `runTool` resolves `{{tool…turn…}}` placeholders against this
@@ -474,8 +476,12 @@ function getAssistantBatchReplayKey(
 ): string {
   return JSON.stringify([
     threadId ?? null,
-    batch.message.id == null || batch.message.id === '' ? null : batch.message.id,
-    createHash('sha256').update(stableStringify(batch.message.tool_calls ?? [])).digest('hex'),
+    batch.message.id == null || batch.message.id === ''
+      ? null
+      : batch.message.id,
+    createHash('sha256')
+      .update(stableStringify(batch.message.tool_calls ?? []))
+      .digest('hex'),
   ]);
 }
 
@@ -647,29 +653,78 @@ function toInjectedFileRef(
 
 /* Stable file identity = `(storage_session_id, id)`. Same name in
  * different storage sessions are distinct files. */
-function fileIdentityKey(file: {
-  storage_session_id?: string;
-  id: string;
-}): string {
-  return `${file.storage_session_id ?? ''}\0${file.id}`;
+function fileIdentityKey(
+  file: {
+    storage_session_id?: string;
+    id: string;
+  },
+  fallbackSessionId = ''
+): string {
+  return `${file.storage_session_id ?? fallbackSessionId}\0${file.id}`;
+}
+
+function codeSessionIdentityByName(
+  context:
+    | t.ToolCallRequest['codeSessionContext']
+    | t.CodeSessionContext
+    | undefined
+): ReadonlyMap<string, string> {
+  return new Map(
+    (context?.files ?? []).map((file) => [
+      file.name,
+      fileIdentityKey(file, context?.session_id),
+    ])
+  );
+}
+
+function storageSessionIdFromIdentity(
+  identity: string | undefined,
+  fileId: string
+): string | undefined {
+  const suffix = `\0${fileId}`;
+  return identity != null && identity.endsWith(suffix)
+    ? identity.slice(0, identity.length - suffix.length)
+    : undefined;
 }
 
 function updateCodeSession(
   sessions: t.ToolSessionMap,
   sessionKey: string,
   execSessionId: string,
-  files: t.FileRefs | undefined
+  files: t.FileRefs | undefined,
+  deletedFiles: readonly string[] | undefined,
+  baselineIdentityByName: ReadonlyMap<string, string> | undefined
 ): void {
   const newFiles = files ?? [];
   const existingSession = sessions.get(sessionKey) as
     | t.CodeSessionContext
     | undefined;
-  const existingFiles = existingSession?.files ?? [];
+  const existingSessionId = existingSession?.session_id ?? '';
+  const existingFiles = (existingSession?.files ?? []).map((file) =>
+    file.storage_session_id == null
+      ? { ...file, storage_session_id: existingSessionId }
+      : file
+  );
+  const deletedIdentityByName = new Map<string, string>();
+  for (const name of deletedFiles ?? []) {
+    const identity = baselineIdentityByName?.get(name);
+    if (identity != null) {
+      deletedIdentityByName.set(name, identity);
+    }
+  }
 
+  /* Absence from `files` is not a deletion signal: older Code API versions
+   * and truncated responses can omit otherwise-live refs. Remove only paths
+   * explicitly attested by `deleted_files`, and only while the current ref
+   * still has the identity injected into that request. */
   if (newFiles.length === 0) {
     sessions.set(sessionKey, {
       session_id: execSessionId,
-      files: existingFiles,
+      files: existingFiles.filter(
+        (file) =>
+          deletedIdentityByName.get(file.name) !==
+          fileIdentityKey(file, existingSessionId)
+      ),
       lastUpdated: Date.now(),
     });
     return;
@@ -681,26 +736,46 @@ function updateCodeSession(
   const filesWithSession: t.FileRefs = [];
   const newFileNames = new Set<string>();
   const incomingByIdentity = new Map<string, number>();
+  const existingIdentities = new Set(
+    existingFiles.map((file) => fileIdentityKey(file, existingSessionId))
+  );
   for (const file of newFiles) {
+    const baselineIdentity = baselineIdentityByName?.get(file.name);
     const withSession = {
       ...file,
-      storage_session_id: file.storage_session_id ?? execSessionId,
+      storage_session_id:
+        file.storage_session_id ??
+        storageSessionIdFromIdentity(baselineIdentity, file.id) ??
+        execSessionId,
     };
-    incomingByIdentity.set(
-      fileIdentityKey(withSession),
-      filesWithSession.length
-    );
+    const identity = fileIdentityKey(withSession);
+    const wasTrackedByRequest = baselineIdentity === identity;
+    /* An inherited echo is evidence that a request observed an existing ref,
+     * not authority to resurrect one removed by a concurrent sibling. Fresh
+     * and modified outputs remain eligible to replace by path below. */
+    if (
+      file.inherited === true &&
+      wasTrackedByRequest &&
+      (!existingIdentities.has(identity) ||
+        deletedIdentityByName.get(file.name) === identity)
+    ) {
+      continue;
+    }
+    incomingByIdentity.set(identity, filesWithSession.length);
     newFileNames.add(withSession.name);
     filesWithSession.push(withSession);
   }
 
   const filteredExisting: t.FileRefs = [];
   for (const e of existingFiles) {
-    const idx = incomingByIdentity.get(fileIdentityKey(e));
+    const idx = incomingByIdentity.get(fileIdentityKey(e, existingSessionId));
     if (idx !== undefined) {
       filesWithSession[idx] = { ...e, ...filesWithSession[idx] };
     }
-    if (!newFileNames.has(e.name)) {
+    const explicitlyDeleted =
+      deletedIdentityByName.get(e.name) ===
+      fileIdentityKey(e, existingSessionId);
+    if (!newFileNames.has(e.name) && !explicitlyDeleted) {
       filteredExisting.push(e);
     }
   }
@@ -712,6 +787,53 @@ function updateCodeSession(
   });
 }
 
+/** Retains host-provisioned inputs even when execution returns no usable artifact. */
+function retainCodeSessionInputs(
+  sessions: t.ToolSessionMap,
+  sessionKey: string,
+  context: t.ToolCallRequest['codeSessionContext']
+): void {
+  if (
+    context?.session_id == null ||
+    context.session_id === '' ||
+    context.files == null
+  )
+    return;
+  const existing = sessions.get(sessionKey) as t.CodeSessionContext | undefined;
+  const existingFiles = existing?.files ?? [];
+  const files = [...existingFiles];
+  const indexByIdentity = new Map<string, number>();
+  const indexByName = new Map<string, number>();
+  for (let i = 0; i < files.length; i++) {
+    indexByIdentity.set(fileIdentityKey(files[i]), i);
+    indexByName.set(files[i].name, i);
+  }
+  let changed = existing?.session_id !== context.session_id;
+  for (const file of context.files) {
+    if (!file.id || !file.name || !file.storage_session_id) continue;
+    const identity = fileIdentityKey(file);
+    if (indexByIdentity.has(identity)) continue;
+    const replacementIndex = indexByName.get(file.name);
+    if (replacementIndex !== undefined) {
+      indexByIdentity.delete(fileIdentityKey(files[replacementIndex]));
+      files[replacementIndex] = { ...file };
+      indexByIdentity.set(identity, replacementIndex);
+      changed = true;
+      continue;
+    }
+    indexByIdentity.set(identity, files.length);
+    indexByName.set(file.name, files.length);
+    files.push({ ...file });
+    changed = true;
+  }
+  if (!changed) return;
+  sessions.set(sessionKey, {
+    session_id: context.session_id,
+    files,
+    lastUpdated: Date.now(),
+  });
+}
+
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 export class ToolNode<T = any> extends RunnableCallable<T, T> {
   private toolMap: Map<string, StructuredToolInterface | RunnableToolLike>;
@@ -787,6 +909,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
    * even where `agentId` (the subagent-scope marker) is undefined.
    */
   private executingAgentId?: string;
+  private readonly executionContext?: t.ToolNodeOptions['executionContext'];
   private executingAgentName?: string;
   private rootAgentId?: string;
   private rootAgentName?: string;
@@ -876,6 +999,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     eagerEventToolSuppressions,
     agentId,
     executingAgentId,
+    executionContext,
     executingAgentName,
     rootAgentId,
     rootAgentName,
@@ -900,65 +1024,124 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       name: name ?? TOOL_NODE_RUN_NAME,
       tags,
       func: async (input, config) => {
-        const messages = Array.isArray(input) ? input as BaseMessage[] :
-          (input as { messages?: BaseMessage[] }).messages;
-        const assistantBatch = messages == null ? undefined : findAssistantBatch(messages);
-        const activeReplayKey = assistantBatch == null || getToolBatchReplayScope(config) == null ? undefined : getAssistantBatchReplayKey(
-          assistantBatch, getToolBatchReplayScope(config)
-        );
+        const messages = Array.isArray(input)
+          ? (input as BaseMessage[])
+          : (input as { messages?: BaseMessage[] }).messages;
+        const assistantBatch =
+          messages == null ? undefined : findAssistantBatch(messages);
+        const activeReplayKey =
+          assistantBatch == null || getToolBatchReplayScope(config) == null
+            ? undefined
+            : getAssistantBatchReplayKey(
+              assistantBatch,
+              getToolBatchReplayScope(config)
+            );
         const replayOwner = getToolBatchReplayOwner(
           config,
           this.executingAgentId ?? this.agentId ?? ''
         );
         let reviewEvidence = getToolApprovalReviewEvidence(config);
         const replayState = getToolBatchReplayState(config.configurable);
-        const ownsParentBatch = replayState?.records.some(
-          (record) => record.owner === replayOwner && record.batch === activeReplayKey
-        ) === true;
+        const ownsParentBatch =
+          replayState?.records.some(
+            (record) =>
+              record.owner === replayOwner && record.batch === activeReplayKey
+          ) === true;
         let isChildReplay = false;
         if (ownsParentBatch) {
-          const calls = this.isSendInput(input) ? [input.lg_tool_call] : assistantBatch?.message.tool_calls ?? [];
-          const trustedParentCallIds = new Set(calls.filter((call) =>
-            call.name === Constants.SUBAGENT &&
-            (this.toolMap.get(call.name) as ReplayableSubagentTool | undefined)?.[SUBAGENT_REPLAY_CONTROLLER] != null
-          ).map((call) => call.id ?? ''));
-          isChildReplay = reviewEvidence?.owner != null
-            ? isChildToolReplayOwner(config, reviewEvidence.owner, trustedParentCallIds)
-            : replayState.records.some((record) => isChildToolReplayOwner(config, record.owner, trustedParentCallIds));
+          const calls = this.isSendInput(input)
+            ? [input.lg_tool_call]
+            : (assistantBatch?.message.tool_calls ?? []);
+          const trustedParentCallIds = new Set(
+            calls
+              .filter(
+                (call) =>
+                  call.name === Constants.SUBAGENT &&
+                  (
+                    this.toolMap.get(call.name) as
+                      | ReplayableSubagentTool
+                      | undefined
+                  )?.[SUBAGENT_REPLAY_CONTROLLER] != null
+              )
+              .map((call) => call.id ?? '')
+          );
+          isChildReplay =
+            reviewEvidence?.owner != null
+              ? isChildToolReplayOwner(
+                config,
+                reviewEvidence.owner,
+                trustedParentCallIds
+              )
+              : replayState.records.some((record) =>
+                isChildToolReplayOwner(
+                  config,
+                  record.owner,
+                  trustedParentCallIds
+                )
+              );
         }
         const hasApprovalAuthorization =
           reviewEvidence?.owner != null || replayState?.approvalOwner != null;
-        const replayResumeStatus = hasApprovalAuthorization || replayState != null
-          ? getToolReplayResumeStatus(
-            config,
-            reviewEvidence?.interruptId ?? replayState?.interruptId,
-            isChildReplay
-          )
-          : undefined;
+        const replayResumeStatus =
+          hasApprovalAuthorization || replayState != null
+            ? getToolReplayResumeStatus(
+              config,
+              reviewEvidence?.interruptId ?? replayState?.interruptId,
+              isChildReplay
+            )
+            : undefined;
         if (replayResumeStatus === 'unverifiable' && hasApprovalAuthorization) {
-          throw new Error('Cannot verify the active tool approval resume — failing closed');
+          throw new Error(
+            'Cannot verify the active tool approval resume — failing closed'
+          );
         }
         if (replayResumeStatus === 'stale') {
           const configurable = { ...config.configurable };
-          clearStaleToolApprovalConfig(configurable, replayOwner, activeReplayKey);
+          clearStaleToolApprovalConfig(
+            configurable,
+            replayOwner,
+            activeReplayKey
+          );
           config = {
             ...config,
             configurable,
           };
           reviewEvidence = undefined;
         }
-        if (reviewEvidence?.owner != null && reviewEvidence.owner !== replayOwner && !isChildReplay) {
-          throw new Error('Tool approval execution owner changed before resume — failing closed');
+        if (
+          reviewEvidence?.owner != null &&
+          reviewEvidence.owner !== replayOwner &&
+          !isChildReplay
+        ) {
+          throw new Error(
+            'Tool approval execution owner changed before resume — failing closed'
+          );
         }
         const referenceReplay: { state?: ToolOutputReferenceState } = {};
-        const restoredBatches = await restoreToolBatchReplayState(config, replayOwner);
-        if (restoredBatches.length > 0 && !restoredBatches.some(({ batch }) => batch === activeReplayKey)) {
-          throw new Error('Tool batch proposal changed before approval replay — failing closed');
+        const restoredBatches = await restoreToolBatchReplayState(
+          config,
+          replayOwner
+        );
+        if (
+          restoredBatches.length > 0 &&
+          !restoredBatches.some(({ batch }) => batch === activeReplayKey)
+        ) {
+          throw new Error(
+            'Tool batch proposal changed before approval replay — failing closed'
+          );
         }
-        if (activeReplayKey != null && getToolBatchReplayState(config.configurable) != null) {
+        if (
+          activeReplayKey != null &&
+          getToolBatchReplayState(config.configurable) != null
+        ) {
           this.settledDirectResultsByBatch.delete(activeReplayKey);
         }
-        for (const { batch, results, turnState, referenceState } of restoredBatches) {
+        for (const {
+          batch,
+          results,
+          turnState,
+          referenceState,
+        } of restoredBatches) {
           if (batch !== activeReplayKey) {
             continue;
           }
@@ -978,27 +1161,45 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             throw error;
           }
           const resumeState = createRunStepResumeState?.();
-          const activeResults = activeReplayKey == null ? undefined : this.settledDirectResultsByBatch.get(activeReplayKey);
-          const activeBatches = activeReplayKey == null
-            ? new Map<string, Map<string, SettledDirectToolResult>>()
-            : new Map([[activeReplayKey, activeResults ?? new Map<string, SettledDirectToolResult>()]]);
+          const activeResults =
+            activeReplayKey == null
+              ? undefined
+              : this.settledDirectResultsByBatch.get(activeReplayKey);
+          const activeBatches =
+            activeReplayKey == null
+              ? new Map<string, Map<string, SettledDirectToolResult>>()
+              : new Map([
+                [
+                  activeReplayKey,
+                  activeResults ?? new Map<string, SettledDirectToolResult>(),
+                ],
+              ]);
           const turnState = this.createSubagentResumeState();
-          const activeCallIds = new Set(assistantBatch?.message.tool_calls?.map((call) => call.id));
-          turnState.directPathTurns = turnState.directPathTurns.filter(({ toolCallId }) => activeCallIds.has(toolCallId));
+          const activeCallIds = new Set(
+            assistantBatch?.message.tool_calls?.map((call) => call.id)
+          );
+          turnState.directPathTurns = turnState.directPathTurns.filter(
+            ({ toolCallId }) => activeCallIds.has(toolCallId)
+          );
           throw new GraphInterrupt(
-            await Promise.all(error.interrupts.map(async (pendingInterrupt) => {
-              const value = await attachToolBatchReplayState(
-                pendingInterrupt.value,
-                replayOwner,
-                activeBatches,
-                turnState,
-                referenceReplay.state
-              );
-              return {
-                ...pendingInterrupt,
-                value: resumeState == null ? value : attachRunStepResumeState(value, resumeState),
-              };
-            }))
+            await Promise.all(
+              error.interrupts.map(async (pendingInterrupt) => {
+                const value = await attachToolBatchReplayState(
+                  pendingInterrupt.value,
+                  replayOwner,
+                  activeBatches,
+                  turnState,
+                  referenceReplay.state
+                );
+                return {
+                  ...pendingInterrupt,
+                  value:
+                    resumeState == null
+                      ? value
+                      : attachRunStepResumeState(value, resumeState),
+                };
+              })
+            )
           );
         }
         if (createRunStepResumeState == null) {
@@ -1040,6 +1241,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     // Default to agentId so callers constructing ToolNode directly (who pass the
     // existing agentId option) still get attribution without knowing the new option.
     this.executingAgentId = executingAgentId ?? agentId;
+    this.executionContext = executionContext;
     this.executingAgentName = executingAgentName;
     this.rootAgentId = rootAgentId;
     this.rootAgentName = rootAgentName;
@@ -1109,8 +1311,13 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         ? options
         : {
           ...options,
+          configurable: {
+            ...options?.configurable,
+            executionContext: this.executionContext,
+          },
           metadata: {
             ...options?.metadata,
+            executionContext: this.executionContext,
             agentId: this.executingAgentId,
             activeAgentId: this.executingAgentId,
             ...(this.executingAgentName == null
@@ -1668,7 +1875,11 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           return next;
         })();
       this.recordToolUsageTurn(call.name, usageCount, call.id);
-      if (batchContext.replayBatchKey != null && call.id != null && call.id !== '') {
+      if (
+        batchContext.replayBatchKey != null &&
+        call.id != null &&
+        call.id !== ''
+      ) {
         this.directPathTurns.set(call.id, usageCount);
       }
       let args = call.args;
@@ -1732,6 +1943,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             runId: (config.configurable?.run_id as string | undefined) ?? '',
             threadId: config.configurable?.thread_id as string | undefined,
             agentId: this.agentId,
+            executionContext: this.executionContext,
             executingAgentId: this.executingAgentId,
           },
         };
@@ -1757,6 +1969,12 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           | t.CodeSessionContext
           | undefined;
         const execSessionId = codeSession?.session_id;
+        if (batchContext.codeSessionBaselineByCallId != null) {
+          batchContext.codeSessionBaselineByCallId.set(
+            call.id ?? '',
+            codeSessionIdentityByName(codeSession)
+          );
+        }
         if (execSessionId != null && execSessionId !== '') {
           invokeParams = {
             ...invokeParams,
@@ -2167,8 +2385,13 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           approvalReplaySessionId,
           approvalReplayKey
         );
-    const approvalReviewEvidence = getToolApprovalReviewEvidence(config,
-      getToolBatchReplayOwner(config, this.executingAgentId ?? this.agentId ?? ''));
+    const approvalReviewEvidence = getToolApprovalReviewEvidence(
+      config,
+      getToolBatchReplayOwner(
+        config,
+        this.executingAgentId ?? this.agentId ?? ''
+      )
+    );
     const reviewedApproval = getReviewedToolApproval(
       approvalReviewEvidence?.payload,
       call.id
@@ -2294,6 +2517,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
               runId,
               threadId,
               agentId: this.agentId,
+              executionContext: this.executionContext,
               executingAgentId: this.executingAgentId,
               toolName: call.name,
               toolInput: resolvedArgs,
@@ -2305,11 +2529,17 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             matchQuery: call.name,
             onceReplayKey: approvalReplayKey,
             onceReplaySessionId: approvalReplaySessionId,
-          }).catch((): AggregatedHookResult | undefined => approvalReviewEvidence == null ? undefined : {
-            decision: 'deny',
-            reason: 'Approval policy could not be evaluated on resume',
-            additionalContexts: [], injectedMessages: [], errors: [],
-          });
+          }).catch((): AggregatedHookResult | undefined =>
+            approvalReviewEvidence == null
+              ? undefined
+              : {
+                decision: 'deny',
+                reason: 'Approval policy could not be evaluated on resume',
+                additionalContexts: [],
+                injectedMessages: [],
+                errors: [],
+              }
+          );
 
       if (preResult != null) {
         // Forward any additionalContext strings hooks returned into
@@ -2651,6 +2881,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           runId,
           threadId,
           agentId: this.agentId,
+          executionContext: this.executionContext,
           executingAgentId: this.executingAgentId,
           toolName: call.name,
           toolInput: effectiveCall.args as Record<string, unknown>,
@@ -2688,6 +2919,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           runId,
           threadId,
           agentId: this.agentId,
+          executionContext: this.executionContext,
           executingAgentId: this.executingAgentId,
           toolName: call.name,
           toolInput: effectiveCall.args as Record<string, unknown>,
@@ -2804,6 +3036,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           runId,
           threadId,
           agentId: this.agentId,
+          executionContext: this.executionContext,
           executingAgentId: this.executingAgentId,
           toolName: call.name,
           toolInput: resolvedArgs,
@@ -2935,18 +3168,26 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
 
   private storeCodeSessionFromResults(
     results: t.ToolExecuteResult[],
-    requestMap: Map<string, t.ToolCallRequest>
+    requestMap: Map<string, t.ToolCallRequest>,
+    baselineByRequestId: ReadonlyMap<
+      string,
+      ReadonlyMap<string, string>
+    > = new Map(),
+    executionBaselineByRequestId: ReadonlyMap<
+      string,
+      ReadonlyMap<string, string>
+    > = baselineByRequestId
   ): void {
     if (!this.sessions) {
       return;
     }
 
+    this.retainCodeSessionInputsFromRequests(
+      requestMap.values(),
+      baselineByRequestId
+    );
     for (let i = 0; i < results.length; i++) {
       const result = results[i];
-      if (result.status !== 'success' || result.artifact == null) {
-        continue;
-      }
-
       const request = requestMap.get(result.toolCallId);
       if (
         request?.name == null ||
@@ -2957,6 +3198,10 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         continue;
       }
 
+      if (result.status !== 'success' || result.artifact == null) {
+        continue;
+      }
+
       const artifact = result.artifact as t.CodeExecutionArtifact | undefined;
       const execSessionId = artifact?.session_id;
       if (execSessionId == null || execSessionId === '') {
@@ -2967,11 +3212,59 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         this.sessions,
         this.codeSessionKey,
         execSessionId,
-        artifact?.files
+        artifact?.files,
+        artifact?.deleted_files,
+        executionBaselineByRequestId.get(request.id)
       );
     }
   }
 
+  private retainCodeSessionInputsFromRequests(
+    requests: Iterable<t.ToolCallRequest>,
+    baselineByRequestId: ReadonlyMap<
+      string,
+      ReadonlyMap<string, string>
+    >
+  ): void {
+    if (!this.sessions) {
+      return;
+    }
+    const refreshedByName = new Map<string, t.CodeEnvFile>();
+    let refreshedSessionId: string | undefined;
+    let fileSessionId: string | undefined;
+    for (const request of requests) {
+      if (
+        request.name === '' ||
+        (!this.participatesInCodeSession(request.name) &&
+          request.name !== Constants.SKILL_TOOL &&
+          request.name !== Constants.READ_FILE &&
+          request.retainCodeSessionInputs !== true)
+      ) {
+        continue;
+      }
+      const context = request.codeSessionContext;
+      if (context?.session_id == null || context.session_id === '') continue;
+      if (context.session_id !== request.codeSessionBaselineId) {
+        refreshedSessionId = context.session_id;
+      }
+      const baselineIdentityByName = baselineByRequestId.get(request.id);
+      for (const file of context.files ?? []) {
+        if (!file.id || !file.name || !file.storage_session_id) continue;
+        if (baselineIdentityByName?.get(file.name) === fileIdentityKey(file)) {
+          continue;
+        }
+        fileSessionId = context.session_id;
+        refreshedByName.set(file.name, file);
+      }
+    }
+    const sessionId = refreshedSessionId ?? fileSessionId;
+    if (sessionId == null) return;
+    retainCodeSessionInputs(this.sessions, this.codeSessionKey, {
+      session_id: sessionId,
+      files: [...refreshedByName.values()],
+    });
+  }
+
   /**
    * Post-processes standard runTool outputs: dispatches ON_RUN_STEP_COMPLETED
    * and stores code session context. Mirrors the completion handling in
@@ -2990,6 +3283,10 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     outputs: (BaseMessage | Command)[],
     config: RunnableConfig,
     resolvedArgsByCallId?: ResolvedArgsByCallId,
+    codeSessionBaselineByCallId?: ReadonlyMap<
+      string,
+      ReadonlyMap<string, string>
+    >,
     errorOwnership?: ToolErrorOwnership
   ): Promise<void> {
     const ownership = errorOwnership ?? this.looseErrorOwnership;
@@ -3025,7 +3322,11 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         ownership.undispatched.delete(toolCallId);
       }
 
-      if (this.sessions && this.participatesInCodeSession(call.name)) {
+      if (
+        toolMessage.status !== 'error' &&
+        this.sessions &&
+        this.participatesInCodeSession(call.name)
+      ) {
         const artifact = toolMessage.artifact as
           | t.CodeExecutionArtifact
           | undefined;
@@ -3035,7 +3336,9 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             this.sessions,
             this.codeSessionKey,
             execSessionId,
-            artifact?.files
+            artifact?.files,
+            artifact?.deleted_files,
+            codeSessionBaselineByCallId?.get(toolCallId)
           );
         }
       }
@@ -3178,10 +3481,13 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           unresolvedByCallId.set(call.id, pre.unresolved);
         }
       } else if (registry != null) {
-        const options = { substituteIntentKey: this.toolDeclaresBusinessIntent(call.name) };
-        const { resolved, unresolved } = preBatchSnapshot != null
-          ? preBatchSnapshot.resolve(originalArgs, options)
-          : registry.resolve(registryRunId, originalArgs, options);
+        const options = {
+          substituteIntentKey: this.toolDeclaresBusinessIntent(call.name),
+        };
+        const { resolved, unresolved } =
+          preBatchSnapshot != null
+            ? preBatchSnapshot.resolve(originalArgs, options)
+            : registry.resolve(registryRunId, originalArgs, options);
         resolvedArgs = resolved as Record<string, unknown>;
         if (unresolved.length > 0 && call.id != null) {
           unresolvedByCallId.set(call.id, unresolved);
@@ -3229,8 +3535,13 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       config,
       runId
     );
-    const approvalReviewEvidence = getToolApprovalReviewEvidence(config,
-      getToolBatchReplayOwner(config, this.executingAgentId ?? this.agentId ?? ''));
+    const approvalReviewEvidence = getToolApprovalReviewEvidence(
+      config,
+      getToolBatchReplayOwner(
+        config,
+        this.executingAgentId ?? this.agentId ?? ''
+      )
+    );
     const hasPendingApproval = preToolCalls.some(
       (entry) =>
         entry.call.id != null &&
@@ -3295,6 +3606,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
                   runId,
                   threadId,
                   agentId: this.agentId,
+                  executionContext: this.executionContext,
                   executingAgentId: this.executingAgentId,
                   toolName: entry.call.name,
                   toolInput: entry.args,
@@ -3397,6 +3709,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
                 runId,
                 threadId,
                 agentId: this.agentId,
+                executionContext: this.executionContext,
                 executingAgentId: this.executingAgentId,
                 toolName: item.toolName,
                 toolInput: item.args,
@@ -3504,9 +3817,15 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           if (hookResult.updatedInput != null) {
             applyInputOverride(entry, hookResult.updatedInput);
           } else {
-            const reviewed = getReviewedToolApproval(approvalReviewEvidence?.payload, entry.call.id);
+            const reviewed = getReviewedToolApproval(
+              approvalReviewEvidence?.payload,
+              entry.call.id
+            );
             if (reviewed != null) {
-              entry.args = this.coerceEventToolArgs(entry.call.name, reviewed.request.arguments);
+              entry.args = this.coerceEventToolArgs(
+                entry.call.name,
+                reviewed.request.arguments
+              );
               if (entry.call.id != null) {
                 unresolvedByCallId.delete(entry.call.id);
               }
@@ -3877,6 +4196,18 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       }
 
       const requestMap = new Map(plan.allRequests.map((r) => [r.id, r]));
+      for (const request of plan.allRequests) {
+        request.codeSessionBaselineId = request.codeSessionContext?.session_id;
+      }
+      const codeSessionBaselineByRequestId = new Map<
+        string,
+        ReadonlyMap<string, string>
+      >(
+        plan.allRequests.map((request) => [
+          request.id,
+          codeSessionIdentityByName(request.codeSessionContext),
+        ])
+      );
       const eagerExecutions: Array<{
         request: t.ToolCallRequest;
         execution: t.EagerEventToolExecution;
@@ -3887,6 +4218,14 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         const eagerExecution = this.takeMatchingEagerEventExecution(request);
         if (eagerExecution != null) {
           eagerExecutions.push({ request, execution: eagerExecution });
+          if (eagerExecution.codeSessionBaselineByName != null) {
+            codeSessionBaselineByRequestId.set(
+              request.id,
+              eagerExecution.codeSessionBaselineByName
+            );
+          }
+          request.codeSessionBaselineId =
+            eagerExecution.request.codeSessionBaselineId;
         } else {
           dispatchRequests.push(request);
         }
@@ -3972,6 +4311,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             };
 
             const batchRequest: t.ToolExecuteBatchRequest = {
+              executionContext: this.executionContext,
               toolCalls: dispatchRequests,
               userId: config.configurable?.user_id as string | undefined,
               // Dispatch attribution, NOT the hook subagent-scope marker:
@@ -4011,24 +4351,46 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
 
       const eagerResultsPromise = Promise.all(
         eagerExecutions.map(async ({ request, execution }) => {
-          const results = await this.resolveEagerEventExecution(
-            request,
-            execution
-          );
-          return {
-            results,
-            completionDispatched:
-              execution.completionDispatched === true &&
-              execution.request.turn === request.turn,
-            toolCallId: request.id,
-          };
+          try {
+            const results = await this.resolveEagerEventExecution(
+              request,
+              execution
+            );
+            return {
+              results,
+              completionDispatched:
+                execution.completionDispatched === true &&
+                execution.request.turn === request.turn,
+              toolCallId: request.id,
+            };
+          } finally {
+            if (execution.request.codeSessionContext != null) {
+              request.codeSessionContext = execution.request.codeSessionContext;
+            }
+          }
         })
       );
 
-      const [eagerResults, dispatchedResults] = await Promise.all([
+      const [eagerOutcome, dispatchedOutcome] = await Promise.allSettled([
         eagerResultsPromise,
         dispatchPromise,
       ]);
+      if (eagerOutcome.status === 'rejected') {
+        this.retainCodeSessionInputsFromRequests(
+          requestMap.values(),
+          codeSessionBaselineByRequestId
+        );
+        throw eagerOutcome.reason;
+      }
+      if (dispatchedOutcome.status === 'rejected') {
+        this.retainCodeSessionInputsFromRequests(
+          requestMap.values(),
+          codeSessionBaselineByRequestId
+        );
+        throw dispatchedOutcome.reason;
+      }
+      const eagerResults = eagerOutcome.value;
+      const dispatchedResults = dispatchedOutcome.value;
       // Settle in-flight early completion dispatches before the batch loop
       // below decides which completions still need emitting.
       await Promise.allSettled(earlyCompletionDispatches);
@@ -4045,8 +4407,26 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         ...flattenedEagerResults,
         ...dispatchedResults,
       ];
+      /* Hosts may lazily refresh codeSessionContext while dispatching. The
+       * pre-dispatch snapshot above remains the authority for detecting those
+       * additions, while deletion reconciliation must use the identities the
+       * execution actually received. */
+      const executionBaselineByRequestId = new Map<
+        string,
+        ReadonlyMap<string, string>
+      >(
+        plan.allRequests.map((request) => [
+          request.id,
+          codeSessionIdentityByName(request.codeSessionContext),
+        ])
+      );
 
-      this.storeCodeSessionFromResults(results, requestMap);
+      this.storeCodeSessionFromResults(
+        results,
+        requestMap,
+        codeSessionBaselineByRequestId,
+        executionBaselineByRequestId
+      );
 
       for (const result of results) {
         if (result.injectedMessages && result.injectedMessages.length > 0) {
@@ -4115,6 +4495,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
                 runId,
                 threadId,
                 agentId: this.agentId,
+                executionContext: this.executionContext,
                 executingAgentId: this.executingAgentId,
                 toolName,
                 toolInput: request?.args ?? {},
@@ -4167,6 +4548,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
                 runId,
                 threadId,
                 agentId: this.agentId,
+                executionContext: this.executionContext,
                 executingAgentId: this.executingAgentId,
                 toolName,
                 toolInput: request?.args ?? {},
@@ -4326,28 +4708,40 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       // effects) every round.
       this.eagerEventToolSuppressions?.add(request.name);
       this.eagerEventToolSuppressions?.add(execution.toolName);
+      if (
+        this.participatesInCodeSession(execution.toolName) ||
+        execution.toolName === Constants.SKILL_TOOL ||
+        execution.toolName === Constants.READ_FILE
+      ) {
+        request.retainCodeSessionInputs = true;
+      }
       // eslint-disable-next-line no-console
       console.warn(
         '[ToolNode] eager prestart args diverged from the final request for ' +
           `tool "${request.name}" (toolCallId=${request.id}); suppressing ` +
           'eager prestart for this tool for the rest of the run'
       );
+      const createMismatchOutcome = (): t.EagerEventToolExecutionOutcome => ({
+        results: [
+          {
+            toolCallId: request.id,
+            status: 'error',
+            content: '',
+            errorMessage:
+              'Tool call changed after eager execution started; refusing to re-run the tool to avoid duplicate side effects.',
+          },
+        ],
+      });
       return {
         toolCallId: request.id,
         toolName: request.name,
         args: request.args,
-        request,
-        promise: Promise.resolve({
-          results: [
-            {
-              toolCallId: request.id,
-              status: 'error',
-              content: '',
-              errorMessage:
-                'Tool call changed after eager execution started; refusing to re-run the tool to avoid duplicate side effects.',
-            },
-          ],
-        }),
+        request: execution.request,
+        codeSessionBaselineByName: execution.codeSessionBaselineByName,
+        promise: execution.promise.then(
+          createMismatchOutcome,
+          createMismatchOutcome
+        ),
       };
     }
 
@@ -4439,6 +4833,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           runId,
           threadId,
           agentId: this.agentId,
+          executionContext: this.executionContext,
           executingAgentId: this.executingAgentId,
           entries: orderedBatchEntries,
         },
@@ -4647,13 +5042,16 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       position: number,
       result: SettledDirectToolResult
     ): void => {
-      if (
-        baseContext.replayBatchKey == null ||
-        settledBatchResults == null
-      ) {
+      if (baseContext.replayBatchKey == null || settledBatchResults == null) {
         return;
       }
-      settledBatchResults.set(directReplayKey(call, baseContext.replayCallPositions?.get(call) ?? batchIndices[position]), result);
+      settledBatchResults.set(
+        directReplayKey(
+          call,
+          baseContext.replayCallPositions?.get(call) ?? batchIndices[position]
+        ),
+        result
+      );
       this.settledDirectResultsByBatch.set(
         baseContext.replayBatchKey,
         settledBatchResults
@@ -4668,14 +5066,26 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       if (turn != null) {
         this.recordToolUsageTurn(call.name, turn, call.id);
       }
-      const refMeta = isBaseMessage(result.output) ? result.output.additional_kwargs as t.ToolMessageRefMetadata : undefined;
+      const refMeta = isBaseMessage(result.output)
+        ? (result.output.additional_kwargs as t.ToolMessageRefMetadata)
+        : undefined;
       if (refMeta?._refKey != null && result.referenceContent != null) {
-        this.toolOutputRegistry?.set(baseContext.batchScopeId, refMeta._refKey, result.referenceContent);
+        this.toolOutputRegistry?.set(
+          baseContext.batchScopeId,
+          refMeta._refKey,
+          result.referenceContent
+        );
         if (isBaseMessage(result.output)) {
-          result = { ...result, output: new ToolMessage({
-            ...(result.output as ToolMessage),
-            additional_kwargs: { ...result.output.additional_kwargs, _refScope: baseContext.batchScopeId },
-          }) };
+          result = {
+            ...result,
+            output: new ToolMessage({
+              ...(result.output as ToolMessage),
+              additional_kwargs: {
+                ...result.output.additional_kwargs,
+                _refScope: baseContext.batchScopeId,
+              },
+            }),
+          };
         }
       }
       if (
@@ -4685,17 +5095,38 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
       ) {
         baseContext.resolvedArgsByCallId.set(call.id, result.resolvedArgs);
       }
+      if (
+        result.codeSessionBaseline != null &&
+        baseContext.codeSessionBaselineByCallId != null
+      ) {
+        baseContext.codeSessionBaselineByCallId.set(
+          call.id ?? '',
+          new Map(result.codeSessionBaseline)
+        );
+      }
       return result;
     };
     const runOne = async (
       call: ToolCall,
       position: number
     ): Promise<SettledDirectToolResult> => {
-      const cachedResult = settledBatchResults?.get(directReplayKey(call, baseContext.replayCallPositions?.get(call) ?? batchIndices[position]));
+      const cachedResult = settledBatchResults?.get(
+        directReplayKey(
+          call,
+          baseContext.replayCallPositions?.get(call) ?? batchIndices[position]
+        )
+      );
       if (cachedResult != null) {
-        if (cachedResult.proposal.name !== call.name ||
-          !recordArgsEqual(cachedResult.proposal.args, call.args as Record<string, unknown>)) {
-          throw new Error('Cannot change an already completed tool call during approval replay');
+        if (
+          cachedResult.proposal.name !== call.name ||
+          !recordArgsEqual(
+            cachedResult.proposal.args,
+            call.args as Record<string, unknown>
+          )
+        ) {
+          throw new Error(
+            'Cannot change an already completed tool call during approval replay'
+          );
         }
         return restoreResult(call, cachedResult);
       }
@@ -4709,16 +5140,33 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         call.id == null
           ? undefined
           : baseContext.resolvedArgsByCallId?.get(call.id);
+      const codeSessionBaseline =
+        baseContext.codeSessionBaselineByCallId?.get(call.id ?? '');
       const result: SettledDirectToolResult = {
-        proposal: structuredClone({ name: call.name, args: call.args as Record<string, unknown> }),
-        turn: call.id == null ? undefined : this.toolCallTurns.get(call.id) ?? this.directPathTurns.get(call.id),
+        proposal: structuredClone({
+          name: call.name,
+          args: call.args as Record<string, unknown>,
+        }),
+        turn:
+          call.id == null
+            ? undefined
+            : (this.toolCallTurns.get(call.id) ??
+              this.directPathTurns.get(call.id)),
         output,
         additionalContexts,
         ...(resolvedArgs == null ? {} : { resolvedArgs }),
+        ...(codeSessionBaseline == null
+          ? {}
+          : { codeSessionBaseline: [...codeSessionBaseline.entries()] }),
       };
-      const refMeta = isBaseMessage(output) ? output.additional_kwargs as t.ToolMessageRefMetadata : undefined;
+      const refMeta = isBaseMessage(output)
+        ? (output.additional_kwargs as t.ToolMessageRefMetadata)
+        : undefined;
       if (refMeta?._refKey != null) {
-        result.referenceContent = this.toolOutputRegistry?.get(refMeta._refScope, refMeta._refKey);
+        result.referenceContent = this.toolOutputRegistry?.get(
+          refMeta._refScope,
+          refMeta._refKey
+        );
       }
       cacheResult(call, position, result);
       return restoreResult(call, result);
@@ -4734,8 +5182,11 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         directCalls.map((call, i) => runOne(call, i))
       );
       this.throwIfBreakerTripped(config);
-      const failed = results.find((result) => result.status === 'rejected' && !isGraphInterrupt(result.reason)) ??
-        results.find((result) => result.status === 'rejected');
+      const failed =
+        results.find(
+          (result) =>
+            result.status === 'rejected' && !isGraphInterrupt(result.reason)
+        ) ?? results.find((result) => result.status === 'rejected');
       if (failed?.status === 'rejected') {
         throw failed.reason;
       }
@@ -4877,8 +5328,11 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     return (Array.isArray(input) ? outputs : { messages: outputs }) as T;
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  protected async run(input: any, config: RunnableConfig, referenceReplay: { state?: ToolOutputReferenceState } = {}): Promise<T> {
+  protected async run(
+    input: T,
+    config: RunnableConfig,
+    referenceReplay: { state?: ToolOutputReferenceState } = {}
+  ): Promise<T> {
     /**
      * Breaker read once at batch entry: every downstream path (direct
      * `tool.invoke` runtimes and the ON_TOOL_EXECUTE dispatch) inherits this
@@ -4922,6 +5376,10 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
      * ToolNode cannot read or wipe each other's entries.
      */
     const resolvedArgsByCallId = new Map<string, Record<string, unknown>>();
+    const codeSessionBaselineByCallId = new Map<
+      string,
+      ReadonlyMap<string, string>
+    >();
     /** Per-invocation error-completion ownership — see `ToolErrorOwnership`. */
     const errorOwnership = createToolErrorOwnership();
     /**
@@ -4938,10 +5396,20 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
     const incomingRunId = config.configurable?.run_id as string | undefined;
     const batchScopeId = incomingRunId ?? `\0anon-${this.anonBatchCounter++}`;
     const resumedReferenceState = referenceReplay.state;
-    const replayInputSnapshot = resumedReferenceState == null ? undefined :
-      this.toolOutputRegistry?.resumeBatch(batchScopeId, resumedReferenceState);
-    referenceReplay.state = resumedReferenceState ?? this.toolOutputRegistry?.snapshotState(batchScopeId);
-    const turn = resumedReferenceState?.turnCounter ?? this.toolOutputRegistry?.nextTurn(batchScopeId) ?? 0;
+    const replayInputSnapshot =
+      resumedReferenceState == null
+        ? undefined
+        : this.toolOutputRegistry?.resumeBatch(
+          batchScopeId,
+          resumedReferenceState
+        );
+    referenceReplay.state =
+      resumedReferenceState ??
+      this.toolOutputRegistry?.snapshotState(batchScopeId);
+    const turn =
+      resumedReferenceState?.turnCounter ??
+      this.toolOutputRegistry?.nextTurn(batchScopeId) ??
+      0;
     let outputs: (BaseMessage | Command)[];
     let replayBatchKey: string | undefined;
     /** Hoisted from the messages-state branch so the Command tail can carry
@@ -4992,6 +5460,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           batchScopeId,
           preBatchSnapshot: replayInputSnapshot,
           resolvedArgsByCallId,
+          codeSessionBaselineByCallId,
           errorOwnership,
           additionalContextsSink: directAdditionalContexts,
           replayBatchKey: sendReplayBatchKey,
@@ -5020,6 +5489,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         [sendOutput],
         config,
         resolvedArgsByCallId,
+        codeSessionBaselineByCallId,
         errorOwnership
       );
     } else {
@@ -5045,10 +5515,13 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         throw new Error('ToolNode only accepts AIMessages as input.');
       }
       const aiMessage = assistantBatch.message;
-      replayBatchKey = getToolBatchReplayScope(config) == null ? undefined : getAssistantBatchReplayKey(
-        assistantBatch,
-        getToolBatchReplayScope(config)
-      );
+      replayBatchKey =
+        getToolBatchReplayScope(config) == null
+          ? undefined
+          : getAssistantBatchReplayKey(
+            assistantBatch,
+            getToolBatchReplayScope(config)
+          );
 
       if (this.loadRuntimeTools) {
         const { tools, toolMap } = this.loadRuntimeTools(
@@ -5059,7 +5532,9 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         this.applyToolExecutionOverrides();
       }
 
-      const replayCallPositions = new Map((aiMessage.tool_calls ?? []).map((call, i) => [call, i]));
+      const replayCallPositions = new Map(
+        (aiMessage.tool_calls ?? []).map((call, i) => [call, i])
+      );
       const filteredCalls =
         aiMessage.tool_calls?.filter((call) => {
           /**
@@ -5069,7 +5544,9 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
            *    which are executed by the provider's API and don't require invocation
            */
           return (
-            (call.id == null || call.id === '' || !toolMessageIds.has(call.id)) &&
+            (call.id == null ||
+              call.id === '' ||
+              !toolMessageIds.has(call.id)) &&
             !(
               call.id?.startsWith(Constants.ANTHROPIC_SERVER_TOOL_PREFIX) ??
               false
@@ -5234,7 +5711,8 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
          * same-turn direct outputs.
          */
         const preBatchSnapshot =
-          replayInputSnapshot ?? this.toolOutputRegistry?.snapshot(batchScopeId);
+          replayInputSnapshot ??
+          this.toolOutputRegistry?.snapshot(batchScopeId);
         if (preBatchSnapshot != null) {
           for (const entry of eventEntries) {
             if (entry.call.id != null) {
@@ -5269,6 +5747,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
                 turn,
                 batchScopeId,
                 resolvedArgsByCallId,
+                codeSessionBaselineByCallId,
                 errorOwnership,
                 preBatchSnapshot,
                 additionalContextsSink: directAdditionalContexts,
@@ -5281,11 +5760,18 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             : [];
 
         const settledDirectResults =
-          replayBatchKey == null ? undefined : this.settledDirectResultsByBatch.get(replayBatchKey);
+          replayBatchKey == null
+            ? undefined
+            : this.settledDirectResultsByBatch.get(replayBatchKey);
         const unhandledDirectCalls: ToolCall[] = [];
         const unhandledDirectOutputs: (BaseMessage | Command)[] = [];
         for (let i = 0; i < directCalls.length; i++) {
-          const settled = settledDirectResults?.get(directReplayKey(directCalls[i], replayCallPositions.get(directCalls[i])!));
+          const settled = settledDirectResults?.get(
+            directReplayKey(
+              directCalls[i],
+              replayCallPositions.get(directCalls[i])!
+            )
+          );
           if (settled?.completionHandled === true) {
             continue;
           }
@@ -5299,10 +5785,16 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             unhandledDirectOutputs,
             config,
             resolvedArgsByCallId,
+            codeSessionBaselineByCallId,
             errorOwnership
           );
           for (let i = 0; i < directCalls.length; i++) {
-            const settled = settledDirectResults?.get(directReplayKey(directCalls[i], replayCallPositions.get(directCalls[i])!));
+            const settled = settledDirectResults?.get(
+              directReplayKey(
+                directCalls[i],
+                replayCallPositions.get(directCalls[i])!
+              )
+            );
             if (settled != null) {
               settled.completionHandled = true;
             }
@@ -5355,7 +5847,8 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
         // leak a sibling's just-registered output into a sister
         // call's args mid-await (Codex P1 #18).
         const preBatchSnapshot =
-          replayInputSnapshot ?? this.toolOutputRegistry?.snapshot(batchScopeId);
+          replayInputSnapshot ??
+          this.toolOutputRegistry?.snapshot(batchScopeId);
         const directAdditionalContexts: string[] = [];
         const approvalScopeCallIds = new Set(
           filteredCalls.flatMap((call) => (call.id == null ? [] : [call.id]))
@@ -5368,6 +5861,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
             turn,
             batchScopeId,
             resolvedArgsByCallId,
+            codeSessionBaselineByCallId,
             errorOwnership,
             preBatchSnapshot,
             additionalContextsSink: directAdditionalContexts,
@@ -5382,6 +5876,7 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
           toolOutputs,
           config,
           resolvedArgsByCallId,
+          codeSessionBaselineByCallId,
           errorOwnership
         );
         // Append accumulated additionalContexts as a single