@librechat/agents 3.5.1 → 3.6.1

package/src/tools/subagent/SubagentExecutor.ts

@@ -46,6 +46,7 @@ import type {
   RunStep,
   RunStepDeltaEvent,
   RunStepClosedEvent,
+  RunStepStatus,
   StandardGraphInput,
   ExecutableSubagentConfigEntry,
   ResolvedSubagentConfig,
@@ -61,7 +62,6 @@ import type {
   ToolApprovalInterruptPayload,
   ToolExecuteBatchRequest,
   ToolCallDelta,
-  ToolSessionContext,
   TokenCounter,
   ToolApprovalDecision,
   ToolApprovalDecisionMap,
@@ -87,6 +87,7 @@ import type {
 import type { GraphFactory } from '@/graphs/graphFactory';
 import type { StandardGraph } from '@/graphs/Graph';
 import type { HandlerRegistry } from '@/events';
+import { stripRunStepResumeState } from '@/tools/runStepResume';
 import {
   getSubagentApprovalExecutionScope,
   SubagentDefinitionBindingError,
@@ -124,6 +125,7 @@ import {
   createChildGraphPlan,
   isGraphSubagentConfig,
 } from './childGraphConfig';
+import { seedAgentInitialSessions } from '@/utils/toolSessions';
 import { stableStringify } from '@/tools/eagerEventExecution';
 import { composeAbortSignals } from '@/utils/misc';
 
@@ -147,71 +149,11 @@ const SUBAGENT_CONFIG_CHANGED_MESSAGE =
 const SUBAGENT_INVOCATION_CHANGED_MESSAGE =
   'Subagent error: Subagent invocation changed for this execution.';
 
-function cloneToolSessionContext(
-  context: ToolSessionContext
-): ToolSessionContext {
-  return {
-    ...context,
-    ...(context.files == null
-      ? {}
-      : {
-        files: context.files.map((file) => ({
-          ...file,
-          storage_session_id: file.storage_session_id ?? context.session_id,
-        })),
-      }),
-  };
-}
-
 function seedChildGraphSessions(
   childGraph: StandardGraph,
   agents: AgentInputs[]
 ): void {
-  const seenFilesByTool = new Map<string, Set<string>>();
-  for (const agent of agents) {
-    if (agent.initialSessions == null) {
-      continue;
-    }
-    for (const [toolName, context] of agent.initialSessions) {
-      const existing = childGraph.sessions.get(toolName);
-      if (existing == null) {
-        const cloned = cloneToolSessionContext(context);
-        childGraph.sessions.set(toolName, cloned);
-        seenFilesByTool.set(
-          toolName,
-          new Set(
-            cloned.files?.map(
-              (file) => `${file.storage_session_id ?? ''}\0${file.id}`
-            ) ?? []
-          )
-        );
-        continue;
-      }
-      if (context.files == null || context.files.length === 0) {
-        continue;
-      }
-      const seenFiles =
-        seenFilesByTool.get(toolName) ??
-        new Set(
-          existing.files?.map(
-            (file) =>
-              `${file.storage_session_id ?? existing.session_id}\0${file.id}`
-          ) ?? []
-        );
-      const files = existing.files == null ? [] : [...existing.files];
-      for (const file of context.files) {
-        const storageSessionId = file.storage_session_id ?? context.session_id;
-        const key = `${storageSessionId}\0${file.id}`;
-        if (seenFiles.has(key)) {
-          continue;
-        }
-        seenFiles.add(key);
-        files.push({ ...file, storage_session_id: storageSessionId });
-      }
-      seenFilesByTool.set(toolName, seenFiles);
-      childGraph.sessions.set(toolName, { ...existing, files });
-    }
-  }
+  seedAgentInitialSessions(childGraph.sessions, agents);
 }
 
 async function dispatchObservationalSubagentUpdate(
@@ -713,11 +655,11 @@ function addSubagentScope(
   resumeManifest?: SubagentResumeManifest
 ): Interrupt[] {
   return interrupts.map((childInterrupt) => {
-    let payload = childInterrupt.value;
+    let payload = stripRunStepResumeState(childInterrupt.value);
     if (isToolApprovalPayload(payload)) {
       payload = {
-        ...childInterrupt.value,
-        subagent: childInterrupt.value.subagent ?? scope,
+        ...payload,
+        subagent: payload.subagent ?? scope,
       };
     }
     return {
@@ -1610,6 +1552,42 @@ export class SubagentExecutor {
     graph.clearHeavyState();
   }
 
+  /**
+   * Terminal sweep for a child graph's run steps. `Run.processStream` sweeps
+   * only the graph it owns, and a child executes through `workflow.invoke()`
+   * outside that loop — so without this, a child's last message step (no
+   * successor step ever opens to close it) and every step of an aborted child
+   * stay `in_progress` forever, leaving hosts with unmatched starts.
+   *
+   * Must run BEFORE `forwarding.drain()`: closures reach the parent only as
+   * child custom events through the forwarder, which stops relaying once
+   * drained.
+   *
+   * Deliberately not called on the interrupt path, mirroring the parent's
+   * `isAwaitingResume` guard: an in-process resume continues into these same
+   * steps, so closing them would strand the resumed run. A resume that
+   * crosses a process boundary is a known gap either way — the child rebuilds
+   * from `SubagentGraphResumeState`, which carries tool-call identity but not
+   * `contentData`, so pre-interrupt steps are unreachable from the new graph
+   * and stay open. That is the child-level instance of the same cross-process
+   * orphan the parent graph has, and it needs persisted lifecycle state to
+   * fix rather than a sweep here.
+   *
+   * `at` is the moment execution actually ended, threaded through so
+   * observational work between then and the sweep cannot inflate the stamps.
+   */
+  private async closeChildRunSteps(
+    graph: StandardGraph,
+    status: Exclude<RunStepStatus, 'in_progress'>,
+    at?: number
+  ): Promise<void> {
+    try {
+      await graph.closeUnfinishedRunSteps(status, at);
+    } catch (_e) {
+      /** A failed sweep must never mask the subagent's own outcome */
+    }
+  }
+
   clearHeavyState(): void {
     this.executions.clear((record) => {
       if (record.activeRun != null) {
@@ -2470,6 +2448,8 @@ export class SubagentExecutor {
         result = childResult;
       }
     } catch (error) {
+      /** Stamped at failure, not after the error-envelope work below. */
+      const childTerminalAt = Date.now();
       if (isGraphInterrupt(error)) {
         const activeChildRun = execution.activeRun;
         if (activeChildRun != null) {
@@ -2505,6 +2485,16 @@ export class SubagentExecutor {
         childBreaker.abort(error);
       }
       const errorMessage = truncateErrorMessage(error);
+      /**
+       * `cancelled` vs `failed` mirrors `Run.resolveSweepStatus`: an aborted
+       * child was stopped on purpose (caller abort, or a breaker trip from a
+       * parallel sibling), anything else died of an unexpected error.
+       */
+      await this.closeChildRunSteps(
+        childGraph,
+        childSignal.aborted ? 'cancelled' : 'failed',
+        childTerminalAt
+      );
       if (forwarding) {
         await forwarding.drain();
         await this.emitSubagentUpdate(parentRegistry!, {
@@ -2537,6 +2527,14 @@ export class SubagentExecutor {
       };
     }
 
+    /**
+     * Child execution ended here. Captured before the observational work
+     * below — an awaited `SubagentStop` hook, then the forwarder drain — so a
+     * slow hook cannot inflate the closure stamps with its own latency.
+     * Mirrors the parent's `terminalAt` capture in `Run.processStream`.
+     */
+    const childTerminalAt = Date.now();
+
     if (result == null) {
       throw new Error('Subagent completed without producing graph state.');
     }
@@ -2573,6 +2571,10 @@ export class SubagentExecutor {
       });
     }
 
+    if (!childAlreadyCompleted) {
+      await this.closeChildRunSteps(childGraph, 'completed', childTerminalAt);
+    }
+
     if (forwarding && !childAlreadyCompleted) {
       await forwarding.drain();
       await this.emitSubagentUpdate(parentRegistry!, {

package/src/tools/subagent/SubagentReplay.ts

@@ -2,7 +2,13 @@ import type { ToolCall, ToolMessage } from '@langchain/core/messages/tool';
 import type { RunnableConfig } from '@langchain/core/runnables';
 import type { ToolOutputReferenceState } from '@/tools/toolOutputReferences';
 import type { ToolApprovalReplaySnapshot } from '@/hooks';
-import type { ToolSessionContext } from '@/types';
+import type { RunStepResumeState, ToolSessionContext } from '@/types';
+import {
+  attachRunStepResumeState,
+  getRunStepResumeState,
+  isRunStepResumeState,
+  stripRunStepResumeState,
+} from '@/tools/runStepResume';
 
 export const SUBAGENT_RESUME_MANIFEST_CONFIG_KEY =
   '__librechat_subagent_resume_manifest';
@@ -57,6 +63,7 @@ export interface SubagentGraphResumeState {
   toolNodes: SubagentToolNodeResumeState[];
   eagerToolUsage: SubagentEagerToolUsageState[];
   eagerToolSuppressions: string[];
+  runStepState?: RunStepResumeState;
   toolOutputReferences?: ToolOutputReferenceState;
 }
 
@@ -291,6 +298,8 @@ function isGraphResumeState(value: unknown): value is SubagentGraphResumeState {
     !state.eagerToolUsage.every(isEagerToolUsageState) ||
     !Array.isArray(state.eagerToolSuppressions) ||
     !state.eagerToolSuppressions.every(isString) ||
+    (state.runStepState != null &&
+      !isRunStepResumeState(state.runStepState)) ||
     (state.toolOutputReferences != null &&
       !isToolOutputReferenceState(state.toolOutputReferences))
   ) {
@@ -412,6 +421,7 @@ function isSubagentResumeManifest(
 export function getSubagentResumeManifest(
   payload: unknown
 ): SubagentResumeManifest | undefined {
+  payload = stripRunStepResumeState(payload);
   if (payload == null || typeof payload !== 'object') {
     return undefined;
   }
@@ -489,6 +499,16 @@ export function attachSubagentResumeManifest(
   payload: unknown,
   manifest: SubagentResumeManifest
 ): object {
+  const runStepState = getRunStepResumeState(payload);
+  if (runStepState != null) {
+    return attachRunStepResumeState(
+      attachSubagentResumeManifest(
+        stripRunStepResumeState(payload),
+        manifest
+      ),
+      runStepState
+    );
+  }
   if (
     isWrappedSubagentResumePayload(payload) ||
     isInlineSubagentResumePayload(payload)
@@ -512,6 +532,7 @@ export function attachSubagentResumeManifest(
 }
 
 export function stripSubagentResumeManifest(payload: unknown): unknown {
+  payload = stripRunStepResumeState(payload);
   if (isWrappedSubagentResumePayload(payload)) {
     return payload[SUBAGENT_RESUME_WRAPPED_PAYLOAD_KEY];
   }

package/src/types/graph.ts

@@ -20,6 +20,7 @@ import type {
 import type {
   RunStep,
   RunStepDeltaEvent,
+  RunStepResumeState,
   RunStepClosedEvent,
   MessageDeltaEvent,
   ReasoningDeltaEvent,
@@ -76,6 +77,7 @@ export type SystemCallbacks = {
 
 export type BaseGraphState = {
   messages: BaseMessage[];
+  runStepState?: RunStepResumeState;
 };
 
 export type AgentSubgraphState = BaseGraphState & {
@@ -803,6 +805,14 @@ export interface LangfuseConfig {
 
 export interface AgentInputs {
   agentId: string;
+  /**
+   * Partition key for transient code-session ids and file refs. Agents with
+   * the same key share those refs; different execution profiles/scopes must
+   * use different keys. Defaults to the legacy shared `execute_code` slot.
+   * Non-default partitions also disable speculative eager execution of code
+   * tools because their factory may target a durable backend.
+   */
+  codeSessionKey?: string;
   /** Human-readable name for the agent (used in handoff context). Defaults to agentId if not provided. */
   name?: string;
   toolEnd?: boolean;

package/src/types/index.ts

@@ -10,3 +10,4 @@ export * from './tools';
 export * from './summarize';
 export * from './activityLabel';
 export * from './assistantPhase';
+export * from './reasoningLabel';

package/src/types/reasoningLabel.ts

@@ -0,0 +1,59 @@
+import type { RunnableConfig } from '@langchain/core/runnables';
+import type { UsageMetadata } from '@langchain/core/messages';
+import type { ClientOptions } from '@/types/llm';
+import type { Providers } from '@/common';
+
+/** Lifecycle state of the visible reasoning snapshot being labeled. */
+export type ReasoningLabelStatus = 'streaming' | 'complete';
+
+/** Result of one reasoning-label revision. */
+export type ReasoningLabelResult = {
+  label?: string;
+  /** Provider-reported usage for host-side billing after a durable commit. */
+  usage?: UsageMetadata;
+};
+
+/** Options for `Run.generateReasoningLabel`. */
+export type RunReasoningLabelOptions = {
+  provider: Providers;
+  clientOptions?: ClientOptions;
+  /**
+   * Complete user-visible reasoning accumulated for this step so far. Hidden
+   * chain-of-thought must never be supplied through this API.
+   */
+  visibleReasoning: string;
+  /** Stable run-step identity shared by every revision of this label. */
+  reasoningStepId: string;
+  /** Monotonically increasing host revision for this reasoning step. */
+  revision: number;
+  /** Whether the snapshot can still grow. Default `streaming`. */
+  status?: ReasoningLabelStatus;
+  /**
+   * Last durably visible label for this step. The model repeats it exactly
+   * when the reasoning direction has not materially changed, allowing hosts
+   * to avoid redundant UI updates.
+   */
+  previousLabel?: string;
+  /**
+   * Agent that emitted the reasoning. Selects its Langfuse overlay and
+   * redaction policy. Unknown agents and omitted multi-agent ownership fail
+   * closed before tracing or generation.
+   */
+  agentId?: string;
+  /** Override for the default reasoning-label system prompt. */
+  prompt?: string;
+  /** Maximum reasoning characters retained in the prompt. Default 6000. */
+  charLimit?: number;
+  /** LangChain runnable config carrier (signal, callbacks, thread/user ids). */
+  chainOptions?: Partial<RunnableConfig> & {
+    configurable?: Record<string, unknown>;
+  };
+  /** Deterministic seed for this reasoning-label revision trace. */
+  traceSeed?: string;
+  /** Stable source run identifier recorded on the observation. */
+  sourceRunId?: string;
+  /** Source Langfuse trace id for linking this detached label trace. */
+  sourceTraceId?: string;
+  /** Host response/message identifier recorded on the observation. */
+  responseId?: string;
+};

package/src/types/stream.ts

@@ -103,6 +103,23 @@ export type RunStep = {
   // };
 };
 
+/** Minimal durable lifecycle state needed to continue open run steps. */
+export interface RunStepResumeEntry {
+  step: RunStep;
+  pendingToolCallIds: string[];
+  latestCompletionAt?: number;
+  openMessageStep: boolean;
+}
+
+/** SDK-private state persisted in LangGraph checkpoints for process-safe resume. */
+export interface RunStepResumeState {
+  version: 1;
+  revision: number;
+  nextIndex: number;
+  toolCallSteps: Array<{ toolCallId: string; stepId: string }>;
+  steps: RunStepResumeEntry[];
+}
+
 /**
  * Represents a run step delta i.e. any changed fields on a run step during
  * streaming.

package/src/types/tools.ts

@@ -1,10 +1,17 @@
 // src/types/tools.ts
 import type { StructuredToolInterface } from '@langchain/core/tools';
-import type { RunnableToolLike } from '@langchain/core/runnables';
+import type {
+  RunnableConfig,
+  RunnableToolLike,
+} from '@langchain/core/runnables';
 import type { ToolCall } from '@langchain/core/messages/tool';
 import type { ToolOutputReferenceRegistry } from '@/tools/toolOutputReferences';
 import type { RunBreakerScope } from '@/llm/streamLimits';
-import type { MessageContentComplex, ToolErrorData } from './stream';
+import type {
+  MessageContentComplex,
+  RunStepResumeState,
+  ToolErrorData,
+} from './stream';
 import type { HumanInTheLoopConfig } from './hitl';
 import type { LangfuseConfig } from './graph';
 import type { HookRegistry } from '@/hooks';
@@ -128,6 +135,8 @@ export type ToolNodeOptions = {
   toolRegistry?: LCToolRegistry;
   /** Reference to Graph's sessions map for automatic session injection */
   sessions?: ToolSessionMap;
+  /** Partition within `sessions` used by this agent's code tools. */
+  codeSessionKey?: string;
   /** When true, dispatches ON_TOOL_EXECUTE events instead of invoking tools directly */
   eventDrivenMode?: boolean;
   /** Tool definitions for event-driven mode (used for context, not invocation) */
@@ -274,6 +283,13 @@ export type ToolNodeOptions = {
    * controller.
    */
   getRunScope?: () => RunBreakerScope;
+  /** SDK-owned checkpoint bridge for open run-step lifecycle state. */
+  restoreRunStepResumeState?: (
+    state?: RunStepResumeState,
+    config?: RunnableConfig
+  ) => void;
+  /** SDK-owned checkpoint snapshot for open run-step lifecycle state. */
+  createRunStepResumeState?: () => RunStepResumeState;
 };
 
 export type ToolNodeConstructorParams = ToolRefs & ToolNodeOptions;
@@ -375,6 +391,19 @@ export type CodeEnvFile =
 export type CodeExecutionToolParams =
   | undefined
   | {
+      /** Trusted Code API endpoint selected by the host for this agent. */
+      baseUrl?: string;
+      /**
+       * Expected Code API execution profile. Sent as a request assertion so
+       * endpoint/configuration drift fails closed instead of silently using
+       * the wrong sandbox backend.
+       */
+      executionProfile?: CodeApiExecutionProfile;
+      /**
+       * Trusted warm-runtime affinity hint selected by the host for this
+       * agent invocation. This overrides the legacy ToolNode-injected hint.
+       */
+      runtimeSessionHint?: string;
       /** Execution session — see `CodeSessionContext.session_id`. */
       session_id?: string;
       user_id?: string;
@@ -382,16 +411,16 @@ export type CodeExecutionToolParams =
       /** Optional host-supplied Code API auth headers. */
       authHeaders?: CodeApiAuthHeaders;
       /**
-       * Advertise best-effort stateful sessions in the tool description
-       * (variables/files may persist between calls, may reset). Prompt text
-       * only, and it must be set here because the description is bound to the
-       * LLM at construction time. Pair it with the run-scoped
-       * `toolExecution.sandbox.statefulSessions` gate, which drives the wire
-       * hint — set both from one flag so the prompt and the backend agree.
+       * Advertise filesystem-tier stateful sessions in the tool description
+       * and allow a trusted runtime affinity hint onto the wire. Resolve this,
+       * `baseUrl`, `executionProfile`, and `runtimeSessionHint` together from
+       * the executing agent so the prompt and backend agree.
        */
       statefulSessions?: boolean;
     };
 
+export type CodeApiExecutionProfile = 'default' | 'stateful';
+
 export type CodeApiAuthHeaderMap = Record<string, string>;
 
 export type CodeApiAuthHeaders =
@@ -1106,19 +1135,17 @@ export type CloudflareSandboxExecutionConfig = {
 
 export type SandboxExecutionConfig = {
   /**
-   * Opt into best-effort stateful runtime sessions on the remote Code API
-   * (its warm per-session MicroVM backend). This gate is run-scoped: it only
-   * controls the wire behavior (ToolNode injecting the session hint on
-   * execute_code/bash calls). The transport is otherwise unchanged.
+   * Legacy run-scoped opt-in for warm runtime hint injection. New hosts that
+   * need mixed execution profiles should resolve `statefulSessions`,
+   * `baseUrl`, `executionProfile`, and `runtimeSessionHint` on each executing
+   * agent's tool factories instead. Direct executors ignore this injected hint
+   * unless their own factory has `statefulSessions: true`.
    *
    * It does NOT change the model-facing tool description. Tool descriptions are
    * bound to the LLM at construction time (`createCodeExecutionTool` /
    * `createBashExecutionTool`), before this run config is applied inside the
    * graph, so they can only be adjusted via the tools' own `statefulSessions`
-   * factory param. Set BOTH from one flag (as LibreChat does): with this on but
-   * the factory param off, the backend runs statefully while the model is still
-   * told the environment is stateless (non-corrupting — the model just won't
-   * exploit persistence).
+   * factory param.
    */
   statefulSessions?: boolean;
   /**
@@ -1294,6 +1321,10 @@ export type BashProgrammaticToolCallingParams = ProgrammaticToolCallingParams;
 export type ProgrammaticToolCallingParams = {
   /** Code API base URL (or use CODE_BASEURL env var) */
   baseUrl?: string;
+  /** Expected Code API execution profile (host-controlled). */
+  executionProfile?: CodeApiExecutionProfile;
+  /** Trusted warm-runtime affinity hint for this agent invocation. */
+  runtimeSessionHint?: string;
   /** Safety limit for round-trips (default: 20) */
   maxRoundTrips?: number;
   /** Maximum per-sandbox-run timeout for PTC's legacy `timeout` field. */
@@ -1304,11 +1335,9 @@ export type ProgrammaticToolCallingParams = {
   debug?: boolean;
   /** Optional host-supplied Code API auth headers. */
   authHeaders?: CodeApiAuthHeaders;
-  /* No `statefulSessions` here: PTC is stateless in v1. The initial
-   * /exec/programmatic request still forwards a ToolNode-injected
-   * `_runtime_session_hint` when present, but there is no factory-level opt-in
-   * to advertise (it would be a no-op). Re-add with real behavior when PTC
-   * stateful prompting lands. */
+  /* No `statefulSessions` here: PTC keeps a stateless runtime prompt in v1.
+   * An explicit stateful execution profile may still use a warm filesystem;
+   * runtime/in-memory state never carries across calls. */
 };
 
 // ============================================================================

package/src/utils/toolSessions.ts

@@ -0,0 +1,113 @@
+import type * as t from '@/types';
+import { Constants } from '@/common';
+
+type CodeSessionAgent = Pick<t.AgentInputs, 'codeSessionKey' | 'initialSessions'>;
+
+function cloneToolSessionContext(
+  context: t.ToolSessionContext
+): t.ToolSessionContext {
+  return {
+    ...context,
+    ...(context.files == null
+      ? {}
+      : {
+        files: context.files.map((file) => ({
+          ...file,
+          storage_session_id:
+            file.storage_session_id ?? context.session_id,
+        })),
+      }),
+  };
+}
+
+function mergeToolSessionContext(
+  sessions: t.ToolSessionMap,
+  key: string,
+  context: t.ToolSessionContext
+): void {
+  const existing = sessions.get(key);
+  if (existing == null) {
+    sessions.set(key, cloneToolSessionContext(context));
+    return;
+  }
+  if (context.files == null || context.files.length === 0) {
+    return;
+  }
+
+  const seenFiles = new Set(
+    existing.files?.map(
+      (file) =>
+        `${file.storage_session_id ?? existing.session_id}\0${file.id}`
+    ) ?? []
+  );
+  const files = existing.files == null ? [] : [...existing.files];
+  for (const file of context.files) {
+    const storageSessionId = file.storage_session_id ?? context.session_id;
+    const fileKey = `${storageSessionId}\0${file.id}`;
+    if (seenFiles.has(fileKey)) {
+      continue;
+    }
+    seenFiles.add(fileKey);
+    files.push({ ...file, storage_session_id: storageSessionId });
+  }
+  sessions.set(key, { ...existing, files });
+}
+
+/**
+ * Seeds a top-level graph while preserving the legacy run-wide session map.
+ * A legacy `execute_code` seed is copied into every custom agent partition
+ * unless the host already supplied an explicit seed for that partition.
+ */
+export function seedRunInitialSessions(args: {
+  sessions: t.ToolSessionMap;
+  initialSessions: t.ToolSessionMap;
+  agents: Iterable<Pick<t.AgentInputs, 'codeSessionKey'>>;
+}): void {
+  const { sessions, initialSessions, agents } = args;
+  for (const [key, context] of initialSessions) {
+    mergeToolSessionContext(sessions, key, context);
+  }
+
+  const legacyCodeSession = initialSessions.get(Constants.EXECUTE_CODE);
+  if (legacyCodeSession == null) {
+    return;
+  }
+  for (const agent of agents) {
+    const key = agent.codeSessionKey ?? Constants.EXECUTE_CODE;
+    if (
+      key === Constants.EXECUTE_CODE ||
+      initialSessions.has(key) ||
+      sessions.has(key)
+    ) {
+      continue;
+    }
+    sessions.set(key, cloneToolSessionContext(legacyCodeSession));
+  }
+}
+
+/**
+ * Seeds an isolated child graph from each selected agent. Legacy code-session
+ * entries are remapped to the owning agent's partition before they are merged.
+ */
+export function seedAgentInitialSessions(
+  sessions: t.ToolSessionMap,
+  agents: Iterable<CodeSessionAgent>
+): void {
+  for (const agent of agents) {
+    if (agent.initialSessions == null) {
+      continue;
+    }
+    const codeSessionKey = agent.codeSessionKey ?? Constants.EXECUTE_CODE;
+    const hasExplicitCodeSession =
+      codeSessionKey !== Constants.EXECUTE_CODE &&
+      agent.initialSessions.has(codeSessionKey);
+    for (const [toolName, context] of agent.initialSessions) {
+      if (toolName === Constants.EXECUTE_CODE && hasExplicitCodeSession) {
+        continue;
+      }
+      const key =
+        toolName === Constants.EXECUTE_CODE ? codeSessionKey : toolName;
+      mergeToolSessionContext(sessions, key, context);
+    }
+  }
+}