🗿 feat: Add Attached Stateful Code Environments

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

@@ -356,6 +356,9 @@ const loadTools = async ({
           resolveCodeExecutionContext({
             statefulSessions,
             environment: agent?.stateful_code_environment,
+            environmentId: agent?.code_environment_id,
+            environments:
+              options.req?.config?.endpoints?.agents?.statefulCodeSessions?.environments,
             userId: user,
             agentId: agent?.id,
             conversationId: options.req?.body?.conversationId,
@@ -365,6 +368,7 @@ const loadTools = async ({
           agentId: agent?.id,
           codeApiBaseUrl: codeExecutionContext.baseUrl,
           executionProfile: codeExecutionContext.executionProfile,
+          executionRouteKey: codeExecutionContext.executionRouteKey,
         });
         if (toolContext) {
           dynamicToolContextMap[tool] = toolContext;

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

@@ -855,6 +855,7 @@ describe('createToolEndCallback', () => {
         codeExecutionContext: {
           baseUrl: 'https://code-stateful.example.com',
           executionProfile: 'stateful',
+          executionRouteKey: `stateful:${'a'.repeat(32)}`,
         },
       });
       await toolEndCallback({ output: event.output }, event.metadata);
@@ -870,6 +871,7 @@ describe('createToolEndCallback', () => {
           conversationId: 'thread789',
           codeApiBaseUrl: 'https://code-stateful.example.com',
           executionProfile: 'stateful',
+          executionRouteKey: `stateful:${'a'.repeat(32)}`,
         }),
       );
       expect(res.write).toHaveBeenCalledTimes(2);

api/server/controllers/agents/callbacks.js

@@ -1047,6 +1047,7 @@ function createToolEndCallback({ req, res, artifactPromises, streamId = null, jo
           session_id: sessionId,
           codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl,
           executionProfile: metadata.codeExecutionContext?.executionProfile,
+          executionRouteKey: metadata.codeExecutionContext?.executionRouteKey,
           preparedBuffer,
           downloadFallback,
         });
@@ -1391,6 +1392,7 @@ function createResponsesToolEndCallback({ req, res, tracker, artifactPromises })
           session_id: sessionId,
           codeApiBaseUrl: metadata.codeExecutionContext?.baseUrl,
           executionProfile: metadata.codeExecutionContext?.executionProfile,
+          executionRouteKey: metadata.codeExecutionContext?.executionRouteKey,
           preparedBuffer,
           downloadFallback,
         });

api/server/controllers/agents/client.js

@@ -402,6 +402,9 @@ class AgentClient extends BaseClient {
             hide_sequential_outputs: agent.hide_sequential_outputs,
             stateful_code_sessions: agent.stateful_code_sessions,
             stateful_code_environment: agent.stateful_code_environment,
+            execution_route_key:
+              agent.codeExecutionContext?.executionRouteKey ??
+              agent.codeExecutionContext?.executionProfile,
             artifacts: agent.artifacts,
             recursion_limit: agent.recursion_limit,
             subagents: agent.subagents,

api/server/controllers/agents/v1.js

@@ -443,7 +443,27 @@ const isCodeInterpreterCapabilityEnabled = (req) => {
 /** Reject a newly selected stateful workspace scope that the deployment owner
  * has excluded. Disabled sessions and unrelated edits remain saveable so an
  * allowlist tightening never silently rewrites or strands an existing agent. */
-const validateStatefulCodeEnvironment = (req, res, enabled, environment) => {
+const validateStatefulCodeEnvironment = (
+  req,
+  res,
+  enabled,
+  environment,
+  environmentId,
+  environmentIdSelected = false,
+) => {
+  if (enabled !== true && !environmentIdSelected) {
+    return true;
+  }
+  if (environmentId != null) {
+    const configuredEnvironments =
+      req.config?.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments ?? [];
+    if (!configuredEnvironments.some((configured) => configured.id === environmentId)) {
+      res.status(400).json({
+        error: `Stateful code environment is not configured: ${environmentId}`,
+      });
+      return false;
+    }
+  }
   if (enabled !== true) {
     return true;
   }
@@ -721,6 +741,8 @@ const createAgentHandler = async (req, res) => {
         res,
         agentData.stateful_code_sessions,
         agentData.stateful_code_environment,
+        agentData.code_environment_id,
+        agentData.code_environment_id != null,
       )
     ) {
       return;
@@ -997,13 +1019,22 @@ const updateAgentHandler = async (req, res) => {
     normalizeToolResourceFiles(req.body?.tool_resources);
     const validatedData = agentUpdateSchema.parse(req.body);
     // Preserve explicit null for avatar to allow resetting the avatar
-    const { avatar: avatarField, _id, ...rest } = validatedData;
+    const {
+      avatar: avatarField,
+      code_environment_id: codeEnvironmentIdField,
+      _id,
+      ...rest
+    } = validatedData;
     const updateData = removeNullishValues(rest);
+    if (codeEnvironmentIdField !== undefined) {
+      updateData.code_environment_id = codeEnvironmentIdField;
+    }
     let existingAgent;
 
     const includesStatefulConfiguration =
       updateData.stateful_code_sessions !== undefined ||
-      updateData.stateful_code_environment !== undefined;
+      updateData.stateful_code_environment !== undefined ||
+      updateData.code_environment_id !== undefined;
     const includesToolsConfiguration = Array.isArray(updateData.tools);
     const includesToolOptionsConfiguration = updateData.tool_options !== undefined;
     if (
@@ -1016,13 +1047,17 @@ const updateAgentHandler = async (req, res) => {
         return res.status(404).json({ error: 'Agent not found' });
       }
 
+      const codeEnvironmentSelectionChanged =
+        updateData.code_environment_id !== undefined &&
+        updateData.code_environment_id !== existingAgent.code_environment_id;
       const statefulConfigurationChanged =
         (updateData.stateful_code_sessions !== undefined &&
           (updateData.stateful_code_sessions === true) !==
             (existingAgent.stateful_code_sessions === true)) ||
         (updateData.stateful_code_environment !== undefined &&
           (updateData.stateful_code_environment ?? 'user') !==
-            (existingAgent.stateful_code_environment ?? 'user'));
+            (existingAgent.stateful_code_environment ?? 'user')) ||
+        codeEnvironmentSelectionChanged;
       const activatesCodeExecution =
         includesToolsConfiguration &&
         updateData.tools.includes(Tools.execute_code) &&
@@ -1032,12 +1067,18 @@ const updateAgentHandler = async (req, res) => {
           updateData.stateful_code_sessions ?? existingAgent.stateful_code_sessions;
         const effectiveStatefulEnvironment =
           updateData.stateful_code_environment ?? existingAgent.stateful_code_environment;
+        const effectiveCodeEnvironmentId =
+          updateData.code_environment_id === null
+            ? undefined
+            : (updateData.code_environment_id ?? existingAgent.code_environment_id);
         if (
           !validateStatefulCodeEnvironment(
             req,
             res,
             effectiveStatefulSessions,
             effectiveStatefulEnvironment,
+            effectiveCodeEnvironmentId,
+            codeEnvironmentSelectionChanged,
           )
         ) {
           return;
@@ -1227,6 +1268,11 @@ const updateAgentHandler = async (req, res) => {
       }
     }
 
+    if (updateData.code_environment_id === null) {
+      delete updateData.code_environment_id;
+      updateData.$unset = { code_environment_id: 1 };
+    }
+
     let updatedAgent =
       Object.keys(updateData).length > 0
         ? await db.updateAgent({ id }, updateData, {
@@ -1333,6 +1379,7 @@ const duplicateAgentHandler = async (req, res) => {
         res,
         newAgentData.stateful_code_sessions,
         newAgentData.stateful_code_environment,
+        newAgentData.code_environment_id,
       )
     ) {
       return;
@@ -1925,6 +1972,7 @@ const revertAgentVersionHandler = async (req, res) => {
         res,
         revertVersion.stateful_code_sessions,
         revertVersion.stateful_code_environment,
+        revertVersion.code_environment_id,
       )
     ) {
       return;

api/server/controllers/agents/v1.spec.js

@@ -237,6 +237,39 @@ describe('Agent Controllers - Mass Assignment Protection', () => {
       expect(await Agent.countDocuments()).toBe(0);
     });
 
+    test('rejects an unconfigured code environment id', async () => {
+      mockReq.config = {
+        endpoints: {
+          agents: {
+            statefulCodeSessions: {
+              allowedEnvironments: ['user'],
+              environments: [
+                {
+                  id: 'configured-vm',
+                  name: 'Configured VM',
+                  type: 'attached',
+                  baseURL: 'https://code.example.com/v1',
+                  default: true,
+                },
+              ],
+            },
+          },
+        },
+      };
+      mockReq.body = {
+        name: 'Invalid Environment Agent',
+        provider: 'openai',
+        model: 'gpt-4',
+        stateful_code_sessions: true,
+        code_environment_id: 'missing-vm',
+      };
+
+      await createAgentHandler(mockReq, mockRes);
+
+      expect(mockRes.status).toHaveBeenCalledWith(400);
+      expect(await Agent.countDocuments()).toBe(0);
+    });
+
     test('should block configured agent instruction content before persistence', async () => {
       mockReq.config = {
         filters: {
@@ -1182,6 +1215,98 @@ describe('Agent Controllers - Mass Assignment Protection', () => {
       expect(agentInDb.stateful_code_sessions).not.toBe(true);
     });
 
+    test('rejects updating an agent to an unconfigured code environment id', async () => {
+      mockReq.user.id = existingAgentAuthorId.toString();
+      mockReq.params.id = existingAgentId;
+      mockReq.config = {
+        endpoints: {
+          agents: {
+            statefulCodeSessions: {
+              allowedEnvironments: ['user'],
+              environments: [],
+            },
+          },
+        },
+      };
+      mockReq.body = { code_environment_id: 'missing-vm' };
+
+      await updateAgentHandler(mockReq, mockRes);
+
+      expect(mockRes.status).toHaveBeenCalledWith(400);
+      const agentInDb = await Agent.findOne({ id: existingAgentId });
+      expect(agentInDb.code_environment_id).toBeUndefined();
+    });
+
+    test('allows disabling stateful sessions after the configured environment is removed', async () => {
+      await Agent.updateOne(
+        { id: existingAgentId },
+        {
+          stateful_code_sessions: true,
+          code_environment_id: 'removed-vm',
+        },
+      );
+      mockReq.user.id = existingAgentAuthorId.toString();
+      mockReq.params.id = existingAgentId;
+      mockReq.config = {
+        endpoints: {
+          agents: {
+            statefulCodeSessions: {
+              allowedEnvironments: ['user'],
+              environments: [],
+            },
+          },
+        },
+      };
+      mockReq.body = {
+        stateful_code_sessions: false,
+        code_environment_id: 'removed-vm',
+      };
+
+      await updateAgentHandler(mockReq, mockRes);
+
+      expect(mockRes.status).not.toHaveBeenCalledWith(400);
+      const agentInDb = await Agent.findOne({ id: existingAgentId });
+      expect(agentInDb.stateful_code_sessions).toBe(false);
+      expect(agentInDb.code_environment_id).toBe('removed-vm');
+    });
+
+    test('restores the deployment-default code environment', async () => {
+      await Agent.updateOne(
+        { id: existingAgentId },
+        {
+          stateful_code_sessions: true,
+          code_environment_id: 'attached-vm',
+        },
+      );
+      mockReq.user.id = existingAgentAuthorId.toString();
+      mockReq.params.id = existingAgentId;
+      mockReq.config = {
+        endpoints: {
+          agents: {
+            statefulCodeSessions: {
+              allowedEnvironments: ['user'],
+              environments: [
+                {
+                  id: 'attached-vm',
+                  name: 'Attached VM',
+                  type: 'attached',
+                  baseURL: 'https://bridge.example.com/v1',
+                  default: true,
+                },
+              ],
+            },
+          },
+        },
+      };
+      mockReq.body = { code_environment_id: null };
+
+      await updateAgentHandler(mockReq, mockRes);
+
+      expect(mockRes.status).not.toHaveBeenCalledWith(400);
+      const agentInDb = await Agent.findOne({ id: existingAgentId });
+      expect(agentInDb.code_environment_id).toBeUndefined();
+    });
+
     test('allows unrelated edits to an existing scope after policy is tightened', async () => {
       await Agent.updateOne(
         { id: existingAgentId },

api/server/routes/files/files.js

@@ -13,6 +13,7 @@ const {
   sendUploadPolicyError,
   resolveUploadErrorMessage,
   verifyAgentUploadPermission,
+  createCodeExecutionRouteKey,
   getCodeExecutionBaseUrl,
   assertUploadContentAllowed,
   hasActiveFilePolicy,
@@ -350,7 +351,29 @@ router.get('/code/download/:session_id/:fileId', async (req, res) => {
       return res.status(400).send('Bad request');
     }
     const executionProfile = requestedProfile ?? 'default';
-    const baseUrl = getCodeExecutionBaseUrl(executionProfile);
+    const requestedRouteKey = req.query.execution_route_key;
+    if (
+      requestedRouteKey != null &&
+      (typeof requestedRouteKey !== 'string' ||
+        executionProfile !== 'stateful' ||
+        !/^stateful:[a-f0-9]{32}$/.test(requestedRouteKey))
+    ) {
+      logger.debug(`${logPrefix} invalid execution_route_key`);
+      return res.status(400).send('Bad request');
+    }
+    const environments =
+      req.config?.endpoints?.[EModelEndpoint.agents]?.statefulCodeSessions?.environments;
+    const configuredEnvironment = requestedRouteKey
+      ? environments?.find(
+          (environment) =>
+            createCodeExecutionRouteKey('stateful', environment) === requestedRouteKey,
+        )
+      : undefined;
+    if (requestedRouteKey && !configuredEnvironment) {
+      logger.debug(`${logPrefix} unknown execution_route_key`);
+      return res.status(404).send('Not found');
+    }
+    const baseUrl = getCodeExecutionBaseUrl(executionProfile, configuredEnvironment);
 
     const { getDownloadStream } = getStrategyFunctions(FileSources.execute_code);
     if (!getDownloadStream) {