Address turn environment review feedback

Co-authored-by: Codex <noreply@openai.com>

codex-rs/codex-mcp/src/runtime.rs

@@ -50,11 +50,11 @@ impl McpRuntimeEnvironment {
         }
     }
 
-    pub fn environment(&self) -> Arc<Environment> {
+    pub(crate) fn environment(&self) -> Arc<Environment> {
         Arc::clone(&self.environment)
     }
 
-    pub fn fallback_cwd(&self) -> PathBuf {
+    pub(crate) fn fallback_cwd(&self) -> PathBuf {
         self.fallback_cwd.clone()
     }
 }

codex-rs/core/src/codex_delegate.rs

@@ -99,6 +99,7 @@ pub(crate) async fn run_codex_thread_interactive(
             .iter()
             .map(TurnEnvironment::selection)
             .collect(),
+        resolved_environments: None,
         analytics_events_client: Some(parent_session.services.analytics_events_client.clone()),
         thread_store: Arc::clone(&parent_session.services.thread_store),
     }))

codex-rs/core/src/environment_selection.rs

@@ -1,13 +1,15 @@
 use std::collections::HashSet;
 use std::sync::Arc;
 
-use codex_exec_server::Environment;
 use codex_exec_server::EnvironmentManager;
+use codex_exec_server::ExecutorFileSystem;
 use codex_protocol::error::CodexErr;
 use codex_protocol::error::Result as CodexResult;
 use codex_protocol::protocol::TurnEnvironmentSelection;
 use codex_utils_absolute_path::AbsolutePathBuf;
 
