mirror of
https://github.com/openai/codex.git
synced 2026-10-16 16:58:46 +00:00
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:
+14
-10
@@ -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,
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user