From ed4ac115c987706affa331a921ccf9371673b88a Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 20 Apr 2026 14:58:01 -0700 Subject: [PATCH] Add sticky thread environments Carry thread-level environment selections through app-server, core turn context, unified exec, and environment manager plumbing. Update generated schema and existing tests for the new environment fields. Co-authored-by: Codex --- .../schema/json/v2/ThreadStartParams.json | 19 ++ codex-rs/app-server/src/lib.rs | 15 +- .../app-server/tests/suite/v2/skills_list.rs | 1 + codex-rs/cli/src/main.rs | 2 +- codex-rs/core/src/config/mod.rs | 123 +++++++++++++ .../core/src/context/environment_context.rs | 82 ++++++++- .../src/context/environment_context_tests.rs | 45 +++++ codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/src/session/tests.rs | 79 ++++++++ codex-rs/core/src/session/turn_context.rs | 88 ++++++++- .../core/src/tools/handlers/unified_exec.rs | 75 ++++++-- codex-rs/core/src/tools/orchestrator.rs | 5 +- .../core/src/tools/runtimes/unified_exec.rs | 36 ++-- codex-rs/core/src/tools/sandboxing.rs | 4 + codex-rs/core/src/unified_exec/mod.rs | 3 + .../core/src/unified_exec/process_manager.rs | 2 + .../tests/suite/collaboration_instructions.rs | 13 ++ codex-rs/core/tests/suite/compact.rs | 1 + codex-rs/core/tests/suite/compact_remote.rs | 2 + .../core/tests/suite/compact_resume_fork.rs | 1 + codex-rs/core/tests/suite/model_overrides.rs | 2 + codex-rs/core/tests/suite/model_switching.rs | 3 + .../core/tests/suite/model_visible_layout.rs | 1 + codex-rs/core/tests/suite/override_updates.rs | 3 + .../core/tests/suite/permissions_messages.rs | 4 + codex-rs/core/tests/suite/personality.rs | 4 + codex-rs/core/tests/suite/prompt_caching.rs | 2 + codex-rs/core/tests/suite/remote_models.rs | 2 + codex-rs/core/tests/suite/resume.rs | 1 + codex-rs/core/tests/suite/review.rs | 1 + codex-rs/core/tests/suite/unified_exec.rs | 94 ++++++++++ codex-rs/exec-server/src/client.rs | 90 ++++++++-- codex-rs/exec-server/src/client_api.rs | 16 ++ codex-rs/exec-server/src/connection.rs | 15 +- codex-rs/exec-server/src/environment.rs | 168 +++++++++++++++--- codex-rs/exec-server/src/lib.rs | 5 + codex-rs/exec-server/src/server/transport.rs | 21 ++- codex-rs/exec/src/lib.rs | 9 +- codex-rs/features/src/lib.rs | 8 + codex-rs/mcp-server/src/lib.rs | 15 +- codex-rs/tools/src/local_tool.rs | 9 + codex-rs/tools/src/local_tool_tests.rs | 2 + codex-rs/tools/src/tool_config.rs | 3 + codex-rs/tools/src/tool_registry_plan.rs | 2 + .../tools/src/tool_registry_plan_tests.rs | 1 + codex-rs/tui/src/app_command.rs | 2 + .../tui/src/chatwidget/tests/permissions.rs | 1 + codex-rs/tui/src/lib.rs | 14 +- 48 files changed, 984 insertions(+), 112 deletions(-) diff --git a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json index 9b1cba6d1a..e40fe65cc7 100644 --- a/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json +++ b/codex-rs/app-server-protocol/schema/json/v2/ThreadStartParams.json @@ -1,6 +1,10 @@ { "$schema": "http://json-schema.org/draft-07/schema#", "definitions": { + "AbsolutePathBuf": { + "description": "A path that is guaranteed to be absolute and normalized (though it is not guaranteed to be canonicalized or exist on the filesystem).\n\nIMPORTANT: When deserializing an `AbsolutePathBuf`, a base path must be set using [AbsolutePathBufGuard::new]. If no base path is set, the deserialization will fail unless the path being deserialized is already absolute.", + "type": "string" + }, "ApprovalsReviewer": { "description": "Configures who approval requests are routed to for review. Examples include sandbox escapes, blocked network access, MCP approval prompts, and ARC escalations. Defaults to `user`. `guardian_subagent` uses a carefully prompted subagent to gather relevant context and apply a risk-based decision framework before approving or denying the request.", "enum": [ @@ -114,6 +118,21 @@ "clear" ], "type": "string" + }, + "TurnEnvironmentParams": { + "properties": { + "cwd": { + "$ref": "#/definitions/AbsolutePathBuf" + }, + "environmentId": { + "type": "string" + } + }, + "required": [ + "cwd", + "environmentId" + ], + "type": "object" } }, "properties": { diff --git a/codex-rs/app-server/src/lib.rs b/codex-rs/app-server/src/lib.rs index d3e874c5c4..10f6d5f1c8 100644 --- a/codex-rs/app-server/src/lib.rs +++ b/codex-rs/app-server/src/lib.rs @@ -362,12 +362,10 @@ pub async fn run_main_with_transport( session_source: SessionSource, auth: AppServerWebsocketAuthSettings, ) -> IoResult<()> { - let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::from_env( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - ))); + let runtime_paths = ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?; let (transport_event_tx, mut transport_event_rx) = mpsc::channel::(CHANNEL_CAPACITY); let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); @@ -444,6 +442,11 @@ pub async fn run_main_with_transport( })? } }; + let exec_server_url = EnvironmentManagerArgs::from_env(runtime_paths.clone()).exec_server_url; + let environment_manager = Arc::new( + EnvironmentManager::try_new(config.environment_manager_args(exec_server_url, runtime_paths)?) + .map_err(|err| std::io::Error::new(ErrorKind::InvalidInput, err))?, + ); if let Ok(Some(err)) = check_execpolicy_for_warnings(&config.config_layer_stack).await { let (path, range) = exec_policy_warning_location(&err); diff --git a/codex-rs/app-server/tests/suite/v2/skills_list.rs b/codex-rs/app-server/tests/suite/v2/skills_list.rs index 8675b3a429..7712b54827 100644 --- a/codex-rs/app-server/tests/suite/v2/skills_list.rs +++ b/codex-rs/app-server/tests/suite/v2/skills_list.rs @@ -322,6 +322,7 @@ async fn skills_changed_notification_is_emitted_after_skill_change() -> Result<( personality: None, ephemeral: None, session_start_source: None, + environments: None, dynamic_tools: None, mock_experimental_field: None, experimental_raw_events: false, diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 9852a2cd5f..e8e9970475 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -424,7 +424,7 @@ struct AppServerCommand { #[derive(Debug, Parser)] struct ExecServerCommand { - /// Transport endpoint URL. Supported values: `ws://IP:PORT` (default). + /// Transport endpoint URL. Supported values: `ws://IP:PORT` (default) and `stdio://`. #[arg( long = "listen", value_name = "URL", diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 7ae710f83c..f2681bafdb 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -49,8 +49,13 @@ use codex_config::types::ToolSuggestDiscoverable; use codex_config::types::TuiNotificationSettings; use codex_config::types::UriBasedFileOpener; use codex_config::types::WindowsSandboxModeToml; +use codex_exec_server::ConfiguredEnvironmentManagerArgs; +use codex_exec_server::ConfiguredEnvironmentSpec; +use codex_exec_server::EnvironmentManagerArgs; +use codex_exec_server::ExecServerRuntimePaths; use codex_exec_server::ExecutorFileSystem; use codex_exec_server::LOCAL_FS; +use codex_exec_server::RemoteExecServerTransport; pub use codex_features::Feature; use codex_features::FeatureConfigSource; use codex_features::FeatureOverrides; @@ -833,6 +838,38 @@ impl Config { } } + pub fn environment_manager_args( + &self, + exec_server_url: Option, + local_runtime_paths: ExecServerRuntimePaths, + ) -> std::io::Result { + let environment_config: EnvironmentConfigToml = self + .config_layer_stack + .effective_config() + .try_into() + .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidData, err))?; + let Some(environments) = environment_config.environments.as_ref() else { + if environment_config.default_environment.is_none() { + return Ok(EnvironmentManagerArgs { + exec_server_url, + local_runtime_paths, + } + .into()); + } + return build_configured_environment_manager_args( + environment_config.default_environment, + &[], + local_runtime_paths, + ); + }; + + build_configured_environment_manager_args( + environment_config.default_environment, + environments, + local_runtime_paths, + ) + } + /// This is the preferred way to create an instance of [Config]. pub async fn load_with_cli_overrides( cli_overrides: Vec<(String, TomlValue)>, @@ -1283,6 +1320,92 @@ fn resolve_tool_suggest_config(config_toml: &ConfigToml) -> ToolSuggestConfig { ToolSuggestConfig { discoverables } } +#[derive(Debug, Deserialize, Default)] +struct EnvironmentConfigToml { + default_environment: Option, + environments: Option>, +} + +#[derive(Debug, Deserialize)] +struct EnvironmentToml { + id: String, + url: Option, + exec_server_command: Option, +} + +fn build_configured_environment_manager_args( + default_environment: Option, + environments: &[EnvironmentToml], + local_runtime_paths: ExecServerRuntimePaths, +) -> std::io::Result { + let mut configured_environments = Vec::with_capacity(environments.len()); + for environment in environments { + let id = environment.id.trim(); + if id.is_empty() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "environment id must not be empty", + )); + } + if id == "local" || id.eq_ignore_ascii_case("none") { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("environment id `{id}` is reserved"), + )); + } + + let url = environment + .url + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()); + let command = environment + .exec_server_command + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()); + let transport = match (url, command) { + (Some(url), None) => { + if !(url.starts_with("ws://") || url.starts_with("wss://")) { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("environment `{id}` url must start with ws:// or wss://"), + )); + } + RemoteExecServerTransport::WebSocket { + url: url.to_string(), + } + } + (None, Some(command)) => RemoteExecServerTransport::Command { + command: command.to_string(), + }, + (Some(_), Some(_)) => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("environment `{id}` must set only one of url or exec_server_command"), + )); + } + _ => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("environment `{id}` must set url or exec_server_command"), + )); + } + }; + configured_environments.push(ConfiguredEnvironmentSpec { + id: id.to_string(), + transport, + }); + } + + Ok(ConfiguredEnvironmentManagerArgs { + default_environment, + environments: configured_environments, + local_runtime_paths, + } + .into()) +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum PermissionConfigSyntax { Legacy, diff --git a/codex-rs/core/src/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index c4e77624f8..d3105a03f2 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -1,5 +1,6 @@ use crate::session::turn_context::TurnContext; use crate::shell::Shell; +use codex_features::Feature; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::TurnContextNetworkItem; use std::path::PathBuf; @@ -9,6 +10,7 @@ use super::ContextualUserFragment; #[derive(Debug, Clone, PartialEq)] pub(crate) struct EnvironmentContext { pub(crate) cwd: Option, + pub(crate) environments: Option>, pub(crate) shell: String, pub(crate) current_date: Option, pub(crate) timezone: Option, @@ -16,6 +18,13 @@ pub(crate) struct EnvironmentContext { pub(crate) subagents: Option, } +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct EnvironmentContextEnvironment { + pub(crate) id: String, + pub(crate) cwd: PathBuf, + pub(crate) primary: bool, +} + #[derive(Debug, Clone, PartialEq, Eq, Default)] pub(crate) struct NetworkContext { allowed_domains: Vec, @@ -39,9 +48,30 @@ impl EnvironmentContext { timezone: Option, network: Option, subagents: Option, + ) -> Self { + Self::new_with_environments( + cwd, + /*environments*/ None, + shell, + current_date, + timezone, + network, + subagents, + ) + } + + fn new_with_environments( + cwd: Option, + environments: Option>, + shell: String, + current_date: Option, + timezone: Option, + network: Option, + subagents: Option, ) -> Self { Self { cwd, + environments, shell, current_date, timezone, @@ -56,6 +86,7 @@ impl EnvironmentContext { pub(crate) fn equals_except_shell(&self, other: &EnvironmentContext) -> bool { let EnvironmentContext { cwd, + environments, current_date, timezone, network, @@ -63,6 +94,7 @@ impl EnvironmentContext { shell: _, } = other; self.cwd == *cwd + && self.environments == *environments && self.current_date == *current_date && self.timezone == *timezone && self.network == *network @@ -83,8 +115,9 @@ impl EnvironmentContext { } else { before_network }; - EnvironmentContext::new( + EnvironmentContext::new_with_environments( cwd, + after.environments.clone(), after.shell.clone(), after.current_date.clone(), after.timezone.clone(), @@ -94,8 +127,9 @@ impl EnvironmentContext { } pub(crate) fn from_turn_context(turn_context: &TurnContext, shell: &Shell) -> Self { - Self::new( + Self::new_with_environments( Some(turn_context.cwd.to_path_buf()), + Self::environments_from_turn_context(turn_context), shell.name().to_string(), turn_context.current_date.clone(), turn_context.timezone.clone(), @@ -159,6 +193,30 @@ impl EnvironmentContext { denied_domains.clone(), )) } + + fn environments_from_turn_context( + turn_context: &TurnContext, + ) -> Option> { + if !turn_context.features.enabled(Feature::MultiEnvironmentTools) { + return None; + } + let environments = turn_context.environments.as_ref()?; + if environments.is_empty() { + return None; + } + + Some( + environments + .iter() + .enumerate() + .map(|(index, environment)| EnvironmentContextEnvironment { + id: environment.environment_id.clone(), + cwd: environment.cwd.to_path_buf(), + primary: index == 0, + }) + .collect(), + ) + } } impl ContextualUserFragment for EnvironmentContext { @@ -171,6 +229,26 @@ impl ContextualUserFragment for EnvironmentContext { if let Some(cwd) = &self.cwd { lines.push(format!(" {}", cwd.to_string_lossy())); } + if let Some(environments) = &self.environments { + lines.push(" ".to_string()); + for environment in environments { + let primary = if environment.primary { + " primary=\"true\"" + } else { + "" + }; + lines.push(format!( + " ", + environment.id, primary + )); + lines.push(format!( + " {}", + environment.cwd.to_string_lossy() + )); + lines.push(" ".to_string()); + } + lines.push(" ".to_string()); + } lines.push(format!(" {}", self.shell)); if let Some(current_date) = &self.current_date { diff --git a/codex-rs/core/src/context/environment_context_tests.rs b/codex-rs/core/src/context/environment_context_tests.rs index 84f8c0d99f..3279797fb3 100644 --- a/codex-rs/core/src/context/environment_context_tests.rs +++ b/codex-rs/core/src/context/environment_context_tests.rs @@ -72,6 +72,51 @@ fn serialize_environment_context_with_network() { assert_eq!(context.render(), expected); } +#[test] +fn serialize_environment_context_with_selected_environments() { + let context = EnvironmentContext::new_with_environments( + Some(test_path_buf("/repo-primary")), + Some(vec![ + EnvironmentContextEnvironment { + id: "remote".to_string(), + cwd: test_path_buf("/repo-primary"), + primary: true, + }, + EnvironmentContextEnvironment { + id: "local".to_string(), + cwd: test_path_buf("/repo-local"), + primary: false, + }, + ]), + fake_shell_name(), + Some("2026-02-26".to_string()), + Some("America/Los_Angeles".to_string()), + /*network*/ None, + /*subagents*/ None, + ); + + let expected = format!( + r#" + {primary} + + + {primary} + + + {local} + + + bash + 2026-02-26 + America/Los_Angeles +"#, + primary = test_path_buf("/repo-primary").display(), + local = test_path_buf("/repo-local").display() + ); + + assert_eq!(context.render(), expected); +} + #[test] fn serialize_read_only_environment_context() { let context = EnvironmentContext::new( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4f27edb5bd..da8daf1f21 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -110,6 +110,7 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnContextItem; use codex_protocol::protocol::TurnContextNetworkItem; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::W3cTraceContext; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -629,7 +630,6 @@ impl Codex { persist_extended_history, inherited_shell_snapshot, user_shell_override, - environment_selections, }; // Generate a unique ID for the lifetime of this Codex session. diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 5855b62141..d921d807b8 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3944,6 +3944,36 @@ async fn turn_environment_selection_sets_primary_environment() { )); assert_eq!(turn_context.cwd.as_path(), selected_cwd.as_path()); assert_eq!(turn_context.config.cwd.as_path(), selected_cwd.as_path()); + + let primary_environment = turn_context + .primary_environment() + .expect("primary environment should resolve"); + assert_eq!(primary_environment.environment_id.as_deref(), Some("local")); + assert!(std::sync::Arc::ptr_eq( + &primary_environment.environment, + &turn_environments[0].environment + )); + assert_eq!(primary_environment.cwd, selected_cwd); + let tool_environment = turn_context + .environment_for_tool(Some("local")) + .expect("tool environment should resolve by id"); + assert_eq!(tool_environment.environment_id.as_deref(), Some("local")); + assert!(std::sync::Arc::ptr_eq( + &tool_environment.environment, + &turn_environments[0].environment + )); + let default_tool_environment = turn_context + .environment_for_tool(/*environment_id*/ None) + .expect("tool environment should default to primary"); + assert_eq!( + default_tool_environment.environment_id.as_deref(), + Some("local") + ); + assert_eq!( + turn_context + .resolve_path_for_environment(&primary_environment, Some("relative/path".to_string())), + selected_cwd.join("relative/path") + ); } #[tokio::test] @@ -3989,6 +4019,15 @@ async fn multiple_turn_environment_selections_use_first_as_primary_environment() )); assert_eq!(turn_context.cwd, first_cwd); assert_eq!(turn_context.config.cwd, first_cwd); + + let primary_environment = turn_context + .primary_environment() + .expect("primary environment should resolve"); + assert_eq!(primary_environment.cwd, first_cwd); + let tool_environment = turn_context + .environment_for_tool(Some("local")) + .expect("tool environment should resolve by id"); + assert_eq!(tool_environment.cwd, first_cwd); } #[tokio::test] @@ -4015,6 +4054,46 @@ async fn empty_turn_environment_selection_clears_primary_environment() { .len(), 0 ); + assert!(turn_context.primary_environment().is_none()); + assert!( + turn_context + .environment_for_tool(/*environment_id*/ None) + .is_none() + ); + assert!(turn_context.environment_for_tool(Some("local")).is_none()); +} + +#[tokio::test] +async fn primary_environment_falls_back_to_compatibility_projection() { + let (_session, turn_context, _rx) = make_session_and_context_with_rx().await; + + let primary_environment = turn_context + .primary_environment() + .expect("primary environment should resolve from compatibility fields"); + assert_eq!(primary_environment.environment_id.as_deref(), None); + assert!(std::sync::Arc::ptr_eq( + &primary_environment.environment, + turn_context + .environment + .as_ref() + .expect("compatibility environment should be set") + )); + assert_eq!(primary_environment.cwd, turn_context.cwd); + + let default_tool_environment = turn_context + .environment_for_tool(/*environment_id*/ None) + .expect("tool environment should resolve from compatibility fields"); + assert_eq!(default_tool_environment.environment_id.as_deref(), None); + assert!(std::sync::Arc::ptr_eq( + &default_tool_environment.environment, + &primary_environment.environment + )); + assert!(turn_context.environment_for_tool(Some("local")).is_none()); + assert_eq!( + turn_context + .resolve_path_for_environment(&primary_environment, Some("relative/path".to_string())), + turn_context.cwd.join("relative/path") + ); } #[tokio::test] diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index b44a416520..243f5b9327 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -33,6 +33,24 @@ pub(crate) struct TurnEnvironment { pub(crate) cwd: AbsolutePathBuf, } +#[derive(Clone, Debug)] +#[allow(dead_code)] +pub(crate) struct SelectedTurnEnvironment { + pub(crate) environment_id: Option, + pub(crate) environment: Arc, + pub(crate) cwd: AbsolutePathBuf, +} + +impl TurnEnvironment { + fn selected(&self) -> SelectedTurnEnvironment { + SelectedTurnEnvironment { + environment_id: Some(self.environment_id.clone()), + environment: Arc::clone(&self.environment), + cwd: self.cwd.clone(), + } + } +} + /// The context needed for a single turn of the thread. #[derive(Debug)] pub(crate) struct TurnContext { @@ -216,14 +234,67 @@ impl TurnContext { .map_or_else(|| self.cwd.clone(), |path| self.cwd.join(path)) } + #[allow(dead_code)] + pub(crate) fn primary_environment(&self) -> Option { + self.environments + .as_ref() + .and_then(|environments| environments.first()) + .map(TurnEnvironment::selected) + .or_else(|| { + self.environment + .as_ref() + .map(|environment| SelectedTurnEnvironment { + environment_id: None, + environment: Arc::clone(environment), + cwd: self.cwd.clone(), + }) + }) + } + + #[allow(dead_code)] + pub(crate) fn environment_for_tool( + &self, + environment_id: Option<&str>, + ) -> Option { + match environment_id { + Some(environment_id) => self + .environments + .as_ref()? + .iter() + .find(|environment| environment.environment_id == environment_id) + .map(TurnEnvironment::selected), + None => self.primary_environment(), + } + } + + #[allow(dead_code)] + pub(crate) fn resolve_path_for_environment( + &self, + environment: &SelectedTurnEnvironment, + path: Option, + ) -> AbsolutePathBuf { + path.as_ref().map_or_else( + || environment.cwd.clone(), + |path| environment.cwd.join(path), + ) + } + pub(crate) fn file_system_sandbox_context( &self, additional_permissions: Option, + ) -> FileSystemSandboxContext { + self.file_system_sandbox_context_for_cwd(&self.cwd, additional_permissions) + } + + pub(crate) fn file_system_sandbox_context_for_cwd( + &self, + cwd: &AbsolutePathBuf, + additional_permissions: Option, ) -> FileSystemSandboxContext { FileSystemSandboxContext { sandbox_policy: self.sandbox_policy.get().clone(), - sandbox_policy_cwd: Some(self.cwd.clone()), - file_system_sandbox_policy: self.non_legacy_file_system_sandbox_policy(), + sandbox_policy_cwd: Some(cwd.clone()), + file_system_sandbox_policy: self.non_legacy_file_system_sandbox_policy(cwd), windows_sandbox_level: self.windows_sandbox_level, windows_sandbox_private_desktop: self .config @@ -234,15 +305,16 @@ impl TurnContext { } } - fn non_legacy_file_system_sandbox_policy(&self) -> Option { + fn non_legacy_file_system_sandbox_policy( + &self, + cwd: &AbsolutePathBuf, + ) -> Option { // Omit the derived split filesystem policy when it is equivalent to // the legacy sandbox policy. This keeps turn-context payloads stable // while both fields exist; once callers consume only the split policy, // this comparison and the legacy projection should go away. - let legacy_file_system_sandbox_policy = FileSystemSandboxPolicy::from_legacy_sandbox_policy( - self.sandbox_policy.get(), - &self.cwd, - ); + let legacy_file_system_sandbox_policy = + FileSystemSandboxPolicy::from_legacy_sandbox_policy(self.sandbox_policy.get(), cwd); (self.file_system_sandbox_policy != legacy_file_system_sandbox_policy) .then(|| self.file_system_sandbox_policy.clone()) } @@ -263,7 +335,7 @@ impl TurnContext { approval_policy: self.approval_policy.value(), sandbox_policy: self.sandbox_policy.get().clone(), network: self.turn_context_network_item(), - file_system_sandbox_policy: self.non_legacy_file_system_sandbox_policy(), + file_system_sandbox_policy: self.non_legacy_file_system_sandbox_policy(&self.cwd), model: self.model_info.slug.clone(), personality: self.personality, collaboration_mode: Some(self.collaboration_mode.clone()), diff --git a/codex-rs/core/src/tools/handlers/unified_exec.rs b/codex-rs/core/src/tools/handlers/unified_exec.rs index fd04d42f2e..33f9bc6bf3 100644 --- a/codex-rs/core/src/tools/handlers/unified_exec.rs +++ b/codex-rs/core/src/tools/handlers/unified_exec.rs @@ -64,6 +64,12 @@ pub(crate) struct ExecCommandArgs { prefix_rule: Option>, } +#[derive(Debug, Deserialize)] +struct ToolEnvironmentArgs { + #[serde(default, alias = "environmentId")] + environment_id: Option, +} + #[derive(Debug, Deserialize)] struct WriteStdinArgs { // The model is trained on `session_id`. @@ -181,21 +187,47 @@ impl ToolHandler for UnifiedExecHandler { } }; - let Some(environment) = turn.environment.as_ref() else { + let Some(default_environment) = turn.environment_for_tool(/*environment_id*/ None) else { return Err(FunctionCallError::RespondToModel( "unified exec is unavailable in this session".to_string(), )); }; - let fs = environment.get_filesystem(); let manager: &UnifiedExecProcessManager = &session.services.unified_exec_manager; let context = UnifiedExecContext::new(session.clone(), turn.clone(), call_id.clone()); let response = match tool_name.name.as_str() { "exec_command" => { - let cwd = resolve_workdir_base_path(&arguments, &context.turn.cwd)?; + let selector: ToolEnvironmentArgs = parse_arguments(&arguments)?; + let requested_environment_id = if context + .turn + .features + .enabled(Feature::MultiEnvironmentTools) + { + selector.environment_id.as_deref() + } else { + None + }; + let should_intercept_apply_patch = requested_environment_id.is_none(); + let tool_environment = context + .turn + .environment_for_tool(requested_environment_id) + .ok_or_else(|| { + FunctionCallError::RespondToModel(match requested_environment_id { + Some(environment_id) => { + format!( + "environment `{environment_id}` is unavailable in this session" + ) + } + None => "unified exec is unavailable in this session".to_string(), + }) + })?; + let fs = tool_environment.environment.get_filesystem(); + let cwd = resolve_workdir_base_path(&arguments, &tool_environment.cwd)?; let args: ExecCommandArgs = parse_arguments_with_base_path(&arguments, &cwd)?; - let workdir = context.turn.resolve_path(args.workdir.clone()); + let workdir = context + .turn + .resolve_path_for_environment(&tool_environment, args.workdir.clone()); maybe_emit_implicit_skill_invocation( session.as_ref(), context.turn.as_ref(), @@ -211,7 +243,8 @@ impl ToolHandler for UnifiedExecHandler { turn.tools_config.allow_login_shell, ) .map_err(FunctionCallError::RespondToModel)?; - let command_for_display = codex_shell_command::parse_command::shlex_join(&command); + let command_for_display = + codex_shell_command::parse_command::shlex_join(&command); let ExecCommandArgs { workdir, @@ -259,7 +292,11 @@ impl ToolHandler for UnifiedExecHandler { let workdir = workdir.filter(|value| !value.is_empty()); - let workdir = workdir.map(|dir| context.turn.resolve_path(Some(dir))); + let workdir = workdir.map(|dir| { + context + .turn + .resolve_path_for_environment(&tool_environment, Some(dir)) + }); let cwd = workdir.clone().unwrap_or(cwd); let normalized_additional_permissions = match implicit_granted_permissions( sandbox_permissions, @@ -286,17 +323,18 @@ impl ToolHandler for UnifiedExecHandler { } }; - if let Some(output) = intercept_apply_patch( - &command, - &cwd, - fs.as_ref(), - context.session.clone(), - context.turn.clone(), - Some(&tracker), - &context.call_id, - &tool_name.name, - ) - .await? + if should_intercept_apply_patch + && let Some(output) = intercept_apply_patch( + &command, + &cwd, + fs.as_ref(), + context.session.clone(), + context.turn.clone(), + Some(&tracker), + &context.call_id, + &tool_name.name, + ) + .await? { manager.release_process_id(process_id).await; return Ok(ExecCommandToolOutput { @@ -320,6 +358,8 @@ impl ToolHandler for UnifiedExecHandler { command, hook_command: args.cmd, process_id, + environment_id: tool_environment.environment_id.clone(), + environment: tool_environment.environment, yield_time_ms, max_output_tokens, workdir, @@ -363,6 +403,7 @@ impl ToolHandler for UnifiedExecHandler { } } "write_stdin" => { + let _ = default_environment; let args: WriteStdinArgs = parse_arguments(&arguments)?; let response = manager .write_stdin(WriteStdinRequest { diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index c1b960bc8d..e57bd09d0c 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -200,6 +200,7 @@ impl ToolOrchestrator { // Platform-specific flag gating is handled by SandboxManager::select_initial. let use_legacy_landlock = turn_ctx.features.use_legacy_landlock(); + let sandbox_cwd = tool.sandbox_cwd(req, turn_ctx); let initial_attempt = SandboxAttempt { sandbox: initial_sandbox, policy: &turn_ctx.sandbox_policy, @@ -207,7 +208,7 @@ impl ToolOrchestrator { network_policy: turn_ctx.network_sandbox_policy, enforce_managed_network: managed_network_active, manager: &self.sandbox, - sandbox_cwd: &turn_ctx.cwd, + sandbox_cwd, codex_linux_sandbox_exe: turn_ctx.codex_linux_sandbox_exe.as_ref(), use_legacy_landlock, windows_sandbox_level: turn_ctx.windows_sandbox_level, @@ -348,7 +349,7 @@ impl ToolOrchestrator { network_policy: turn_ctx.network_sandbox_policy, enforce_managed_network: managed_network_active, manager: &self.sandbox, - sandbox_cwd: &turn_ctx.cwd, + sandbox_cwd, codex_linux_sandbox_exe: None, use_legacy_landlock, windows_sandbox_level: turn_ctx.windows_sandbox_level, diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index fdd434e542..c0f1da7aed 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -34,6 +34,7 @@ use crate::unified_exec::NoopSpawnLifecycle; use crate::unified_exec::UnifiedExecError; use crate::unified_exec::UnifiedExecProcess; use crate::unified_exec::UnifiedExecProcessManager; +use codex_exec_server::Environment; use codex_network_proxy::NetworkProxy; use codex_protocol::error::CodexErr; use codex_protocol::error::SandboxErr; @@ -45,6 +46,7 @@ use codex_tools::UnifiedExecShellMode; use codex_utils_absolute_path::AbsolutePathBuf; use futures::future::BoxFuture; use std::collections::HashMap; +use std::sync::Arc; /// Request payload used by the unified-exec runtime after approvals and /// sandbox preferences have been resolved for the current turn. @@ -53,6 +55,8 @@ pub struct UnifiedExecRequest { pub command: Vec, pub hook_command: String, pub process_id: i32, + pub environment_id: Option, + pub environment: Arc, pub cwd: AbsolutePathBuf, pub env: HashMap, pub exec_server_env_config: Option, @@ -72,6 +76,7 @@ pub struct UnifiedExecRequest { #[derive(serde::Serialize, Clone, Debug, Eq, PartialEq, Hash)] pub struct UnifiedExecApprovalKey { pub command: Vec, + pub environment_id: Option, pub cwd: AbsolutePathBuf, pub tty: bool, pub sandbox_permissions: SandboxPermissions, @@ -111,6 +116,7 @@ impl Approvable for UnifiedExecRuntime<'_> { fn approval_keys(&self, req: &UnifiedExecRequest) -> Vec { vec![UnifiedExecApprovalKey { command: canonicalize_command_for_approval(&req.command), + environment_id: req.environment_id.clone(), cwd: req.cwd.clone(), tty: req.tty, sandbox_permissions: req.sandbox_permissions, @@ -198,6 +204,14 @@ impl Approvable for UnifiedExecRuntime<'_> { } impl<'a> ToolRuntime for UnifiedExecRuntime<'a> { + fn sandbox_cwd<'b>( + &self, + req: &'b UnifiedExecRequest, + _turn_ctx: &'b crate::session::turn_context::TurnContext, + ) -> &'b AbsolutePathBuf { + &req.cwd + } + fn network_approval_spec( &self, req: &UnifiedExecRequest, @@ -219,11 +233,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt ) -> Result { let base_command = &req.command; let session_shell = ctx.session.user_shell(); - let environment_is_remote = ctx - .turn - .environment - .as_ref() - .is_some_and(|environment| environment.is_remote()); + let environment_is_remote = req.environment.is_remote(); let command = if environment_is_remote { base_command.to_vec() } else { @@ -267,12 +277,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt .await? { Some(prepared) => { - let Some(environment) = ctx.turn.environment.as_ref() else { - return Err(ToolError::Rejected( - "exec_command is unavailable in this session".to_string(), - )); - }; - if environment.is_remote() { + if req.environment.is_remote() { return Err(ToolError::Rejected( "unified_exec zsh-fork is not supported when exec_server_url is configured".to_string(), )); @@ -284,7 +289,7 @@ impl<'a> ToolRuntime for UnifiedExecRunt &prepared.exec_request, req.tty, prepared.spawn_lifecycle, - environment.as_ref(), + req.environment.as_ref(), ) .await .map_err(|err| match err { @@ -315,18 +320,13 @@ impl<'a> ToolRuntime for UnifiedExecRunt .env_for(command, options, req.network.as_ref()) .map_err(|err| ToolError::Codex(err.into()))?; exec_env.exec_server_env_config = req.exec_server_env_config.clone(); - let Some(environment) = ctx.turn.environment.as_ref() else { - return Err(ToolError::Rejected( - "exec_command is unavailable in this session".to_string(), - )); - }; self.manager .open_session_with_exec_env( req.process_id, &exec_env, req.tty, Box::new(NoopSpawnLifecycle), - environment.as_ref(), + req.environment.as_ref(), ) .await .map_err(|err| match err { diff --git a/codex-rs/core/src/tools/sandboxing.rs b/codex-rs/core/src/tools/sandboxing.rs index 87bef4617c..d6c17118e5 100644 --- a/codex-rs/core/src/tools/sandboxing.rs +++ b/codex-rs/core/src/tools/sandboxing.rs @@ -329,6 +329,10 @@ pub(crate) trait ToolRuntime: Approvable + Sandboxable { None } + fn sandbox_cwd<'a>(&self, _req: &'a Req, turn_ctx: &'a TurnContext) -> &'a AbsolutePathBuf { + &turn_ctx.cwd + } + async fn run( &mut self, req: &Req, diff --git a/codex-rs/core/src/unified_exec/mod.rs b/codex-rs/core/src/unified_exec/mod.rs index 01c8914bc2..06dd6fb48a 100644 --- a/codex-rs/core/src/unified_exec/mod.rs +++ b/codex-rs/core/src/unified_exec/mod.rs @@ -27,6 +27,7 @@ use std::collections::HashSet; use std::sync::Arc; use std::sync::Weak; +use codex_exec_server::Environment; use codex_network_proxy::NetworkProxy; use codex_protocol::models::PermissionProfile; use codex_utils_absolute_path::AbsolutePathBuf; @@ -90,6 +91,8 @@ pub(crate) struct ExecCommandRequest { pub command: Vec, pub hook_command: String, pub process_id: i32, + pub environment_id: Option, + pub environment: Arc, pub yield_time_ms: u64, pub max_output_tokens: Option, pub workdir: Option, diff --git a/codex-rs/core/src/unified_exec/process_manager.rs b/codex-rs/core/src/unified_exec/process_manager.rs index 5181359698..d401946819 100644 --- a/codex-rs/core/src/unified_exec/process_manager.rs +++ b/codex-rs/core/src/unified_exec/process_manager.rs @@ -753,6 +753,8 @@ impl UnifiedExecProcessManager { command: request.command.clone(), hook_command: request.hook_command.clone(), process_id: request.process_id, + environment_id: request.environment_id.clone(), + environment: Arc::clone(&request.environment), cwd, env, exec_server_env_config: Some(exec_server_env_config), diff --git a/codex-rs/core/tests/suite/collaboration_instructions.rs b/codex-rs/core/tests/suite/collaboration_instructions.rs index 3f8d2378ad..610671561b 100644 --- a/codex-rs/core/tests/suite/collaboration_instructions.rs +++ b/codex-rs/core/tests/suite/collaboration_instructions.rs @@ -135,6 +135,7 @@ async fn user_input_includes_collaboration_instructions_after_override() -> Resu service_tier: None, collaboration_mode: Some(collaboration_mode), personality: None, + environments: None, }) .await?; @@ -236,6 +237,7 @@ async fn override_then_next_turn_uses_updated_collaboration_instructions() -> Re service_tier: None, collaboration_mode: Some(collaboration_mode), personality: None, + environments: None, }) .await?; @@ -290,6 +292,7 @@ async fn user_turn_overrides_collaboration_instructions_after_override() -> Resu service_tier: None, collaboration_mode: Some(base_mode), personality: None, + environments: None, }) .await?; @@ -362,6 +365,7 @@ async fn collaboration_mode_update_emits_new_instruction_message() -> Result<()> service_tier: None, collaboration_mode: Some(collab_mode_with_instructions(Some(first_text))), personality: None, + environments: None, }) .await?; @@ -391,6 +395,7 @@ async fn collaboration_mode_update_emits_new_instruction_message() -> Result<()> service_tier: None, collaboration_mode: Some(collab_mode_with_instructions(Some(second_text))), personality: None, + environments: None, }) .await?; @@ -449,6 +454,7 @@ async fn collaboration_mode_update_noop_does_not_append() -> Result<()> { service_tier: None, collaboration_mode: Some(collab_mode_with_instructions(Some(collab_text))), personality: None, + environments: None, }) .await?; @@ -478,6 +484,7 @@ async fn collaboration_mode_update_noop_does_not_append() -> Result<()> { service_tier: None, collaboration_mode: Some(collab_mode_with_instructions(Some(collab_text))), personality: None, + environments: None, }) .await?; @@ -538,6 +545,7 @@ async fn collaboration_mode_update_emits_new_instruction_message_when_mode_chang Some(default_text), )), personality: None, + environments: None, }) .await?; @@ -570,6 +578,7 @@ async fn collaboration_mode_update_emits_new_instruction_message_when_mode_chang Some(plan_text), )), personality: None, + environments: None, }) .await?; @@ -631,6 +640,7 @@ async fn collaboration_mode_update_noop_does_not_append_when_mode_is_unchanged() Some(collab_text), )), personality: None, + environments: None, }) .await?; @@ -663,6 +673,7 @@ async fn collaboration_mode_update_noop_does_not_append_when_mode_is_unchanged() Some(collab_text), )), personality: None, + environments: None, }) .await?; @@ -727,6 +738,7 @@ async fn resume_replays_collaboration_instructions() -> Result<()> { service_tier: None, collaboration_mode: Some(collab_mode_with_instructions(Some(collab_text))), personality: None, + environments: None, }) .await?; @@ -801,6 +813,7 @@ async fn empty_collaboration_instructions_are_ignored() -> Result<()> { }, }), personality: None, + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index c37326ffc9..f37e6f307e 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -3090,6 +3090,7 @@ async fn snapshot_request_shape_pre_turn_compaction_including_incoming_user_mess service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await .expect("override turn context"); diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 58ed42c288..98ef96b199 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -2110,6 +2110,7 @@ async fn snapshot_request_shape_remote_pre_turn_compaction_including_incoming_us service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; } @@ -2225,6 +2226,7 @@ async fn snapshot_request_shape_remote_pre_turn_compaction_strips_incoming_model service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; codex diff --git a/codex-rs/core/tests/suite/compact_resume_fork.rs b/codex-rs/core/tests/suite/compact_resume_fork.rs index 48d359529b..01a7105264 100644 --- a/codex-rs/core/tests/suite/compact_resume_fork.rs +++ b/codex-rs/core/tests/suite/compact_resume_fork.rs @@ -589,6 +589,7 @@ async fn snapshot_rollback_followup_turn_trims_context_updates() -> Result<()> { }, }), personality: None, + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/model_overrides.rs b/codex-rs/core/tests/suite/model_overrides.rs index f7fe7cc322..aba3ed3bfb 100644 --- a/codex-rs/core/tests/suite/model_overrides.rs +++ b/codex-rs/core/tests/suite/model_overrides.rs @@ -37,6 +37,7 @@ async fn override_turn_context_does_not_persist_when_config_exists() { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await .expect("submit override"); @@ -75,6 +76,7 @@ async fn override_turn_context_does_not_create_config_file() { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await .expect("submit override"); diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs index 54c2d1035f..793b0f5e7e 100644 --- a/codex-rs/core/tests/suite/model_switching.rs +++ b/codex-rs/core/tests/suite/model_switching.rs @@ -154,6 +154,7 @@ async fn model_change_appends_model_instructions_developer_message() -> Result<( service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -253,6 +254,7 @@ async fn model_and_personality_change_only_appends_model_instructions() -> Resul service_tier: None, collaboration_mode: None, personality: Some(Personality::Pragmatic), + environments: None, }) .await?; @@ -1045,6 +1047,7 @@ async fn model_switch_to_smaller_model_updates_token_context_window() -> Result< service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/model_visible_layout.rs b/codex-rs/core/tests/suite/model_visible_layout.rs index c93294121c..175979725a 100644 --- a/codex-rs/core/tests/suite/model_visible_layout.rs +++ b/codex-rs/core/tests/suite/model_visible_layout.rs @@ -465,6 +465,7 @@ async fn snapshot_model_visible_layout_resume_override_matches_rollout_model() - service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; resumed diff --git a/codex-rs/core/tests/suite/override_updates.rs b/codex-rs/core/tests/suite/override_updates.rs index 13ace1fe6d..193accbccb 100644 --- a/codex-rs/core/tests/suite/override_updates.rs +++ b/codex-rs/core/tests/suite/override_updates.rs @@ -125,6 +125,7 @@ async fn override_turn_context_without_user_turn_does_not_record_permissions_upd service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -167,6 +168,7 @@ async fn override_turn_context_without_user_turn_does_not_record_environment_upd service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -206,6 +208,7 @@ async fn override_turn_context_without_user_turn_does_not_record_collaboration_u service_tier: None, collaboration_mode: Some(collaboration_mode), personality: None, + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/permissions_messages.rs b/codex-rs/core/tests/suite/permissions_messages.rs index fb0acf7519..8c7bc22e8a 100644 --- a/codex-rs/core/tests/suite/permissions_messages.rs +++ b/codex-rs/core/tests/suite/permissions_messages.rs @@ -112,6 +112,7 @@ async fn permissions_message_added_on_override_change() -> Result<()> { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -244,6 +245,7 @@ async fn permissions_message_omitted_when_disabled() -> Result<()> { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -332,6 +334,7 @@ async fn resume_replays_permissions_messages() -> Result<()> { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -437,6 +440,7 @@ async fn resume_and_fork_append_permissions_messages() -> Result<()> { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/personality.rs b/codex-rs/core/tests/suite/personality.rs index 7c72478ace..944eba13af 100644 --- a/codex-rs/core/tests/suite/personality.rs +++ b/codex-rs/core/tests/suite/personality.rs @@ -359,6 +359,7 @@ async fn user_turn_personality_some_adds_update_message() -> anyhow::Result<()> service_tier: None, collaboration_mode: None, personality: Some(Personality::Friendly), + environments: None, }) .await?; @@ -466,6 +467,7 @@ async fn user_turn_personality_same_value_does_not_add_update_message() -> anyho service_tier: None, collaboration_mode: None, personality: Some(Personality::Pragmatic), + environments: None, }) .await?; @@ -586,6 +588,7 @@ async fn user_turn_personality_skips_if_feature_disabled() -> anyhow::Result<()> service_tier: None, collaboration_mode: None, personality: Some(Personality::Pragmatic), + environments: None, }) .await?; @@ -867,6 +870,7 @@ async fn user_turn_personality_remote_model_template_includes_update_message() - service_tier: None, collaboration_mode: None, personality: Some(Personality::Friendly), + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/prompt_caching.rs b/codex-rs/core/tests/suite/prompt_caching.rs index cee43e9820..7f2d090e30 100644 --- a/codex-rs/core/tests/suite/prompt_caching.rs +++ b/codex-rs/core/tests/suite/prompt_caching.rs @@ -451,6 +451,7 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() -> an service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -536,6 +537,7 @@ async fn override_before_first_turn_emits_environment_context() -> anyhow::Resul service_tier: None, collaboration_mode: Some(collaboration_mode), personality: None, + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/remote_models.rs b/codex-rs/core/tests/suite/remote_models.rs index e62fa0c323..d54e4b8b6e 100644 --- a/codex-rs/core/tests/suite/remote_models.rs +++ b/codex-rs/core/tests/suite/remote_models.rs @@ -602,6 +602,7 @@ async fn remote_models_remote_model_uses_unified_exec() -> Result<()> { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; @@ -847,6 +848,7 @@ async fn remote_models_apply_remote_base_instructions() -> Result<()> { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; diff --git a/codex-rs/core/tests/suite/resume.rs b/codex-rs/core/tests/suite/resume.rs index febeb926b1..6b6a30fd30 100644 --- a/codex-rs/core/tests/suite/resume.rs +++ b/codex-rs/core/tests/suite/resume.rs @@ -435,6 +435,7 @@ async fn resume_model_switch_is_not_duplicated_after_pre_turn_override() -> Resu service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await?; resumed diff --git a/codex-rs/core/tests/suite/review.rs b/codex-rs/core/tests/suite/review.rs index faf469c981..9baae7c48f 100644 --- a/codex-rs/core/tests/suite/review.rs +++ b/codex-rs/core/tests/suite/review.rs @@ -846,6 +846,7 @@ async fn review_uses_overridden_cwd_for_base_branch_merge_base() { service_tier: None, collaboration_mode: None, personality: None, + environments: None, }) .await .unwrap(); diff --git a/codex-rs/core/tests/suite/unified_exec.rs b/codex-rs/core/tests/suite/unified_exec.rs index 58dbad7b1d..d4612b6d40 100644 --- a/codex-rs/core/tests/suite/unified_exec.rs +++ b/codex-rs/core/tests/suite/unified_exec.rs @@ -12,6 +12,7 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ExecCommandSource; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use core_test_support::assert_regex_match; use core_test_support::process::process_is_alive; @@ -339,6 +340,99 @@ async fn unified_exec_intercepts_apply_patch_exec_command() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unified_exec_environment_id_targets_non_primary_environment() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + skip_if_windows!(Ok(())); + + let server = start_mock_server().await; + let mut builder = test_codex() + .with_exec_server_url("ws://127.0.0.1:9") + .with_config(|config| { + config + .features + .enable(Feature::UnifiedExec) + .expect("test config should allow unified exec"); + config + .features + .enable(Feature::MultiEnvironmentTools) + .expect("test config should allow multi-environment tools"); + }); + let test = builder.build(&server).await?; + let selected_local_cwd = create_workspace_directory(&test, "selected-local").await?; + let primary_remote_cwd = test.config.cwd.join("primary-remote"); + + let call_id = "uexec-env-local"; + let args = json!({ + "cmd": "pwd", + "environment_id": "local", + "yield_time_ms": 5_000, + }); + let responses = vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "exec_command", &serde_json::to_string(&args)?), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_assistant_message("msg-1", "done"), + ev_completed("resp-2"), + ]), + ]; + let request_log = mount_sse_sequence(&server, responses).await; + + test.submit_turn_with_environments( + "print the selected local cwd", + Some(vec![ + TurnEnvironmentSelection { + environment_id: "remote".to_string(), + cwd: primary_remote_cwd, + }, + TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: selected_local_cwd.clone().try_into()?, + }, + ]), + ) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let bodies = request_log + .requests() + .into_iter() + .map(|request| request.body_json()) + .collect::>(); + let request_text = serde_json::to_string(&bodies)?; + assert!( + request_text.contains(""), + "environment list should be visible to the model: {request_text}" + ); + assert!( + request_text.contains("id=\"local\""), + "local environment id should be visible to the model: {request_text}" + ); + assert!( + request_text.contains("\"environment_id\""), + "exec_command schema should expose environment_id when gated: {request_text}" + ); + + let outputs = collect_tool_outputs(&bodies)?; + let output = outputs + .get(call_id) + .expect("missing exec_command environment-id output"); + assert_eq!( + output.output.trim(), + selected_local_cwd.display().to_string() + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn unified_exec_emits_exec_command_begin_event() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 375571b779..a79171f438 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -14,13 +14,17 @@ use tokio::sync::OnceCell; use tokio::sync::mpsc; use tokio::sync::watch; +use tokio::process::Child; +use tokio::process::Command; use tokio::time::timeout; use tokio_tungstenite::connect_async; use tracing::debug; use crate::ProcessId; +use crate::client_api::CommandExecServerConnectArgs; use crate::client_api::ExecServerClientConnectOptions; use crate::client_api::RemoteExecServerConnectArgs; +use crate::client_api::RemoteExecServerTransport; use crate::connection::JsonRpcConnection; use crate::process::ExecProcessEvent; use crate::process::ExecProcessEventLog; @@ -102,6 +106,16 @@ impl From for ExecServerClientConnectOptions { } } +impl From for ExecServerClientConnectOptions { + fn from(value: CommandExecServerConnectArgs) -> Self { + Self { + client_name: value.client_name, + initialize_timeout: value.initialize_timeout, + resume_session_id: value.resume_session_id, + } + } +} + impl RemoteExecServerConnectArgs { pub fn new(websocket_url: String, client_name: String) -> Self { Self { @@ -162,11 +176,15 @@ struct Inner { http_body_stream_next_id: AtomicU64, session_id: std::sync::RwLock>, reader_task: tokio::task::JoinHandle<()>, + child: Option, } impl Drop for Inner { fn drop(&mut self) { self.reader_task.abort(); + if let Some(child) = &mut self.child { + let _ = child.start_kill(); + } } } @@ -177,14 +195,14 @@ pub struct ExecServerClient { #[derive(Clone)] pub(crate) struct LazyRemoteExecServerClient { - websocket_url: String, + transport: RemoteExecServerTransport, client: Arc>, } impl LazyRemoteExecServerClient { - pub(crate) fn new(websocket_url: String) -> Self { + pub(crate) fn new(transport: RemoteExecServerTransport) -> Self { Self { - websocket_url, + transport, client: Arc::new(OnceCell::new()), } } @@ -192,14 +210,27 @@ impl LazyRemoteExecServerClient { 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 + match &self.transport { + RemoteExecServerTransport::WebSocket { url } => { + ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { + websocket_url: 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 + } + RemoteExecServerTransport::Command { command } => { + ExecServerClient::connect_command(CommandExecServerConnectArgs { + command: command.clone(), + client_name: "codex-environment".to_string(), + initialize_timeout: Duration::from_secs(5), + resume_session_id: None, + }) + .await + } + } }) .await .cloned() @@ -255,6 +286,39 @@ impl ExecServerClient { format!("exec-server websocket {websocket_url}"), ), args.into(), + /*child*/ None, + ) + .await + } + + pub async fn connect_command( + args: CommandExecServerConnectArgs, + ) -> Result { + let command = args.command.clone(); + let mut child = Command::new("sh") + .arg("-lc") + .arg(&command) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::inherit()) + .kill_on_drop(true) + .spawn() + .map_err(ExecServerError::Spawn)?; + let stdout = child.stdout.take().ok_or_else(|| { + ExecServerError::Protocol(format!("exec-server command `{command}` has no stdout")) + })?; + let stdin = child.stdin.take().ok_or_else(|| { + ExecServerError::Protocol(format!("exec-server command `{command}` has no stdin")) + })?; + + Self::connect( + JsonRpcConnection::from_stdio( + stdout, + stdin, + format!("exec-server command `{command}`"), + ), + args.into(), + Some(child), ) .await } @@ -410,6 +474,7 @@ impl ExecServerClient { async fn connect( connection: JsonRpcConnection, options: ExecServerClientConnectOptions, + child: Option, ) -> Result { let (rpc_client, mut events_rx) = RpcClient::new(connection); let inner = Arc::new_cyclic(|weak| { @@ -455,6 +520,7 @@ impl ExecServerClient { http_body_stream_next_id: AtomicU64::new(1), session_id: std::sync::RwLock::new(None), reader_task, + child, } }); @@ -949,6 +1015,7 @@ mod tests { "test-exec-server-client".to_string(), ), ExecServerClientConnectOptions::default(), + /*child*/ None, ) .await .expect("client should connect"); @@ -1092,6 +1159,7 @@ mod tests { "test-exec-server-client".to_string(), ), ExecServerClientConnectOptions::default(), + /*child*/ None, ) .await .expect("client should connect"); diff --git a/codex-rs/exec-server/src/client_api.rs b/codex-rs/exec-server/src/client_api.rs index ac4371e2ea..e2aa614e07 100644 --- a/codex-rs/exec-server/src/client_api.rs +++ b/codex-rs/exec-server/src/client_api.rs @@ -8,6 +8,13 @@ pub struct ExecServerClientConnectOptions { pub resume_session_id: Option, } +/// Transport used to connect to a remote exec-server. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RemoteExecServerTransport { + WebSocket { url: String }, + Command { command: String }, +} + /// WebSocket connection arguments for a remote exec-server. #[derive(Debug, Clone, PartialEq, Eq)] pub struct RemoteExecServerConnectArgs { @@ -17,3 +24,12 @@ pub struct RemoteExecServerConnectArgs { pub initialize_timeout: Duration, pub resume_session_id: Option, } + +/// Command-spawn connection arguments for a remote exec-server over stdio. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CommandExecServerConnectArgs { + pub command: String, + pub client_name: String, + pub initialize_timeout: Duration, + pub resume_session_id: Option, +} diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 21eac6b4c5..b8e3610583 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -1,22 +1,17 @@ use codex_app_server_protocol::JSONRPCMessage; use futures::SinkExt; use futures::StreamExt; +use tokio::io::AsyncBufReadExt; use tokio::io::AsyncRead; use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::io::BufWriter; use tokio::sync::mpsc; use tokio::sync::watch; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; -#[cfg(test)] -use tokio::io::AsyncBufReadExt; -#[cfg(test)] -use tokio::io::AsyncWriteExt; -#[cfg(test)] -use tokio::io::BufReader; -#[cfg(test)] -use tokio::io::BufWriter; - pub(crate) const CHANNEL_CAPACITY: usize = 128; #[derive(Debug)] @@ -34,7 +29,6 @@ pub(crate) struct JsonRpcConnection { } impl JsonRpcConnection { - #[cfg(test)] pub(crate) fn from_stdio(reader: R, writer: W, connection_label: String) -> Self where R: AsyncRead + Unpin + Send + 'static, @@ -298,7 +292,6 @@ async fn send_malformed_message( .await; } -#[cfg(test)] async fn write_jsonrpc_line_message( writer: &mut BufWriter, message: &JSONRPCMessage, diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 8e10c4a34e..922c65f89f 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::ExecServerError; use crate::ExecServerRuntimePaths; +use crate::RemoteExecServerTransport; use crate::client::LazyRemoteExecServerClient; use crate::file_system::ExecutorFileSystem; use crate::local_file_system::LocalFileSystem; @@ -16,10 +17,13 @@ pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; /// Owns the execution/filesystem environments available to the Codex runtime. /// /// `EnvironmentManager` is a shared registry for concrete environments. It -/// always creates a local environment under [`LOCAL_ENVIRONMENT_ID`]. When -/// `CODEX_EXEC_SERVER_URL` is set to a websocket URL, it also creates a remote -/// environment under [`REMOTE_ENVIRONMENT_ID`] and makes that the default -/// environment. Otherwise the local environment is the default. +/// always creates a local environment under [`LOCAL_ENVIRONMENT_ID`]. Additional +/// environments are remote exec-server endpoints that can connect over either a +/// websocket URL or a command-backed stdio JSON-RPC transport. +/// +/// In legacy mode, when `CODEX_EXEC_SERVER_URL` is set to a websocket URL, it +/// also creates a remote environment under [`REMOTE_ENVIRONMENT_ID`] and makes +/// that the default environment. Otherwise the local environment is the default. /// /// Setting `CODEX_EXEC_SERVER_URL=none` disables environment access by leaving /// the default environment unset while still keeping the local environment @@ -61,6 +65,37 @@ impl EnvironmentManagerArgs { } } +#[derive(Clone, Debug)] +pub struct ConfiguredEnvironmentManagerArgs { + pub default_environment: Option, + pub environments: Vec, + pub local_runtime_paths: ExecServerRuntimePaths, +} + +#[derive(Clone, Debug)] +pub struct ConfiguredEnvironmentSpec { + pub id: String, + pub transport: RemoteExecServerTransport, +} + +#[derive(Clone, Debug)] +pub enum EnvironmentManagerConfig { + ExecServerUrl(EnvironmentManagerArgs), + Configured(ConfiguredEnvironmentManagerArgs), +} + +impl From for EnvironmentManagerConfig { + fn from(args: EnvironmentManagerArgs) -> Self { + Self::ExecServerUrl(args) + } +} + +impl From for EnvironmentManagerConfig { + fn from(args: ConfiguredEnvironmentManagerArgs) -> Self { + Self::Configured(args) + } +} + impl EnvironmentManager { /// Builds a test-only manager without configured sandbox helper paths. pub fn default_for_tests() -> Self { @@ -75,7 +110,15 @@ impl EnvironmentManager { /// Builds a manager from the raw `CODEX_EXEC_SERVER_URL` value and local /// runtime paths used when creating local filesystem helpers. - pub fn new(args: EnvironmentManagerArgs) -> Self { + pub fn new(config: impl Into) -> Self { + match config.into() { + EnvironmentManagerConfig::ExecServerUrl(args) => Self::from_exec_server_url_args(args), + EnvironmentManagerConfig::Configured(args) => Self::from_configured_args(args) + .unwrap_or_else(|err| panic!("invalid environment manager config: {err}")), + } + } + + fn from_exec_server_url_args(args: EnvironmentManagerArgs) -> Self { let EnvironmentManagerArgs { exec_server_url, local_runtime_paths, @@ -92,7 +135,12 @@ impl EnvironmentManager { Some(exec_server_url) => { environments.insert( REMOTE_ENVIRONMENT_ID.to_string(), - Arc::new(Environment::remote(exec_server_url, local_runtime_paths)), + Arc::new(Environment::remote_with_runtime_paths( + RemoteExecServerTransport::WebSocket { + url: exec_server_url, + }, + Some(local_runtime_paths), + )), ); Some(REMOTE_ENVIRONMENT_ID.to_string()) } @@ -106,6 +154,73 @@ impl EnvironmentManager { } } + pub fn try_new(config: impl Into) -> Result { + match config.into() { + EnvironmentManagerConfig::ExecServerUrl(args) => { + Ok(Self::from_exec_server_url_args(args)) + } + EnvironmentManagerConfig::Configured(args) => Self::from_configured_args(args), + } + } + + fn from_configured_args( + args: ConfiguredEnvironmentManagerArgs, + ) -> Result { + let ConfiguredEnvironmentManagerArgs { + default_environment, + environments: configured_environments, + local_runtime_paths, + } = args; + let mut environments = HashMap::from([( + LOCAL_ENVIRONMENT_ID.to_string(), + Arc::new(Environment::local(local_runtime_paths.clone())), + )]); + + for environment in configured_environments { + let id = environment.id.trim(); + if id.is_empty() { + return Err(ExecServerError::Protocol( + "environment id must not be empty".to_string(), + )); + } + if id == LOCAL_ENVIRONMENT_ID || id.eq_ignore_ascii_case("none") { + return Err(ExecServerError::Protocol(format!( + "environment id `{id}` is reserved" + ))); + } + if environments.contains_key(id) { + return Err(ExecServerError::Protocol(format!( + "duplicate environment id `{id}`" + ))); + } + environments.insert( + id.to_string(), + Arc::new(Environment::remote_with_runtime_paths( + environment.transport, + Some(local_runtime_paths.clone()), + )), + ); + } + + let default_environment = match default_environment.as_deref().map(str::trim) { + None | Some("") => Some(LOCAL_ENVIRONMENT_ID.to_string()), + Some(default_environment) if default_environment.eq_ignore_ascii_case("none") => None, + Some(default_environment) => { + if !environments.contains_key(default_environment) { + return Err(ExecServerError::Protocol(format!( + "default_environment `{default_environment}` is not a configured environment id" + ))); + } + Some(default_environment.to_string()) + } + }; + + Ok(Self { + default_environment, + environments, + }) + } + /// Returns the default environment instance. pub fn default_environment(&self) -> Option> { self.default_environment @@ -133,7 +248,7 @@ impl EnvironmentManager { /// paths used by filesystem helpers. #[derive(Clone)] pub struct Environment { - exec_server_url: Option, + remote_transport: Option, exec_backend: Arc, filesystem: Arc, local_runtime_paths: Option, @@ -143,7 +258,7 @@ impl Environment { /// Builds a test-only local environment without configured sandbox helper paths. 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()), local_runtime_paths: None, @@ -154,7 +269,7 @@ impl Environment { impl std::fmt::Debug for Environment { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("Environment") - .field("exec_server_url", &self.exec_server_url) + .field("remote_transport", &self.remote_transport) .finish_non_exhaustive() } } @@ -187,7 +302,12 @@ impl Environment { } Ok(match exec_server_url { - Some(exec_server_url) => Self::remote_inner(exec_server_url, local_runtime_paths), + Some(exec_server_url) => Self::remote_with_runtime_paths( + RemoteExecServerTransport::WebSocket { + url: exec_server_url, + }, + local_runtime_paths, + ), None => match local_runtime_paths { Some(local_runtime_paths) => Self::local(local_runtime_paths), None => Self::default_for_tests(), @@ -197,7 +317,7 @@ impl Environment { 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(), @@ -206,20 +326,16 @@ impl Environment { } } - fn remote(exec_server_url: String, local_runtime_paths: ExecServerRuntimePaths) -> Self { - Self::remote_inner(exec_server_url, Some(local_runtime_paths)) - } - - fn remote_inner( - exec_server_url: String, + fn remote_with_runtime_paths( + transport: RemoteExecServerTransport, local_runtime_paths: Option, ) -> Self { - let client = LazyRemoteExecServerClient::new(exec_server_url.clone()); + let client = LazyRemoteExecServerClient::new(transport.clone()); let exec_backend: Arc = Arc::new(RemoteProcess::new(client.clone())); let filesystem: Arc = Arc::new(RemoteFileSystem::new(client)); Self { - exec_server_url: Some(exec_server_url), + remote_transport: Some(transport), exec_backend, filesystem, local_runtime_paths, @@ -227,12 +343,22 @@ 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. pub fn exec_server_url(&self) -> Option<&str> { - self.exec_server_url.as_deref() + match self.remote_transport.as_ref() { + Some(RemoteExecServerTransport::WebSocket { url }) => Some(url.as_str()), + Some(RemoteExecServerTransport::Command { .. }) | None => None, + } + } + + pub fn exec_server_command(&self) -> Option<&str> { + match self.remote_transport.as_ref() { + Some(RemoteExecServerTransport::Command { command }) => Some(command.as_str()), + Some(RemoteExecServerTransport::WebSocket { .. }) | None => None, + } } pub fn local_runtime_paths(&self) -> Option<&ExecServerRuntimePaths> { diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index fc6a86f508..4fde4de4f7 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -20,12 +20,17 @@ mod server; pub use client::ExecServerClient; pub use client::ExecServerError; +pub use client_api::CommandExecServerConnectArgs; pub use client_api::ExecServerClientConnectOptions; pub use client_api::RemoteExecServerConnectArgs; +pub use client_api::RemoteExecServerTransport; pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR; +pub use environment::ConfiguredEnvironmentManagerArgs; +pub use environment::ConfiguredEnvironmentSpec; pub use environment::Environment; pub use environment::EnvironmentManager; pub use environment::EnvironmentManagerArgs; +pub use environment::EnvironmentManagerConfig; pub use file_system::CopyOptions; pub use file_system::CreateDirectoryOptions; pub use file_system::ExecutorFileSystem; diff --git a/codex-rs/exec-server/src/server/transport.rs b/codex-rs/exec-server/src/server/transport.rs index b8a5a086b6..9b0da965d9 100644 --- a/codex-rs/exec-server/src/server/transport.rs +++ b/codex-rs/exec-server/src/server/transport.rs @@ -1,5 +1,6 @@ use std::io::Write as _; use std::net::SocketAddr; +use tokio::io; use tokio::net::TcpListener; use tokio_tungstenite::accept_async; use tracing::warn; @@ -21,7 +22,7 @@ impl std::fmt::Display for ExecServerListenUrlParseError { match self { ExecServerListenUrlParseError::UnsupportedListenUrl(listen_url) => write!( f, - "unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT`" + "unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT` or `stdio://`" ), ExecServerListenUrlParseError::InvalidWebSocketListenUrl(listen_url) => write!( f, @@ -51,10 +52,28 @@ pub(crate) async fn run_transport( listen_url: &str, runtime_paths: ExecServerRuntimePaths, ) -> Result<(), Box> { + if listen_url == "stdio://" { + return run_stdio(runtime_paths).await; + } + let bind_address = parse_listen_url(listen_url)?; run_websocket_listener(bind_address, runtime_paths).await } +async fn run_stdio( + runtime_paths: ExecServerRuntimePaths, +) -> Result<(), Box> { + let processor = ConnectionProcessor::new(runtime_paths); + processor + .run_connection(JsonRpcConnection::from_stdio( + io::stdin(), + io::stdout(), + "exec-server stdio".to_string(), + )) + .await; + Ok(()) +} + async fn run_websocket_listener( bind_address: SocketAddr, runtime_paths: ExecServerRuntimePaths, diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 11bf1b7a60..d5cd4769a4 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -498,9 +498,12 @@ pub async fn run_main(cli: Cli, arg0_paths: Arg0DispatchPaths) -> anyhow::Result cloud_requirements: run_cloud_requirements, feedback: CodexFeedback::new(), log_db: None, - environment_manager: std::sync::Arc::new(EnvironmentManager::new( - EnvironmentManagerArgs::from_env(local_runtime_paths), - )), + environment_manager: std::sync::Arc::new(EnvironmentManager::try_new( + config.environment_manager_args( + EnvironmentManagerArgs::from_env(local_runtime_paths.clone()).exec_server_url, + local_runtime_paths, + )?, + )?), config_warnings, session_source: SessionSource::Exec, enable_codex_api_key_env: true, diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 046063a554..7d836406af 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -91,6 +91,8 @@ pub enum Feature { ShellZshFork, /// Include the freeform apply_patch tool. ApplyPatchFreeform, + /// Expose selected turn environments to the model and allow tools to target them by id. + MultiEnvironmentTools, /// Stream structured progress while apply_patch input is being generated. ApplyPatchStreamingEvents, /// Allow exec tools to request additional permissions while staying sandboxed. @@ -728,6 +730,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::MultiEnvironmentTools, + key: "multi_environment_tools", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::ApplyPatchStreamingEvents, key: "apply_patch_streaming_events", diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index 1d904e4577..c027350d23 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -60,12 +60,10 @@ pub async fn run_main( arg0_paths: Arg0DispatchPaths, cli_config_overrides: CliConfigOverrides, ) -> IoResult<()> { - let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::from_env( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, - ))); + let runtime_paths = ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?; // Parse CLI overrides once and derive the base Config eagerly so later // components do not need to work with raw TOML values. let cli_kv_overrides = cli_config_overrides.parse_overrides().map_err(|e| { @@ -79,6 +77,11 @@ pub async fn run_main( .map_err(|e| { std::io::Error::new(ErrorKind::InvalidData, format!("error loading config: {e}")) })?; + let exec_server_url = EnvironmentManagerArgs::from_env(runtime_paths.clone()).exec_server_url; + let environment_manager = Arc::new( + EnvironmentManager::try_new(config.environment_manager_args(exec_server_url, runtime_paths)?) + .map_err(|err| std::io::Error::new(ErrorKind::InvalidInput, err))?, + ); set_default_client_residency_requirement(config.enforce_residency.value()); let otel = codex_core::otel_init::build_provider( diff --git a/codex-rs/tools/src/local_tool.rs b/codex-rs/tools/src/local_tool.rs index 3e369ab1e3..8148ad60dd 100644 --- a/codex-rs/tools/src/local_tool.rs +++ b/codex-rs/tools/src/local_tool.rs @@ -9,6 +9,7 @@ use std::collections::BTreeMap; pub struct CommandToolOptions { pub allow_login_shell: bool, pub exec_permission_approvals_enabled: bool, + pub multi_environment_tools: bool, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -63,6 +64,14 @@ pub fn create_exec_command_tool(options: CommandToolOptions) -> ToolSpec { )), ); } + if options.multi_environment_tools { + properties.insert( + "environment_id".to_string(), + JsonSchema::string(Some( + "Optional selected environment id to run the command in; omit to use the primary environment.".to_string(), + )), + ); + } properties.extend(create_approval_parameters( options.exec_permission_approvals_enabled, )); diff --git a/codex-rs/tools/src/local_tool_tests.rs b/codex-rs/tools/src/local_tool_tests.rs index b751545b3a..a14cfd1d62 100644 --- a/codex-rs/tools/src/local_tool_tests.rs +++ b/codex-rs/tools/src/local_tool_tests.rs @@ -96,6 +96,7 @@ fn exec_command_tool_matches_expected_spec() { let tool = create_exec_command_tool(CommandToolOptions { allow_login_shell: true, exec_permission_approvals_enabled: false, + multi_environment_tools: false, }); let description = if cfg!(windows) { @@ -332,6 +333,7 @@ fn shell_command_tool_matches_expected_spec() { let tool = create_shell_command_tool(CommandToolOptions { allow_login_shell: true, exec_permission_approvals_enabled: false, + multi_environment_tools: false, }); let description = if cfg!(windows) { diff --git a/codex-rs/tools/src/tool_config.rs b/codex-rs/tools/src/tool_config.rs index 20c812ce96..2ff0e064fe 100644 --- a/codex-rs/tools/src/tool_config.rs +++ b/codex-rs/tools/src/tool_config.rs @@ -97,6 +97,7 @@ pub struct ToolsConfig { pub tool_suggest: bool, pub exec_permission_approvals_enabled: bool, pub request_permissions_tool_enabled: bool, + pub multi_environment_tools: bool, pub code_mode_enabled: bool, pub code_mode_only_enabled: bool, pub js_repl_enabled: bool, @@ -161,6 +162,7 @@ impl ToolsConfig { && supports_image_generation(model_info); let exec_permission_approvals_enabled = features.enabled(Feature::ExecPermissionApprovals); let request_permissions_tool_enabled = features.enabled(Feature::RequestPermissionsTool); + let multi_environment_tools = features.enabled(Feature::MultiEnvironmentTools); let shell_command_backend = if features.enabled(Feature::ShellTool) && features.enabled(Feature::ShellZshFork) { ShellCommandBackendConfig::ZshFork @@ -218,6 +220,7 @@ impl ToolsConfig { tool_suggest: include_tool_suggest, exec_permission_approvals_enabled, request_permissions_tool_enabled, + multi_environment_tools, code_mode_enabled: include_code_mode, code_mode_only_enabled: include_code_mode_only, js_repl_enabled: include_js_repl, diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index d3b075f5ab..3a4b25679e 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -157,6 +157,7 @@ pub fn build_tool_registry_plan( create_exec_command_tool(CommandToolOptions { allow_login_shell: config.allow_login_shell, exec_permission_approvals_enabled, + multi_environment_tools: config.multi_environment_tools, }), /*supports_parallel_tool_calls*/ true, config.code_mode_enabled, @@ -175,6 +176,7 @@ pub fn build_tool_registry_plan( create_shell_command_tool(CommandToolOptions { allow_login_shell: config.allow_login_shell, exec_permission_approvals_enabled, + multi_environment_tools: false, }), /*supports_parallel_tool_calls*/ true, config.code_mode_enabled, diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index 0b94ef64ca..2b81dee155 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -85,6 +85,7 @@ fn test_full_toolset_specs_for_gpt5_codex_unified_exec_web_search() { create_exec_command_tool(CommandToolOptions { allow_login_shell: true, exec_permission_approvals_enabled: false, + multi_environment_tools: false, }), create_write_stdin_tool(), create_update_plan_tool(), diff --git a/codex-rs/tui/src/app_command.rs b/codex-rs/tui/src/app_command.rs index 45425c3494..4893f7a9fe 100644 --- a/codex-rs/tui/src/app_command.rs +++ b/codex-rs/tui/src/app_command.rs @@ -181,6 +181,7 @@ impl AppCommand { ) -> Self { Self(Op::OverrideTurnContext { cwd, + environments: None, approval_policy, approvals_reviewer, sandbox_policy, @@ -314,6 +315,7 @@ impl AppCommand { }, Op::OverrideTurnContext { cwd, + environments: _, approval_policy, approvals_reviewer, sandbox_policy, diff --git a/codex-rs/tui/src/chatwidget/tests/permissions.rs b/codex-rs/tui/src/chatwidget/tests/permissions.rs index c1e142c4e1..8b19bc7538 100644 --- a/codex-rs/tui/src/chatwidget/tests/permissions.rs +++ b/codex-rs/tui/src/chatwidget/tests/permissions.rs @@ -629,6 +629,7 @@ async fn permissions_selection_sends_approvals_reviewer_in_override_turn_context op, Op::OverrideTurnContext { cwd: None, + environments: None, approval_policy: Some(AskForApproval::OnRequest), approvals_reviewer: Some(ApprovalsReviewer::GuardianSubagent), sandbox_policy: Some(SandboxPolicy::new_workspace_write_policy()), diff --git a/codex-rs/tui/src/lib.rs b/codex-rs/tui/src/lib.rs index 7e33f2e8a3..96cba25785 100644 --- a/codex-rs/tui/src/lib.rs +++ b/codex-rs/tui/src/lib.rs @@ -729,11 +729,12 @@ pub async fn run_main( } }; + let runtime_paths = ExecServerRuntimePaths::from_optional_paths( + arg0_paths.codex_self_exe.clone(), + arg0_paths.codex_linux_sandbox_exe.clone(), + )?; let environment_manager = Arc::new(EnvironmentManager::new(EnvironmentManagerArgs::from_env( - ExecServerRuntimePaths::from_optional_paths( - arg0_paths.codex_self_exe.clone(), - arg0_paths.codex_linux_sandbox_exe.clone(), - )?, + runtime_paths.clone(), ))); let cwd = cli.cwd.clone(); let config_cwd = @@ -848,6 +849,11 @@ pub async fn run_main( cloud_requirements.clone(), ) .await; + let exec_server_url = EnvironmentManagerArgs::from_env(runtime_paths.clone()).exec_server_url; + let environment_manager = Arc::new( + EnvironmentManager::try_new(config.environment_manager_args(exec_server_url, runtime_paths)?) + .map_err(|err| std::io::Error::new(std::io::ErrorKind::InvalidInput, err))?, + ); #[allow(clippy::print_stderr)] match check_execpolicy_for_warnings(&config.config_layer_stack).await {