diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 4796147c69..55c36f77b0 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -165,17 +165,21 @@ impl McpManager { ready_selected_capability_roots: &[SelectedCapabilityRoot], executor_capability_discovery: Option<&ExecutorCapabilityDiscoverySnapshot>, ) -> McpRuntimeProjection { + let mut context = McpServerContributionContext::for_step( + config, + thread_init, + thread_store, + identity.originator, + ready_selected_capability_roots, + executor_capability_discovery, + ) + .with_session_source(identity.session_source) + .with_auth_changed(identity.auth_changed); + if let McpEnvironmentScope::Selected(environments) = &identity.environments { + context = context.with_selected_environments(environments); + } self.runtime_config_with_context( - McpServerContributionContext::for_step( - config, - thread_init, - thread_store, - identity.originator, - ready_selected_capability_roots, - executor_capability_discovery, - ) - .with_session_source(identity.session_source) - .with_auth_changed(identity.auth_changed), + context, Some(identity.originator), identity.environments, identity.disabled_plugin_ids, diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 579ac76406..e7832e1e9f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3823,7 +3823,8 @@ impl Session { &ready_selected_capability_roots, executor_capability_discovery.as_deref(), ) - .with_session_source(&turn_context.session_source), + .with_session_source(&turn_context.session_source) + .with_selected_environments(&environments.all_selections()), &turn_context.disabled_plugin_ids, ) .await; diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index d240d6b4e3..c9c6fdea1b 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1619,6 +1619,11 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() ) -> codex_extension_api::ExtensionFuture<'a, Vec>> { Box::pin(async move { + assert_eq!( + context.selected_environments(), + Some([].as_slice()), + "thread MCP projection must preserve explicitly empty selections" + ); let thread_init = context .thread_init() .expect("initial MCP resolution should be thread-scoped"); diff --git a/codex-rs/core/tests/suite/mcp_executor_context_tests.rs b/codex-rs/core/tests/suite/mcp_executor_context_tests.rs new file mode 100644 index 0000000000..642e369052 --- /dev/null +++ b/codex-rs/core/tests/suite/mcp_executor_context_tests.rs @@ -0,0 +1,194 @@ +//! Hosted MCP contributors observe executor selection without filtering unavailable entries. + +use std::sync::Arc; +use std::sync::Mutex; + +use anyhow::Result; +use codex_core::StartThreadOptions; +use codex_core::config::Config; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::McpServerContribution; +use codex_extension_api::McpServerContributionContext; +use codex_extension_api::McpServerContributor; +use codex_extension_api::SelectedPlugin; +use codex_features::Feature; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::protocol::EnvironmentConfigState; +use codex_protocol::protocol::TurnEnvironmentRequest; +use codex_protocol::protocol::TurnEnvironmentSelection; +use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::environment_config_for_selection; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; + +#[derive(Default)] +struct ExecutorSelectionRecorder { + observed: Mutex>>, + plugin_contexts: Mutex>>>, +} + +impl McpServerContributor for ExecutorSelectionRecorder { + fn id(&self) -> &'static str { + "executor_selection_recorder" + } + + fn selected_plugins<'a>( + &'a self, + context: McpServerContributionContext<'a, Config>, + ) -> ExtensionFuture<'a, Vec>> { + Box::pin(async move { + self.plugin_contexts + .lock() + .expect("plugin context recorder lock") + .push(context.selected_environments().map(<[_]>::to_vec)); + Vec::new() + }) + } + + fn contribute<'a>( + &'a self, + context: McpServerContributionContext<'a, Config>, + ) -> ExtensionFuture<'a, Vec> { + Box::pin(async move { + if let Some(environments) = context.selected_environments() { + self.observed + .lock() + .expect("executor selection recorder lock") + .push(environments.to_vec()); + } + Vec::new() + }) + } +} + +/// A pending or failed primary must reach hosted contributors ahead of a ready secondary. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn thread_projection_preserves_executor_order_and_unavailable_selections() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let recorder = Arc::new(ExecutorSelectionRecorder::default()); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.mcp_server_contributor(recorder.clone()); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config + .features + .enable(Feature::DeferredExecutor) + .expect("enable selected executors"); + }) + .build_with_auto_env(&server) + .await?; + let ready = test.executor_environment().request(); + test.thread_manager + .environment_manager() + .upsert_environment( + "unavailable-executor".to_owned(), + "ws://127.0.0.1:1".to_owned(), + /*connect_timeout*/ None, + )?; + for state in [ + EnvironmentConfigState::Pending, + EnvironmentConfigState::Failed("configuration unavailable".to_owned()), + ] { + let unavailable = TurnEnvironmentRequest { + environment_id: "unavailable-executor".to_owned(), + config: state, + ..ready.clone() + }; + for requests in [ + vec![unavailable.clone(), ready.clone()], + vec![ready.clone(), unavailable.clone()], + ] { + let selections = requests + .iter() + .cloned() + .map(TurnEnvironmentSelection::new) + .collect::>(); + recorder + .observed + .lock() + .expect("executor selection recorder lock") + .clear(); + let thread = test + .thread_manager + .start_thread(StartThreadOptions { + environments: Some(requests), + ..StartThreadOptions::new(test.config.clone()) + }) + .await?; + thread.thread.current_mcp_config_and_runtime_context().await; + assert_eq!(thread.thread.environment_selections().await, selections); + { + let observed = recorder + .observed + .lock() + .expect("executor selection recorder lock"); + assert!( + observed.contains(&selections), + "MCP projection must receive the complete ordered selection: {observed:?}" + ); + } + thread.thread.shutdown_and_wait().await?; + } + } + Ok(()) +} + +#[tokio::test] +async fn model_step_plugin_projection_preserves_selected_environments() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let recorder = Arc::new(ExecutorSelectionRecorder::default()); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.mcp_server_contributor(recorder.clone()); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .build_with_auto_env(&server) + .await?; + let selection = test.codex.environment_selections().await.remove(0); + let mut environment_config = environment_config_for_selection(&test.config, &selection); + // Step-local plugin projection returns early without a ready selected root. + environment_config.selected_capability_roots = vec![SelectedCapabilityRoot { + id: "selected-plugin-root".to_owned(), + location: CapabilityRootLocation::Environment { + environment_id: selection.environment_id.clone(), + path: selection.cwd.clone(), + }, + }]; + test.codex + .environment_ready(&selection, environment_config) + .await?; + let selections = test.codex.environment_selections().await; + recorder + .plugin_contexts + .lock() + .expect("plugin context recorder lock") + .clear(); + core_test_support::responses::mount_sse_once( + &server, + core_test_support::responses::sse(vec![core_test_support::responses::ev_completed("done")]), + ) + .await; + + test.submit_turn("Respond without using tools.").await?; + + let contexts = recorder + .plugin_contexts + .lock() + .expect("plugin context recorder lock"); + assert!( + !contexts.is_empty(), + "the model step must project selected plugins" + ); + assert!( + contexts + .iter() + .all(|context| context.as_ref() == Some(&selections)), + "every step-local plugin projection must receive its selected environments: {contexts:?}" + ); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index a65fbf2d88..dada42382a 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -129,6 +129,8 @@ mod managed_threads; mod mcp_auth_elicitation; mod mcp_auth_refresh; mod mcp_ema_config; +#[path = "mcp_executor_context_tests.rs"] +mod mcp_executor_context; mod mcp_extension_protocol; mod mcp_optional_startup_grace; mod mcp_refresh_cleanup; diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 138005a5af..699048ddee 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -2,6 +2,7 @@ use codex_config::McpServerConfig; use codex_exec_server_protocol::ExecutorCapabilityDiscoverySnapshot; use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::TurnEnvironmentSelection; use crate::ExtensionData; use crate::ExtensionDataInit; @@ -25,6 +26,8 @@ pub struct McpServerContributionContext<'a, C> { session_source: Option<&'a SessionSource>, /// Effective request originator for the active thread, when resolution is thread-scoped. originator: Option<&'a str>, + /// All selected environments captured for this projection, in priority order. + selected_environments: Option<&'a [TurnEnvironmentSelection]>, /// Selected roots resolved against ready environments for this exact step. ready_selected_capability_roots: Option<&'a [SelectedCapabilityRoot]>, /// Executor-materialized capability files shared by all consumers in this exact step. @@ -49,6 +52,7 @@ impl<'a, C> McpServerContributionContext<'a, C> { thread_init: None, session_source: None, originator: None, + selected_environments: None, ready_selected_capability_roots: None, executor_capability_discovery: None, } @@ -70,11 +74,30 @@ impl<'a, C> McpServerContributionContext<'a, C> { thread_init: Some(thread_init), session_source: None, originator: Some(originator), + selected_environments: None, ready_selected_capability_roots: Some(ready_selected_capability_roots), executor_capability_discovery, } } + /// Attaches the same environment snapshot used to project MCP authority. + /// + /// Preserve pending and failed selections: an unavailable primary executor + /// must not silently select a different ready executor for hosted tools. + pub fn with_selected_environments( + mut self, + environments: &'a [TurnEnvironmentSelection], + ) -> Self { + self.selected_environments = Some(environments); + self + } + + /// Returns the selected executor snapshot, or `None` for threadless discovery. + /// An empty slice means the thread explicitly has no selected environments. + pub fn selected_environments(&self) -> Option<&'a [TurnEnvironmentSelection]> { + self.selected_environments + } + /// Marks whether this projection replaces the published MCP authentication. pub fn with_auth_changed(mut self, auth_changed: bool) -> Self { self.auth_changed = auth_changed; diff --git a/codex-rs/ext/extension-api/tests/registry.rs b/codex-rs/ext/extension-api/tests/registry.rs index e9a86b9cb3..5f959ad866 100644 --- a/codex-rs/ext/extension-api/tests/registry.rs +++ b/codex-rs/ext/extension-api/tests/registry.rs @@ -33,11 +33,14 @@ use codex_extension_api::TurnItemContributor; use codex_extension_api::TurnLifecycleContributor; use codex_protocol::items::HookPromptItem; use codex_protocol::items::TurnItem; +use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; +use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::protocol::WarningEvent; +use codex_utils_path_uri::PathUri; use pretty_assertions::assert_eq; struct AllContributors; @@ -66,6 +69,39 @@ fn mcp_contribution_context_identifies_the_running_thread() { ); } +#[test] +fn mcp_context_preserves_selected_executor_order_and_unavailability() { + let config = (); + let environments = vec![ + TurnEnvironmentSelection { + environment_id: "starting-primary".to_owned(), + cwd: PathUri::parse("file:///workspace").expect("workspace URI"), + workspace_roots: Vec::new(), + config: EnvironmentConfigState::Pending, + }, + TurnEnvironmentSelection { + environment_id: "ready-secondary".to_owned(), + cwd: PathUri::parse("file:///other").expect("workspace URI"), + workspace_roots: Vec::new(), + config: EnvironmentConfigState::FromThread, + }, + ]; + let context = McpServerContributionContext::global(&config); + assert_eq!(context.selected_environments(), None); + assert_eq!( + context + .with_selected_environments(&[]) + .selected_environments(), + Some([].as_slice()) + ); + assert_eq!( + context + .with_selected_environments(&environments) + .selected_environments(), + Some(environments.as_slice()) + ); +} + impl ContextContributor for AllContributors { fn contribute_thread_context<'a>( &'a self,