From 8436adb8e31772800fc8f84ce69bfd59d0b742e4 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Wed, 6 May 2026 21:12:37 -0700 Subject: [PATCH] Attach configured environments on thread startup Preserve provider environment order and derive startup selections from the configured default plus remaining environments. This lets multi-environment CODEX_HOME setups attach every configured environment by default while keeping the default first as primary. Co-authored-by: Codex --- codex-rs/core/src/environment_selection.rs | 21 +++-- codex-rs/core/src/thread_manager_tests.rs | 92 +++++++++++++++++++ codex-rs/exec-server/src/client_transport.rs | 31 +++++-- codex-rs/exec-server/src/environment.rs | 69 ++++++++++---- .../exec-server/src/environment_provider.rs | 62 ++++++++----- codex-rs/exec-server/src/environment_toml.rs | 36 +++++--- 6 files changed, 243 insertions(+), 68 deletions(-) diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index f14361bdb5..50371e86a0 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -15,12 +15,12 @@ pub(crate) fn default_thread_environment_selections( cwd: &AbsolutePathBuf, ) -> Vec { environment_manager - .default_environment_id() + .default_environment_ids() + .into_iter() .map(|environment_id| TurnEnvironmentSelection { - environment_id: environment_id.to_string(), + environment_id, cwd: cwd.clone(), }) - .into_iter() .collect() } @@ -85,6 +85,7 @@ pub(crate) fn resolve_environment_selections( #[cfg(test)] mod tests { use codex_exec_server::ExecServerRuntimePaths; + use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; @@ -110,10 +111,16 @@ mod tests { assert_eq!( default_thread_environment_selections(&manager, &cwd), - vec![TurnEnvironmentSelection { - environment_id: REMOTE_ENVIRONMENT_ID.to_string(), - cwd, - }] + vec![ + TurnEnvironmentSelection { + environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + cwd: cwd.clone(), + }, + TurnEnvironmentSelection { + environment_id: LOCAL_ENVIRONMENT_ID.to_string(), + cwd, + }, + ] ); } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 7ad54565ff..aa01ce14d8 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -19,6 +19,7 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::UserMessageEvent; +use codex_protocol::user_input::UserInput; use core_test_support::PathBufExt; use core_test_support::PathExt; use core_test_support::responses::mount_models_once; @@ -351,6 +352,97 @@ async fn start_thread_accepts_explicit_environment_when_default_environment_is_d assert_eq!(manager.list_thread_ids().await, vec![thread.thread_id]); } +#[tokio::test] +async fn start_thread_uses_all_default_environments_from_codex_home() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + std::fs::write( + config.codex_home.join("environments.toml"), + r#" +default = "local" + +[[environments]] +id = "dev" +program = "ssh" +args = ["dev", "cd /tmp && true"] +"#, + ) + .expect("write environments.toml"); + + let runtime_paths = codex_exec_server::ExecServerRuntimePaths::new( + std::env::current_exe().expect("current exe path"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths"); + let environment_manager = Arc::new( + codex_exec_server::EnvironmentManager::from_codex_home( + config.codex_home.clone(), + runtime_paths, + ) + .await + .expect("environment manager"), + ); + assert_eq!( + environment_manager.default_environment_ids(), + vec!["local".to_string(), "dev".to_string()] + ); + + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + environment_manager, + ) + .await; + + let thread = manager + .start_thread(config) + .await + .expect("thread should start"); + + let selections = thread.thread.environment_selections().await; + assert_eq!( + selections, + vec![ + TurnEnvironmentSelection { + environment_id: "local".to_string(), + cwd: thread.session_configured.cwd.abs(), + }, + TurnEnvironmentSelection { + environment_id: "dev".to_string(), + cwd: thread.session_configured.cwd.abs(), + }, + ] + ); + + let prompt_items = crate::prompt_debug::build_prompt_input_from_session( + thread.thread.codex.session.as_ref(), + Vec::::new(), + ) + .await + .expect("prompt input"); + let environment_context = prompt_items + .iter() + .filter_map(|item| match item { + ResponseItem::Message { content, .. } => Some(content), + _ => None, + }) + .flatten() + .find_map(|content| match content { + ContentItem::InputText { text } if text.contains("") => { + Some(text.as_str()) + } + _ => None, + }) + .expect("environment context prompt item"); + assert!(environment_context.contains("")); + assert!(environment_context.contains("")); + assert!(environment_context.contains("")); +} + #[tokio::test] async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { let temp_dir = tempdir().expect("tempdir"); diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index f7fffaa822..3fccfa25c5 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -24,16 +24,27 @@ impl ExecServerClient { pub(crate) async fn connect_for_transport( transport_params: crate::client_api::ExecServerTransportParams, ) -> Result { - let crate::client_api::ExecServerTransportParams::WebSocketUrl(websocket_url) = - transport_params; - Self::connect_websocket(RemoteExecServerConnectArgs { - websocket_url, - client_name: ENVIRONMENT_CLIENT_NAME.to_string(), - connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, - initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, - resume_session_id: None, - }) - .await + match transport_params { + crate::client_api::ExecServerTransportParams::WebSocketUrl(websocket_url) => { + Self::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + connect_timeout: ENVIRONMENT_CONNECT_TIMEOUT, + initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, + resume_session_id: None, + }) + .await + } + crate::client_api::ExecServerTransportParams::StdioCommand(command) => { + Self::connect_stdio_command(StdioExecServerConnectArgs { + command, + client_name: ENVIRONMENT_CLIENT_NAME.to_string(), + initialize_timeout: ENVIRONMENT_INITIALIZE_TIMEOUT, + resume_session_id: None, + }) + .await + } + } } pub async fn connect_websocket( diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 524314337e..60d20b3cdd 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -41,6 +41,7 @@ pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL"; #[derive(Debug)] pub struct EnvironmentManager { default_environment: Option, + configured_environment_ids: Vec, environments: HashMap>, local_environment: Arc, } @@ -53,6 +54,7 @@ impl EnvironmentManager { pub fn default_for_tests() -> Self { Self { default_environment: Some(LOCAL_ENVIRONMENT_ID.to_string()), + configured_environment_ids: vec![LOCAL_ENVIRONMENT_ID.to_string()], environments: HashMap::from([( LOCAL_ENVIRONMENT_ID.to_string(), Arc::new(Environment::default_for_tests()), @@ -65,6 +67,7 @@ impl EnvironmentManager { pub fn disabled_for_tests(local_runtime_paths: ExecServerRuntimePaths) -> Self { Self { default_environment: None, + configured_environment_ids: Vec::new(), environments: HashMap::new(), local_environment: Arc::new(Environment::local(local_runtime_paths)), } @@ -135,18 +138,28 @@ impl EnvironmentManager { environments, default, } = snapshot; - for id in environments.keys() { + let mut configured_environment_ids = Vec::with_capacity(environments.len()); + let mut environment_map = HashMap::with_capacity(environments.len()); + for (id, environment) in environments { if id.is_empty() { return Err(ExecServerError::Protocol( "environment id cannot be empty".to_string(), )); } + if environment_map + .insert(id.clone(), Arc::new(environment)) + .is_some() + { + return Err(ExecServerError::Protocol(format!( + "environment id `{id}` is duplicated" + ))); + } + configured_environment_ids.push(id); } - let default_environment = match default { EnvironmentDefault::Disabled => None, EnvironmentDefault::EnvironmentId(environment_id) => { - if !environments.contains_key(&environment_id) { + if !environment_map.contains_key(&environment_id) { return Err(ExecServerError::Protocol(format!( "default environment `{environment_id}` is not configured" ))); @@ -155,14 +168,11 @@ impl EnvironmentManager { } }; let local_environment = Arc::new(Environment::local(local_runtime_paths)); - let environments = environments - .into_iter() - .map(|(id, environment)| (id, Arc::new(environment))) - .collect(); Ok(Self { default_environment, - environments, + configured_environment_ids, + environments: environment_map, local_environment, }) } @@ -179,6 +189,26 @@ impl EnvironmentManager { self.default_environment.as_deref() } + /// Returns the ordered environment ids used for new thread startup. + pub fn default_environment_ids(&self) -> Vec { + let Some(default_environment_id) = self.default_environment.as_deref() else { + return Vec::new(); + }; + if self.configured_environment_ids.len() <= 1 { + return vec![default_environment_id.to_string()]; + } + + let mut environment_ids = Vec::with_capacity(self.configured_environment_ids.len()); + environment_ids.push(default_environment_id.to_string()); + environment_ids.extend( + self.configured_environment_ids + .iter() + .filter(|environment_id| environment_id.as_str() != default_environment_id) + .cloned(), + ); + environment_ids + } + /// Returns the local environment instance used for internal runtime work. pub fn local_environment(&self) -> Arc { Arc::clone(&self.local_environment) @@ -331,7 +361,6 @@ impl Environment { #[cfg(test)] mod tests { - use std::collections::HashMap; use std::sync::Arc; use super::Environment; @@ -453,11 +482,11 @@ mod tests { async fn environment_manager_builds_from_provider() { let provider = TestEnvironmentProvider { snapshot: EnvironmentProviderSnapshot { - environments: HashMap::from([( + environments: vec![( REMOTE_ENVIRONMENT_ID.to_string(), Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), - )]), + )], default: EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()), }, }; @@ -483,7 +512,7 @@ mod tests { async fn environment_manager_rejects_empty_environment_id() { let provider = TestEnvironmentProvider { snapshot: EnvironmentProviderSnapshot { - environments: HashMap::from([("".to_string(), Environment::default_for_tests())]), + environments: vec![("".to_string(), Environment::default_for_tests())], default: EnvironmentDefault::Disabled, }, }; @@ -501,7 +530,7 @@ mod tests { async fn environment_manager_uses_explicit_provider_default() { let provider = TestEnvironmentProvider { snapshot: EnvironmentProviderSnapshot { - environments: HashMap::from([ + environments: vec![ ( LOCAL_ENVIRONMENT_ID.to_string(), Environment::default_for_tests(), @@ -511,7 +540,7 @@ mod tests { Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) .expect("remote environment"), ), - ]), + ], default: EnvironmentDefault::EnvironmentId("devbox".to_string()), }, }; @@ -520,6 +549,10 @@ mod tests { .expect("manager"); assert_eq!(manager.default_environment_id(), Some("devbox")); + assert_eq!( + manager.default_environment_ids(), + vec!["devbox".to_string(), LOCAL_ENVIRONMENT_ID.to_string()] + ); assert!(manager.default_environment().expect("default").is_remote()); } @@ -527,10 +560,10 @@ mod tests { async fn environment_manager_disables_provider_default() { let provider = TestEnvironmentProvider { snapshot: EnvironmentProviderSnapshot { - environments: HashMap::from([( + environments: vec![( LOCAL_ENVIRONMENT_ID.to_string(), Environment::default_for_tests(), - )]), + )], default: EnvironmentDefault::Disabled, }, }; @@ -547,10 +580,10 @@ mod tests { async fn environment_manager_rejects_unknown_provider_default() { let provider = TestEnvironmentProvider { snapshot: EnvironmentProviderSnapshot { - environments: HashMap::from([( + environments: vec![( LOCAL_ENVIRONMENT_ID.to_string(), Environment::default_for_tests(), - )]), + )], default: EnvironmentDefault::EnvironmentId("missing".to_string()), }, }; diff --git a/codex-rs/exec-server/src/environment_provider.rs b/codex-rs/exec-server/src/environment_provider.rs index 0e4bcc5191..bced67db55 100644 --- a/codex-rs/exec-server/src/environment_provider.rs +++ b/codex-rs/exec-server/src/environment_provider.rs @@ -1,5 +1,3 @@ -use std::collections::HashMap; - use async_trait::async_trait; use crate::Environment; @@ -12,9 +10,9 @@ use crate::environment::REMOTE_ENVIRONMENT_ID; /// Lists the concrete environments available to Codex. /// /// Implementations own a startup snapshot containing both the available -/// environment list and default environment selection. Providers that want the -/// local environment to be addressable by id should include it explicitly in -/// the returned map. +/// environment list in configured order and the default environment +/// selection. Providers that want the local environment to be addressable by +/// id should include it explicitly in the returned list. #[async_trait] pub trait EnvironmentProvider: Send + Sync { /// Returns the provider-owned environment startup snapshot. @@ -26,7 +24,7 @@ pub trait EnvironmentProvider: Send + Sync { #[derive(Clone, Debug)] pub struct EnvironmentProviderSnapshot { - pub environments: HashMap, + pub environments: Vec<(String, Environment)>, pub default: EnvironmentDefault, } @@ -57,22 +55,25 @@ impl DefaultEnvironmentProvider { &self, local_runtime_paths: &ExecServerRuntimePaths, ) -> EnvironmentProviderSnapshot { - let mut environments = HashMap::from([( + let mut environments = vec![( LOCAL_ENVIRONMENT_ID.to_string(), Environment::local(local_runtime_paths.clone()), - )]); + )]; let (exec_server_url, disabled) = normalize_exec_server_url(self.exec_server_url.clone()); if let Some(exec_server_url) = exec_server_url { - environments.insert( + environments.push(( REMOTE_ENVIRONMENT_ID.to_string(), Environment::remote_inner(exec_server_url, Some(local_runtime_paths.clone())), - ); + )); } + let has_remote = environments + .iter() + .any(|(id, _environment)| id == REMOTE_ENVIRONMENT_ID); let default = if disabled { EnvironmentDefault::Disabled - } else if environments.contains_key(REMOTE_ENVIRONMENT_ID) { + } else if has_remote { EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()) } else { EnvironmentDefault::EnvironmentId(LOCAL_ENVIRONMENT_ID.to_string()) @@ -105,6 +106,8 @@ pub(crate) fn normalize_exec_server_url(exec_server_url: Option) -> (Opt #[cfg(test)] mod tests { + use std::collections::HashMap; + use pretty_assertions::assert_eq; use super::*; @@ -126,7 +129,11 @@ mod tests { .snapshot(&runtime_paths) .await .expect("environments"); - let environments = snapshot.environments; + let EnvironmentProviderSnapshot { + environments, + default, + } = snapshot; + let environments: HashMap<_, _> = environments.into_iter().collect(); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert_eq!( @@ -135,7 +142,7 @@ mod tests { ); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); assert_eq!( - snapshot.default, + default, EnvironmentDefault::EnvironmentId(LOCAL_ENVIRONMENT_ID.to_string()) ); } @@ -148,12 +155,16 @@ mod tests { .snapshot(&runtime_paths) .await .expect("environments"); - let environments = snapshot.environments; + let EnvironmentProviderSnapshot { + environments, + default, + } = snapshot; + let environments: HashMap<_, _> = environments.into_iter().collect(); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); assert_eq!( - snapshot.default, + default, EnvironmentDefault::EnvironmentId(LOCAL_ENVIRONMENT_ID.to_string()) ); } @@ -166,11 +177,15 @@ mod tests { .snapshot(&runtime_paths) .await .expect("environments"); - let environments = snapshot.environments; + let EnvironmentProviderSnapshot { + environments, + default, + } = snapshot; + let environments: HashMap<_, _> = environments.into_iter().collect(); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert!(!environments.contains_key(REMOTE_ENVIRONMENT_ID)); - assert_eq!(snapshot.default, EnvironmentDefault::Disabled); + assert_eq!(default, EnvironmentDefault::Disabled); } #[tokio::test] @@ -181,7 +196,11 @@ mod tests { .snapshot(&runtime_paths) .await .expect("environments"); - let environments = snapshot.environments; + let EnvironmentProviderSnapshot { + environments, + default, + } = snapshot; + let environments: HashMap<_, _> = environments.into_iter().collect(); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); let remote_environment = &environments[REMOTE_ENVIRONMENT_ID]; @@ -191,7 +210,7 @@ mod tests { Some("ws://127.0.0.1:8765") ); assert_eq!( - snapshot.default, + default, EnvironmentDefault::EnvironmentId(REMOTE_ENVIRONMENT_ID.to_string()) ); } @@ -200,13 +219,14 @@ mod tests { async fn default_provider_normalizes_exec_server_url() { let provider = DefaultEnvironmentProvider::new(Some(" ws://127.0.0.1:8765 ".to_string())); let runtime_paths = test_runtime_paths(); - let environments = provider + let snapshot = provider .snapshot(&runtime_paths) .await .expect("environments"); + let environments: HashMap<_, _> = snapshot.environments.into_iter().collect(); assert_eq!( - environments.environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), + environments[REMOTE_ENVIRONMENT_ID].exec_server_url(), Some("ws://127.0.0.1:8765") ); } diff --git a/codex-rs/exec-server/src/environment_toml.rs b/codex-rs/exec-server/src/environment_toml.rs index 99808d7896..5ad236c1a0 100644 --- a/codex-rs/exec-server/src/environment_toml.rs +++ b/codex-rs/exec-server/src/environment_toml.rs @@ -44,7 +44,7 @@ struct EnvironmentToml { #[derive(Clone, Debug, PartialEq, Eq)] struct TomlEnvironmentProvider { default: EnvironmentDefault, - environments: HashMap, + environments: Vec<(String, ExecServerTransportParams)>, } impl TomlEnvironmentProvider { @@ -58,7 +58,7 @@ impl TomlEnvironmentProvider { config_dir: Option<&Path>, ) -> Result { let mut ids = HashSet::from([LOCAL_ENVIRONMENT_ID.to_string()]); - let mut environments = HashMap::with_capacity(config.environments.len()); + let mut environments = Vec::with_capacity(config.environments.len()); for item in config.environments { let (id, transport) = parse_environment_toml(item, config_dir)?; if !ids.insert(id.clone()) { @@ -66,7 +66,7 @@ impl TomlEnvironmentProvider { "environment id `{id}` is duplicated" ))); } - environments.insert(id, transport); + environments.push((id, transport)); } let default = normalize_default_environment_id(config.default.as_deref(), &ids)?; Ok(Self { @@ -82,19 +82,19 @@ impl EnvironmentProvider for TomlEnvironmentProvider { &self, local_runtime_paths: &ExecServerRuntimePaths, ) -> Result { - let mut environments = HashMap::from([( + let mut environments = Vec::with_capacity(self.environments.len() + 1); + environments.push(( LOCAL_ENVIRONMENT_ID.to_string(), Environment::local(local_runtime_paths.clone()), - )]); - + )); for (id, transport_params) in &self.environments { - environments.insert( + environments.push(( id.clone(), Environment::remote_with_transport( transport_params.clone(), Some(local_runtime_paths.clone()), ), - ); + )); } Ok(EnvironmentProviderSnapshot { @@ -345,13 +345,15 @@ mod tests { environments, default, } = snapshot; + let environments: HashMap<_, _> = environments.into_iter().collect(); assert!(!environments[LOCAL_ENVIRONMENT_ID].is_remote()); assert_eq!( environments["devbox"].exec_server_url(), Some("ws://127.0.0.1:8765") ); - assert_eq!(provider.environments["ssh-dev"], ssh_transport); + assert_eq!(provider.environments[1].0, "ssh-dev"); + assert_eq!(provider.environments[1].1, ssh_transport); assert!(environments["ssh-dev"].is_remote()); assert_eq!(environments["ssh-dev"].exec_server_url(), None); assert_eq!( @@ -483,7 +485,7 @@ mod tests { .expect("provider"); assert_eq!( - provider.environments["ssh-dev"], + provider.environments[0].1, ExecServerTransportParams::StdioCommand(StdioExecServerCommand { program: "ssh".to_string(), args: Vec::new(), @@ -686,8 +688,13 @@ default = "none" .snapshot(&test_runtime_paths()) .await .expect("environments"); + let environment_ids: Vec<_> = snapshot + .environments + .into_iter() + .map(|(id, _environment)| id) + .collect(); - assert!(snapshot.environments.contains_key(LOCAL_ENVIRONMENT_ID)); + assert!(environment_ids.contains(&LOCAL_ENVIRONMENT_ID.to_string())); assert_eq!(snapshot.default, EnvironmentDefault::Disabled); } @@ -702,7 +709,12 @@ default = "none" .snapshot(&test_runtime_paths()) .await .expect("environments"); + let environment_ids: Vec<_> = snapshot + .environments + .into_iter() + .map(|(id, _environment)| id) + .collect(); - assert!(snapshot.environments.contains_key(LOCAL_ENVIRONMENT_ID)); + assert!(environment_ids.contains(&LOCAL_ENVIRONMENT_ID.to_string())); } }