From 6c13b113a3d28b16ce6d00cdbd18ede5bbc1f283 Mon Sep 17 00:00:00 2001 From: TAFOYA-OAI Date: Wed, 29 Jul 2026 02:57:25 +0000 Subject: [PATCH] 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 | 3 +- codex-rs/exec-server/src/environment.rs | 76 +++++--- .../exec-server/tests/deferred_environment.rs | 162 ++++++++++++++++++ 3 files changed, 219 insertions(+), 22 deletions(-) diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 21c1293aae..0b90712229 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/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 diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 5f7b45ef05..1618a301a7 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/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>, environment_id: String, - ready_info: Arc>, + ready_info: Arc>, } /// 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 { #[derive(Clone)] pub struct Environment { remote_client: Option, - ready_info: Option>>, + ready_info: Arc>, // Dropping the environment stops unfinished background startup work. startup_task: Arc>>>, exec_backend: Arc, @@ -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 { 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() }) } diff --git a/codex-rs/exec-server/tests/deferred_environment.rs b/codex-rs/exec-server/tests/deferred_environment.rs index 899a7605a6..5398e5247d 100644 --- a/codex-rs/exec-server/tests/deferred_environment.rs +++ b/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, ¤t_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, ¤t_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, ¤t_environment)); + Ok(()) +} + #[tokio::test] async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> { let manager = environment_manager_without_environments();