+use crate::session::turn_context::TurnEnvironment;
+
 pub(crate) fn default_thread_environment_selections(
     environment_manager: &EnvironmentManager,
     cwd: &AbsolutePathBuf,
@@ -22,59 +24,87 @@ pub(crate) fn default_thread_environment_selections(
         .collect()
 }
 
-pub(crate) fn validate_environment_selections(
+#[derive(Clone)]
+pub(crate) struct ResolvedEnvironmentSelections {
+    pub(crate) turn_environments: Vec<TurnEnvironment>,
+}
+
+impl ResolvedEnvironmentSelections {
+    pub(crate) fn primary_environment(&self) -> Option<&TurnEnvironment> {
+        self.turn_environments.first()
+    }
+
+    pub(crate) fn primary_cwd_or_fallback(
+        &self,
+        fallback_cwd: &AbsolutePathBuf,
+    ) -> AbsolutePathBuf {
+        self.primary_environment()
+            .map(|environment| environment.cwd.clone())
+            .unwrap_or_else(|| fallback_cwd.clone())
+    }
+
+    pub(crate) fn primary_environment_backend(
+        &self,
+    ) -> Option<Arc<codex_exec_server::Environment>> {
+        self.primary_environment()
+            .map(|environment| Arc::clone(&environment.environment))
+    }
+
+    pub(crate) fn primary_filesystem(&self) -> Option<Arc<dyn ExecutorFileSystem>> {
+        self.primary_environment()
+            .map(|environment| environment.environment.get_filesystem())
+    }
+}
+
+pub(crate) fn resolve_environment_selections(
     environment_manager: &EnvironmentManager,
     environments: &[TurnEnvironmentSelection],
-) -> CodexResult<()> {
+) -> CodexResult<ResolvedEnvironmentSelections> {
+    resolve_environment_selections_with_primary_cwd(
+        environment_manager,
+        environments,
+        /*primary_cwd_override*/ None,
+    )
+}
+
+pub(crate) fn resolve_environment_selections_with_primary_cwd(
+    environment_manager: &EnvironmentManager,
+    environments: &[TurnEnvironmentSelection],
+    primary_cwd_override: Option<&AbsolutePathBuf>,
+) -> CodexResult<ResolvedEnvironmentSelections> {
     let mut seen_environment_ids = HashSet::with_capacity(environments.len());
-    for selected_environment in environments {
+    let mut turn_environments = Vec::with_capacity(environments.len());
+    for (index, selected_environment) in environments.iter().enumerate() {
         if !seen_environment_ids.insert(selected_environment.environment_id.as_str()) {
             return Err(CodexErr::InvalidRequest(format!(
                 "duplicate turn environment id `{}`",
                 selected_environment.environment_id
             )));
         }
-        if environment_manager
-            .get_environment(&selected_environment.environment_id)
-            .is_none()
-        {
-            return Err(CodexErr::InvalidRequest(format!(
-                "unknown turn environment id `{}`",
-                selected_environment.environment_id
-            )));
-        }
+        let environment_id = selected_environment.environment_id.clone();
+        let environment = environment_manager
+            .get_environment(&environment_id)
+            .ok_or_else(|| {
+                CodexErr::InvalidRequest(format!("unknown turn environment id `{environment_id}`"))
+            })?;
+        turn_environments.push(TurnEnvironment {
+            environment_id,
+            environment,
+            cwd: primary_cwd_override
+                .filter(|_| index == 0)
+                .cloned()
+                .unwrap_or_else(|| selected_environment.cwd.clone()),
+        });
     }
 
-    Ok(())
+    Ok(ResolvedEnvironmentSelections { turn_environments })
 }
 
-pub(crate) fn selected_primary_environment(
+pub(crate) fn validate_environment_selections(
     environment_manager: &EnvironmentManager,
     environments: &[TurnEnvironmentSelection],
-) -> CodexResult<Option<Arc<Environment>>> {
-    environments
-        .first()
-        .map(|selected_environment| {
-            environment_manager
-                .get_environment(&selected_environment.environment_id)
-                .ok_or_else(|| {
-                    CodexErr::InvalidRequest(format!(
-                        "unknown turn environment id `{}`",
-                        selected_environment.environment_id
-                    ))
-                })
-        })
-        .transpose()
-}
-
-pub(crate) fn primary_selected_cwd_or_fallback(
-    environments: &[TurnEnvironmentSelection],
-    fallback_cwd: &AbsolutePathBuf,
-) -> AbsolutePathBuf {
-    environments
-        .first()
-        .map(|selected_environment| selected_environment.cwd.clone())
-        .unwrap_or_else(|| fallback_cwd.clone())
+) -> CodexResult<()> {
+    resolve_environment_selections(environment_manager, environments).map(|_| ())
 }
 
 #[cfg(test)]
@@ -148,26 +178,58 @@ mod tests {
     }
 
     #[test]
-    fn primary_selected_cwd_or_fallback_uses_first_selection() {
+    fn resolved_environment_selections_use_first_selection_as_primary() {
         let cwd = AbsolutePathBuf::current_dir().expect("cwd");
         let selected_cwd = cwd.join("selected");
+        let manager = EnvironmentManager::default_for_tests();
 
+        let resolved = resolve_environment_selections(
+            &manager,
+            &[TurnEnvironmentSelection {
+                environment_id: "local".to_string(),
+                cwd: selected_cwd.clone(),
+            }],
+        )
+        .expect("environment selections should resolve");
+
+        assert_eq!(resolved.primary_cwd_or_fallback(&cwd), selected_cwd);
         assert_eq!(
-            primary_selected_cwd_or_fallback(
-                &[TurnEnvironmentSelection {
-                    environment_id: "local".to_string(),
-                    cwd: selected_cwd.clone(),
-                }],
-                &cwd,
-            ),
-            selected_cwd
+            resolved
+                .primary_environment()
+                .expect("primary environment")
+                .environment_id,
+            "local"
         );
     }
 
     #[test]
-    fn primary_selected_cwd_or_fallback_uses_fallback_without_selections() {
+    fn resolved_environment_selections_can_override_primary_cwd() {
         let cwd = AbsolutePathBuf::current_dir().expect("cwd");
+        let selected_cwd = cwd.join("selected");
+        let override_cwd = cwd.join("override");
+        let manager = EnvironmentManager::default_for_tests();
 
-        assert_eq!(primary_selected_cwd_or_fallback(&[], &cwd), cwd);
+        let resolved = resolve_environment_selections_with_primary_cwd(
+            &manager,
+            &[TurnEnvironmentSelection {
+                environment_id: "local".to_string(),
+                cwd: selected_cwd,
+            }],
+            Some(&override_cwd),
+        )
+        .expect("environment selections should resolve");
+
+        assert_eq!(resolved.primary_cwd_or_fallback(&cwd), override_cwd);
+    }
+
+    #[test]
+    fn resolved_environment_selections_use_fallback_without_selections() {
+        let cwd = AbsolutePathBuf::current_dir().expect("cwd");
+        let manager = EnvironmentManager::default_for_tests();
+
+        let resolved =
+            resolve_environment_selections(&manager, &[]).expect("empty selections are valid");
+
+        assert_eq!(resolved.primary_cwd_or_fallback(&cwd), cwd);
     }
 }

codex-rs/core/src/session/mcp.rs

@@ -1,40 +1,24 @@
 use super::*;
 
 impl Session {
-    pub(crate) fn mcp_runtime_environment_for_turn_context(
-        &self,
-        turn_context: &TurnContext,
-    ) -> McpRuntimeEnvironment {
-        if let Some(turn_environment) = turn_context.primary_environment() {
-            return McpRuntimeEnvironment::new(
-                Arc::clone(&turn_environment.environment),
-                turn_environment.cwd.to_path_buf(),
-            );
-        }
-
-        McpRuntimeEnvironment::new(
-            self.services.environment_manager.local_environment(),
-            turn_context.cwd.to_path_buf(),
-        )
-    }
-
     pub(crate) fn mcp_runtime_environment_for_configuration(
         &self,
         session_configuration: &SessionConfiguration,
     ) -> CodexResult<McpRuntimeEnvironment> {
-        if let Some(turn_environment) = session_configuration.environments.first() {
-            let environment = self
-                .services
-                .environment_manager
-                .get_environment(&turn_environment.environment_id)
-                .ok_or_else(|| {
-                    CodexErr::InvalidRequest(format!(
-                        "unknown stored MCP environment id `{}`",
-                        turn_environment.environment_id
-                    ))
-                })?;
+        let resolved_environments = crate::environment_selection::resolve_environment_selections(
+            self.services.environment_manager.as_ref(),
+            &session_configuration.environments,
+        )
+        .map_err(|err| {
+            CodexErr::InvalidRequest(err.to_string().replace(
+                "unknown turn environment id",
+                "unknown stored MCP environment id",
+            ))
+        })?;
+
+        if let Some(turn_environment) = resolved_environments.primary_environment() {
             return Ok(McpRuntimeEnvironment::new(
-                environment,
+                Arc::clone(&turn_environment.environment),
                 turn_environment.cwd.to_path_buf(),
             ));
         }
@@ -265,6 +249,18 @@ impl Session {
         let mcp_servers = with_codex_apps_mcp(mcp_servers, auth.as_ref(), &mcp_config);
         let auth_statuses =
             compute_auth_statuses(mcp_servers.iter(), store_mode, auth.as_ref()).await;
+        let session_configuration = {
+            let state = self.state.lock().await;
+            state.session_configuration.clone()
+        };
+        let mcp_runtime_environment =
+            match self.mcp_runtime_environment_for_configuration(&session_configuration) {
+                Ok(runtime_environment) => runtime_environment,
+                Err(err) => {
+                    warn!("failed to resolve MCP runtime environment for refresh: {err}");
+                    return;
+                }
+            };
         {
             let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
             guard.cancel();
@@ -278,7 +274,7 @@ impl Session {
             turn_context.sub_id.clone(),
             self.get_tx_event(),
             turn_context.permission_profile(),
-            self.mcp_runtime_environment_for_turn_context(turn_context),
+            mcp_runtime_environment,
             config.codex_home.to_path_buf(),
             codex_apps_tools_cache_key(auth.as_ref()),
             tool_plugin_provenance,

codex-rs/core/src/session/mod.rs

@@ -30,9 +30,8 @@ use crate::context::NetworkRuleSaved;
 use crate::context::PermissionsInstructions;
 use crate::context::PersonalitySpecInstructions;
 use crate::default_skill_metadata_budget;
-use crate::environment_selection::primary_selected_cwd_or_fallback;
-use crate::environment_selection::selected_primary_environment;
-use crate::environment_selection::validate_environment_selections;
+use crate::environment_selection::ResolvedEnvironmentSelections;
+use crate::environment_selection::resolve_environment_selections;
 use crate::exec_policy::ExecPolicyManager;
 use crate::installation_id::resolve_installation_id;
 use crate::parse_turn_item;
@@ -408,6 +407,7 @@ pub(crate) struct CodexSpawnArgs {
     pub(crate) user_shell_override: Option<shell::Shell>,
     pub(crate) parent_trace: Option<W3cTraceContext>,
     pub(crate) environments: Vec<TurnEnvironmentSelection>,
+    pub(crate) resolved_environments: Option<ResolvedEnvironmentSelections>,
     pub(crate) analytics_events_client: Option<AnalyticsEventsClient>,
     pub(crate) thread_store: Arc<dyn ThreadStore>,
 }
@@ -465,19 +465,19 @@ impl Codex {
             parent_rollout_thread_trace,
             parent_trace: _,
             environments,
+            resolved_environments,
             analytics_events_client,
             thread_store,
         } = args;
         let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
         let (tx_event, rx_event) = async_channel::unbounded();
-        validate_environment_selections(environment_manager.as_ref(), &environments)?;
-        let environment =
-            selected_primary_environment(environment_manager.as_ref(), &environments)?;
-        let mut load_config = config.clone();
-        load_config.cwd = primary_selected_cwd_or_fallback(&environments, &config.cwd);
-        let fs = environment
-            .as_ref()
-            .map(|environment| environment.get_filesystem());
+        let resolved_environments = match resolved_environments {
+            Some(resolved_environments) => resolved_environments,
+            None => resolve_environment_selections(environment_manager.as_ref(), &environments)?,
+        };
+        config.cwd = resolved_environments.primary_cwd_or_fallback(&config.cwd);
+        let load_config = config.clone();
+        let fs = resolved_environments.primary_filesystem();
         let plugins_input = load_config.plugins_config_input();
         let plugin_outcome = plugins_manager.plugins_for_config(&plugins_input).await;
         let effective_skill_roots = plugin_outcome.effective_skill_roots();
@@ -500,8 +500,9 @@ impl Codex {
             let _ = config.features.disable(Feature::Collab);
         }
 
+        let primary_environment = resolved_environments.primary_environment_backend();
         let user_instructions = AgentsMdManager::new(&load_config)
-            .user_instructions(environment.as_deref())
+            .user_instructions(primary_environment.as_deref())
             .await;
 
         let exec_policy = if crate::guardian::is_guardian_reviewer_source(&session_source) {

codex-rs/core/src/session/session.rs

@@ -199,11 +199,11 @@ impl SessionConfiguration {
             .cwd
             .as_ref()
             .map(|cwd| {
-                crate::environment_selection::primary_selected_cwd_or_fallback(
-                    effective_environments,
-                    &self.cwd,
-                )
-                .join(normalize_for_native_workdir(cwd.as_path()))
+                let base_cwd = effective_environments
+                    .first()
+                    .map(|environment| &environment.cwd)
+                    .unwrap_or(&self.cwd);
+                base_cwd.join(normalize_for_native_workdir(cwd.as_path()))
             })
             .unwrap_or_else(|| self.cwd.clone());
 

codex-rs/core/src/session/tests.rs

@@ -4621,98 +4621,6 @@ async fn duplicate_turn_environment_returns_error_without_mutating_session() {
     );
 }
 
-#[tokio::test]
-async fn mcp_runtime_environment_uses_primary_selected_environment() {
-    let (session, _turn_context) = make_session_and_context().await;
-    let selected_cwd = session
-        .get_config()
-        .await
-        .cwd
-        .join("mcp-selected-environment");
-    {
-        let mut state = session.state.lock().await;
-        state.session_configuration.environments = vec![TurnEnvironmentSelection {
-            environment_id: codex_exec_server::LOCAL_ENVIRONMENT_ID.to_string(),
-            cwd: selected_cwd.clone(),
-        }];
-    }
-
-    let session_configuration = {
-        let state = session.state.lock().await;
-        state.session_configuration.clone()
-    };
-    let runtime_environment = session
-        .mcp_runtime_environment_for_configuration(&session_configuration)
-        .expect("runtime environment");
-    let turn_environment = session
-        .services
-        .environment_manager
-        .get_environment(codex_exec_server::LOCAL_ENVIRONMENT_ID)
-        .expect("local environment");
-
-    assert!(Arc::ptr_eq(
-        &runtime_environment.environment(),
-        &turn_environment
-    ));
-    assert_eq!(
-        runtime_environment.fallback_cwd(),
-        selected_cwd.to_path_buf()
-    );
-}
-
-#[tokio::test]
-async fn mcp_runtime_environment_for_turn_context_uses_turn_environment() {
-    let (session, mut turn_context) = make_session_and_context().await;
-    let selected_cwd = turn_context.cwd.join("mcp-turn-environment");
-    turn_context.environments[0].cwd = selected_cwd.clone();
-    {
-        let mut state = session.state.lock().await;
-        state.session_configuration.environments = Vec::new();
-        state.session_configuration.cwd = turn_context.cwd.join("stored-session-cwd");
-    }
-
-    let runtime_environment = session.mcp_runtime_environment_for_turn_context(&turn_context);
-    let turn_environment = turn_context
-        .primary_environment()
-        .expect("primary turn environment");
-
-    assert!(Arc::ptr_eq(
-        &runtime_environment.environment(),
-        &turn_environment.environment
-    ));
-    assert_eq!(
-        runtime_environment.fallback_cwd(),
-        selected_cwd.to_path_buf()
-    );
-}
-
-#[tokio::test]
-async fn mcp_runtime_environment_without_selection_uses_local_fallback() {
-    let (session, _turn_context) = make_session_and_context().await;
-    {
-        let mut state = session.state.lock().await;
-        state.session_configuration.environments = Vec::new();
-    }
-
-    let session_configuration = {
-        let state = session.state.lock().await;
-        state.session_configuration.clone()
-    };
-    let runtime_environment = session
-        .mcp_runtime_environment_for_configuration(&session_configuration)
-        .expect("runtime environment");
-    let local_environment = session.services.environment_manager.local_environment();
-
-    assert!(Arc::ptr_eq(
-        &runtime_environment.environment(),
-        &local_environment
-    ));
-    assert_eq!(
-        runtime_environment.fallback_cwd(),
-        session_configuration.cwd.to_path_buf()
-    );
-}
-
 #[tokio::test]
 async fn mcp_runtime_environment_for_configuration_errors_on_unknown_environment() {
     let (session, _turn_context) = make_session_and_context().await;

codex-rs/core/src/session/tests/guardian_tests.rs

@@ -755,6 +755,7 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() {
         user_shell_override: None,
         parent_trace: None,
         environments: Vec::new(),
+        resolved_environments: None,
         analytics_events_client: None,
         thread_store,
     })

codex-rs/core/src/session/turn_context.rs

@@ -564,18 +564,15 @@ impl Session {
             let mut state = self.state.lock().await;
             match state.session_configuration.clone().apply(&updates) {
                 Ok(next) => {
-                    let mut effective_environments = updates
+                    let effective_environments = updates
                         .environments
                         .clone()
                         .unwrap_or_else(|| next.environments.clone());
-                    if updates.environments.is_some()
-                        && updates.cwd.is_some()
-                        && let Some(turn_environment) = effective_environments.first_mut()
-                    {
-                        turn_environment.cwd = next.cwd.clone();
-                    }
-                    let turn_environments =
-                        self.resolve_turn_environments(&effective_environments)?;
+                    let primary_cwd_override = (updates.environments.is_some()
+                        && updates.cwd.is_some())
+                    .then_some(&next.cwd);
+                    let turn_environments = self
+                        .resolve_turn_environments(&effective_environments, primary_cwd_override)?;
                     let previous_cwd = state.session_configuration.cwd.clone();
                     let previous_permission_profile =
                         state.session_configuration.permission_profile();
@@ -646,32 +643,14 @@ impl Session {
     fn resolve_turn_environments(
         &self,
         environments: &[TurnEnvironmentSelection],
+        primary_cwd_override: Option<&AbsolutePathBuf>,
     ) -> CodexResult<Vec<TurnEnvironment>> {
-        crate::environment_selection::validate_environment_selections(
+        crate::environment_selection::resolve_environment_selections_with_primary_cwd(
             self.services.environment_manager.as_ref(),
             environments,
-        )?;
-        let mut turn_environments = Vec::with_capacity(environments.len());
-        for selected_environment in environments {
-            let environment_id = selected_environment.environment_id.clone();
-            let environment = self
-                .services
-                .environment_manager
-                .get_environment(&environment_id)
-                .ok_or_else(|| {
-                    CodexErr::InvalidRequest(format!(
-                        "unknown turn environment id `{environment_id}`"
-                    ))
-                })?;
-            let cwd = selected_environment.cwd.clone();
-            turn_environments.push(TurnEnvironment {
-                environment_id,
-                environment,
-                cwd,
-            });
-        }
-
-        Ok(turn_environments)
+            primary_cwd_override,
+        )
+        .map(|resolved| resolved.turn_environments)
     }
 
     async fn new_turn_from_configuration(
@@ -779,14 +758,16 @@ impl Session {
             let state = self.state.lock().await;
             state.session_configuration.clone()
         };
-        let turn_environments =
-            match self.resolve_turn_environments(&session_configuration.environments) {
-                Ok(turn_environments) => turn_environments,
-                Err(err) => {
-                    warn!("failed to resolve stored session environments: {err}");
-                    Vec::new()
-                }
-            };
+        let turn_environments = match self.resolve_turn_environments(
+            &session_configuration.environments,
+            /*primary_cwd_override*/ None,
+        ) {
+            Ok(turn_environments) => turn_environments,
+            Err(err) => {
+                warn!("failed to resolve stored session environments: {err}");
+                Vec::new()
+            }
+        };
 
         self.new_turn_from_configuration(
             sub_id,

codex-rs/core/src/thread_manager.rs

@@ -4,7 +4,7 @@ use crate::codex_thread::CodexThread;
 use crate::config::Config;
 use crate::config::ThreadStoreConfig;
 use crate::environment_selection::default_thread_environment_selections;
-use crate::environment_selection::selected_primary_environment;
+use crate::environment_selection::resolve_environment_selections;
 use crate::environment_selection::validate_environment_selections;
 use crate::file_watcher::FileWatcher;
 use crate::mcp::McpManager;
@@ -1072,23 +1072,19 @@ impl ThreadManagerState {
         user_shell_override: Option<crate::shell::Shell>,
     ) -> CodexResult<NewThread> {
         let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_));
-        validate_environment_selections(self.environment_manager.as_ref(), &environments)?;
-        let environment =
-            selected_primary_environment(self.environment_manager.as_ref(), &environments)?;
-        let effective_cwd = environments
-            .first()
-            .map(|environment| environment.cwd.clone())
-            .unwrap_or_else(|| config.cwd.clone());
+        let resolved_environments =
+            resolve_environment_selections(self.environment_manager.as_ref(), &environments)?;
+        let effective_cwd = resolved_environments.primary_cwd_or_fallback(&config.cwd);
         let mut load_config = config.clone();
         load_config.cwd = effective_cwd;
-        let watch_registration = match environment.as_ref() {
-            Some(environment) if !environment.is_remote() => {
+        let watch_registration = match resolved_environments.primary_environment() {
+            Some(turn_environment) if !turn_environment.environment.is_remote() => {
                 self.skills_watcher
                     .register_config(
                         &load_config,
                         self.skills_manager.as_ref(),
                         self.plugins_manager.as_ref(),
-                        Some(environment.get_filesystem()),
+                        Some(turn_environment.environment.get_filesystem()),
                     )
                     .await
             }
@@ -1121,6 +1117,7 @@ impl ThreadManagerState {
             user_shell_override,
             parent_trace,
             environments,
+            resolved_environments: Some(resolved_environments),
             analytics_events_client: self.analytics_events_client.clone(),
             thread_store,
         })

codex-rs/core/tests/suite/rmcp_client.rs

@@ -20,6 +20,7 @@ use std::time::UNIX_EPOCH;
 use codex_config::types::McpServerConfig;
 use codex_config::types::McpServerEnvVar;
 use codex_config::types::McpServerTransportConfig;
+use codex_config::types::OAuthCredentialsStoreMode;
 use codex_core::config::Config;
 use codex_exec_server::CreateDirectoryOptions;
 use codex_exec_server::Environment;
@@ -40,8 +41,10 @@ use codex_protocol::openai_models::TruncationPolicyConfig;
 use codex_protocol::protocol::AskForApproval;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::McpInvocation;
+use codex_protocol::protocol::McpServerRefreshConfig;
 use codex_protocol::protocol::McpToolCallBeginEvent;
 use codex_protocol::protocol::Op;
+use codex_protocol::protocol::TurnEnvironmentSelection;
 use codex_protocol::user_input::UserInput;
 use codex_utils_cargo_bin::cargo_bin;
 use core_test_support::assert_regex_match;
@@ -335,6 +338,23 @@ async fn call_cwd_tool(
     fixture: &TestCodex,
     server_name: &str,
     call_id: &str,
+) -> anyhow::Result<Value> {
+    call_cwd_tool_with_turn(
+        server,
+        fixture,
+        server_name,
+        call_id,
+        read_only_user_turn(fixture, "call the rmcp cwd tool"),
+    )
+    .await
+}
+
+async fn call_cwd_tool_with_turn(
+    server: &MockServer,
+    fixture: &TestCodex,
+    server_name: &str,
+    call_id: &str,
+    turn: Op,
 ) -> anyhow::Result<Value> {
     let namespace = format!("mcp__{server_name}__");
     mount_sse_once(
@@ -355,10 +375,7 @@ async fn call_cwd_tool(
     )
     .await;
 
-    fixture
-        .codex
-        .submit(read_only_user_turn(fixture, "call the rmcp cwd tool"))
-        .await?;
+    fixture.codex.submit(turn).await?;
 
     wait_for_event(&fixture.codex, |ev| {
         matches!(ev, EventMsg::McpToolCallBegin(_))
@@ -547,6 +564,129 @@ async fn stdio_server_round_trip() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+#[serial(mcp_cwd)]
+async fn refreshed_stdio_server_uses_stored_environment_cwd() -> anyhow::Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = responses::start_mock_server().await;
+    let server_name = "rmcp_stored_environment_fallback_cwd";
+    let selected_dir = "mcp-turn-environment-cwd";
+    let selected_cwd = Arc::new(Mutex::new(None::<PathBuf>));
+    let selected_cwd_for_setup = Arc::clone(&selected_cwd);
+    let rmcp_test_server_bin = remote_aware_stdio_server_bin()?;
+
+    let fixture = test_codex()
+        .with_workspace_setup(move |cwd, fs| {
+            let selected_cwd_for_setup = Arc::clone(&selected_cwd_for_setup);
+            async move {
+                let selected = cwd.join(selected_dir);
+                fs.create_directory(
+                    &selected,
+                    CreateDirectoryOptions { recursive: true },
+                    /*sandbox*/ None,
+                )
+                .await?;
+                *selected_cwd_for_setup
+                    .lock()
+                    .expect("selected cwd lock should not be poisoned") =
+                    Some(selected.to_path_buf());
+                Ok(())
+            }
+        })
+        .build_remote_aware(&server)
+        .await?;
+
+    let selected_cwd = selected_cwd
+        .lock()
+        .expect("selected cwd lock should not be poisoned")
+        .clone()
+        .expect("workspace setup should record selected cwd");
+    let environment_id = if std::env::var_os(remote_env_env_var()).is_some() {
+        codex_exec_server::REMOTE_ENVIRONMENT_ID
+    } else {
+        codex_exec_server::LOCAL_ENVIRONMENT_ID
+    };
+    let refreshed_servers = HashMap::from([(
+        server_name.to_string(),
+        McpServerConfig {
+            transport: stdio_transport(
+                rmcp_test_server_bin,
+                Some(HashMap::from([(
+                    "MCP_TEST_VALUE".to_string(),
+                    "stored-environment-fallback-cwd".to_string(),
+                )])),
+                Vec::new(),
+            ),
+            experimental_environment: remote_aware_experimental_environment(),
+            enabled: true,
+            required: false,
+            startup_timeout_sec: None,
+            tool_timeout_sec: None,
+            default_tools_approval_mode: None,
+            enabled_tools: None,
+            disabled_tools: None,
+            scopes: None,
+            oauth_resource: None,
+            tools: HashMap::new(),
+        },
+    )]);
+
+    mount_sse_once(
+        &server,
+        responses::sse(vec![
+            responses::ev_response_created("resp-store-environment"),
+            responses::ev_assistant_message(
+                "msg-store-environment",
+                "stored environment selected.",
+            ),
+            responses::ev_completed("resp-store-environment"),
+        ]),
+    )
+    .await;
+    fixture
+        .codex
+        .submit(Op::UserInput {
+            environments: Some(vec![TurnEnvironmentSelection {
+                environment_id: environment_id.to_string(),
+                cwd: selected_cwd.clone().abs(),
+            }]),
+            items: vec![UserInput::Text {
+                text: "store the selected environment".to_string(),
+                text_elements: Vec::new(),
+            }],
+            final_output_json_schema: None,
+            responsesapi_client_metadata: None,
+        })
+        .await?;
+    wait_for_event(&fixture.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await;
+
+    fixture
+        .codex
+        .submit(Op::RefreshMcpServers {
+            config: McpServerRefreshConfig {
+                mcp_servers: serde_json::to_value(refreshed_servers)?,
+                mcp_oauth_credentials_store_mode: serde_json::to_value(
+                    OAuthCredentialsStoreMode::Auto,
+                )?,
+            },
+        })
+        .await?;
+
+    let structured = call_cwd_tool_with_turn(
+        &server,
+        &fixture,
+        server_name,
+        "call-stored-environment-fallback-cwd",
+        read_only_user_turn(&fixture, "call the rmcp cwd tool"),
+    )
+    .await?;
+
+    assert_cwd_tool_output(&structured, &selected_cwd);
+    server.verify().await;
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 #[serial(mcp_cwd)]
 async fn stdio_server_uses_configured_cwd_before_runtime_fallback() -> anyhow::Result<()> {