From a7fa3679ba9389eaf113c54ae4c0507236defd3a Mon Sep 17 00:00:00 2001 From: starr-openai Date: Fri, 17 Apr 2026 15:02:29 -0700 Subject: [PATCH] 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 --- codex-rs/app-server-client/src/lib.rs | 21 +- .../app-server/src/bespoke_event_handling.rs | 4 +- .../app-server/src/codex_message_processor.rs | 80 ++--- codex-rs/app-server/src/in_process.rs | 4 +- .../src/message_processor/tracing_tests.rs | 4 +- .../app-server/tests/suite/v2/mcp_resource.rs | 4 +- codex-rs/core/src/agent/control_tests.rs | 28 +- codex-rs/core/src/codex_delegate.rs | 5 +- codex-rs/core/src/memories/tests.rs | 4 +- codex-rs/core/src/session/mod.rs | 7 +- codex-rs/core/src/session/review.rs | 2 +- codex-rs/core/src/session/session.rs | 8 +- codex-rs/core/src/session/tests.rs | 20 +- .../core/src/session/tests/guardian_tests.rs | 4 +- codex-rs/core/src/session/turn_context.rs | 4 +- codex-rs/core/src/state/service.rs | 3 + codex-rs/core/src/thread_manager.rs | 10 +- codex-rs/core/src/thread_manager_tests.rs | 20 +- codex-rs/core/tests/common/test_codex.rs | 14 +- codex-rs/core/tests/suite/client.rs | 4 +- codex-rs/core/tests/suite/skills.rs | 7 +- codex-rs/exec-server/src/environment.rs | 334 +++++++++++------- .../exec-server/src/remote_file_system.rs | 34 +- codex-rs/exec-server/src/remote_process.rs | 11 +- codex-rs/exec-server/tests/exec_process.rs | 4 +- codex-rs/exec-server/tests/file_system.rs | 4 +- codex-rs/tui/src/app/test_support.rs | 4 +- codex-rs/tui/src/app/tests.rs | 8 +- codex-rs/tui/src/lib.rs | 34 +- codex-rs/tui/src/onboarding/auth.rs | 8 +- 30 files changed, 399 insertions(+), 299 deletions(-) diff --git a/codex-rs/app-server-client/src/lib.rs b/codex-rs/app-server-client/src/lib.rs index c6f678c2aa..31184063e9 100644 --- a/codex-rs/app-server-client/src/lib.rs +++ b/codex-rs/app-server-client/src/lib.rs @@ -46,6 +46,7 @@ use codex_core::config::Config; use codex_core::config_loader::CloudRequirementsLoader; use codex_core::config_loader::LoaderOverrides; pub use codex_exec_server::EnvironmentManager; +pub use codex_exec_server::EnvironmentManagerArgs; pub use codex_exec_server::ExecServerRuntimePaths; use codex_feedback::CodexFeedback; use codex_protocol::protocol::SessionSource; @@ -968,7 +969,9 @@ mod tests { cloud_requirements: CloudRequirementsLoader::default(), feedback: CodexFeedback::new(), log_db: None, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), config_warnings: Vec::new(), session_source, enable_codex_api_key_env: false, @@ -1969,9 +1972,12 @@ mod tests { #[tokio::test] async fn runtime_start_args_forward_environment_manager() { let config = Arc::new(build_test_config().await); - let environment_manager = Arc::new(EnvironmentManager::new(Some( - "ws://127.0.0.1:8765".to_string(), - ))); + let environment_manager = Arc::new(EnvironmentManager::from_exec_server_url( + EnvironmentManagerArgs { + exec_server_url: Some("ws://127.0.0.1:8765".to_string()), + local_runtime_paths: None, + }, + )); let runtime_args = InProcessClientStartArgs { arg0_paths: Arg0DispatchPaths::default(), @@ -1998,7 +2004,12 @@ mod tests { &runtime_args.environment_manager, &environment_manager )); - assert!(runtime_args.environment_manager.is_remote()); + assert!( + runtime_args + .environment_manager + .default_environment() + .is_some_and(|environment| environment.is_remote()) + ); } #[tokio::test] diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index a2ea77900a..43a75df678 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -3497,8 +3497,8 @@ mod tests { CodexAuth::create_dummy_chatgpt_auth_for_testing(), config.model_provider.clone(), config.codex_home.to_path_buf(), - Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ), ); diff --git a/codex-rs/app-server/src/codex_message_processor.rs b/codex-rs/app-server/src/codex_message_processor.rs index 3b92d4199f..08fb559d28 100644 --- a/codex-rs/app-server/src/codex_message_processor.rs +++ b/codex-rs/app-server/src/codex_message_processor.rs @@ -5694,27 +5694,15 @@ impl CodexMessageProcessor { .await; let auth_manager = Arc::clone(&self.auth_manager); let auth = auth_manager.auth().await; - let runtime_environment = match self.thread_manager.environment_manager().current().await { - Ok(Some(environment)) => { - // Status listing has no turn cwd. This fallback is used only - // by executor-backed stdio MCPs whose config omits `cwd`. - McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()) - } - Ok(None) => McpRuntimeEnvironment::new( - Arc::new(codex_exec_server::Environment::default()), - config.cwd.to_path_buf(), - ), - Err(err) => { - // TODO(aibrahim): Investigate degrading MCP status listing when - // executor environment creation fails. - let error = JSONRPCErrorError { - code: INTERNAL_ERROR_CODE, - message: format!("failed to create environment: {err}"), - data: None, - }; - self.outgoing.send_error(request, error).await; - return; - } + let runtime_environment = { + let environment = self + .thread_manager + .environment_manager() + .default_environment() + .unwrap_or_else(|| Arc::new(codex_exec_server::Environment::default())); + // Status listing has no turn cwd. This fallback is used only + // by executor-backed stdio MCPs whose config omits `cwd`. + McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()) }; tokio::spawn(async move { @@ -5864,25 +5852,15 @@ impl CodexMessageProcessor { .to_mcp_config(self.thread_manager.plugins_manager().as_ref()) .await; let auth = self.auth_manager.auth().await; - let runtime_environment = match self.thread_manager.environment_manager().current().await { - Ok(Some(environment)) => { - // Resource reads without a thread have no turn cwd. This fallback - // is used only by executor-backed stdio MCPs whose config omits `cwd`. - McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()) - } - Ok(None) => McpRuntimeEnvironment::new( - Arc::new(codex_exec_server::Environment::default()), - config.cwd.to_path_buf(), - ), - Err(err) => { - let error = JSONRPCErrorError { - code: INTERNAL_ERROR_CODE, - message: format!("failed to create environment: {err}"), - data: None, - }; - self.outgoing.send_error(request_id, error).await; - return; - } + let runtime_environment = { + let environment = self + .thread_manager + .environment_manager() + .default_environment() + .unwrap_or_else(|| Arc::new(codex_exec_server::Environment::default())); + // Resource reads without a thread have no turn cwd. This fallback + // is used only by executor-backed stdio MCPs whose config omits `cwd`. + McpRuntimeEnvironment::new(environment, config.cwd.to_path_buf()) }; tokio::spawn(async move { @@ -6469,23 +6447,11 @@ impl CodexMessageProcessor { }; let skills_manager = self.thread_manager.skills_manager(); let plugins_manager = self.thread_manager.plugins_manager(); - let fs = match self.thread_manager.environment_manager().current().await { - Ok(Some(environment)) => Some(environment.get_filesystem()), - Ok(None) => None, - Err(err) => { - self.outgoing - .send_error( - request_id, - JSONRPCErrorError { - code: INTERNAL_ERROR_CODE, - message: format!("failed to create environment: {err}"), - data: None, - }, - ) - .await; - return; - } - }; + let fs = self + .thread_manager + .environment_manager() + .default_environment() + .map(|environment| environment.get_filesystem()); let mut data = Vec::new(); for cwd in cwds { let extra_roots = extra_roots_by_cwd diff --git a/codex-rs/app-server/src/in_process.rs b/codex-rs/app-server/src/in_process.rs index e73124c0d3..c4544a7de7 100644 --- a/codex-rs/app-server/src/in_process.rs +++ b/codex-rs/app-server/src/in_process.rs @@ -738,7 +738,9 @@ mod tests { thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), config_warnings: Vec::new(), session_source, enable_codex_api_key_env: false, diff --git a/codex-rs/app-server/src/message_processor/tracing_tests.rs b/codex-rs/app-server/src/message_processor/tracing_tests.rs index 83f8bc98d7..a6fffce95a 100644 --- a/codex-rs/app-server/src/message_processor/tracing_tests.rs +++ b/codex-rs/app-server/src/message_processor/tracing_tests.rs @@ -241,7 +241,9 @@ fn build_test_processor( outgoing, arg0_paths: Arg0DispatchPaths::default(), config, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), cli_overrides: Vec::new(), loader_overrides: LoaderOverrides::default(), cloud_requirements: CloudRequirementsLoader::default(), diff --git a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs index c26b456fa9..34b7a55d6a 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_resource.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_resource.rs @@ -204,7 +204,9 @@ async fn mcp_resource_read_returns_error_for_unknown_thread() -> Result<()> { thread_config_loader: Arc::new(codex_config::NoopThreadConfigLoader), feedback: CodexFeedback::new(), log_db: None, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), config_warnings: Vec::new(), session_source: SessionSource::Cli, enable_codex_api_key_env: false, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 10850ef8c7..473180caa3 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -95,8 +95,8 @@ impl AgentControlHarness { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let control = manager.agent_control(); @@ -911,8 +911,8 @@ async fn spawn_agent_respects_max_threads_limit() { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let control = manager.agent_control(); @@ -965,8 +965,8 @@ async fn spawn_agent_releases_slot_after_shutdown() { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let control = manager.agent_control(); @@ -1010,8 +1010,8 @@ async fn spawn_agent_limit_shared_across_clones() { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let control = manager.agent_control(); @@ -1057,8 +1057,8 @@ async fn resume_agent_respects_max_threads_limit() { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let control = manager.agent_control(); @@ -1115,8 +1115,8 @@ async fn resume_agent_releases_slot_after_resume_failure() { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let control = manager.agent_control(); @@ -1512,8 +1512,8 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let control = manager.agent_control(); diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 7247c601f4..4f4ced4101 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -4,7 +4,6 @@ use std::sync::Arc; use async_channel::Receiver; use async_channel::Sender; use codex_async_utils::OrCancelExt; -use codex_exec_server::EnvironmentManager; use codex_protocol::protocol::ApplyPatchApprovalRequestEvent; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; @@ -78,9 +77,7 @@ pub(crate) async fn run_codex_thread_interactive( config, auth_manager, models_manager, - environment_manager: Arc::new(EnvironmentManager::from_environment( - parent_ctx.environment.as_deref(), - )), + environment_manager: Arc::clone(&parent_session.services.environment_manager), skills_manager: Arc::clone(&parent_session.services.skills_manager), plugins_manager: Arc::clone(&parent_session.services.plugins_manager), mcp_manager: Arc::clone(&parent_session.services.mcp_manager), diff --git a/codex-rs/core/src/memories/tests.rs b/codex-rs/core/src/memories/tests.rs index af698e1eaa..2b68596fd4 100644 --- a/codex-rs/core/src/memories/tests.rs +++ b/codex-rs/core/src/memories/tests.rs @@ -491,8 +491,8 @@ mod phase2 { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + std::sync::Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let (mut session, _turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4b919800cd..ad392efda0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -460,10 +460,7 @@ impl Codex { let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY); let (tx_event, rx_event) = async_channel::unbounded(); - let environment = environment_manager - .current() - .await - .map_err(|err| CodexErr::Fatal(format!("failed to create environment: {err}")))?; + let environment = environment_manager.default_environment(); let fs = environment .as_ref() .map(|environment| environment.get_filesystem()); @@ -650,7 +647,7 @@ impl Codex { mcp_manager.clone(), skills_watcher, agent_control, - environment, + environment_manager, analytics_events_client, ) .await diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index 94de4617d5..af1028686d 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -46,7 +46,7 @@ pub(super) async fn spawn_review_thread( ) .with_web_search_config(/*web_search_config*/ None) .with_allow_login_shell(config.permissions.allow_login_shell) - .with_has_environment(parent_turn_context.environment.is_some()) + .with_has_environment(parent_turn_context.tools_config.has_environment) .with_spawn_agent_usage_hint(config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(config.multi_agent_v2.hide_spawn_agent_metadata) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 16e86e3aea..94cf8ae3f8 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -228,7 +228,7 @@ impl Session { mcp_manager: Arc, skills_watcher: Arc, agent_control: AgentControl, - environment: Option>, + environment_manager: Arc, analytics_events_client: Option, ) -> anyhow::Result> { debug!( @@ -641,6 +641,8 @@ impl Session { Arc::clone(&auth_manager), session_configuration.session_source.clone(), )); + let environment = environment_manager.default_environment(); + let allows_agent_environment_access = environment_manager.allows_agent_environment_access(); let services = SessionServices { // Initialize the MCP connection manager with an uninitialized // instance. It will be replaced with one created via @@ -695,7 +697,9 @@ impl Session { code_mode_service: crate::tools::code_mode::CodeModeService::new( config.js_repl_node_path.clone(), ), - environment: environment.clone(), + environment_manager, + environment, + allows_agent_environment_access, }; services .model_client diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index f1a11bfcb6..1da7ad13e7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3082,11 +3082,7 @@ async fn session_new_fails_when_zsh_fork_enabled_without_zsh_path() { mcp_manager, Arc::new(SkillsWatcher::noop()), AgentControl::default(), - Some(Arc::new( - codex_exec_server::Environment::create(/*exec_server_url*/ None) - .await - .expect("create environment"), - )), + Arc::new(codex_exec_server::EnvironmentManager::default()), /*analytics_events_client*/ None, ) .await; @@ -3184,7 +3180,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { let network_approval = Arc::new(NetworkApprovalService::default()); let environment = Arc::new( codex_exec_server::Environment::create(/*exec_server_url*/ None) - .await .expect("create environment"), ); @@ -3249,7 +3244,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { code_mode_service: crate::tools::code_mode::CodeModeService::new( config.js_repl_node_path.clone(), ), + environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default()), environment: Some(Arc::clone(&environment)), + allows_agent_environment_access: true, }; let js_repl = Arc::new(JsReplHandle::with_node_path( config.js_repl_node_path.clone(), @@ -3283,6 +3280,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { model_info, &models_manager, /*network*/ None, + /*allows_agent_environment_access*/ true, Some(environment), "turn_id".to_string(), Arc::clone(&js_repl), @@ -3405,11 +3403,7 @@ async fn make_session_with_config_and_rx( mcp_manager, Arc::new(SkillsWatcher::noop()), AgentControl::default(), - Some(Arc::new( - codex_exec_server::Environment::create(/*exec_server_url*/ None) - .await - .expect("create environment"), - )), + Arc::new(codex_exec_server::EnvironmentManager::default()), /*analytics_events_client*/ None, ) .await?; @@ -4287,7 +4281,6 @@ where let network_approval = Arc::new(NetworkApprovalService::default()); let environment = Arc::new( codex_exec_server::Environment::create(/*exec_server_url*/ None) - .await .expect("create environment"), ); @@ -4352,7 +4345,9 @@ where code_mode_service: crate::tools::code_mode::CodeModeService::new( config.js_repl_node_path.clone(), ), + environment_manager: Arc::new(codex_exec_server::EnvironmentManager::default()), environment: Some(Arc::clone(&environment)), + allows_agent_environment_access: true, }; let js_repl = Arc::new(JsReplHandle::with_node_path( config.js_repl_node_path.clone(), @@ -4386,6 +4381,7 @@ where model_info, &models_manager, /*network*/ None, + /*allows_agent_environment_access*/ true, Some(environment), "turn_id".to_string(), Arc::clone(&js_repl), diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 9b1172d578..9ee27bd202 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -634,7 +634,9 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { config, auth_manager, models_manager, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), skills_manager, plugins_manager, mcp_manager, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index dd86804ee5..58fcfc73b8 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -345,6 +345,7 @@ impl Session { model_info: ModelInfo, models_manager: &ModelsManager, network: Option, + allows_agent_environment_access: bool, environment: Option>, sub_id: String, js_repl: Arc, @@ -381,7 +382,7 @@ impl Session { ) .with_web_search_config(per_turn_config.web_search_config.clone()) .with_allow_login_shell(per_turn_config.permissions.allow_login_shell) - .with_has_environment(environment.is_some()) + .with_has_environment(allows_agent_environment_access) .with_spawn_agent_usage_hint(per_turn_config.multi_agent_v2.usage_hint_enabled) .with_spawn_agent_usage_hint_text(per_turn_config.multi_agent_v2.usage_hint_text.clone()) .with_hide_spawn_agent_metadata(per_turn_config.multi_agent_v2.hide_spawn_agent_metadata) @@ -576,6 +577,7 @@ impl Session { ) .then(|| started_proxy.proxy()) }), + self.services.allows_agent_environment_access, self.services.environment.clone(), sub_id, Arc::clone(&self.js_repl), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 5db38f7b72..6e42227b6b 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -18,6 +18,7 @@ use crate::tools::sandboxing::ApprovalStore; use crate::unified_exec::UnifiedExecProcessManager; use codex_analytics::AnalyticsEventsClient; use codex_exec_server::Environment; +use codex_exec_server::EnvironmentManager; use codex_hooks::Hooks; use codex_login::AuthManager; use codex_mcp::McpConnectionManager; @@ -64,5 +65,7 @@ pub(crate) struct SessionServices { /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, pub(crate) code_mode_service: CodeModeService, + pub(crate) environment_manager: Arc, pub(crate) environment: Option>, + pub(crate) allows_agent_environment_access: bool, } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index e4da99bb55..293423defe 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -301,7 +301,9 @@ impl ThreadManager { auth, provider, codex_home.clone(), - Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), ); manager._test_codex_home_guard = Some(TempCodexHomeGuard { path: codex_home }); manager @@ -920,11 +922,7 @@ impl ThreadManagerState { parent_trace: Option, user_shell_override: Option, ) -> CodexResult { - let environment = self - .environment_manager - .current() - .await - .map_err(|err| CodexErr::Fatal(format!("failed to create environment: {err}")))?; + let environment = self.environment_manager.default_environment(); let watch_registration = match environment.as_ref() { Some(environment) if !environment.is_remote() => { self.skills_watcher diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index fe2039e89b..8c0fbe5e38 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -246,8 +246,8 @@ async fn shutdown_all_threads_bounded_submits_shutdown_to_every_thread() { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.to_path_buf(), - Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), ); let thread_1 = manager @@ -297,8 +297,8 @@ async fn new_uses_configured_openai_provider_for_model_refresh() { auth_manager, SessionSource::Exec, CollaborationModesConfig::default(), - Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), /*analytics_events_client*/ None, ); @@ -434,8 +434,8 @@ async fn interrupted_fork_snapshot_does_not_synthesize_turn_id_for_legacy_histor auth_manager.clone(), SessionSource::Exec, CollaborationModesConfig::default(), - Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), /*analytics_events_client*/ None, ); @@ -537,8 +537,8 @@ async fn interrupted_fork_snapshot_preserves_explicit_turn_id() { auth_manager.clone(), SessionSource::Exec, CollaborationModesConfig::default(), - Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), /*analytics_events_client*/ None, ); @@ -630,8 +630,8 @@ async fn interrupted_fork_snapshot_uses_persisted_mid_turn_history_without_live_ auth_manager.clone(), SessionSource::Exec, CollaborationModesConfig::default(), - Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), /*analytics_events_client*/ None, ); diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 5075f91620..32f5657b34 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -76,7 +76,7 @@ impl TestEnv { pub async fn local() -> Result { let local_cwd_temp_dir = Arc::new(TempDir::new()?); let cwd = local_cwd_temp_dir.abs(); - let environment = codex_exec_server::Environment::create(/*exec_server_url*/ None).await?; + let environment = codex_exec_server::Environment::create(/*exec_server_url*/ None)?; Ok(Self { environment, cwd, @@ -115,7 +115,7 @@ pub async fn test_env() -> Result { match get_remote_test_env() { Some(remote_env) => { let websocket_url = remote_exec_server_url()?; - let environment = codex_exec_server::Environment::create(Some(websocket_url)).await?; + let environment = codex_exec_server::Environment::create(Some(websocket_url))?; let cwd = remote_aware_cwd_path(); environment .get_filesystem() @@ -350,9 +350,13 @@ impl TestCodexBuilder { let (config, fallback_cwd) = self .prepare_config(base_url, &home, test_env.cwd().clone()) .await?; - let environment_manager = Arc::new(codex_exec_server::EnvironmentManager::new( - test_env.exec_server_url().map(str::to_owned), - )); + let environment_manager = + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs { + exec_server_url: test_env.exec_server_url().map(str::to_owned), + local_runtime_paths: None, + }, + )); let file_system = test_env.environment().get_filesystem(); let mut workspace_setups = vec![]; swap(&mut self.workspace_setups, &mut workspace_setups); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 2086367b21..3bebd2192c 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -1103,8 +1103,8 @@ async fn prefers_apikey_when_config_prefers_apikey_even_with_chatgpt_tokens() { .features .enabled(Feature::DefaultModeRequestUserInput), }, - Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, + Arc::new(codex_exec_server::EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), )), /*analytics_events_client*/ None, ); diff --git a/codex-rs/core/tests/suite/skills.rs b/codex-rs/core/tests/suite/skills.rs index 1b2ac71643..4ff6da00d1 100644 --- a/codex-rs/core/tests/suite/skills.rs +++ b/codex-rs/core/tests/suite/skills.rs @@ -234,7 +234,12 @@ async fn list_skills_skips_cwd_roots_when_environment_disabled() -> Result<()> { codex_core::test_support::auth_manager_from_auth(CodexAuth::from_api_key("dummy")), SessionSource::Exec, CollaborationModesConfig::default(), - Arc::new(EnvironmentManager::new(Some("none".to_string()))), + Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs { + exec_server_url: Some("none".to_string()), + local_runtime_paths: None, + }, + )), /*analytics_events_client*/ None, ); let new_thread = thread_manager.start_thread(config.clone()).await?; diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index afe0720196..d43d2ba2c4 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/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, - local_runtime_paths: Option, - disabled: bool, - current_environment: OnceCell>>, + default_environment: Option, + environment_disabled_for_agent: bool, + environments: HashMap>, +} + +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, + pub local_runtime_paths: Option, +} + +#[derive(Clone, Debug)] +pub(crate) struct LazyRemoteExecServerClient { + websocket_url: String, + client: Arc>, +} + +impl LazyRemoteExecServerClient { + fn new(websocket_url: String) -> Self { + Self { + websocket_url, + client: Arc::new(OnceCell::new()), + } + } + + pub(crate) async fn get(&self) -> Result { + 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) -> 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, - local_runtime_paths: Option, - ) -> 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, ) -> 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> { + 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>, 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>::cloned) + /// Returns the local environment instance. + pub fn local_environment(&self) -> Option> { + self.get_environment(LOCAL_ENVIRONMENT_ID) + } + + /// Returns a named environment instance. + pub fn get_environment(&self, environment_id: &str) -> Option> { + self.environments.get(environment_id).cloned() } } @@ -128,7 +170,7 @@ impl EnvironmentManager { #[derive(Clone)] pub struct Environment { exec_server_url: Option, - remote_exec_server_client: Option, + remote_exec_server_client: Option, exec_backend: Arc, local_runtime_paths: Option, } @@ -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) -> Result { - Self::create_with_runtime_paths(exec_server_url, /*local_runtime_paths*/ None).await + pub fn create(exec_server_url: Option) -> Result { + 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, local_runtime_paths: Option, ) -> Result { @@ -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] diff --git a/codex-rs/exec-server/src/remote_file_system.rs b/codex-rs/exec-server/src/remote_file_system.rs index d6a32ba4d5..1b4150fe5e 100644 --- a/codex-rs/exec-server/src/remote_file_system.rs +++ b/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 { + self.client.get().await.map_err(map_remote_error) + } } #[async_trait] @@ -46,8 +50,8 @@ impl ExecutorFileSystem for RemoteFileSystem { sandbox: Option<&FileSystemSandboxContext>, ) -> FileSystemResult> { 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 { 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> { 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(), diff --git a/codex-rs/exec-server/src/remote_process.rs b/codex-rs/exec-server/src/remote_process.rs index 86786a54f7..19828d9d69 100644 --- a/codex-rs/exec-server/src/remote_process.rs +++ b/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 { 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); } diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index d449315c8d..4887a0be4d 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -49,13 +49,13 @@ enum ProcessEventSnapshot { async fn create_process_context(use_remote: bool) -> Result { if use_remote { let server = exec_server().await?; - let environment = Environment::create(Some(server.websocket_url().to_string())).await?; + let environment = Environment::create(Some(server.websocket_url().to_string()))?; Ok(ProcessContext { backend: environment.get_exec_backend(), server: Some(server), }) } else { - let environment = Environment::create(/*exec_server_url*/ None).await?; + let environment = Environment::create(/*exec_server_url*/ None)?; Ok(ProcessContext { backend: environment.get_exec_backend(), server: None, diff --git a/codex-rs/exec-server/tests/file_system.rs b/codex-rs/exec-server/tests/file_system.rs index d4f94c7e44..4bb654198f 100644 --- a/codex-rs/exec-server/tests/file_system.rs +++ b/codex-rs/exec-server/tests/file_system.rs @@ -46,7 +46,7 @@ struct FileSystemContext { async fn create_file_system_context(use_remote: bool) -> Result { if use_remote { let server = exec_server().await?; - let environment = Environment::create(Some(server.websocket_url().to_string())).await?; + let environment = Environment::create(Some(server.websocket_url().to_string()))?; Ok(FileSystemContext { file_system: environment.get_filesystem(), _helper_paths: None, @@ -214,7 +214,7 @@ async fn sandboxed_file_system_helper_finds_bwrap_on_preserved_path() -> Result< let helper_path = std::env::join_paths(path_entries)?; let server = exec_server_with_env([("PATH", helper_path.as_os_str())]).await?; - let environment = Environment::create(Some(server.websocket_url().to_string())).await?; + let environment = Environment::create(Some(server.websocket_url().to_string()))?; let file_system = environment.get_filesystem(); let workspace = tmp.path().join("workspace"); std::fs::create_dir_all(&workspace)?; diff --git a/codex-rs/tui/src/app/test_support.rs b/codex-rs/tui/src/app/test_support.rs index 8b2dfe8d47..254fc4cfa7 100644 --- a/codex-rs/tui/src/app/test_support.rs +++ b/codex-rs/tui/src/app/test_support.rs @@ -38,7 +38,9 @@ pub(super) async fn make_test_app() -> App { backtrack_render_pending: false, feedback: codex_feedback::CodexFeedback::new(), feedback_audience: FeedbackAudience::External, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), remote_app_server_url: None, remote_app_server_auth_token: None, pending_update_action: None, diff --git a/codex-rs/tui/src/app/tests.rs b/codex-rs/tui/src/app/tests.rs index 1749ee9f1e..6f4ff6d1aa 100644 --- a/codex-rs/tui/src/app/tests.rs +++ b/codex-rs/tui/src/app/tests.rs @@ -3576,7 +3576,9 @@ async fn make_test_app() -> App { backtrack_render_pending: false, feedback: codex_feedback::CodexFeedback::new(), feedback_audience: FeedbackAudience::External, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), remote_app_server_url: None, remote_app_server_auth_token: None, pending_update_action: None, @@ -3633,7 +3635,9 @@ async fn make_test_app_with_channels() -> ( backtrack_render_pending: false, feedback: codex_feedback::CodexFeedback::new(), feedback_audience: FeedbackAudience::External, - environment_manager: Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + environment_manager: Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), remote_app_server_url: None, remote_app_server_auth_token: None, pending_update_action: None, diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 7c60b2e38a..2fb6ace80a 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -425,7 +425,9 @@ pub(crate) async fn start_embedded_app_server_for_picker( start_app_server_for_picker( config, &AppServerTarget::Embedded, - Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), ) .await } @@ -623,7 +625,9 @@ fn config_cwd_for_app_server_target( app_server_target: &AppServerTarget, environment_manager: &EnvironmentManager, ) -> std::io::Result> { - if environment_manager.is_remote() + if environment_manager + .default_environment() + .is_some_and(|environment| environment.is_remote()) || matches!(app_server_target, AppServerTarget::Remote { .. }) { return Ok(None); @@ -1771,7 +1775,9 @@ mod tests { CloudRequirementsLoader::default(), codex_feedback::CodexFeedback::new(), /*log_db*/ None, - Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), ) .await } @@ -1930,7 +1936,9 @@ mod tests { websocket_url: "ws://127.0.0.1:1234/".to_string(), auth_token: None, }; - let environment_manager = EnvironmentManager::new(/*exec_server_url*/ None); + let environment_manager = EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + ); let config_cwd = config_cwd_for_app_server_target(Some(remote_only_cwd), &target, &environment_manager)?; @@ -1943,7 +1951,9 @@ mod tests { fn config_cwd_for_app_server_target_canonicalizes_embedded_cli_cwd() -> std::io::Result<()> { let temp_dir = TempDir::new()?; let target = AppServerTarget::Embedded; - let environment_manager = EnvironmentManager::new(/*exec_server_url*/ None); + let environment_manager = EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + ); let config_cwd = config_cwd_for_app_server_target(Some(temp_dir.path()), &target, &environment_manager)?; @@ -1963,7 +1973,9 @@ mod tests { let temp_dir = TempDir::new()?; let missing = temp_dir.path().join("missing"); let target = AppServerTarget::Embedded; - let environment_manager = EnvironmentManager::new(/*exec_server_url*/ None); + let environment_manager = EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + ); let err = config_cwd_for_app_server_target(Some(&missing), &target, &environment_manager) .expect_err("missing embedded cwd should fail"); @@ -1980,7 +1992,11 @@ mod tests { Path::new("/definitely/not/local/to/this/test") }; let target = AppServerTarget::Embedded; - let environment_manager = EnvironmentManager::new(Some("ws://127.0.0.1:8765".to_string())); + let environment_manager = + EnvironmentManager::from_exec_server_url(codex_exec_server::EnvironmentManagerArgs { + exec_server_url: Some("ws://127.0.0.1:8765".to_string()), + local_runtime_paths: None, + }); let config_cwd = config_cwd_for_app_server_target(Some(remote_only_cwd), &target, &environment_manager)?; @@ -2107,7 +2123,9 @@ mod tests { CloudRequirementsLoader::default(), codex_feedback::CodexFeedback::new(), /*log_db*/ None, - Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + Arc::new(EnvironmentManager::from_exec_server_url( + codex_exec_server::EnvironmentManagerArgs::default(), + )), |_args| async { Err(std::io::Error::other("boom")) }, ) .await; diff --git a/codex-rs/tui/src/onboarding/auth.rs b/codex-rs/tui/src/onboarding/auth.rs index aee909cea1..ee6f843c8c 100644 --- a/codex-rs/tui/src/onboarding/auth.rs +++ b/codex-rs/tui/src/onboarding/auth.rs @@ -989,9 +989,11 @@ mod tests { ), feedback: codex_feedback::CodexFeedback::new(), log_db: None, - environment_manager: Arc::new(codex_app_server_client::EnvironmentManager::new( - /*exec_server_url*/ None, - )), + environment_manager: Arc::new( + codex_app_server_client::EnvironmentManager::from_exec_server_url( + codex_app_server_client::EnvironmentManagerArgs::default(), + ), + ), config_warnings: Vec::new(), session_source: SessionSource::Cli, enable_codex_api_key_env: false,