From dfa3c95375da7a5fdcabb5a7cd8b4bebd839f72b Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 25 Jun 2026 00:59:25 +0100 Subject: [PATCH] Track selected capability readiness --- codex-rs/Cargo.lock | 1 + codex-rs/core-plugins/Cargo.toml | 3 +- codex-rs/core-plugins/src/lib.rs | 7 + codex-rs/core-plugins/src/provider.rs | 121 ++++++- .../core-plugins/src/selected_capabilities.rs | 335 ++++++++++++++++++ .../src/selected_capabilities_tests.rs | 290 +++++++++++++++ codex-rs/core/src/thread_manager.rs | 13 + codex-rs/core/src/thread_manager_tests.rs | 5 + 8 files changed, 768 insertions(+), 7 deletions(-) create mode 100644 codex-rs/core-plugins/src/selected_capabilities.rs create mode 100644 codex-rs/core-plugins/src/selected_capabilities_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5448821756..c21d639b5d 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2793,6 +2793,7 @@ dependencies = [ "codex-utils-plugins", "dirs", "flate2", + "futures", "libc", "pretty_assertions", "regex", diff --git a/codex-rs/core-plugins/Cargo.toml b/codex-rs/core-plugins/Cargo.toml index 81cfe2f89c..5b2b744be6 100644 --- a/codex-rs/core-plugins/Cargo.toml +++ b/codex-rs/core-plugins/Cargo.toml @@ -36,6 +36,7 @@ codex-utils-plugins = { workspace = true } chrono = { workspace = true } dirs = { workspace = true } flate2 = { workspace = true } +futures = { workspace = true } reqwest = { workspace = true } regex = { workspace = true } semver = { workspace = true } @@ -44,7 +45,7 @@ serde_json = { workspace = true } tar = { workspace = true } tempfile = { workspace = true } thiserror = { workspace = true } -tokio = { workspace = true, features = ["fs", "macros", "rt", "time"] } +tokio = { workspace = true, features = ["fs", "macros", "rt", "sync", "time"] } toml = { workspace = true } tracing = { workspace = true } url = { workspace = true } diff --git a/codex-rs/core-plugins/src/lib.rs b/codex-rs/core-plugins/src/lib.rs index 7ee5ddd117..be557936d7 100644 --- a/codex-rs/core-plugins/src/lib.rs +++ b/codex-rs/core-plugins/src/lib.rs @@ -14,6 +14,7 @@ mod provider; pub mod remote; pub mod remote_bundle; pub mod remote_legacy; +mod selected_capabilities; pub mod startup_sync; pub mod store; #[cfg(test)] @@ -59,5 +60,11 @@ pub use marketplace_upgrade::ConfiguredMarketplaceUpgradeOutcome as PluginMarket pub use provider::ExecutorPluginProvider; pub use provider::ExecutorPluginProviderError; pub use provider::ResolvedExecutorPlugin; +pub use provider::ResolvedSelectedCapabilityRoot; pub use remote::RecommendedPlugin; pub use remote::RecommendedPluginsMode; +pub use selected_capabilities::SelectedCapabilityBindingSnapshot; +pub use selected_capabilities::SelectedCapabilityBindingStatus; +pub use selected_capabilities::SelectedCapabilityBindings; +pub use selected_capabilities::SelectedCapabilityFailure; +pub use selected_capabilities::SelectedCapabilitySnapshot; diff --git a/codex-rs/core-plugins/src/provider.rs b/codex-rs/core-plugins/src/provider.rs index 5a371b3d2e..fea0fe9a1c 100644 --- a/codex-rs/core-plugins/src/provider.rs +++ b/codex-rs/core-plugins/src/provider.rs @@ -1,5 +1,7 @@ use crate::manifest::parse_plugin_manifest_uri; +use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; +use codex_exec_server::ExecServerError; use codex_exec_server::ExecutorFileSystem; use codex_plugin::PluginProvider; use codex_plugin::ResolvedPlugin; @@ -12,6 +14,9 @@ use codex_utils_plugins::DISCOVERABLE_PLUGIN_MANIFEST_PATHS; use std::io; use std::sync::Arc; use thiserror::Error; +use tokio::sync::Semaphore; + +const MAX_CONCURRENT_PLUGIN_INSPECTIONS: usize = 4; /// Failure to resolve an environment-owned capability root as a plugin package. #[derive(Debug, Error)] @@ -23,6 +28,15 @@ pub enum ExecutorPluginProviderError { root_id: String, environment_id: String, }, + #[error( + "selected capability root `{root_id}` environment `{environment_id}` failed to start: {source}" + )] + EnvironmentStartup { + root_id: String, + environment_id: String, + #[source] + source: ExecServerError, + }, #[error("failed to inspect selected capability root `{root_id}` at {path}: {source}")] InspectRoot { root_id: String, @@ -75,6 +89,7 @@ pub enum ExecutorPluginProviderError { #[derive(Clone, Debug)] pub struct ExecutorPluginProvider { environment_manager: Arc, + inspection_permits: Arc, } /// A resolved plugin paired with the concrete filesystem used to read it. @@ -84,6 +99,62 @@ pub struct ResolvedExecutorPlugin { file_system: Arc, } +/// One selected capability root bound to its owning execution environment. +/// +/// The optional plugin descriptor is absent when the root is valid but does +/// not contain a plugin manifest. Consumers such as skills may still use the +/// exact executor filesystem in that case. +#[derive(Clone)] +pub struct ResolvedSelectedCapabilityRoot { + selection_order: usize, + selected_root: SelectedCapabilityRoot, + environment: Arc, + plugin: Option, +} + +impl ResolvedSelectedCapabilityRoot { + pub(crate) fn new( + selection_order: usize, + selected_root: SelectedCapabilityRoot, + environment: Arc, + plugin: Option, + ) -> Self { + Self { + selection_order, + selected_root, + environment, + plugin, + } + } + + /// Returns this root's position in the caller-provided selection. + pub fn selection_order(&self) -> usize { + self.selection_order + } + + /// Returns the original environment-qualified selection. + pub fn selected_root(&self) -> &SelectedCapabilityRoot { + &self.selected_root + } + + /// Returns the plugin descriptor when the selected root declares one. + pub fn plugin(&self) -> Option<&ResolvedPlugin> { + self.plugin.as_ref() + } + + /// Returns the filesystem owned by the selected root's exact executor. + pub fn file_system(&self) -> Arc { + self.environment.get_filesystem() + } + + fn into_plugin(self) -> Option { + self.plugin.map(|plugin| ResolvedExecutorPlugin { + plugin, + file_system: self.environment.get_filesystem(), + }) + } +} + impl ResolvedExecutorPlugin { /// Returns the source-neutral plugin descriptor. pub fn plugin(&self) -> &ResolvedPlugin { @@ -101,6 +172,7 @@ impl ExecutorPluginProvider { pub fn new(environment_manager: Arc) -> Self { Self { environment_manager, + inspection_permits: Arc::new(Semaphore::new(MAX_CONCURRENT_PLUGIN_INSPECTIONS)), } } @@ -109,8 +181,18 @@ impl ExecutorPluginProvider { &self, selected_root: &SelectedCapabilityRoot, ) -> Result, ExecutorPluginProviderError> { + self.resolve_selected_root(/*selection_order*/ 0, selected_root.clone()) + .await + .map(ResolvedSelectedCapabilityRoot::into_plugin) + } + + /// Resolves one selected root after its owning environment becomes ready. + pub async fn resolve_selected_root( + &self, + selection_order: usize, + selected_root: SelectedCapabilityRoot, + ) -> Result { let root_id = &selected_root.id; - let plugin_root = selected_plugin_root(selected_root); let CapabilityRootLocation::Environment { environment_id, .. } = &selected_root.location; let environment = self .environment_manager @@ -119,13 +201,40 @@ impl ExecutorPluginProvider { root_id: root_id.clone(), environment_id: environment_id.clone(), })?; - let file_system = environment.get_filesystem(); - let plugin = resolve_plugin_root(selected_root, plugin_root, file_system.as_ref()).await?; + self.resolve_selected_root_with_environment(selection_order, selected_root, environment) + .await + } - Ok(plugin.map(|plugin| ResolvedExecutorPlugin { + pub(crate) async fn resolve_selected_root_with_environment( + &self, + selection_order: usize, + selected_root: SelectedCapabilityRoot, + environment: Arc, + ) -> Result { + let root_id = &selected_root.id; + let plugin_root = selected_plugin_root(&selected_root); + let CapabilityRootLocation::Environment { environment_id, .. } = &selected_root.location; + environment.wait_until_ready().await.map_err(|source| { + ExecutorPluginProviderError::EnvironmentStartup { + root_id: root_id.clone(), + environment_id: environment_id.clone(), + source, + } + })?; + let _inspection_permit = self + .inspection_permits + .acquire() + .await + .expect("plugin inspection semaphore should remain open"); + let file_system = environment.get_filesystem(); + let plugin = resolve_plugin_root(&selected_root, plugin_root, file_system.as_ref()).await?; + + Ok(ResolvedSelectedCapabilityRoot::new( + selection_order, + selected_root, + environment, plugin, - file_system, - })) + )) } } diff --git a/codex-rs/core-plugins/src/selected_capabilities.rs b/codex-rs/core-plugins/src/selected_capabilities.rs new file mode 100644 index 0000000000..8b0b9f7e9e --- /dev/null +++ b/codex-rs/core-plugins/src/selected_capabilities.rs @@ -0,0 +1,335 @@ +use std::collections::HashMap; +use std::future::Future; +use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; + +use codex_exec_server::EnvironmentManager; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use futures::FutureExt; +use futures::StreamExt; +use futures::stream::FuturesUnordered; +use tokio::sync::Notify; + +use crate::ExecutorPluginProvider; +use crate::ResolvedSelectedCapabilityRoot; + +type ResolutionFuture = Pin< + Box< + dyn Future> + + Send + + 'static, + >, +>; +type IndexedResolutionFuture = Pin< + Box< + dyn Future< + Output = ( + usize, + Result, + ), + > + Send + + 'static, + >, +>; + +/// Terminal failure while binding one selected root to its executor. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SelectedCapabilityFailure { + message: String, +} + +impl SelectedCapabilityFailure { + /// Returns the stable diagnostic captured for this failed binding. + pub fn message(&self) -> &str { + &self.message + } +} + +/// Resolution state for one environment-qualified selected capability root. +#[derive(Clone)] +pub enum SelectedCapabilityBindingStatus { + /// The owning executor has not completed capability discovery yet. + Pending, + /// The root is bound to its exact executor and optional plugin descriptor. + Ready(Arc), + /// Discovery failed permanently for this thread. + Failed(Arc), +} + +/// One selected root and its state in an immutable binding snapshot. +#[derive(Clone)] +pub struct SelectedCapabilityBindingSnapshot { + selected_root: SelectedCapabilityRoot, + status: SelectedCapabilityBindingStatus, +} + +impl SelectedCapabilityBindingSnapshot { + /// Returns the original root in caller-provided selection order. + pub fn selected_root(&self) -> &SelectedCapabilityRoot { + &self.selected_root + } + + /// Returns this root's state at the captured generation. + pub fn status(&self) -> &SelectedCapabilityBindingStatus { + &self.status + } +} + +/// Immutable selected-capability view captured at one generation. +#[derive(Clone)] +pub struct SelectedCapabilitySnapshot { + generation: u64, + entries: Vec, +} + +impl SelectedCapabilitySnapshot { + /// Returns the monotonically increasing binding generation. + pub fn generation(&self) -> u64 { + self.generation + } + + /// Returns entries in the original selected-root order. + pub fn entries(&self) -> &[SelectedCapabilityBindingSnapshot] { + &self.entries + } + + /// Returns whether every selected root is ready or permanently failed. + pub fn is_terminal(&self) -> bool { + self.entries + .iter() + .all(|entry| !matches!(entry.status(), SelectedCapabilityBindingStatus::Pending)) + } + + /// Iterates roots that are ready in caller-provided selection order. + pub fn ready(&self) -> impl Iterator { + self.entries.iter().filter_map(|entry| match &entry.status { + SelectedCapabilityBindingStatus::Ready(resolved) => Some(resolved.as_ref()), + SelectedCapabilityBindingStatus::Pending + | SelectedCapabilityBindingStatus::Failed(_) => None, + }) + } +} + +/// Thread-owned, nonblocking selected-capability resolution state. +/// +/// Each selected root resolves once against its original environment. Callers +/// may capture an immutable snapshot without waiting, or await the terminal +/// snapshot to preserve legacy startup behavior. +#[derive(Clone)] +pub struct SelectedCapabilityBindings { + inner: Arc, +} + +struct SelectedCapabilityBindingsInner { + roots: Vec, + state: Mutex, + changed: Notify, + resolution_task: Mutex>, +} + +struct SelectedCapabilityBindingsState { + generation: u64, + statuses: Vec, +} + +impl SelectedCapabilityBindings { + /// Starts background binding for the supplied environment-qualified roots. + /// + /// # Panics + /// + /// Panics when called outside a Tokio runtime if `selected_roots` is not + /// empty. + pub fn new( + selected_roots: Vec, + environment_manager: Arc, + ) -> Self { + let provider = ExecutorPluginProvider::new(Arc::clone(&environment_manager)); + let mut environments = HashMap::new(); + let resolutions = selected_roots + .iter() + .cloned() + .enumerate() + .map(|(selection_order, selected_root)| { + let provider = provider.clone(); + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + let environment = environments + .entry(environment_id.clone()) + .or_insert_with(|| environment_manager.get_environment(environment_id)) + .clone(); + let Some(environment) = environment else { + let message = format!( + "selected capability root `{}` references unavailable environment `{environment_id}`", + selected_root.id + ); + return Box::pin(async move { Err(SelectedCapabilityFailure { message }) }) + as ResolutionFuture; + }; + Box::pin(async move { + provider + .resolve_selected_root_with_environment( + selection_order, + selected_root, + environment, + ) + .await + .map_err(|err| SelectedCapabilityFailure { + message: err.to_string(), + }) + }) as ResolutionFuture + }) + .collect(); + Self::from_resolutions(selected_roots, resolutions) + } + + fn from_resolutions( + selected_roots: Vec, + resolutions: Vec, + ) -> Self { + assert_eq!(selected_roots.len(), resolutions.len()); + let inner = Arc::new(SelectedCapabilityBindingsInner { + state: Mutex::new(SelectedCapabilityBindingsState { + generation: 0, + statuses: vec![SelectedCapabilityBindingStatus::Pending; selected_roots.len()], + }), + roots: selected_roots, + changed: Notify::new(), + resolution_task: Mutex::new(None), + }); + + if !resolutions.is_empty() { + let weak_inner = Arc::downgrade(&inner); + let resolution_task = tokio::spawn(async move { + let mut active = FuturesUnordered::::new(); + for (selection_order, resolution) in resolutions.into_iter().enumerate() { + active.push(index_resolution(selection_order, resolution)); + } + while let Some((selection_order, resolution)) = active.next().await { + let Some(inner) = weak_inner.upgrade() else { + return; + }; + let status = match resolution { + Ok(resolved) => SelectedCapabilityBindingStatus::Ready(Arc::new(resolved)), + Err(failure) => { + tracing::warn!( + selected_root = %inner.roots[selection_order].id, + error = failure.message(), + "failed to bind selected capability root" + ); + SelectedCapabilityBindingStatus::Failed(Arc::new(failure)) + } + }; + { + let mut state = inner.state(); + if !matches!( + &state.statuses[selection_order], + SelectedCapabilityBindingStatus::Pending + ) { + return; + } + state.statuses[selection_order] = status; + state.generation = state.generation.saturating_add(1); + } + inner.changed.notify_waiters(); + } + }); + *inner + .resolution_task + .lock() + .unwrap_or_else(PoisonError::into_inner) = Some(resolution_task.abort_handle()); + } + + Self { inner } + } + + /// Captures current states without waiting for pending executors. + pub fn snapshot(&self) -> SelectedCapabilitySnapshot { + let state = self.inner.state(); + SelectedCapabilitySnapshot { + generation: state.generation, + entries: self + .inner + .roots + .iter() + .cloned() + .zip(state.statuses.iter().cloned()) + .map( + |(selected_root, status)| SelectedCapabilityBindingSnapshot { + selected_root, + status, + }, + ) + .collect(), + } + } + + /// Waits until the binding generation differs from `generation`. + pub async fn wait_for_change(&self, generation: u64) -> SelectedCapabilitySnapshot { + loop { + let changed = self.inner.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + let snapshot = self.snapshot(); + if snapshot.generation != generation { + return snapshot; + } + changed.await; + } + } + + /// Waits for every selected root to become ready or permanently fail. + pub async fn resolve_all(&self) -> SelectedCapabilitySnapshot { + loop { + let changed = self.inner.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + let snapshot = self.snapshot(); + if snapshot.is_terminal() { + return snapshot; + } + changed.await; + } + } +} + +fn index_resolution( + selection_order: usize, + resolution: ResolutionFuture, +) -> IndexedResolutionFuture { + Box::pin(async move { + let resolution = AssertUnwindSafe(resolution).catch_unwind().await; + let resolution = resolution.unwrap_or_else(|_| { + Err(SelectedCapabilityFailure { + message: "selected capability resolution panicked".to_string(), + }) + }); + (selection_order, resolution) + }) +} + +impl SelectedCapabilityBindingsInner { + fn state(&self) -> std::sync::MutexGuard<'_, SelectedCapabilityBindingsState> { + self.state.lock().unwrap_or_else(PoisonError::into_inner) + } +} + +impl Drop for SelectedCapabilityBindingsInner { + fn drop(&mut self) { + if let Some(resolution_task) = self + .resolution_task + .get_mut() + .unwrap_or_else(PoisonError::into_inner) + .take() + { + resolution_task.abort(); + } + } +} + +#[cfg(test)] +#[path = "selected_capabilities_tests.rs"] +mod tests; diff --git a/codex-rs/core-plugins/src/selected_capabilities_tests.rs b/codex-rs/core-plugins/src/selected_capabilities_tests.rs new file mode 100644 index 0000000000..b671713a7b --- /dev/null +++ b/codex-rs/core-plugins/src/selected_capabilities_tests.rs @@ -0,0 +1,290 @@ +use std::sync::Arc; +use std::time::Duration; + +use codex_exec_server::EnvironmentManager; +use codex_exec_server::LOCAL_ENVIRONMENT_ID; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_utils_path_uri::PathUri; +use pretty_assertions::assert_eq; +use tokio::sync::oneshot; +use tokio::time::timeout; + +use super::*; + +struct DropSignal(Option>); + +impl Drop for DropSignal { + fn drop(&mut self) { + if let Some(sender) = self.0.take() { + let _ = sender.send(()); + } + } +} + +#[tokio::test] +async fn snapshots_preserve_root_order_and_advance_on_terminal_transitions() { + let roots = vec![ + selected_root("root-a", "local"), + selected_root("root-b", "local"), + ]; + let environment = EnvironmentManager::default_for_tests() + .get_environment(LOCAL_ENVIRONMENT_ID) + .expect("local environment"); + let (first_tx, first_rx) = oneshot::channel(); + let (second_tx, second_rx) = oneshot::channel(); + let bindings = SelectedCapabilityBindings::from_resolutions( + roots.clone(), + vec![ + Box::pin(async move { first_rx.await.expect("first resolution") }), + Box::pin(async move { second_rx.await.expect("second resolution") }), + ], + ); + + let initial = bindings.snapshot(); + assert_eq!(initial.generation(), 0); + assert_eq!( + initial + .entries() + .iter() + .map(|entry| entry.selected_root().id.as_str()) + .collect::>(), + vec!["root-a", "root-b"] + ); + assert!( + initial + .entries() + .iter() + .all(|entry| matches!(entry.status(), SelectedCapabilityBindingStatus::Pending)) + ); + + assert!( + second_tx + .send(Ok(ResolvedSelectedCapabilityRoot::new( + /*selection_order*/ 1, + roots[1].clone(), + Arc::clone(&environment), + None, + ))) + .is_ok() + ); + let after_second = timeout( + Duration::from_secs(1), + bindings.wait_for_change(/*generation*/ 0), + ) + .await + .expect("second root should resolve"); + assert_eq!(after_second.generation(), 1); + assert!(matches!( + after_second.entries()[0].status(), + SelectedCapabilityBindingStatus::Pending + )); + let SelectedCapabilityBindingStatus::Ready(second) = after_second.entries()[1].status() else { + panic!("second root should be ready"); + }; + assert_eq!(second.selection_order(), 1); + assert_eq!(second.selected_root(), &roots[1]); + assert!(Arc::ptr_eq( + &second.file_system(), + &environment.get_filesystem() + )); + + assert!( + first_tx + .send(Err(SelectedCapabilityFailure { + message: "root-a failed".to_string(), + })) + .is_ok() + ); + let terminal = timeout( + Duration::from_secs(1), + bindings.wait_for_change(/*generation*/ 1), + ) + .await + .expect("first root should fail"); + assert_eq!(terminal.generation(), 2); + let SelectedCapabilityBindingStatus::Failed(first) = terminal.entries()[0].status() else { + panic!("first root should have failed"); + }; + assert_eq!(first.message(), "root-a failed"); + assert_eq!( + bindings.resolve_all().await.generation(), + terminal.generation() + ); +} + +#[tokio::test] +async fn unavailable_environment_becomes_one_terminal_failure() { + let bindings = SelectedCapabilityBindings::new( + vec![selected_root("root-a", "missing")], + Arc::new(EnvironmentManager::without_environments()), + ); + + let terminal = timeout(Duration::from_secs(1), bindings.resolve_all()) + .await + .expect("missing environment should resolve as a failure"); + + assert_eq!(terminal.generation(), 1); + let SelectedCapabilityBindingStatus::Failed(failure) = terminal.entries()[0].status() else { + panic!("missing environment should fail"); + }; + assert!( + failure + .message() + .contains("unavailable environment `missing`") + ); +} + +#[tokio::test] +async fn one_transition_wakes_all_generation_waiters() { + let root = selected_root("root-a", "local"); + let (resolution_tx, resolution_rx) = oneshot::channel(); + let bindings = SelectedCapabilityBindings::from_resolutions( + vec![root], + vec![Box::pin( + async move { resolution_rx.await.expect("resolution") }, + )], + ); + let first = { + let bindings = bindings.clone(); + tokio::spawn(async move { + bindings + .wait_for_change(/*generation*/ 0) + .await + .generation() + }) + }; + let second = { + let bindings = bindings.clone(); + tokio::spawn(async move { + bindings + .wait_for_change(/*generation*/ 0) + .await + .generation() + }) + }; + tokio::task::yield_now().await; + + assert!( + resolution_tx + .send(Err(SelectedCapabilityFailure { + message: "failed".to_string(), + })) + .is_ok() + ); + + assert_eq!(first.await.expect("first waiter"), 1); + assert_eq!(second.await.expect("second waiter"), 1); +} + +#[tokio::test] +async fn later_ready_root_is_not_blocked_by_pending_roots() { + let roots = (0..5) + .map(|index| selected_root(&format!("root-{index}"), "local")) + .collect::>(); + let environment = EnvironmentManager::default_for_tests() + .get_environment(LOCAL_ENVIRONMENT_ID) + .expect("local environment"); + let mut senders = Vec::new(); + let mut resolutions = Vec::new(); + for _ in &roots { + let (sender, receiver) = oneshot::channel(); + senders.push(sender); + resolutions + .push(Box::pin(async move { receiver.await.expect("resolution") }) as ResolutionFuture); + } + let bindings = SelectedCapabilityBindings::from_resolutions(roots.clone(), resolutions); + + assert!( + senders + .pop() + .expect("fifth sender") + .send(Ok(ResolvedSelectedCapabilityRoot::new( + /*selection_order*/ 4, + roots[4].clone(), + environment, + None, + ))) + .is_ok() + ); + + let snapshot = timeout( + Duration::from_secs(1), + bindings.wait_for_change(/*generation*/ 0), + ) + .await + .expect("fifth root should resolve independently"); + assert_eq!(snapshot.generation(), 1); + assert!( + snapshot.entries()[..4] + .iter() + .all(|entry| matches!(entry.status(), SelectedCapabilityBindingStatus::Pending)) + ); + assert!(matches!( + snapshot.entries()[4].status(), + SelectedCapabilityBindingStatus::Ready(_) + )); +} + +#[tokio::test] +async fn resolution_panic_becomes_terminal_failure() { + let bindings = SelectedCapabilityBindings::from_resolutions( + vec![selected_root("root-a", "local")], + vec![Box::pin(async { panic!("resolution panic") })], + ); + + let terminal = timeout(Duration::from_secs(1), bindings.resolve_all()) + .await + .expect("panicked resolution should become terminal"); + + let SelectedCapabilityBindingStatus::Failed(failure) = terminal.entries()[0].status() else { + panic!("panicked resolution should fail"); + }; + assert_eq!(failure.message(), "selected capability resolution panicked"); +} + +#[tokio::test] +async fn dropping_bindings_cancels_pending_resolution() { + let (started_tx, started_rx) = oneshot::channel(); + let (dropped_tx, dropped_rx) = oneshot::channel(); + let bindings = SelectedCapabilityBindings::from_resolutions( + vec![selected_root("root-a", "local")], + vec![Box::pin(async move { + let _drop_signal = DropSignal(Some(dropped_tx)); + let _ = started_tx.send(()); + std::future::pending().await + })], + ); + started_rx.await.expect("resolution should start"); + + drop(bindings); + + timeout(Duration::from_secs(1), dropped_rx) + .await + .expect("resolution should be canceled") + .expect("drop signal should be sent"); +} + +#[tokio::test] +async fn empty_bindings_are_immediately_terminal() { + let bindings = SelectedCapabilityBindings::new( + Vec::new(), + Arc::new(EnvironmentManager::without_environments()), + ); + + let snapshot = bindings.resolve_all().await; + + assert_eq!(snapshot.generation(), 0); + assert!(snapshot.entries().is_empty()); + assert!(snapshot.is_terminal()); +} + +fn selected_root(id: &str, environment_id: &str) -> SelectedCapabilityRoot { + SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: environment_id.to_string(), + path: PathUri::parse(&format!("file:///plugins/{id}")).expect("plugin root URI"), + }, + } +} diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 907de03705..203ec695e1 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -22,6 +22,7 @@ use codex_analytics::AnalyticsEventsClient; use codex_app_server_protocol::ThreadHistoryBuilder; use codex_app_server_protocol::TurnStatus; use codex_core_plugins::PluginsManager; +use codex_core_plugins::SelectedCapabilityBindings; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; @@ -1549,6 +1550,18 @@ impl ThreadManagerState { thread_extension_init.insert(selected_capability_roots); } } + if thread_extension_init + .get::() + .is_none() + && let Some(selected_capability_roots) = + thread_extension_init.get::>() + && !selected_capability_roots.is_empty() + { + thread_extension_init.insert(SelectedCapabilityBindings::new( + selected_capability_roots.as_ref().clone(), + Arc::clone(&self.environment_manager), + )); + } self.extensions .initialize_thread_data(&mut thread_extension_init) .await; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 75f0bdf08a..dba379f624 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -439,6 +439,11 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() thread_init: &'a mut codex_extension_api::ExtensionDataInit, ) -> codex_extension_api::ExtensionFuture<'a, ()> { Box::pin(async move { + assert!( + thread_init + .get::() + .is_some() + ); let selected_root = thread_init .get::>() .and_then(|roots| roots.first().cloned())