Allow environment readiness updates in place (#35875)

## What changed

- Add `EnvironmentManager::publish_ready_info` to update the selected capability
  roots for an existing environment without replacing it.
- Store readiness as an atomically replaceable snapshot so repeated publications
  expose the latest roots while deferred completion still controls connection
  readiness.
- Validate published roots and reject updates for missing environments.

## Testing

- Cover publication, replacement, repeated and invalid updates, and interaction
  with deferred environment completion.

GitOrigin-RevId: 1477f60e1319f840e6b33703ef2af052713779fd

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

@@ -323,6 +323,7 @@ impl Session {
             .services
             .selected_capability_roots
             .iter()
+            .cloned()
             .chain(
                 environments
                     .turn_environments()
@@ -352,7 +353,7 @@ impl Session {
                 ready_environment_root_count += 1;
             }
             root_locations_by_id.insert(root.id.clone(), root.location.clone());
-            selected_capability_roots.push(root.clone());
+            selected_capability_roots.push(root);
         }
         self.services
             .turn_environments

codex-rs/exec-server/src/environment.rs

@@ -2,9 +2,9 @@ use std::collections::HashMap;
 use std::collections::HashSet;
 use std::sync::Arc;
 use std::sync::Mutex;
-use std::sync::OnceLock;
 use std::sync::RwLock;
 
+use arc_swap::ArcSwapOption;
 use codex_http_client::HttpClientFactory;
 use codex_http_client::OutboundProxyPolicy;
 use codex_protocol::capabilities::CapabilityRootLocation;
@@ -95,7 +95,7 @@ pub struct EnvironmentReadyInfo {
 pub struct DeferredEnvironmentRegistration {
     completion: oneshot::Sender<Result<(), String>>,
     environment_id: String,
-    ready_info: Arc<OnceLock<EnvironmentReadyInfo>>,
+    ready_info: Arc<ArcSwapOption<EnvironmentReadyInfo>>,
 }
 
 /// Maximum capability roots accepted from deferred environment ready information.
@@ -382,6 +382,47 @@ impl EnvironmentManager {
             .cloned()
     }
 
+    /// Publishes readiness to a registered environment without replacing it.
+    pub fn publish_ready_info(
+        &self,
+        environment_id: &str,
+        ready_info: EnvironmentReadyInfo,
+    ) -> Result<(), ExecServerError> {
+        validate_environment_id(environment_id)?;
+        if ready_info.selected_capability_roots.len() > MAX_SELECTED_CAPABILITY_ROOTS {
+            return Err(ExecServerError::Protocol(format!(
+                "environment ready info contains more than {MAX_SELECTED_CAPABILITY_ROOTS} selected capability roots"
+            )));
+        }
+
+        let mut root_ids = HashSet::with_capacity(ready_info.selected_capability_roots.len());
+        for root in &ready_info.selected_capability_roots {
+            let CapabilityRootLocation::Environment {
+                environment_id: root_environment_id,
+                ..
+            } = &root.location;
+            if root.id.trim().is_empty()
+                || root_environment_id != environment_id
+                || !root_ids.insert(root.id.as_str())
+            {
+                return Err(ExecServerError::Protocol(format!(
+                    "selected capability roots must have unique non-empty IDs and belong to environment `{environment_id}`"
+                )));
+            }
+        }
+
+        let environments = self
+            .environments
+            .read()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let environment = environments.get(environment_id).ok_or_else(|| {
+            ExecServerError::Protocol(format!("environment `{environment_id}` is not registered"))
+        })?;
+
+        environment.ready_info.store(Some(Arc::new(ready_info)));
+        Ok(())
+    }
+
     /// Returns the outbound HTTP policy carried by this manager.
     pub fn http_client_factory(&self) -> &HttpClientFactory {
         &self.http_client_factory
@@ -428,8 +469,7 @@ impl EnvironmentManager {
         validate_environment_id(&environment_id)?;
         let identity = noise_channel_identity()?;
         let (completion, readiness) = oneshot::channel();
-        let ready_info = Arc::new(OnceLock::new());
-        let mut environment = Environment::remote_with_transport(
+        let environment = Environment::remote_with_transport(
             ExecServerTransportParams::Deferred(Box::new(crate::client_api::Deferred {
                 readiness: readiness.shared(),
                 transport: ExecServerTransportParams::NoiseRendezvous { provider, identity },
@@ -437,7 +477,7 @@ impl EnvironmentManager {
             self.local_runtime_paths.clone(),
             self.http_client_factory.clone(),
         );
-        environment.ready_info = Some(Arc::clone(&ready_info));
+        let ready_info = Arc::clone(&environment.ready_info);
         let environment = Arc::new(environment);
         self.insert_environment(environment_id.clone(), environment);
         Ok(DeferredEnvironmentRegistration {
@@ -508,13 +548,7 @@ impl DeferredEnvironmentRegistration {
                         return Err(error);
                     }
                 }
-                if self.ready_info.set(ready_info).is_err() {
-                    let error = ExecServerError::Protocol(
-                        "deferred environment ready info was already set".to_string(),
-                    );
-                    let _ = self.completion.send(Err(error.to_string()));
-                    return Err(error);
-                }
+                self.ready_info.store(Some(Arc::new(ready_info)));
                 Ok(())
             }
             Err(message) => Err(message),
@@ -613,7 +647,7 @@ fn optional_environment_value(name: &str) -> Option<String> {
 #[derive(Clone)]
 pub struct Environment {
     remote_client: Option<LazyRemoteExecServerClient>,
-    ready_info: Option<Arc<OnceLock<EnvironmentReadyInfo>>>,
+    ready_info: Arc<ArcSwapOption<EnvironmentReadyInfo>>,
     // Dropping the environment stops unfinished background startup work.
     startup_task: Arc<Mutex<Option<AbortOnDropHandle<()>>>>,
     exec_backend: Arc<dyn ExecBackend>,
@@ -627,7 +661,7 @@ impl Environment {
     pub fn default_for_tests() -> Self {
         Self {
             remote_client: None,
-            ready_info: None,
+            ready_info: Arc::new(ArcSwapOption::empty()),
             startup_task: Arc::new(Mutex::new(None)),
             exec_backend: Arc::new(LocalProcess::default()),
             filesystem: Arc::new(LocalFileSystem::unsandboxed()),
@@ -705,7 +739,7 @@ impl Environment {
     ) -> Self {
         Self {
             remote_client: None,
-            ready_info: None,
+            ready_info: Arc::new(ArcSwapOption::empty()),
             startup_task: Arc::new(Mutex::new(None)),
             exec_backend: Arc::new(LocalProcess::with_local_runtime_paths(
                 local_runtime_paths.clone(),
@@ -730,7 +764,7 @@ impl Environment {
 
         Self {
             remote_client: Some(client.clone()),
-            ready_info: None,
+            ready_info: Arc::new(ArcSwapOption::empty()),
             startup_task: Arc::new(Mutex::new(None)),
             exec_backend,
             filesystem,
@@ -743,13 +777,13 @@ impl Environment {
         self.remote_client.is_some()
     }
 
-    /// Returns capability roots supplied with the deferred environment's ready signal.
-    pub fn selected_capability_roots(&self) -> &[SelectedCapabilityRoot] {
+    /// Returns the capability roots most recently published for this environment.
+    pub fn selected_capability_roots(&self) -> Vec<SelectedCapabilityRoot> {
         self.ready_info
+            .load()
             .as_ref()
-            .and_then(|ready_info| ready_info.get())
-            .map_or(&[], |ready_info| {
-                ready_info.selected_capability_roots.as_slice()
+            .map_or_else(Vec::new, |ready_info| {
+                ready_info.selected_capability_roots.clone()
             })
     }
 

codex-rs/exec-server/tests/deferred_environment.rs

@@ -83,6 +83,168 @@ async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn existing_environment_publishes_readiness_without_replacement() -> anyhow::Result<()> {
+    let manager = environment_manager_without_environments();
+    let existing_provider = Arc::new(FailingNoiseConnectProvider::default());
+    manager.upsert_noise_environment("tools".to_string(), existing_provider)?;
+    let existing_environment = manager
+        .get_environment("tools")
+        .expect("existing environment");
+
+    let ready_info = ready_info("selected-root", "tools")?;
+    manager.publish_ready_info("tools", ready_info.clone())?;
+
+    let current_environment = manager
+        .get_environment("tools")
+        .expect("current environment");
+    assert!(Arc::ptr_eq(&existing_environment, &current_environment));
+    assert_eq!(
+        existing_environment.selected_capability_roots(),
+        ready_info.selected_capability_roots
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn readiness_updates_the_current_environment_after_replacement() -> anyhow::Result<()> {
+    let manager = environment_manager_without_environments();
+    manager.upsert_noise_environment(
+        "tools".to_string(),
+        Arc::new(FailingNoiseConnectProvider::default()),
+    )?;
+    let captured_environment = manager
+        .get_environment("tools")
+        .expect("captured environment");
+    manager.upsert_noise_environment(
+        "tools".to_string(),
+        Arc::new(FailingNoiseConnectProvider::default()),
+    )?;
+    let current_environment = manager
+        .get_environment("tools")
+        .expect("replacement environment");
+    let ready_info = ready_info("selected-root", "tools")?;
+
+    manager.publish_ready_info("tools", ready_info.clone())?;
+
+    assert!(!Arc::ptr_eq(&captured_environment, &current_environment));
+    assert!(captured_environment.selected_capability_roots().is_empty());
+    assert_eq!(
+        current_environment.selected_capability_roots(),
+        ready_info.selected_capability_roots
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn publishing_readiness_requires_existing_environment() -> anyhow::Result<()> {
+    let manager = environment_manager_without_environments();
+
+    let error = manager
+        .publish_ready_info("tools", ready_info("selected-root", "tools")?)
+        .unwrap_err();
+
+    assert!(matches!(error, ExecServerError::Protocol(_)));
+    assert!(manager.get_environment("tools").is_none());
+    Ok(())
+}
+
+#[tokio::test]
+async fn existing_environment_accepts_matching_readiness() -> anyhow::Result<()> {
+    let manager = environment_manager_without_environments();
+    manager.upsert_noise_environment(
+        "tools".to_string(),
+        Arc::new(FailingNoiseConnectProvider::default()),
+    )?;
+    let environment = manager.get_environment("tools").expect("environment");
+    let ready_info = ready_info("selected-root", "tools")?;
+
+    for _ in 0..2 {
+        manager.publish_ready_info("tools", ready_info.clone())?;
+    }
+    assert_eq!(
+        environment.selected_capability_roots(),
+        ready_info.selected_capability_roots
+    );
+    Ok(())
+}
+
+#[tokio::test]
+async fn existing_environment_overwrites_published_readiness() -> anyhow::Result<()> {
+    let manager = environment_manager_without_environments();
+    manager.upsert_noise_environment(
+        "tools".to_string(),
+        Arc::new(FailingNoiseConnectProvider::default()),
+    )?;
+    let environment = manager.get_environment("tools").expect("environment");
+    let selected_ready_info = ready_info("selected-root", "tools")?;
+    manager.publish_ready_info("tools", selected_ready_info)?;
+
+    let updated_ready_info = ready_info("different-root", "tools")?;
+    manager.publish_ready_info("tools", updated_ready_info.clone())?;
+    assert_eq!(
+        environment.selected_capability_roots(),
+        updated_ready_info.selected_capability_roots
+    );
+    assert!(Arc::ptr_eq(
+        &environment,
+        &manager.get_environment("tools").expect("environment")
+    ));
+
+    manager.publish_ready_info("tools", EnvironmentReadyInfo::default())?;
+    assert!(environment.selected_capability_roots().is_empty());
+    Ok(())
+}
+
+#[tokio::test]
+async fn publishing_readiness_before_deferred_completion_preserves_the_gate() -> anyhow::Result<()>
+{
+    let manager = environment_manager_without_environments();
+    let provider = Arc::new(FailingNoiseConnectProvider::default());
+    let registration =
+        manager.register_deferred_noise_environment("tools".to_string(), provider.clone())?;
+    let environment = manager.get_environment("tools").expect("environment");
+    let mut readiness = Box::pin(environment.wait_until_ready());
+
+    manager.publish_ready_info("tools", ready_info("published-root", "tools")?)?;
+    assert!(poll!(&mut readiness).is_pending());
+    assert_eq!(provider.calls(), 0);
+
+    let completed_ready_info = ready_info("completed-root", "tools")?;
+    registration.complete(Ok(completed_ready_info.clone()))?;
+    assert_eq!(
+        environment.selected_capability_roots(),
+        completed_ready_info.selected_capability_roots
+    );
+    let error = readiness.await.unwrap_err();
+    assert!(error.to_string().contains("test Noise provider called"));
+    assert_eq!(provider.calls(), 1);
+    Ok(())
+}
+
+#[tokio::test]
+async fn existing_environment_rejects_invalid_readiness() -> anyhow::Result<()> {
+    let manager = environment_manager_without_environments();
+    manager.upsert_noise_environment(
+        "tools".to_string(),
+        Arc::new(FailingNoiseConnectProvider::default()),
+    )?;
+    let existing_environment = manager
+        .get_environment("tools")
+        .expect("existing environment");
+    let error = manager
+        .publish_ready_info("tools", ready_info("selected-root", "other")?)
+        .unwrap_err();
+
+    assert!(matches!(error, ExecServerError::Protocol(_)));
+    assert!(existing_environment.selected_capability_roots().is_empty());
+    let current_environment = manager
+        .get_environment("tools")
+        .expect("current environment");
+    assert!(Arc::ptr_eq(&existing_environment, &current_environment));
+    Ok(())
+}
+
 #[tokio::test]
 async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> {
     let manager = environment_manager_without_environments();