feat: route code execution per agent profile

src/__tests__/stream.eagerEventExecution.test.ts

@@ -4945,6 +4945,50 @@ describe('ChatModelStreamHandler eager event tool execution', () => {
     expect(graph.eagerEventToolExecutions.has('call_file')).toBe(false);
   });
 
+  it('does not prestart code execution for a custom agent session partition', async () => {
+    const graph = createGraph({
+      eagerEventToolExecution: { enabled: true },
+      getAgentContext: jest.fn(() => ({
+        provider: Providers.ANTHROPIC,
+        reasoningKey: 'reasoning',
+        toolDefinitions: [{ name: Constants.EXECUTE_CODE }],
+        graphTools: [],
+        agentId: 'stateful-agent',
+        codeSessionKey: 'execute_code:stateful:user-1',
+      })) as unknown as StandardGraph['getAgentContext'],
+    });
+    const toolExecuteCalls: t.ToolExecuteBatchRequest[] = [];
+    jest
+      .spyOn(events, 'safeDispatchCustomEvent')
+      .mockImplementation(async (event, data): Promise<void> => {
+        if (event === GraphEvents.ON_TOOL_EXECUTE) {
+          toolExecuteCalls.push(data as t.ToolExecuteBatchRequest);
+        }
+      });
+
+    await new ChatModelStreamHandler().handle(
+      GraphEvents.CHAT_MODEL_STREAM,
+      {
+        chunk: {
+          content: '',
+          tool_calls: [
+            {
+              id: 'call_code',
+              name: Constants.EXECUTE_CODE,
+              args: { lang: 'python', code: 'print(1)' },
+            },
+          ],
+          response_metadata: finalToolCallResponseMetadata,
+        } as unknown as t.StreamChunk,
+      },
+      { langgraph_node: 'agent' },
+      graph
+    );
+
+    expect(toolExecuteCalls).toHaveLength(0);
+    expect(graph.eagerEventToolExecutions.has('call_code')).toBe(false);
+  });
+
   it('does not prestart codeSessionToolNames tools even without excludeToolNames', async () => {
     // A declared session-writing host tool is side-effecting, so it must not be
     // eagerly prestarted even when the host didn't also list it in excludeToolNames.

src/agents/AgentContext.ts

@@ -63,6 +63,7 @@ export class AgentContext {
   ): AgentContext {
     const {
       agentId,
+      codeSessionKey,
       name,
       provider,
       clientOptions,
@@ -92,6 +93,7 @@ export class AgentContext {
 
     const agentContext = new AgentContext({
       agentId,
+      codeSessionKey,
       name: name ?? agentId,
       provider,
       clientOptions,
@@ -167,6 +169,8 @@ export class AgentContext {
 
   /** Agent identifier */
   agentId: string;
+  /** Partition for this agent's transient code session and file refs. */
+  codeSessionKey?: string;
   /** Human-readable name for this agent (used in handoff context). Falls back to agentId if not provided. */
   name?: string;
   /** Provider for this specific agent */
@@ -370,6 +374,7 @@ export class AgentContext {
 
   constructor({
     agentId,
+    codeSessionKey,
     name,
     provider,
     clientOptions,
@@ -394,6 +399,7 @@ export class AgentContext {
     maxToolResultChars,
   }: {
     agentId: string;
+    codeSessionKey?: string;
     name?: string;
     provider: Providers;
     clientOptions?: t.ClientOptions;
@@ -418,6 +424,7 @@ export class AgentContext {
     maxToolResultChars?: number;
   }) {
     this.agentId = agentId;
+    this.codeSessionKey = codeSessionKey;
     this.name = name;
     this.provider = provider;
     this.clientOptions = clientOptions;

src/graphs/Graph.ts

@@ -2391,6 +2391,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         agentLangfuse: agentContext?.langfuse,
         eventDrivenMode: true,
         sessions: this.sessions,
+        codeSessionKey: agentContext?.codeSessionKey,
         toolDefinitions: toolDefMap,
         // `agentId` is the subagent-scope marker — set ONLY for child-run
         // graphs (hooks fire for child scopes too, via the inherited
@@ -2471,6 +2472,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         StandardGraph.handleToolCallErrorStatic(this, data, metadata),
       toolRegistry: agentContext?.toolRegistry,
       sessions: this.sessions,
+      codeSessionKey: agentContext?.codeSessionKey,
       toolExecution: this.toolExecution,
       codeSessionToolNames: this.codeSessionToolNames,
       interruptingToolNames: effectiveInterruptingToolNames,

src/run.ts

@@ -72,6 +72,7 @@ import {
 import { applyGraphRuntimeConfig } from '@/graphs/applyGraphRuntimeConfig';
 import { createTokenCounter, encodingForModel } from '@/utils/tokens';
 import { initializeLangfuseTracing } from './instrumentation';
+import { seedRunInitialSessions } from '@/utils/toolSessions';
 import { getTraceIdSeed } from '@/langfuseRuntimeContext';
 import { createGraph } from '@/graphs/createGraph';
 import { resolveMaxSeals } from '@/llm/preempt';
@@ -351,9 +352,25 @@ export class Run<_T extends t.BaseGraphState> {
     }
 
     if (config.initialSessions && this.Graph) {
-      for (const [key, value] of config.initialSessions) {
-        this.Graph.sessions.set(key, value);
-      }
+      const configuredAgents =
+        'agents' in config.graphConfig &&
+        Array.isArray(config.graphConfig.agents)
+          ? config.graphConfig.agents
+          : undefined;
+      const agents: Array<Pick<t.AgentInputs, 'codeSessionKey'>> =
+        configuredAgents ?? [
+          {
+            codeSessionKey:
+              'codeSessionKey' in config.graphConfig
+                ? config.graphConfig.codeSessionKey
+                : undefined,
+          },
+        ];
+      seedRunInitialSessions({
+        sessions: this.Graph.sessions,
+        initialSessions: config.initialSessions,
+        agents,
+      });
     }
 
     this.returnContent = config.returnContent ?? false;

src/stream.ts

@@ -169,7 +169,8 @@ function hasToolOutputReference(value: unknown): boolean {
 
 function isEagerExecutionExcludedTool(
   name: string,
-  graph: StandardGraph
+  graph: StandardGraph,
+  agentContext?: AgentContext
 ): boolean {
   if (name === '') {
     return false;
@@ -200,9 +201,21 @@ function isEagerExecutionExcludedTool(
   // args, ToolNode discards the eager result but the mutation has already
   // landed in the session workspace, corrupting later runs. Stateless mode
   // uses a throwaway VM per call, so eager prestart stays safe there.
+  if (!CODE_EXECUTION_TOOLS.has(name)) {
+    return false;
+  }
+  if (graph.toolExecution?.sandbox?.statefulSessions === true) {
+    return true;
+  }
+  // A non-default code-session partition is the trusted per-agent signal that
+  // this call must remain isolated from the graph-wide stateless session. The
+  // actual tool factory may route it to a durable backend, which the stream
+  // layer cannot inspect in event-driven mode. Conservatively avoid speculative
+  // execution for these calls so discarded eager results cannot mutate that
+  // agent's workspace.
   return (
-    graph.toolExecution?.sandbox?.statefulSessions === true &&
-    CODE_EXECUTION_TOOLS.has(name)
+    agentContext?.codeSessionKey != null &&
+    agentContext.codeSessionKey !== Constants.EXECUTE_CODE
   );
 }
 
@@ -258,7 +271,8 @@ function toCodeEnvFile(file: t.FileRef, execSessionId: string): t.CodeEnvFile {
 
 function getCodeSessionContext(
   graph: StandardGraph,
-  name: string
+  name: string,
+  agentContext?: AgentContext
 ): t.ToolCallRequest['codeSessionContext'] | undefined {
   if (
     !CODE_EXECUTION_TOOLS.has(name) &&
@@ -269,9 +283,9 @@ function getCodeSessionContext(
     return undefined;
   }
 
-  const codeSession = graph.sessions.get(Constants.EXECUTE_CODE) as
-    | t.CodeSessionContext
-    | undefined;
+  const codeSession = graph.sessions.get(
+    agentContext?.codeSessionKey ?? Constants.EXECUTE_CODE
+  ) as t.CodeSessionContext | undefined;
   if (codeSession?.session_id == null || codeSession.session_id === '') {
     return undefined;
   }
@@ -716,7 +730,8 @@ function createEagerToolExecutionPlan(args: {
   // tool from `hasDirectToolCallInBatch`. Excluded calls fall through to normal
   // ToolNode execution; siblings may still eager-execute.
   const candidateToolCalls = unstartedToolCalls.filter(
-    (toolCall) => !isEagerExecutionExcludedTool(toolCall.name, graph)
+    (toolCall) =>
+      !isEagerExecutionExcludedTool(toolCall.name, graph, agentContext)
   );
   if (candidateToolCalls.length === 0) {
     return [];
@@ -748,7 +763,11 @@ function createEagerToolExecutionPlan(args: {
       name: toolCall.name,
       args: toolCall.args,
       stepId: graph.toolCallStepIds.get(toolCall.id!) ?? '',
-      codeSessionContext: getCodeSessionContext(graph, toolCall.name),
+      codeSessionContext: getCodeSessionContext(
+        graph,
+        toolCall.name,
+        agentContext
+      ),
     })),
     usageCount: graph.getEagerEventToolUsageCount(agentContext?.agentId),
   });

src/tools/BashExecutor.ts

@@ -8,22 +8,22 @@ import {
   appendFailedExecutionFileReminder,
   appendTmpScratchReminder,
   appendCodeSessionFileSummary,
+  addCodeApiExecutionProfileHeader,
   emptyOutputMessage,
   buildCodeApiHttpErrorMessage,
   CodeApiRequestError,
+  buildCodeApiEndpoint,
   getCodeBaseURL,
   normalizeCodeApiRequestError,
   resolveCodeApiAuthHeaders,
+  selectRuntimeSessionHint,
 } from './CodeExecutor';
 import { resolveFetchProxyAgent } from '@/utils/proxy';
 import { INTENT_PROPERTY } from '@/tools/intentArg';
 import { Constants } from '@/common';
 
 config();
 
-const baseEndpoint = getCodeBaseURL();
-const EXEC_ENDPOINT = `${baseEndpoint}/exec`;
-
 export const BashExecutionToolSchema = {
   type: 'object',
   properties: {
@@ -177,15 +177,24 @@ export const BashExecutionToolDefinition = {
 function createBashExecutionTool(
   params: t.BashExecutionToolParams | null = {}
 ): DynamicStructuredTool {
+  const execEndpoint = buildCodeApiEndpoint(
+    params?.baseUrl ?? getCodeBaseURL(),
+    'exec'
+  );
+
   return tool(
     async (rawInput, config) => {
-      /* `statefulSessions` is prompt-only — keep it out of the wire body. */
+      /* `statefulSessions` drives the prompt and gates runtime affinity hints;
+       * keep the flag itself out of the wire body. */
       const {
         authHeaders,
-        statefulSessions: _statefulSessions,
+        baseUrl: _baseUrl,
+        executionProfile,
+        runtimeSessionHint,
+        statefulSessions,
         ...executionParams
       } = params ?? {};
-      void _statefulSessions;
+      void _baseUrl;
       /* Drop any model-supplied `runtime_session_hint` from the raw args: the
        * hint must only come from ToolNode's injected `_runtime_session_hint`
        * (below), never from the tool call itself. */
@@ -217,11 +226,17 @@ function createBashExecutionTool(
         ...executionParams,
       };
 
+      const effectiveRuntimeSessionHint = selectRuntimeSessionHint(
+        runtimeSessionHint,
+        _runtime_session_hint
+      );
       if (
-        typeof _runtime_session_hint === 'string' &&
-        _runtime_session_hint !== ''
+        statefulSessions === true &&
+        executionProfile !== 'default' &&
+        typeof effectiveRuntimeSessionHint === 'string' &&
+        effectiveRuntimeSessionHint !== ''
       ) {
-        postData.runtime_session_hint = _runtime_session_hint;
+        postData.runtime_session_hint = effectiveRuntimeSessionHint;
       }
 
       /* See `CodeExecutor.ts` for the rationale — `/files/<session_id>`
@@ -248,19 +263,22 @@ function createBashExecutionTool(
           headers: {
             'Content-Type': 'application/json',
             'User-Agent': 'LibreChat/1.0',
-            ...resolvedAuthHeaders,
+            ...addCodeApiExecutionProfileHeader(
+              resolvedAuthHeaders,
+              executionProfile
+            ),
           },
           body: JSON.stringify(postData),
         };
 
-        const proxyAgent = resolveFetchProxyAgent(EXEC_ENDPOINT);
+        const proxyAgent = resolveFetchProxyAgent(execEndpoint);
         if (proxyAgent) {
           fetchOptions.agent = proxyAgent;
         }
-        const response = await fetch(EXEC_ENDPOINT, fetchOptions);
+        const response = await fetch(execEndpoint, fetchOptions);
         if (!response.ok) {
           throw new CodeApiRequestError(
-            await buildCodeApiHttpErrorMessage('POST', EXEC_ENDPOINT, response)
+            await buildCodeApiHttpErrorMessage('POST', execEndpoint, response)
           );
         }