Support multiple managed environments

Refactor EnvironmentManager to own a keyed environment registry with explicit default and local lookups. Keep remote exec-server connections lazy at environment use sites and preserve disabled agent environment access separately from internal local environment access.

Co-authored-by: Codex <noreply@openai.com>

codex-rs/exec-server/src/environment.rs

@@ -1,4 +1,6 @@
+use std::collections::HashMap;
 use std::sync::Arc;
+use std::time::Duration;
 
 use tokio::sync::OnceCell;
 
@@ -15,45 +17,64 @@ use crate::remote_process::RemoteProcess;
 
 pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL";
 
-/// Lazily creates and caches the active environment for a session.
+/// Owns the execution/filesystem environments available to a session.
 ///
-/// The manager keeps the session's environment selection stable so subagents
-/// and follow-up turns preserve an explicit disabled state.
+/// The manager keeps the session's default environment selection stable while
+/// separately tracking whether model-facing tools may access environments.
 #[derive(Debug)]
 pub struct EnvironmentManager {
-    exec_server_url: Option<String>,
-    local_runtime_paths: Option<ExecServerRuntimePaths>,
-    disabled: bool,
-    current_environment: OnceCell<Option<Arc<Environment>>>,
+    default_environment: Option<String>,
+    environment_disabled_for_agent: bool,
+    environments: HashMap<String, Arc<Environment>>,
+}
+
+pub const LOCAL_ENVIRONMENT_ID: &str = "local";
+pub const REMOTE_ENVIRONMENT_ID: &str = "remote";
+
+#[derive(Clone, Debug, Default)]
+pub struct EnvironmentManagerArgs {
+    pub exec_server_url: Option<String>,
+    pub local_runtime_paths: Option<ExecServerRuntimePaths>,
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct LazyRemoteExecServerClient {
+    websocket_url: String,
+    client: Arc<OnceCell<ExecServerClient>>,
+}
+
+impl LazyRemoteExecServerClient {
+    fn new(websocket_url: String) -> Self {
+        Self {
+            websocket_url,
+            client: Arc::new(OnceCell::new()),
+        }
+    }
+
+    pub(crate) async fn get(&self) -> Result<ExecServerClient, ExecServerError> {
+        self.client
+            .get_or_try_init(|| async {
+                ExecServerClient::connect_websocket(RemoteExecServerConnectArgs {
+                    websocket_url: self.websocket_url.clone(),
+                    client_name: "codex-environment".to_string(),
+                    connect_timeout: Duration::from_secs(5),
+                    initialize_timeout: Duration::from_secs(5),
+                    resume_session_id: None,
+                })
+                .await
+            })
+            .await
+            .cloned()
+    }
 }
 
 impl Default for EnvironmentManager {
     fn default() -> Self {
-        Self::new(/*exec_server_url*/ None)
+        Self::from_exec_server_url(EnvironmentManagerArgs::default())
     }
 }
 
 impl EnvironmentManager {
-    /// Builds a manager from the raw `CODEX_EXEC_SERVER_URL` value.
-    pub fn new(exec_server_url: Option<String>) -> Self {
-        Self::new_with_runtime_paths(exec_server_url, /*local_runtime_paths*/ None)
-    }
-
-    /// Builds a manager from the raw `CODEX_EXEC_SERVER_URL` value and local
-    /// runtime paths used when creating local filesystem helpers.
-    pub fn new_with_runtime_paths(
-        exec_server_url: Option<String>,
-        local_runtime_paths: Option<ExecServerRuntimePaths>,
-    ) -> Self {
-        let (exec_server_url, disabled) = normalize_exec_server_url(exec_server_url);
-        Self {
-            exec_server_url,
-            local_runtime_paths,
-            disabled,
-            current_environment: OnceCell::new(),
-        }
-    }
-
     /// Builds a manager from process environment variables.
     pub fn from_env() -> Self {
         Self::from_env_with_runtime_paths(/*local_runtime_paths*/ None)
@@ -64,60 +85,81 @@ impl EnvironmentManager {
     pub fn from_env_with_runtime_paths(
         local_runtime_paths: Option<ExecServerRuntimePaths>,
     ) -> Self {
-        Self::new_with_runtime_paths(
-            std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok(),
+        Self::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok(),
             local_runtime_paths,
-        )
+        })
     }
 
-    /// Builds a manager from the currently selected environment, or from the
-    /// disabled mode when no environment is available.
-    pub fn from_environment(environment: Option<&Environment>) -> Self {
-        match environment {
-            Some(environment) => Self {
-                exec_server_url: environment.exec_server_url().map(str::to_owned),
-                local_runtime_paths: environment.local_runtime_paths().cloned(),
-                disabled: false,
-                current_environment: OnceCell::new(),
-            },
-            None => Self {
-                exec_server_url: None,
-                local_runtime_paths: None,
-                disabled: true,
-                current_environment: OnceCell::new(),
-            },
+    /// Builds a manager from the raw `CODEX_EXEC_SERVER_URL` value and local
+    /// runtime paths used when creating local filesystem helpers.
+    pub fn from_exec_server_url(args: EnvironmentManagerArgs) -> Self {
+        let EnvironmentManagerArgs {
+            exec_server_url,
+            local_runtime_paths,
+        } = args;
+        let (exec_server_url, environment_disabled_for_agent) =
+            normalize_exec_server_url(exec_server_url);
+        let mut environments = HashMap::new();
+        environments.insert(
+            LOCAL_ENVIRONMENT_ID.to_string(),
+            Arc::new(
+                Environment::create_with_runtime_paths(
+                    /*exec_server_url*/ None,
+                    local_runtime_paths.clone(),
+                )
+                .expect("valid local environment"),
+            ),
+        );
+
+        let default_environment = match exec_server_url {
+            Some(exec_server_url) => {
+                environments.insert(
+                    REMOTE_ENVIRONMENT_ID.to_string(),
+                    Arc::new(
+                        Environment::create_with_runtime_paths(
+                            Some(exec_server_url),
+                            local_runtime_paths,
+                        )
+                        .expect("valid remote environment"),
+                    ),
+                );
+                Some(REMOTE_ENVIRONMENT_ID.to_string())
+            }
+            None => Some(LOCAL_ENVIRONMENT_ID.to_string()),
+        };
+
+        Self {
+            default_environment,
+            environment_disabled_for_agent,
+            environments,
         }
     }
 
-    /// Returns the remote exec-server URL when one is configured.
-    pub fn exec_server_url(&self) -> Option<&str> {
-        self.exec_server_url.as_deref()
+    /// Returns true when model-facing tools may access an environment.
+    pub fn allows_agent_environment_access(&self) -> bool {
+        !self.environment_disabled_for_agent
+            && self
+                .default_environment
+                .as_deref()
+                .is_some_and(|environment_id| self.environments.contains_key(environment_id))
     }
 
-    /// Returns true when this manager is configured to use a remote exec server.
-    pub fn is_remote(&self) -> bool {
-        self.exec_server_url.is_some()
+    /// Returns the default environment instance.
+    pub fn default_environment(&self) -> Option<Arc<Environment>> {
+        self.default_environment
+            .as_deref()
+            .and_then(|environment_id| self.get_environment(environment_id))
     }
 
-    /// Returns the cached environment, creating it on first access.
-    pub async fn current(&self) -> Result<Option<Arc<Environment>>, ExecServerError> {
-        self.current_environment
-            .get_or_try_init(|| async {
-                if self.disabled {
-                    Ok(None)
-                } else {
-                    Ok(Some(Arc::new(
-                        Environment::create_with_runtime_paths(
-                            self.exec_server_url.clone(),
-                            self.local_runtime_paths.clone(),
-                        )
-                        .await?,
-                    )))
-                }
-            })
-            .await
-            .map(Option::as_ref)
-            .map(std::option::Option::<&Arc<Environment>>::cloned)
+    /// Returns the local environment instance.
+    pub fn local_environment(&self) -> Option<Arc<Environment>> {
+        self.get_environment(LOCAL_ENVIRONMENT_ID)
+    }
+
+    /// Returns a named environment instance.
+    pub fn get_environment(&self, environment_id: &str) -> Option<Arc<Environment>> {
+        self.environments.get(environment_id).cloned()
     }
 }
 
@@ -128,7 +170,7 @@ impl EnvironmentManager {
 #[derive(Clone)]
 pub struct Environment {
     exec_server_url: Option<String>,
-    remote_exec_server_client: Option<ExecServerClient>,
+    remote_exec_server_client: Option<LazyRemoteExecServerClient>,
     exec_backend: Arc<dyn ExecBackend>,
     local_runtime_paths: Option<ExecServerRuntimePaths>,
 }
@@ -154,13 +196,13 @@ impl std::fmt::Debug for Environment {
 
 impl Environment {
     /// Builds an environment from the raw `CODEX_EXEC_SERVER_URL` value.
-    pub async fn create(exec_server_url: Option<String>) -> Result<Self, ExecServerError> {
-        Self::create_with_runtime_paths(exec_server_url, /*local_runtime_paths*/ None).await
+    pub fn create(exec_server_url: Option<String>) -> Result<Self, ExecServerError> {
+        Self::create_with_runtime_paths(exec_server_url, /*local_runtime_paths*/ None)
     }
 
     /// Builds an environment from the raw `CODEX_EXEC_SERVER_URL` value and
     /// local runtime paths used when creating local filesystem helpers.
-    pub async fn create_with_runtime_paths(
+    pub(crate) fn create_with_runtime_paths(
         exec_server_url: Option<String>,
         local_runtime_paths: Option<ExecServerRuntimePaths>,
     ) -> Result<Self, ExecServerError> {
@@ -171,17 +213,8 @@ impl Environment {
             ));
         }
 
-        let remote_exec_server_client = if let Some(exec_server_url) = &exec_server_url {
-            Some(
-                ExecServerClient::connect_websocket(RemoteExecServerConnectArgs {
-                    websocket_url: exec_server_url.clone(),
-                    client_name: "codex-environment".to_string(),
-                    connect_timeout: std::time::Duration::from_secs(5),
-                    initialize_timeout: std::time::Duration::from_secs(5),
-                    resume_session_id: None,
-                })
-                .await?,
-            )
+        let remote_exec_server_client = if let Some(exec_server_url) = exec_server_url.clone() {
+            Some(LazyRemoteExecServerClient::new(exec_server_url))
         } else {
             None
         };
@@ -242,15 +275,16 @@ mod tests {
 
     use super::Environment;
     use super::EnvironmentManager;
+    use super::EnvironmentManagerArgs;
+    use super::REMOTE_ENVIRONMENT_ID;
     use crate::ExecServerRuntimePaths;
     use crate::ProcessId;
     use pretty_assertions::assert_eq;
 
     #[tokio::test]
     async fn create_local_environment_does_not_connect() {
-        let environment = Environment::create(/*exec_server_url*/ None)
-            .await
-            .expect("create environment");
+        let environment =
+            Environment::create(/*exec_server_url*/ None).expect("create environment");
 
         assert_eq!(environment.exec_server_url(), None);
         assert!(environment.remote_exec_server_client.is_none());
@@ -258,39 +292,63 @@ mod tests {
 
     #[test]
     fn environment_manager_normalizes_empty_url() {
-        let manager = EnvironmentManager::new(Some(String::new()));
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: Some(String::new()),
+            local_runtime_paths: None,
+        });
 
-        assert!(!manager.disabled);
-        assert_eq!(manager.exec_server_url(), None);
-        assert!(!manager.is_remote());
+        let environment = manager.default_environment().expect("local environment");
+        assert!(!environment.is_remote());
+        assert!(manager.allows_agent_environment_access());
+        assert!(manager.local_environment().is_some());
+        assert!(manager.get_environment(REMOTE_ENVIRONMENT_ID).is_none());
     }
 
     #[test]
     fn environment_manager_treats_none_value_as_disabled() {
-        let manager = EnvironmentManager::new(Some("none".to_string()));
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: Some("none".to_string()),
+            local_runtime_paths: None,
+        });
 
-        assert!(manager.disabled);
-        assert_eq!(manager.exec_server_url(), None);
-        assert!(!manager.is_remote());
+        assert!(!manager.allows_agent_environment_access());
+        assert!(manager.default_environment().is_some());
+        assert!(manager.local_environment().is_some());
+        assert!(manager.get_environment(REMOTE_ENVIRONMENT_ID).is_none());
     }
 
     #[test]
     fn environment_manager_reports_remote_url() {
-        let manager = EnvironmentManager::new(Some("ws://127.0.0.1:8765".to_string()));
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: Some("ws://127.0.0.1:8765".to_string()),
+            local_runtime_paths: None,
+        });
 
-        assert!(manager.is_remote());
-        assert_eq!(manager.exec_server_url(), Some("ws://127.0.0.1:8765"));
+        let environment = manager.default_environment().expect("remote environment");
+        assert!(environment.is_remote());
+        assert!(manager.allows_agent_environment_access());
+        assert_eq!(environment.exec_server_url(), Some("ws://127.0.0.1:8765"));
+        assert!(
+            !manager
+                .local_environment()
+                .expect("local environment")
+                .is_remote()
+        );
+        assert_eq!(
+            manager
+                .get_environment(REMOTE_ENVIRONMENT_ID)
+                .expect("remote environment")
+                .exec_server_url(),
+            Some("ws://127.0.0.1:8765")
+        );
     }
 
     #[tokio::test]
-    async fn environment_manager_current_caches_environment() {
-        let manager = EnvironmentManager::new(/*exec_server_url*/ None);
+    async fn environment_manager_default_environment_caches_environment() {
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs::default());
 
-        let first = manager.current().await.expect("get current environment");
-        let second = manager.current().await.expect("get current environment");
-
-        let first = first.expect("local environment");
-        let second = second.expect("local environment");
+        let first = manager.default_environment().expect("local environment");
+        let second = manager.default_environment().expect("local environment");
 
         assert!(Arc::ptr_eq(&first, &second));
     }
@@ -302,35 +360,51 @@ mod tests {
             /*codex_linux_sandbox_exe*/ None,
         )
         .expect("runtime paths");
-        let manager = EnvironmentManager::new_with_runtime_paths(
-            /*exec_server_url*/ None,
-            Some(runtime_paths.clone()),
-        );
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: None,
+            local_runtime_paths: Some(runtime_paths.clone()),
+        });
 
-        let environment = manager
-            .current()
-            .await
-            .expect("get current environment")
-            .expect("local environment");
+        let environment = manager.default_environment().expect("local environment");
 
         assert_eq!(environment.local_runtime_paths(), Some(&runtime_paths));
-        assert_eq!(
-            EnvironmentManager::from_environment(Some(&environment)).local_runtime_paths,
-            Some(runtime_paths)
-        );
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: environment.exec_server_url().map(str::to_owned),
+            local_runtime_paths: environment.local_runtime_paths().cloned(),
+        });
+        let environment = manager.default_environment().expect("local environment");
+        assert_eq!(environment.local_runtime_paths(), Some(&runtime_paths));
     }
 
     #[tokio::test]
-    async fn disabled_environment_manager_has_no_current_environment() {
-        let manager = EnvironmentManager::new(Some("none".to_string()));
+    async fn disabled_environment_manager_has_default_environment_but_no_tool_environment() {
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: Some("none".to_string()),
+            local_runtime_paths: None,
+        });
 
-        assert!(
-            manager
-                .current()
-                .await
-                .expect("get current environment")
-                .is_none()
-        );
+        assert!(manager.default_environment().is_some());
+        assert!(!manager.allows_agent_environment_access());
+    }
+
+    #[tokio::test]
+    async fn environment_manager_allows_local_lookup_when_disabled() {
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs {
+            exec_server_url: Some("none".to_string()),
+            local_runtime_paths: None,
+        });
+
+        assert!(manager.default_environment().is_some());
+        assert!(!manager.allows_agent_environment_access());
+        assert!(manager.local_environment().is_some());
+        assert!(manager.get_environment(REMOTE_ENVIRONMENT_ID).is_none());
+    }
+
+    #[tokio::test]
+    async fn get_environment_returns_none_for_unknown_id() {
+        let manager = EnvironmentManager::from_exec_server_url(EnvironmentManagerArgs::default());
+
+        assert!(manager.get_environment("does-not-exist").is_none());
     }
 
     #[tokio::test]

codex-rs/exec-server/src/remote_file_system.rs

@@ -7,7 +7,6 @@ use tracing::trace;
 
 use crate::CopyOptions;
 use crate::CreateDirectoryOptions;
-use crate::ExecServerClient;
 use crate::ExecServerError;
 use crate::ExecutorFileSystem;
 use crate::FileMetadata;
@@ -15,6 +14,7 @@ use crate::FileSystemResult;
 use crate::FileSystemSandboxContext;
 use crate::ReadDirectoryEntry;
 use crate::RemoveOptions;
+use crate::environment::LazyRemoteExecServerClient;
 use crate::protocol::FsCopyParams;
 use crate::protocol::FsCreateDirectoryParams;
 use crate::protocol::FsGetMetadataParams;
@@ -28,14 +28,18 @@ const NOT_FOUND_ERROR_CODE: i64 = -32004;
 
 #[derive(Clone)]
 pub(crate) struct RemoteFileSystem {
-    client: ExecServerClient,
+    client: LazyRemoteExecServerClient,
 }
 
 impl RemoteFileSystem {
-    pub(crate) fn new(client: ExecServerClient) -> Self {
+    pub(crate) fn new(client: LazyRemoteExecServerClient) -> Self {
         trace!("remote fs new");
         Self { client }
     }
+
+    async fn client(&self) -> FileSystemResult<crate::ExecServerClient> {
+        self.client.get().await.map_err(map_remote_error)
+    }
 }
 
 #[async_trait]
@@ -46,8 +50,8 @@ impl ExecutorFileSystem for RemoteFileSystem {
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<Vec<u8>> {
         trace!("remote fs read_file");
-        let response = self
-            .client
+        let client = self.client().await?;
+        let response = client
             .fs_read_file(FsReadFileParams {
                 path: path.clone(),
                 sandbox: sandbox.cloned(),
@@ -69,7 +73,8 @@ impl ExecutorFileSystem for RemoteFileSystem {
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<()> {
         trace!("remote fs write_file");
-        self.client
+        let client = self.client().await?;
+        client
             .fs_write_file(FsWriteFileParams {
                 path: path.clone(),
                 data_base64: STANDARD.encode(contents),
@@ -87,7 +92,8 @@ impl ExecutorFileSystem for RemoteFileSystem {
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<()> {
         trace!("remote fs create_directory");
-        self.client
+        let client = self.client().await?;
+        client
             .fs_create_directory(FsCreateDirectoryParams {
                 path: path.clone(),
                 recursive: Some(options.recursive),
@@ -104,8 +110,8 @@ impl ExecutorFileSystem for RemoteFileSystem {
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<FileMetadata> {
         trace!("remote fs get_metadata");
-        let response = self
-            .client
+        let client = self.client().await?;
+        let response = client
             .fs_get_metadata(FsGetMetadataParams {
                 path: path.clone(),
                 sandbox: sandbox.cloned(),
@@ -127,8 +133,8 @@ impl ExecutorFileSystem for RemoteFileSystem {
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<Vec<ReadDirectoryEntry>> {
         trace!("remote fs read_directory");
-        let response = self
-            .client
+        let client = self.client().await?;
+        let response = client
             .fs_read_directory(FsReadDirectoryParams {
                 path: path.clone(),
                 sandbox: sandbox.cloned(),
@@ -153,7 +159,8 @@ impl ExecutorFileSystem for RemoteFileSystem {
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<()> {
         trace!("remote fs remove");
-        self.client
+        let client = self.client().await?;
+        client
             .fs_remove(FsRemoveParams {
                 path: path.clone(),
                 recursive: Some(options.recursive),
@@ -173,7 +180,8 @@ impl ExecutorFileSystem for RemoteFileSystem {
         sandbox: Option<&FileSystemSandboxContext>,
     ) -> FileSystemResult<()> {
         trace!("remote fs copy");
-        self.client
+        let client = self.client().await?;
+        client
             .fs_copy(FsCopyParams {
                 source_path: source_path.clone(),
                 destination_path: destination_path.clone(),

codex-rs/exec-server/src/remote_process.rs

@@ -9,15 +9,15 @@ use crate::ExecProcess;
 use crate::ExecProcessEventReceiver;
 use crate::ExecServerError;
 use crate::StartedExecProcess;
-use crate::client::ExecServerClient;
 use crate::client::Session;
+use crate::environment::LazyRemoteExecServerClient;
 use crate::protocol::ExecParams;
 use crate::protocol::ReadResponse;
 use crate::protocol::WriteResponse;
 
 #[derive(Clone)]
 pub(crate) struct RemoteProcess {
-    client: ExecServerClient,
+    client: LazyRemoteExecServerClient,
 }
 
 struct RemoteExecProcess {
@@ -25,7 +25,7 @@ struct RemoteExecProcess {
 }
 
 impl RemoteProcess {
-    pub(crate) fn new(client: ExecServerClient) -> Self {
+    pub(crate) fn new(client: LazyRemoteExecServerClient) -> Self {
         trace!("remote process new");
         Self { client }
     }
@@ -35,8 +35,9 @@ impl RemoteProcess {
 impl ExecBackend for RemoteProcess {
     async fn start(&self, params: ExecParams) -> Result<StartedExecProcess, ExecServerError> {
         let process_id = params.process_id.clone();
-        let session = self.client.register_session(&process_id).await?;
-        if let Err(err) = self.client.exec(params).await {
+        let client = self.client.get().await?;
+        let session = client.register_session(&process_id).await?;
+        if let Err(err) = client.exec(params).await {
             session.unregister().await;
             return Err(err);
         }