From 740c8d8fccff57456a862c183520511d6a007df7 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 12:56:25 +0100 Subject: [PATCH] Resolve selected roots through current environments --- codex-rs/core/src/environment_selection.rs | 82 +++---------------- codex-rs/core/src/session/mod.rs | 11 ++- codex-rs/exec-server/src/environment.rs | 21 ++++- .../exec-server/src/resolved_capability.rs | 71 +++++++++++++++- .../tests/selected_capability_roots.rs | 71 ++++++++++++++++ 5 files changed, 183 insertions(+), 73 deletions(-) create mode 100644 codex-rs/exec-server/tests/selected_capability_roots.rs diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index fcab147725..6d8de052a7 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/codex-rs/core/src/environment_selection.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::collections::HashSet; use std::fmt; use std::sync::Arc; @@ -7,9 +8,6 @@ use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerError; use codex_exec_server::ExecutorFileSystem; -use codex_exec_server::ResolvedSelectedCapabilityRoot; -use codex_protocol::capabilities::CapabilityRootLocation; -use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::TurnEnvironmentSelection; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_path_uri::PathUri; @@ -230,24 +228,20 @@ pub(crate) struct TurnEnvironmentSnapshot { } impl TurnEnvironmentSnapshot { - pub(crate) fn resolve_selected_capability_roots( - &self, - selected_roots: &[SelectedCapabilityRoot], - ) -> Vec { - selected_roots + pub(crate) fn captured_environments(&self) -> HashMap>> { + self.turn_environments .iter() - .filter_map(|selected_root| { - let CapabilityRootLocation::Environment { environment_id, .. } = - &selected_root.location; - let environment = self - .turn_environments - .iter() - .find(|environment| environment.environment_id == *environment_id)?; - Some(ResolvedSelectedCapabilityRoot::new( - selected_root.clone(), - Arc::clone(&environment.environment), - )) + .map(|environment| { + ( + environment.environment_id.clone(), + Some(Arc::clone(&environment.environment)), + ) }) + .chain( + self.starting + .iter() + .map(|environment| (environment.selection.environment_id.clone(), None)), + ) .collect() } @@ -584,56 +578,6 @@ url = "ws://127.0.0.1:8765" assert_eq!(resolved.snapshot().await.to_selections(), vec![local]); } - #[tokio::test] - async fn selected_capability_roots_bind_only_to_their_ready_environment() { - let cwd = PathUri::parse("file:///workspace").expect("cwd URI"); - let primary = Arc::new(Environment::default_for_tests()); - let selected = Arc::new(Environment::default_for_tests()); - let snapshot = TurnEnvironmentSnapshot { - turn_environments: vec![ - TurnEnvironment::new( - "primary".to_string(), - Arc::clone(&primary), - cwd.clone(), - /*shell*/ None, - ), - TurnEnvironment::new( - "selected".to_string(), - Arc::clone(&selected), - cwd, - /*shell*/ None, - ), - ], - starting: Vec::new(), - }; - let plugin_path = PathUri::parse("file:///plugins/demo").expect("plugin URI"); - let roots = [ - SelectedCapabilityRoot { - id: "demo@1".to_string(), - location: CapabilityRootLocation::Environment { - environment_id: "selected".to_string(), - path: plugin_path.clone(), - }, - }, - SelectedCapabilityRoot { - id: "pending@1".to_string(), - location: CapabilityRootLocation::Environment { - environment_id: "pending".to_string(), - path: plugin_path, - }, - }, - ]; - - let resolved = snapshot.resolve_selected_capability_roots(&roots); - - let [resolved] = resolved.as_slice() else { - panic!("only the ready selected environment should resolve"); - }; - assert_eq!(resolved.selected_root(), &roots[0]); - assert!(Arc::ptr_eq(resolved.environment(), &selected)); - assert!(!Arc::ptr_eq(resolved.environment(), &primary)); - } - #[tokio::test] async fn blocking_snapshot_waits_for_starting_environment() { let listener = TcpListener::bind("127.0.0.1:0") diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index e006bd8c34..94bef646aa 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2828,8 +2828,15 @@ impl Session { .await; } let loaded_agents_md = self.services.agents_md_manager.get_loaded().await; - let selected_capability_roots = environments - .resolve_selected_capability_roots(&self.services.selected_capability_roots); + let selected_capability_roots = self + .services + .turn_environments + .environment_manager() + .resolve_selected_capability_roots( + &self.services.selected_capability_roots, + &environments.captured_environments(), + ) + .await; Arc::new(StepContext::new( turn_context, environments, diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 664ef730d0..66f7381103 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/codex-rs/exec-server/src/environment.rs @@ -54,7 +54,7 @@ pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str = #[derive(Debug)] pub struct EnvironmentManager { default_environment: Option, - environments: RwLock>>, + pub(super) environments: RwLock>>, local_environment: Option>, local_runtime_paths: Option, } @@ -574,6 +574,25 @@ impl Environment { } } + /// Starts the initial connection after an environment is actually selected for use. + pub(crate) fn start_connecting_for_use(environment: &Arc) { + if environment.remote_client.is_none() { + return; + } + let mut startup_task = environment + .startup_task + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if startup_task.is_none() { + let environment = Arc::clone(environment); + *startup_task = Some(AbortOnDropHandle::new(tokio::spawn(async move { + if let Err(error) = environment.wait_until_ready().await { + tracing::debug!(%error, "exec-server environment startup failed"); + } + }))); + } + } + /// Returns whether initial startup has either succeeded or permanently failed. pub fn startup_finished(&self) -> bool { self.remote_client diff --git a/codex-rs/exec-server/src/resolved_capability.rs b/codex-rs/exec-server/src/resolved_capability.rs index 5e72b23c5a..666657e264 100644 --- a/codex-rs/exec-server/src/resolved_capability.rs +++ b/codex-rs/exec-server/src/resolved_capability.rs @@ -1,9 +1,12 @@ +use std::collections::HashMap; use std::fmt; use std::sync::Arc; +use codex_protocol::capabilities::CapabilityRootLocation; use codex_protocol::capabilities::SelectedCapabilityRoot; use crate::Environment; +use crate::EnvironmentManager; use crate::ExecutorFileSystem; /// A selected capability root pinned to the exact environment instance that owns it. @@ -17,7 +20,7 @@ pub struct ResolvedSelectedCapabilityRoot { } impl ResolvedSelectedCapabilityRoot { - pub fn new(selected_root: SelectedCapabilityRoot, environment: Arc) -> Self { + fn new(selected_root: SelectedCapabilityRoot, environment: Arc) -> Self { Self { selected_root, environment, @@ -37,6 +40,72 @@ impl ResolvedSelectedCapabilityRoot { } } +impl EnvironmentManager { + /// Binds selected roots to the ready environment instances currently registered for them. + /// + /// The registry and readiness are each captured once per environment ID, so every root owned + /// by one environment uses the same executor instance and readiness result. Missing, starting, + /// or failed environments are omitted. A lazy environment is started for a later step. + pub async fn resolve_selected_capability_roots( + &self, + selected_roots: &[SelectedCapabilityRoot], + captured_environments: &HashMap>>, + ) -> Vec { + let candidates = { + let environments = self + .environments + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + selected_roots + .iter() + .filter_map(|selected_root| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + let (environment, already_ready) = + match captured_environments.get(environment_id) { + Some(Some(environment)) => (Arc::clone(environment), true), + Some(None) => return None, + None => (Arc::clone(environments.get(environment_id)?), false), + }; + Some(( + ResolvedSelectedCapabilityRoot::new(selected_root.clone(), environment), + already_ready, + )) + }) + .collect::>() + }; + + let mut readiness = HashMap::new(); + for (candidate, already_ready) in &candidates { + let CapabilityRootLocation::Environment { environment_id, .. } = + &candidate.selected_root().location; + if readiness.contains_key(environment_id) { + continue; + } + let environment = candidate.environment(); + let ready = if *already_ready { + true + } else if environment.startup_finished() { + environment.wait_until_ready().await.is_ok() + } else { + Environment::start_connecting_for_use(environment); + false + }; + readiness.insert(environment_id.clone(), ready); + } + + candidates + .into_iter() + .map(|(candidate, _)| candidate) + .filter(|candidate| { + let CapabilityRootLocation::Environment { environment_id, .. } = + &candidate.selected_root().location; + readiness.get(environment_id).copied().unwrap_or(false) + }) + .collect() + } +} + impl fmt::Debug for ResolvedSelectedCapabilityRoot { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter diff --git a/codex-rs/exec-server/tests/selected_capability_roots.rs b/codex-rs/exec-server/tests/selected_capability_roots.rs new file mode 100644 index 0000000000..e20fc5b68a --- /dev/null +++ b/codex-rs/exec-server/tests/selected_capability_roots.rs @@ -0,0 +1,71 @@ +#![cfg(unix)] + +mod common; + +use std::collections::HashMap; +use std::sync::Arc; + +use codex_exec_server::EnvironmentManager; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_utils_path_uri::PathUri; +use common::exec_server::exec_server; +use pretty_assertions::assert_eq; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn selected_capability_roots_follow_same_id_environment_replacement() -> anyhow::Result<()> { + let mut executor_a = exec_server().await?; + let mut executor_b = exec_server().await?; + let manager = EnvironmentManager::without_environments(); + let selected_root = SelectedCapabilityRoot { + id: "demo@1".to_string(), + location: CapabilityRootLocation::Environment { + environment_id: "tools".to_string(), + path: PathUri::parse("file:///plugins/demo")?, + }, + }; + + manager.upsert_environment( + "tools".to_string(), + executor_a.websocket_url().to_string(), + /*connect_timeout*/ None, + )?; + let environment_a = manager + .get_environment("tools") + .expect("executor A should be registered"); + environment_a.wait_until_ready().await?; + let resolved_a = manager + .resolve_selected_capability_roots(std::slice::from_ref(&selected_root), &HashMap::new()) + .await; + let [resolved_a] = resolved_a.as_slice() else { + anyhow::bail!("selected root should resolve through executor A"); + }; + + manager.upsert_environment( + "tools".to_string(), + executor_b.websocket_url().to_string(), + /*connect_timeout*/ None, + )?; + let environment_b = manager + .get_environment("tools") + .expect("executor B should be registered"); + environment_b.wait_until_ready().await?; + let resolved_b = manager + .resolve_selected_capability_roots(std::slice::from_ref(&selected_root), &HashMap::new()) + .await; + let [resolved_b] = resolved_b.as_slice() else { + anyhow::bail!("selected root should resolve through executor B"); + }; + + assert_eq!(resolved_b.selected_root(), &selected_root); + assert!(Arc::ptr_eq(resolved_a.environment(), &environment_a)); + assert!(Arc::ptr_eq(resolved_b.environment(), &environment_b)); + assert!(!Arc::ptr_eq( + resolved_a.environment(), + resolved_b.environment() + )); + + executor_a.shutdown().await?; + executor_b.shutdown().await?; + Ok(()) +}