From be338ee9a28ce5a1c75455343e9712aded82c70f Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 11 Jun 2026 20:35:28 -0700 Subject: [PATCH] [codex] resolve environment shell metadata eagerly (#27709) ## Why Turn construction passed resolved environments through several layers while leaving the environment shell unresolved. As a result, model-visible environment context could fall back to the session shell instead of reporting the selected remote environment's shell. Resolve environment metadata at the turn-context boundary so each turn carries the shell that belongs to its selected environment. Keep request validation in app-server, where invalid selections can be returned as straightforward JSON-RPC errors without coupling core turn construction to that policy. ## What changed - resolve environment selections eagerly in `new_turn_context_from_configuration` - store the full resolved `Shell` on each `TurnEnvironment` - simplify the now-redundant resolved-environment constructor plumbing - keep duplicate and unknown-environment validation as a small app-server preflight - add a remote-environment integration test that runs a full `test_codex` turn and verifies the model-visible environment message reports `bash` ## Testing - `cargo check -p codex-core --test all -p codex-app-server` - `remote_test_env_exposes_bash_shell_to_model` on the Linux remote-executor harness --- .../src/request_processors/request_errors.rs | 6 +- .../request_processors/thread_processor.rs | 2 +- .../src/request_processors/turn_processor.rs | 2 +- .../core/src/context/environment_context.rs | 3 +- codex-rs/core/src/environment_selection.rs | 83 ++++++++++++------- codex-rs/core/src/session/session.rs | 2 + codex-rs/core/src/session/tests.rs | 83 ++----------------- codex-rs/core/src/session/turn_context.rs | 65 ++++----------- codex-rs/core/src/shell.rs | 20 +++++ codex-rs/core/src/thread_manager.rs | 24 +++++- codex-rs/core/tests/suite/remote_env.rs | 37 +++++++++ 11 files changed, 162 insertions(+), 165 deletions(-) diff --git a/codex-rs/app-server/src/request_processors/request_errors.rs b/codex-rs/app-server/src/request_processors/request_errors.rs index 18082aebe..48bae43af 100644 --- a/codex-rs/app-server/src/request_processors/request_errors.rs +++ b/codex-rs/app-server/src/request_processors/request_errors.rs @@ -1,8 +1,8 @@ use super::*; -pub(super) fn environment_selection_error_message(err: CodexErr) -> String { +pub(super) fn environment_selection_error(err: CodexErr) -> JSONRPCErrorError { match err { - CodexErr::InvalidRequest(message) => message, - err => err.to_string(), + CodexErr::InvalidRequest(message) => invalid_request(message), + err => internal_error(format!("failed to validate environment selections: {err}")), } } diff --git a/codex-rs/app-server/src/request_processors/thread_processor.rs b/codex-rs/app-server/src/request_processors/thread_processor.rs index 58719d55b..813a85fb4 100644 --- a/codex-rs/app-server/src/request_processors/thread_processor.rs +++ b/codex-rs/app-server/src/request_processors/thread_processor.rs @@ -1299,7 +1299,7 @@ impl ThreadRequestProcessor { if let Some(environment_selections) = environment_selections.as_ref() { self.thread_manager .validate_environment_selections(environment_selections) - .map_err(|err| invalid_request(environment_selection_error_message(err)))?; + .map_err(environment_selection_error)?; } Ok(environment_selections) } diff --git a/codex-rs/app-server/src/request_processors/turn_processor.rs b/codex-rs/app-server/src/request_processors/turn_processor.rs index b2e6ddc0b..fe5d3278e 100644 --- a/codex-rs/app-server/src/request_processors/turn_processor.rs +++ b/codex-rs/app-server/src/request_processors/turn_processor.rs @@ -346,7 +346,7 @@ impl TurnRequestProcessor { if let Some(environment_selections) = environment_selections.as_ref() { self.thread_manager .validate_environment_selections(environment_selections) - .map_err(|err| invalid_request(environment_selection_error_message(err)))?; + .map_err(environment_selection_error)?; } Ok(environment_selections) } diff --git a/codex-rs/core/src/context/environment_context.rs b/codex-rs/core/src/context/environment_context.rs index aaaa2b62b..dd11ebd39 100644 --- a/codex-rs/core/src/context/environment_context.rs +++ b/codex-rs/core/src/context/environment_context.rs @@ -49,7 +49,8 @@ impl EnvironmentContextEnvironment { cwd: environment.cwd.clone(), shell: environment .shell - .clone() + .as_ref() + .map(|shell| shell.name().to_string()) .unwrap_or_else(|| shell.name().to_string()), }) .collect() diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index 479132925..3ac1e43a3 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -9,6 +9,7 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use crate::session::turn_context::TurnEnvironment; +use crate::shell::Shell; pub(crate) fn default_thread_environment_selections( environment_manager: &EnvironmentManager, @@ -60,7 +61,7 @@ impl ResolvedTurnEnvironments { } } -pub(crate) fn resolve_environment_selections( +pub(crate) async fn resolve_environment_selections( environment_manager: &EnvironmentManager, environments: &[TurnEnvironmentSelection], ) -> CodexResult { @@ -79,19 +80,34 @@ pub(crate) fn resolve_environment_selections( .ok_or_else(|| { CodexErr::InvalidRequest(format!("unknown turn environment id `{environment_id}`")) })?; + let shell = match environment.info().await { + Ok(info) => match Shell::from_environment_shell_info(info.shell) { + Ok(shell) => Some(shell), + Err(err) => { + tracing::warn!( + "failed to resolve shell for environment `{environment_id}`: {err}" + ); + None + } + }, + Err(err) => { + tracing::warn!("failed to get info for environment `{environment_id}`: {err}"); + None + } + }; turn_environments.push(TurnEnvironment { environment_id, environment, cwd: selected_environment.cwd.clone(), - shell: None, + shell, }); } - Ok(ResolvedTurnEnvironments { turn_environments }) } #[cfg(test)] mod tests { + use codex_exec_server::Environment; use codex_exec_server::ExecServerRuntimePaths; use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_exec_server::REMOTE_ENVIRONMENT_ID; @@ -189,6 +205,7 @@ url = "ws://127.0.0.1:8765" }, ], ) + .await .expect_err("duplicate environment id should fail"); assert!(err.to_string().contains("duplicate")); @@ -207,6 +224,7 @@ url = "ws://127.0.0.1:8765" cwd: selected_cwd, }], ) + .await .expect("environment selections should resolve"); assert_eq!( @@ -216,7 +234,21 @@ url = "ws://127.0.0.1:8765" .environment_id, "local" ); - assert_eq!(resolved.primary().expect("primary environment").shell, None); + assert_eq!( + resolved.primary().expect("primary environment").shell, + Some( + Shell::from_environment_shell_info( + manager + .get_environment("local") + .expect("local environment") + .info() + .await + .expect("local environment info") + .shell + ) + .expect("resolved shell") + ) + ); } #[tokio::test] @@ -230,40 +262,31 @@ url = "ws://127.0.0.1:8765" cwd: cwd.clone(), }], ) + .await .expect("local environment should resolve"); - let remote_manager = EnvironmentManager::create_for_tests( - Some("ws://127.0.0.1:8765".to_string()), - Some(test_runtime_paths()), - ) - .await; - let remote = resolve_environment_selections( - &remote_manager, - &[TurnEnvironmentSelection { + let remote_environment = Arc::new( + Environment::create_for_tests(Some("ws://127.0.0.1:8765".to_string())) + .expect("remote environment"), + ); + let remote = ResolvedTurnEnvironments { + turn_environments: vec![TurnEnvironment { environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + environment: remote_environment.clone(), cwd: cwd.clone(), + shell: None, }], - ) - .expect("remote environment should resolve"); - local_manager - .upsert_environment( - REMOTE_ENVIRONMENT_ID.to_string(), - "ws://127.0.0.1:8765".to_string(), - ) - .expect("remote environment should register"); - let multiple = resolve_environment_selections( - &local_manager, - &[ - TurnEnvironmentSelection { - environment_id: LOCAL_ENVIRONMENT_ID.to_string(), - cwd: cwd.clone(), - }, - TurnEnvironmentSelection { + }; + let multiple = ResolvedTurnEnvironments { + turn_environments: vec![ + local.primary().expect("local environment").clone(), + TurnEnvironment { environment_id: REMOTE_ENVIRONMENT_ID.to_string(), + environment: remote_environment, cwd: cwd.clone(), + shell: None, }, ], - ) - .expect("multiple environments should resolve"); + }; assert_eq!(local.single_local_environment_cwd(), Some(&cwd)); assert_eq!(remote.single_local_environment_cwd(), None); diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index b8d6d7a1b..df60870a6 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -444,6 +444,7 @@ async fn warm_plugins_and_skills_for_session_init( environment_manager.as_ref(), &environments, ) + .await .ok() .and_then(|resolved| resolved.primary_filesystem()); let plugins_input = config.plugins_config_input(); @@ -1117,6 +1118,7 @@ impl Session { sess.services.environment_manager.as_ref(), session_configuration.environment_selections(), ) + .await .map_err(|err| { CodexErr::InvalidRequest(err.to_string().replace( "unknown turn environment id", diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0dbb13fcd..214489a18 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -55,7 +55,6 @@ use codex_protocol::permissions::FileSystemSandboxPolicy; use codex_protocol::permissions::FileSystemSpecialPath; use codex_protocol::protocol::NonSteerableTurnKind; use codex_protocol::protocol::SandboxPolicy; -use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::TurnEnvironmentSelections; use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; @@ -6307,82 +6306,6 @@ async fn empty_turn_environments_clear_primary_environment() { assert_eq!(turn_context.config.cwd, session.get_config().await.cwd); } -#[tokio::test] -async fn unknown_turn_environment_returns_error() { - let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; - let original_configuration = { - let state = session.state.lock().await; - state.session_configuration.clone() - }; - - let err = session - .new_turn_with_sub_id( - "sub-1".to_string(), - SessionSettingsUpdate { - environments: Some(TurnEnvironmentSelections::new( - original_configuration.cwd().clone(), - vec![TurnEnvironmentSelection { - environment_id: "missing".to_string(), - cwd: original_configuration.cwd().clone(), - }], - )), - ..Default::default() - }, - ) - .await - .expect_err("unknown environment should fail"); - - let current_configuration = { - let state = session.state.lock().await; - state.session_configuration.clone() - }; - assert!(matches!(err, CodexErr::InvalidRequest(_))); - assert!(err.to_string().contains("missing")); - assert_eq!(current_configuration.cwd(), original_configuration.cwd()); - assert_eq!( - current_configuration.environment_selections(), - original_configuration.environment_selections() - ); -} - -#[tokio::test] -async fn duplicate_turn_environment_returns_error_without_mutating_session() { - let (session, _turn_context, _rx) = make_session_and_context_with_rx().await; - let original_configuration = { - let state = session.state.lock().await; - state.session_configuration.clone() - }; - - let err = session - .new_turn_with_sub_id( - "sub-1".to_string(), - SessionSettingsUpdate { - environments: Some(TurnEnvironmentSelections::new( - original_configuration.cwd().clone(), - vec![ - local(original_configuration.cwd().clone()), - local(original_configuration.cwd().join("second")), - ], - )), - ..Default::default() - }, - ) - .await - .expect_err("duplicate environment should fail"); - - let current_configuration = { - let state = session.state.lock().await; - state.session_configuration.clone() - }; - assert!(matches!(err, CodexErr::InvalidRequest(_))); - assert!(err.to_string().contains("duplicate")); - assert_eq!(current_configuration.cwd(), original_configuration.cwd()); - assert_eq!( - current_configuration.environment_selections(), - original_configuration.environment_selections() - ); -} - #[tokio::test] async fn spawn_task_turn_span_inherits_dispatch_trace_context() { struct TraceCaptureTask { @@ -7323,7 +7246,11 @@ async fn environment_context_uses_session_shell_when_environment_shell_is_absent .turn_environments .first_mut() .expect("primary environment"); - primary_environment.shell = Some("cmd".to_string()); + primary_environment.shell = Some(crate::shell::Shell { + shell_type: crate::shell::ShellType::Cmd, + shell_path: PathBuf::from("cmd"), + shell_snapshot: crate::shell::empty_shell_snapshot_receiver(), + }); let environment_context = crate::context::EnvironmentContext::from_turn_context( &turn_context, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 5e5d5b760..c249d713a 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -40,7 +40,7 @@ pub(crate) struct TurnEnvironment { pub(crate) environment_id: String, pub(crate) environment: Arc, pub(crate) cwd: AbsolutePathBuf, - pub(crate) shell: Option, + pub(crate) shell: Option, } impl TurnEnvironment { @@ -590,8 +590,6 @@ impl Session { let mut state = self.state.lock().await; match state.session_configuration.clone().apply(&updates) { Ok(next) => { - let turn_environments = - self.resolve_turn_environments(next.environment_selections())?; let previous_cwd = state.session_configuration.cwd().clone(); let previous_permission_profile = state.session_configuration.permission_profile(); @@ -608,7 +606,6 @@ impl Session { state.session_configuration = next.clone(); Ok(( next, - turn_environments, permission_profile_changed, previous_cwd, codex_home, @@ -623,7 +620,6 @@ impl Session { let ( session_configuration, - turn_environments, permission_profile_changed, previous_cwd, codex_home, @@ -664,33 +660,20 @@ impl Session { sub_id, session_configuration, updates.final_output_json_schema, - turn_environments, ) .await) } - fn resolve_turn_environments( - &self, - environments: &[TurnEnvironmentSelection], - ) -> CodexResult { - crate::environment_selection::resolve_environment_selections( - self.services.environment_manager.as_ref(), - environments, - ) - } - async fn new_turn_from_configuration( &self, sub_id: String, session_configuration: SessionConfiguration, final_output_json_schema: Option>, - turn_environments: ResolvedTurnEnvironments, ) -> Arc { self.new_turn_context_from_configuration( sub_id, session_configuration, final_output_json_schema, - turn_environments, TurnMultiAgentRuntime::ResolveAndStore, ) .await @@ -700,13 +683,11 @@ impl Session { &self, sub_id: String, session_configuration: SessionConfiguration, - turn_environments: ResolvedTurnEnvironments, ) -> Arc { self.new_turn_context_from_configuration( sub_id, session_configuration, /*final_output_json_schema*/ None, - turn_environments, TurnMultiAgentRuntime::Preview, ) .await @@ -717,9 +698,17 @@ impl Session { sub_id: String, session_configuration: SessionConfiguration, final_output_json_schema: Option>, - turn_environments: ResolvedTurnEnvironments, multi_agent_runtime: TurnMultiAgentRuntime, ) -> Arc { + let turn_environments = crate::environment_selection::resolve_environment_selections( + self.services.environment_manager.as_ref(), + session_configuration.environment_selections(), + ) + .await + .unwrap_or_else(|err| { + warn!("failed to resolve turn environments: {err}"); + ResolvedTurnEnvironments::default() + }); let primary_turn_environment = turn_environments.primary().cloned(); let cwd = primary_turn_environment .as_ref() @@ -832,13 +821,11 @@ impl Session { } pub(crate) async fn new_default_turn_with_sub_id(&self, sub_id: String) -> Arc { - let (session_configuration, turn_environments) = - self.default_turn_configuration_and_environments().await; + let session_configuration = self.default_turn_configuration().await; self.new_turn_from_configuration( sub_id, session_configuration, /*final_output_json_schema*/ None, - turn_environments, ) .await } @@ -847,31 +834,13 @@ impl Session { &self, sub_id: String, ) -> Arc { - let (session_configuration, turn_environments) = - self.default_turn_configuration_and_environments().await; - self.new_startup_prewarm_turn_from_configuration( - sub_id, - session_configuration, - turn_environments, - ) - .await + let session_configuration = self.default_turn_configuration().await; + self.new_startup_prewarm_turn_from_configuration(sub_id, session_configuration) + .await } - async fn default_turn_configuration_and_environments( - &self, - ) -> (SessionConfiguration, ResolvedTurnEnvironments) { - let session_configuration = { - let state = self.state.lock().await; - state.session_configuration.clone() - }; - let turn_environments = - match self.resolve_turn_environments(session_configuration.environment_selections()) { - Ok(turn_environments) => turn_environments, - Err(err) => { - warn!("failed to resolve stored session environments: {err}"); - ResolvedTurnEnvironments::default() - } - }; - (session_configuration, turn_environments) + async fn default_turn_configuration(&self) -> SessionConfiguration { + let state = self.state.lock().await; + state.session_configuration.clone() } } diff --git a/codex-rs/core/src/shell.rs b/codex-rs/core/src/shell.rs index 9843d4cea..b1049ef1d 100644 --- a/codex-rs/core/src/shell.rs +++ b/codex-rs/core/src/shell.rs @@ -1,4 +1,5 @@ use crate::shell_snapshot::ShellSnapshot; +use codex_exec_server::ShellInfo; use codex_shell_command::shell_detect::DetectedShell; use serde::Deserialize; use serde::Serialize; @@ -85,6 +86,25 @@ impl From for Shell { } } +impl Shell { + pub(crate) fn from_environment_shell_info(shell_info: ShellInfo) -> anyhow::Result { + let shell_type = match shell_info.name.as_str() { + "zsh" => ShellType::Zsh, + "bash" => ShellType::Bash, + "powershell" => ShellType::PowerShell, + "sh" => ShellType::Sh, + "cmd" => ShellType::Cmd, + name => anyhow::bail!("unknown environment shell `{name}`"), + }; + + Ok(Self { + shell_type, + shell_path: PathBuf::from(shell_info.path), + shell_snapshot: empty_shell_snapshot_receiver(), + }) + } +} + #[cfg(all(test, unix))] fn ultimate_fallback_shell() -> Shell { codex_shell_command::shell_detect::ultimate_fallback_shell().into() diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b82e63947..c332622ac 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -451,8 +451,25 @@ impl ThreadManager { &self, environments: &[TurnEnvironmentSelection], ) -> CodexResult<()> { - resolve_environment_selections(self.state.environment_manager.as_ref(), environments) - .map(|_| ()) + let mut environment_ids = HashSet::with_capacity(environments.len()); + for environment in environments { + if !environment_ids.insert(environment.environment_id.as_str()) { + return Err(CodexErr::InvalidRequest(format!( + "duplicate turn environment id `{}`", + environment.environment_id + ))); + } + self.state + .environment_manager + .get_environment(&environment.environment_id) + .ok_or_else(|| { + CodexErr::InvalidRequest(format!( + "unknown turn environment id `{}`", + environment.environment_id + )) + })?; + } + Ok(()) } pub fn get_models_manager(&self) -> SharedModelsManager { @@ -1365,7 +1382,8 @@ impl ThreadManagerState { } } let environment_selections = - resolve_environment_selections(self.environment_manager.as_ref(), &environments)?; + resolve_environment_selections(self.environment_manager.as_ref(), &environments) + .await?; let user_instructions = self .user_instructions_for_spawn(&session_source, parent_thread_id, forked_from_thread_id) .await; diff --git a/codex-rs/core/tests/suite/remote_env.rs b/codex-rs/core/tests/suite/remote_env.rs index 64d5c7081..2f18425c2 100644 --- a/codex-rs/core/tests/suite/remote_env.rs +++ b/codex-rs/core/tests/suite/remote_env.rs @@ -37,6 +37,7 @@ use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; @@ -181,6 +182,42 @@ async fn remote_test_env_can_connect_and_use_filesystem() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn remote_test_env_exposes_bash_shell_to_model() -> Result<()> { + let Some(_remote_env) = get_remote_test_env() else { + return Ok(()); + }; + + let server = start_mock_server().await; + let response_mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + let test = test_codex().build_with_remote_env(&server).await?; + + test.submit_turn("report remote environment").await?; + + let request = response_mock.single_request(); + let environment_context = request + .message_input_texts("user") + .into_iter() + .find(|text| text.starts_with("")) + .context("environment context should be model visible")?; + assert_eq!( + environment_context + .lines() + .find(|line| line.trim_start().starts_with("")), + Some(" bash"), + ); + + Ok(()) +} + fn absolute_path(path: PathBuf) -> AbsolutePathBuf { match AbsolutePathBuf::try_from(path) { Ok(path) => path,