diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 395cac9e63..4372f69c1d 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -9,6 +9,7 @@ use crate::client::LazyRemoteExecServerClient; use crate::client::http_client::ReqwestHttpClient; use crate::client_api::ExecServerTransportParams; use crate::environment_provider::DefaultEnvironmentProvider; +use crate::environment_provider::DefaultEnvironmentSelection; use crate::environment_provider::EnvironmentProvider; use crate::environment_provider::normalize_exec_server_url; use crate::local_file_system::LocalFileSystem; @@ -32,8 +33,8 @@ pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; /// shell/filesystem tool availability. /// /// Remote environments create remote filesystem and execution backends that -/// lazy-connect to the configured exec-server on first use. The websocket is -/// not opened when the manager or environment is constructed. +/// lazy-connect to the configured exec-server on first use. The remote +/// transport is not opened when the manager or environment is constructed. #[derive(Debug)] pub struct EnvironmentManager { default_environment: Option, @@ -72,9 +73,12 @@ impl EnvironmentManager { /// Builds a test-only manager with environment access disabled. pub fn disabled_for_tests(local_runtime_paths: ExecServerRuntimePaths) -> Self { - let mut manager = Self::from_environments(HashMap::new(), local_runtime_paths); - manager.default_environment = None; - manager + Self::from_environments( + HashMap::new(), + local_runtime_paths, + DefaultEnvironmentSelection::Disabled, + ) + .expect("disabled test environment manager") } /// Builds a test-only manager from a raw exec-server URL value. @@ -99,16 +103,14 @@ impl EnvironmentManager { exec_server_url: Option, local_runtime_paths: ExecServerRuntimePaths, ) -> Self { - let environment_disabled = normalize_exec_server_url(exec_server_url.clone()).1; let provider = DefaultEnvironmentProvider::new(exec_server_url); let provider_environments = provider.environments(&local_runtime_paths); - let mut manager = Self::from_environments(provider_environments, local_runtime_paths); - if environment_disabled { - // TODO: Remove this legacy `CODEX_EXEC_SERVER_URL=none` crutch once - // environment attachment defaulting moves out of EnvironmentManager. - manager.default_environment = None; - } - manager + Self::from_environments( + provider_environments, + local_runtime_paths, + provider.default_environment_selection(), + ) + .expect("default provider should create valid environments") } /// Builds a manager from a provider-supplied startup snapshot. @@ -122,12 +124,14 @@ impl EnvironmentManager { Self::from_provider_environments( provider.get_environments(&local_runtime_paths).await?, local_runtime_paths, + provider.default_environment_selection(), ) } fn from_provider_environments( environments: HashMap, local_runtime_paths: ExecServerRuntimePaths, + default_selection: DefaultEnvironmentSelection, ) -> Result { for id in environments.keys() { if id.is_empty() { @@ -137,21 +141,35 @@ impl EnvironmentManager { } } - Ok(Self::from_environments(environments, local_runtime_paths)) + Self::from_environments(environments, local_runtime_paths, default_selection) } fn from_environments( environments: HashMap, local_runtime_paths: ExecServerRuntimePaths, - ) -> Self { + default_selection: DefaultEnvironmentSelection, + ) -> Result { // TODO: Stop deriving a default environment here once omitted // environment attachment is owned by thread/session setup. - let default_environment = if environments.contains_key(REMOTE_ENVIRONMENT_ID) { - Some(REMOTE_ENVIRONMENT_ID.to_string()) - } else if environments.contains_key(LOCAL_ENVIRONMENT_ID) { - Some(LOCAL_ENVIRONMENT_ID.to_string()) - } else { - None + let default_environment = match default_selection { + DefaultEnvironmentSelection::Derived => { + if environments.contains_key(REMOTE_ENVIRONMENT_ID) { + Some(REMOTE_ENVIRONMENT_ID.to_string()) + } else if environments.contains_key(LOCAL_ENVIRONMENT_ID) { + Some(LOCAL_ENVIRONMENT_ID.to_string()) + } else { + None + } + } + DefaultEnvironmentSelection::Environment(environment_id) => { + if !environments.contains_key(&environment_id) { + return Err(ExecServerError::Protocol(format!( + "default environment `{environment_id}` is not configured" + ))); + } + Some(environment_id) + } + DefaultEnvironmentSelection::Disabled => None, }; let local_environment = Arc::new(Environment::local(local_runtime_paths)); let environments = environments @@ -159,11 +177,11 @@ impl EnvironmentManager { .map(|(id, environment)| (id, Arc::new(environment))) .collect(); - Self { + Ok(Self { default_environment, environments, local_environment, - } + }) } /// Returns the default environment instance. @@ -196,6 +214,7 @@ impl EnvironmentManager { #[derive(Clone)] pub struct Environment { exec_server_url: Option, + remote_transport: Option, exec_backend: Arc, filesystem: Arc, http_client: Arc, @@ -207,6 +226,7 @@ impl Environment { pub fn default_for_tests() -> Self { Self { exec_server_url: None, + remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::unsandboxed()), http_client: Arc::new(ReqwestHttpClient), @@ -262,6 +282,7 @@ impl Environment { pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { exec_server_url: None, + remote_transport: None, exec_backend: Arc::new(LocalProcess::default()), filesystem: Arc::new(LocalFileSystem::with_runtime_paths( local_runtime_paths.clone(), @@ -275,15 +296,28 @@ impl Environment { exec_server_url: String, local_runtime_paths: Option, ) -> Self { - let client = LazyRemoteExecServerClient::new(ExecServerTransportParams::WebSocketUrl( - exec_server_url.clone(), - )); + Self::remote_with_transport( + ExecServerTransportParams::WebSocketUrl(exec_server_url), + local_runtime_paths, + ) + } + + fn remote_with_transport( + transport_params: ExecServerTransportParams, + local_runtime_paths: Option, + ) -> Self { + let exec_server_url = match &transport_params { + ExecServerTransportParams::WebSocketUrl(url) => Some(url.clone()), + ExecServerTransportParams::StdioCommand(_) => None, + }; + let client = LazyRemoteExecServerClient::new(transport_params.clone()); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client.clone())); Self { - exec_server_url: Some(exec_server_url), + exec_server_url, + remote_transport: Some(transport_params), exec_backend, filesystem, http_client: Arc::new(client), @@ -292,7 +326,7 @@ impl Environment { } pub fn is_remote(&self) -> bool { - self.exec_server_url.is_some() + self.remote_transport.is_some() } /// Returns the remote exec-server URL when this environment is remote. @@ -322,6 +356,7 @@ mod tests { use std::collections::HashMap; use std::sync::Arc; + use super::DefaultEnvironmentSelection; use super::Environment; use super::EnvironmentManager; use super::LOCAL_ENVIRONMENT_ID; @@ -428,7 +463,9 @@ mod tests { .expect("remote environment"), )]), test_runtime_paths(), - ); + DefaultEnvironmentSelection::Derived, + ) + .expect("environment manager"); assert_eq!( manager.default_environment_id(), @@ -449,6 +486,7 @@ mod tests { let err = EnvironmentManager::from_provider_environments( HashMap::from([("".to_string(), Environment::default_for_tests())]), test_runtime_paths(), + DefaultEnvironmentSelection::Derived, ) .expect_err("empty id should fail"); @@ -458,6 +496,64 @@ mod tests { ); } + #[tokio::test] + async fn environment_manager_uses_explicit_provider_default() { + let manager = EnvironmentManager::from_provider_environments( + HashMap::from([ + ( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + ), + ( + "devbox".to_string(), + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + ), + ]), + test_runtime_paths(), + DefaultEnvironmentSelection::Environment("devbox".to_string()), + ) + .expect("manager"); + + assert_eq!(manager.default_environment_id(), Some("devbox")); + assert!(manager.default_environment().expect("default").is_remote()); + } + + #[tokio::test] + async fn environment_manager_disables_provider_default() { + let manager = EnvironmentManager::from_provider_environments( + HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + test_runtime_paths(), + DefaultEnvironmentSelection::Disabled, + ) + .expect("manager"); + + assert_eq!(manager.default_environment_id(), None); + assert!(manager.default_environment().is_none()); + assert!(manager.get_environment(LOCAL_ENVIRONMENT_ID).is_some()); + } + + #[tokio::test] + async fn environment_manager_rejects_unknown_provider_default() { + let err = EnvironmentManager::from_provider_environments( + HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Environment::default_for_tests(), + )]), + test_runtime_paths(), + DefaultEnvironmentSelection::Environment("missing".to_string()), + ) + .expect_err("unknown default should fail"); + + assert_eq!( + err.to_string(), + "exec-server protocol error: default environment `missing` is not configured" + ); + } + #[tokio::test] async fn environment_manager_uses_provider_supplied_local_environment() { let manager = EnvironmentManager::create_for_tests( diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 7c8db07e85..eebfff6f13 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -21,6 +21,17 @@ pub trait EnvironmentProvider: Send + Sync { &self, local_runtime_paths: &ExecServerRuntimePaths, ) -> Result, ExecServerError>; + + fn default_environment_selection(&self) -> DefaultEnvironmentSelection { + DefaultEnvironmentSelection::Derived + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DefaultEnvironmentSelection { + Derived, + Environment(String), + Disabled, } /// Default provider backed by `CODEX_EXEC_SERVER_URL`. @@ -69,6 +80,14 @@ impl EnvironmentProvider for DefaultEnvironmentProvider { ) -> Result, ExecServerError> { Ok(self.environments(local_runtime_paths)) } + + fn default_environment_selection(&self) -> DefaultEnvironmentSelection { + if normalize_exec_server_url(self.exec_server_url.clone()).1 { + DefaultEnvironmentSelection::Disabled + } else { + DefaultEnvironmentSelection::Derived + } + } } pub(crate) fn normalize_exec_server_url(exec_server_url: Option) -> (Option, bool) { @@ -135,6 +154,10 @@ mod tests { assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); + assert_eq!( + provider.default_environment_selection(), + DefaultEnvironmentSelection::Disabled + ); } #[tokio::test] diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index b36ab39d01..1fd1ea9c1d 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -42,6 +42,7 @@ pub use environment::EnvironmentManagerArgs; pub use environment::LOCAL_ENVIRONMENT_ID; pub use environment::REMOTE_ENVIRONMENT_ID; pub use environment_provider::DefaultEnvironmentProvider; +pub use environment_provider::DefaultEnvironmentSelection; pub use environment_provider::EnvironmentProvider; pub use fs_helper::CODEX_FS_HELPER_ARG1; pub use fs_helper_main::main as run_fs_helper_main;