🛤️ feat: Per-Agent Code Execution Routing With Stateful Session Scopes

.env.example

@@ -621,7 +621,11 @@ TTS_API_KEY=
 
 # LIBRECHAT_CODE_API_KEY=
 # LIBRECHAT_CODE_BASEURL=
-# Prewarm stateful per-conversation sandboxes in parallel with model generation (default: true).
+# Optional dedicated Code API deployment for agents with Stateful code sessions enabled.
+# When configured, stateless agents continue using LIBRECHAT_CODE_BASEURL while stateful
+# agents fail closed onto this endpoint. The endpoint must advertise the `stateful` profile.
+# LIBRECHAT_CODE_BASEURL_STATEFUL=
+# Prewarm selected stateful sandboxes in parallel with model generation (default: true).
 # CODE_SANDBOX_PREWARM=true
 # Time in milliseconds before LibreChat treats a tracked sandbox as cold (default: 2100000 / 35 minutes).
 # CODE_SANDBOX_COLD_AFTER_MS=2100000

api/app/clients/BaseClient.js

@@ -1444,6 +1444,7 @@ class BaseClient {
       if (
         file.embedded === true ||
         file.metadata?.codeEnvRef != null ||
+        file.metadata?.codeEnvRefs != null ||
         file.metadata?.fileIdentifier != null
       ) {
         allFiles.push(file);

api/app/clients/tools/util/handleTools.js

@@ -20,6 +20,7 @@ const {
   ASK_USER_QUESTION_TOOL_NAME,
   resolveWebSearchSSRFAgents,
   buildWebSearchDynamicContext,
+  resolveCodeExecutionContext,
 } = require('@librechat/api');
 const {
   Tools,
@@ -331,28 +332,35 @@ const loadTools = async ({
   for (const tool of tools) {
     if (tool === Tools.execute_code) {
       requestedTools[tool] = async () => {
+        const statefulSessions =
+          agent?.stateful_code_sessions === true &&
+          (await checkCapability(options.req, AgentCapabilities.stateful_code_sessions));
+        const codeExecutionContext =
+          options.codeExecutionContext ??
+          resolveCodeExecutionContext({
+            statefulSessions,
+            environment: agent?.stateful_code_environment,
+            userId: user,
+            agentId: agent?.id,
+            conversationId: options.req?.body?.conversationId,
+          });
         const { files, toolContext } = await primeCodeFiles({
           ...options,
           agentId: agent?.id,
+          codeApiBaseUrl: codeExecutionContext.baseUrl,
+          executionProfile: codeExecutionContext.executionProfile,
         });
         if (toolContext) {
           dynamicToolContextMap[tool] = toolContext;
         }
         if (files?.length) {
           primedCodeFiles = files;
         }
-        /* Hedge the execute_code description toward persistence only when the
-         * admin `stateful_code_sessions` capability is on AND the agent opted
-         * in via the builder (off by default); the matching wire hint is set
-         * in the run config. Older @librechat/agents ignore the param. */
-        const statefulSessions =
-          agent?.stateful_code_sessions === true &&
-          (await checkCapability(options.req, AgentCapabilities.stateful_code_sessions));
         return createCodeExecutionTool({
           user_id: user,
           files,
           authHeaders: () => getCodeApiAuthHeaders(options.req),
-          statefulSessions,
+          ...codeExecutionContext,
         });
       };
       continue;

api/server/controllers/agents/__tests__/callbacks.spec.js

@@ -444,6 +444,7 @@ describe('createToolEndCallback', () => {
       name,
       toolName = 'execute_code',
       hostFileAuthoring = false,
+      codeExecutionContext,
     }) {
       return {
         output: {
@@ -455,7 +456,7 @@ describe('createToolEndCallback', () => {
             files: [{ id: fileId, name, session_id: 'sess-1' }],
           },
         },
-        metadata: { run_id: runId, thread_id: threadId },
+        metadata: { run_id: runId, thread_id: threadId, codeExecutionContext },
       };
     }
 
@@ -679,6 +680,10 @@ describe('createToolEndCallback', () => {
         name: 'created.txt',
         toolName: 'create_file',
         hostFileAuthoring: true,
+        codeExecutionContext: {
+          baseUrl: 'https://code-stateful.example.com',
+          executionProfile: 'stateful',
+        },
       });
       await toolEndCallback({ output: event.output }, event.metadata);
       await Promise.all(artifactPromises);
@@ -690,6 +695,8 @@ describe('createToolEndCallback', () => {
           messageId: 'run-create',
           toolCallId: 'tool-create',
           conversationId: 'thread789',
+          codeApiBaseUrl: 'https://code-stateful.example.com',
+          executionProfile: 'stateful',
         }),
       );
       expect(res.write).toHaveBeenCalledTimes(1);

api/server/controllers/agents/callbacks.js

@@ -926,6 +926,8 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo
              * ids.
              */
             session_id: file.storage_session_id ?? output.artifact.session_id,
+            codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl,
+            executionProfile: metadata.codeExecutionContext?.executionProfile,
           });
           const fileMetadata = result?.file ?? null;
           const finalize = result?.finalize;
@@ -1248,6 +1250,8 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises })
              * ids.
              */
             session_id: file.storage_session_id ?? output.artifact.session_id,
+            codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl,
+            executionProfile: metadata.codeExecutionContext?.executionProfile,
           });
           const fileMetadata = result?.file ?? null;
           const finalize = result?.finalize;

api/server/controllers/agents/client.js

@@ -1401,7 +1401,7 @@ class AgentClient extends BaseClient {
             this.contextHandlers?.processFile(file);
             continue;
           }
-          if (file.metadata?.codeEnvRef) {
+          if (file.metadata?.codeEnvRef || file.metadata?.codeEnvRefs) {
             continue;
           }
         }
@@ -2571,7 +2571,7 @@ class AgentClient extends BaseClient {
         abortController = new AbortController();
       }
 
-      /** Fire-and-forget: boot the per-conversation stateful sandbox in
+      /** Fire-and-forget: boot each selected stateful environment in
        *  parallel with generation so the first execute_code/bash call lands
        *  on a warm VM. No-op unless a reachable agent resolved
        *  `statefulCodeSessions`. */
@@ -2617,13 +2617,7 @@ class AgentClient extends BaseClient {
         ? await this.options.primeInvokedSkills(payload)
         : undefined;
 
-      /**
-       * Seed `Graph.sessions` with code-env files primed across every
-       * reachable agent (primary, handoff/addedConvo, and nested
-       * subagents) plus skill-priming output. The merge logic and its
-       * run-wide semantics live in `buildInitialToolSessions`; see that
-       * helper's doc for why this is intentionally NOT per-agent.
-       */
+      /** Seed each reachable agent's trusted code-session partition. */
       const initialSessions = buildInitialToolSessions({
         skillSessions: skillPrimeResult?.initialSessions,
         agents: [this.options.agent, ...(this.agentConfigs ? this.agentConfigs.values() : [])],

api/server/controllers/agents/openai.js

@@ -87,6 +87,7 @@ function createToolLoader(signal, definitionsOnly = true) {
     provider,
     tool_options,
     tool_resources,
+    codeExecutionContext,
     accessibleMcpServerNames,
   }) {
     const agent = { id: agentId, tools, provider, model, tool_options };
@@ -97,6 +98,7 @@ function createToolLoader(signal, definitionsOnly = true) {
         agent,
         signal,
         tool_resources,
+        codeExecutionContext,
         agentResourceType: ResourceType.REMOTE_AGENT,
         definitionsOnly,
         accessibleMcpServerNames,
@@ -508,6 +510,7 @@ const executeOpenAIChatCompletion = async (envelope, { req, res }) => {
           req,
           res,
           agentResourceType: ResourceType.REMOTE_AGENT,
+          conversationId,
           toolNames,
           agent: ctx.agent ?? agent,
           signal: abortController.signal,

api/server/controllers/agents/responses.js

@@ -101,6 +101,7 @@ function createToolLoader(signal, definitionsOnly = true) {
     provider,
     tool_options,
     tool_resources,
+    codeExecutionContext,
     accessibleMcpServerNames,
   }) {
     const agent = { id: agentId, tools, provider, model, tool_options };
@@ -111,6 +112,7 @@ function createToolLoader(signal, definitionsOnly = true) {
         agent,
         signal,
         tool_resources,
+        codeExecutionContext,
         agentResourceType: ResourceType.REMOTE_AGENT,
         definitionsOnly,
         accessibleMcpServerNames,
@@ -735,6 +737,7 @@ const executeResponse = async (envelope, { req, res }) => {
             req,
             res,
             agentResourceType: ResourceType.REMOTE_AGENT,
+            conversationId,
             toolNames,
             agent: ctx.agent ?? agent,
             signal: abortController.signal,
@@ -919,6 +922,7 @@ const executeResponse = async (envelope, { req, res }) => {
             req,
             res,
             agentResourceType: ResourceType.REMOTE_AGENT,
+            conversationId,
             toolNames,
             agent: ctx.agent ?? agent,
             signal: abortController.signal,

api/server/routes/files/files.js

@@ -10,6 +10,7 @@ const {
   startUploadSseStream,
   resolveUploadErrorMessage,
   verifyAgentUploadPermission,
+  getCodeExecutionBaseUrl,
 } = require('@librechat/api');
 const {
   Time,
@@ -329,6 +330,18 @@ router.get('/code/download/:session_id/:fileId', async (req, res) => {
       return res.status(400).send('Bad request');
     }
 
+    const requestedProfile = req.query.execution_profile;
+    if (
+      requestedProfile != null &&
+      requestedProfile !== 'default' &&
+      requestedProfile !== 'stateful'
+    ) {
+      logger.debug(`${logPrefix} invalid execution_profile`);
+      return res.status(400).send('Bad request');
+    }
+    const executionProfile = requestedProfile ?? 'default';
+    const baseUrl = getCodeExecutionBaseUrl(executionProfile);
+
     const { getDownloadStream } = getStrategyFunctions(FileSources.execute_code);
     if (!getDownloadStream) {
       logger.warn(
@@ -352,6 +365,7 @@ router.get('/code/download/:session_id/:fileId', async (req, res) => {
         id: req.user.id,
       },
       req,
+      { baseUrl, executionProfile },
     );
     res.set(response.headers);
     response.data.pipe(res);

api/server/routes/files/files.test.js

@@ -45,6 +45,11 @@ jest.mock('sharp', () =>
 jest.mock('@librechat/api', () => ({
   ...jest.requireActual('@librechat/api'),
   refreshS3FileUrls: jest.fn(),
+  getCodeExecutionBaseUrl: jest.fn((profile) =>
+    profile === 'stateful'
+      ? process.env.LIBRECHAT_CODE_BASEURL_STATEFUL
+      : 'https://code-default.example.com/v1',
+  ),
 }));
 
 jest.mock('~/cache', () => ({
@@ -1088,4 +1093,43 @@ describe('File Routes - Delete with Agent Access', () => {
       expect(response.status).toBe(401);
     });
   });
+
+  describe('GET /files/code/download/:session_id/:fileId', () => {
+    it('routes a persisted stateful fallback through the stateful Code API', async () => {
+      const getDownloadStream = jest.fn().mockResolvedValue({
+        headers: { 'content-type': 'text/plain' },
+        data: Readable.from(['stateful output']),
+      });
+      getStrategyFunctions.mockReturnValue({ getDownloadStream });
+      process.env.LIBRECHAT_CODE_BASEURL_STATEFUL = 'https://code-stateful.example.com/v1';
+
+      try {
+        const sessionId = 's'.repeat(21);
+        const codeFileId = 'f'.repeat(21);
+        const response = await request(app).get(
+          `/files/code/download/${sessionId}/${codeFileId}?execution_profile=stateful`,
+        );
+
+        expect(response.status).toBe(200);
+        expect(response.text).toBe('stateful output');
+        expect(getDownloadStream).toHaveBeenCalledWith(
+          `${sessionId}/${codeFileId}`,
+          { kind: 'user', id: otherUserId.toString() },
+          expect.any(Object),
+          { baseUrl: 'https://code-stateful.example.com/v1', executionProfile: 'stateful' },
+        );
+      } finally {
+        delete process.env.LIBRECHAT_CODE_BASEURL_STATEFUL;
+      }
+    });
+
+    it('rejects an unknown execution profile', async () => {
+      const response = await request(app).get(
+        `/files/code/download/${'s'.repeat(21)}/${'f'.repeat(21)}?execution_profile=attacker`,
+      );
+
+      expect(response.status).toBe(400);
+      expect(getStrategyFunctions).not.toHaveBeenCalled();
+    });
+  });
 });