Expose selected environments to MCP contributors (#51503)

## Why

MCP contributors need the complete executor selection to distinguish an unavailable primary executor from a ready secondary. Ready capability roots alone do not preserve pending or failed selections.

## What changed

Add `with_selected_environments` and `selected_environments` to `McpServerContributionContext`, and pass the captured selections during thread MCP resolution and model-step plugin selection. Preserve priority order and pending or failed entries. Use `None` for threadless discovery and `Some(&[])` for an explicitly empty selection.

## Testing

Add regression coverage for executor ordering with pending and failed selections, environment context in model-step plugin selection, and the distinction between absent and explicitly empty selections.

GitOrigin-RevId: 2c1222fae72c7e2c75b2f40abfad3b2048cf9316
This commit is contained in:
mtsui-oai
2026-10-06 23:07:48 +00:00
committed by copyberry
parent ddabe594e6
commit 044da98b4a
7 changed files with 276 additions and 11 deletions
+14 -10
View File
@@ -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,
+2 -1
View File
@@ -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;
@@ -1619,6 +1619,11 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
) -> codex_extension_api::ExtensionFuture<'a, Vec<codex_extension_api::SelectedPlugin<'a>>>
{
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");
@@ -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<Vec<Vec<TurnEnvironmentSelection>>>,
plugin_contexts: Mutex<Vec<Option<Vec<TurnEnvironmentSelection>>>>,
}
impl McpServerContributor<Config> for ExecutorSelectionRecorder {
fn id(&self) -> &'static str {
"executor_selection_recorder"
}
fn selected_plugins<'a>(
&'a self,
context: McpServerContributionContext<'a, Config>,
) -> ExtensionFuture<'a, Vec<SelectedPlugin<'a>>> {
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<McpServerContribution>> {
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::<Vec<_>>();
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(())
}
+2
View File
@@ -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;
@@ -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;
@@ -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,