openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Propagate deferred environment capability roots to MCP (#33427)

Browse Source 
## What changed

- Let deferred environments provide selected capability roots with their ready signal.
- Validate that those roots have unique, non-empty IDs, belong to the registering environment, and stay within the root limit.
- Include roots from ready turn environments when resolving MCP contributions, and refresh the MCP runtime when the selected root set changes.
- Expose the exact ready root set to MCP contributors so executor plugins become available with their environment.

## Testing

- Cover ready-root propagation, validation failures, replacement isolation, reconnection, and MCP plugin availability refresh.

GitOrigin-RevId: ec3498aab1164824025094e96a9b1063b7b731ad
This commit is contained in:
jif
2026-07-15 21:07:43 +00:00
committed by copyberry
parent 1c4af96394
commit 79177c3e20
16 changed files with 333 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/mcp.rs
+3 -2
View File
@@ -24,6 +24,7 @@ use codex_mcp::codex_apps_mcp_server_config;
use codex_mcp::configured_mcp_servers;
use codex_mcp::effective_mcp_servers;
use codex_plugin::AppConnectorId;
use codex_protocol::capabilities::SelectedCapabilityRoot;
const LEGACY_CODEX_APPS_REGISTRATION_ID: &str = "legacy_codex_apps";
@@ -107,7 +108,7 @@ impl McpManager {
thread_init: &ExtensionDataInit,
thread_store: &ExtensionData,
originator: &str,
available_environment_ids: &[String],
ready_selected_capability_roots: &[SelectedCapabilityRoot],
) -> McpRuntimeProjection {
self.runtime_config_with_context(
McpServerContributionContext::for_step(
@@ -115,7 +116,7 @@ impl McpManager {
thread_init,
thread_store,
originator,
available_environment_ids,
ready_selected_capability_roots,
),
Some(originator),
)
codex-rs/core/src/session/mcp.rs
+69 -29
View File
@@ -1,10 +1,12 @@
use super::*;
use crate::mcp::McpRuntimeProjection;
use codex_exec_server::MAX_SELECTED_CAPABILITY_ROOTS;
use codex_exec_server::ResolvedSelectedCapabilityRoot;
use codex_mcp::ElicitationReviewRequest;
use codex_mcp::ElicitationReviewer;
use codex_mcp::ElicitationReviewerHandle;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::mcp_approval_meta::APPROVAL_KIND_KEY as MCP_ELICITATION_APPROVAL_KIND_KEY;
use codex_protocol::mcp_approval_meta::APPROVAL_KIND_MCP_TOOL_CALL as MCP_ELICITATION_APPROVAL_KIND_MCP_TOOL_CALL;
@@ -83,8 +85,8 @@ impl Session {
let selected_capability_roots = self
.resolve_selected_capability_roots_for_step(&environments)
.await;
let available_environment_ids =
Self::available_selected_environment_ids(&selected_capability_roots);
let ready_selected_capability_roots =
Self::ready_selected_capability_roots(&selected_capability_roots);
self.services
.mcp_manager
.runtime_config_for_step(
@@ -92,7 +94,7 @@ impl Session {
&self.services.mcp_thread_init,
&self.services.thread_extension_data,
&originator,
&available_environment_ids,
&ready_selected_capability_roots,
)
.await
.config
@@ -116,16 +118,18 @@ impl Session {
environments: &TurnEnvironmentSnapshot,
selected_capability_roots: &[ResolvedSelectedCapabilityRoot],
) -> Arc<McpRuntimeSnapshot> {
let ready_selected_capability_roots =
Self::ready_selected_capability_roots(selected_capability_roots);
let available_environment_ids =
Self::available_selected_environment_ids(selected_capability_roots);
let current = self.services.latest_mcp_runtime();
if current.available_environment_ids() == available_environment_ids {
if current.ready_selected_capability_roots() == ready_selected_capability_roots {
return current;
}
let _guard = self.services.mcp_projection_lock.lock().await;
let current = self.services.latest_mcp_runtime();
if current.available_environment_ids() == available_environment_ids {
if current.ready_selected_capability_roots() == ready_selected_capability_roots {
return current;
}
let mcp_projection = self
@@ -136,7 +140,7 @@ impl Session {
&self.services.mcp_thread_init,
&self.services.thread_extension_data,
&turn_context.originator,
&available_environment_ids,
&ready_selected_capability_roots,
)
.await;
let mcp_config = &mcp_projection.config;
@@ -146,8 +150,13 @@ impl Session {
.values()
.any(|server| {
let was_available = current
.available_environment_ids()
.contains(&server.environment_id);
.ready_selected_capability_roots()
.iter()
.any(|root| {
let CapabilityRootLocation::Environment { environment_id, .. } =
&root.location;
environment_id == &server.environment_id
});
let is_available = available_environment_ids.contains(&server.environment_id);
server.enabled && was_available != is_available
});
@@ -158,15 +167,15 @@ impl Session {
.has_same_servers(&mcp_config.mcp_server_catalog)
&& current.config().connector_snapshot == mcp_config.connector_snapshot
{
// Availability is only an input to the MCP projection. When that input changes but
// the projected servers and connectors do not, advance the input key without
// Selected roots are only an input to the MCP projection. When they change but the
// projected servers and connectors do not, advance the input key without
// replacing the live manager and restarting its processes.
let runtime = Arc::new(McpRuntimeSnapshot::new(
Arc::new(current.config().clone()),
mcp_projection.plugins_available,
current.manager_arc(),
current.runtime_context().clone(),
available_environment_ids,
ready_selected_capability_roots,
));
self.services.mcp_runtime.store(Some(Arc::clone(&runtime)));
return runtime;
@@ -175,7 +184,7 @@ impl Session {
turn_context,
mcp_projection,
environments,
&available_environment_ids,
&ready_selected_capability_roots,
Some(self.mcp_elicitation_reviewer()),
)
.await
@@ -185,11 +194,37 @@ impl Session {
&self,
environments: &TurnEnvironmentSnapshot,
) -> Vec<ResolvedSelectedCapabilityRoot> {
let mut seen_root_ids = self
.services
.selected_capability_roots
.iter()
.map(|root| root.id.as_str())
.collect::<HashSet<_>>();
let mut selected_capability_roots = self.services.selected_capability_roots.clone();
let mut ready_environment_root_count = 0;
for root in environments
.turn_environments
.iter()
.flat_map(|environment| environment.environment.selected_capability_roots())
{
if !seen_root_ids.insert(root.id.as_str()) {
continue;
}
if ready_environment_root_count == MAX_SELECTED_CAPABILITY_ROOTS {
tracing::warn!(
max_root_count = MAX_SELECTED_CAPABILITY_ROOTS,
"ignoring excess selected capability roots from ready environments"
);
break;
}
selected_capability_roots.push(root.clone());
ready_environment_root_count += 1;
}
self.services
.turn_environments
.environment_manager()
.resolve_selected_capability_roots(
&self.services.selected_capability_roots,
&selected_capability_roots,
&environments.captured_environments(),
)
.await
@@ -326,7 +361,7 @@ impl Session {
turn_context: &TurnContext,
mcp_projection: McpRuntimeProjection,
environments: &TurnEnvironmentSnapshot,
available_environment_ids: &[String],
ready_selected_capability_roots: &[SelectedCapabilityRoot],
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) -> Arc<McpRuntimeSnapshot> {
let auth = self.services.auth_manager.auth().await;
@@ -394,7 +429,7 @@ impl Session {
mcp_config,
plugins_available,
mcp_runtime_context,
available_environment_ids.to_vec(),
ready_selected_capability_roots.to_vec(),
refreshed_manager,
)
}
@@ -460,11 +495,9 @@ impl Session {
}
let _guard = self.services.mcp_projection_lock.lock().await;
let available_environment_ids = self
.services
.latest_mcp_runtime()
.available_environment_ids()
.to_vec();
let current_runtime = self.services.latest_mcp_runtime();
let ready_selected_capability_roots =
current_runtime.ready_selected_capability_roots().to_vec();
let mut mcp_projection = self
.services
.mcp_manager
@@ -473,7 +506,7 @@ impl Session {
&self.services.mcp_thread_init,
&self.services.thread_extension_data,
&turn_context.originator,
&available_environment_ids,
&ready_selected_capability_roots,
)
.await;
mcp_projection.config.mcp_server_catalog = mcp_projection
@@ -484,7 +517,7 @@ impl Session {
turn_context,
mcp_projection,
&turn_context.environments,
&available_environment_ids,
&ready_selected_capability_roots,
elicitation_reviewer,
)
.await;
@@ -529,11 +562,9 @@ impl Session {
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) {
let _guard = self.services.mcp_projection_lock.lock().await;
let available_environment_ids = self
.services
.latest_mcp_runtime()
.available_environment_ids()
.to_vec();
let current_runtime = self.services.latest_mcp_runtime();
let ready_selected_capability_roots =
current_runtime.ready_selected_capability_roots().to_vec();
let mcp_projection = self
.services
.mcp_manager
@@ -542,14 +573,14 @@ impl Session {
&self.services.mcp_thread_init,
&self.services.thread_extension_data,
&turn_context.originator,
&available_environment_ids,
&ready_selected_capability_roots,
)
.await;
self.refresh_mcp_servers_inner(
turn_context,
mcp_projection,
&turn_context.environments,
&available_environment_ids,
&ready_selected_capability_roots,
elicitation_reviewer,
)
.await;
@@ -569,6 +600,15 @@ impl Session {
available
}
fn ready_selected_capability_roots(
selected_capability_roots: &[ResolvedSelectedCapabilityRoot],
) -> Vec<SelectedCapabilityRoot> {
selected_capability_roots
.iter()
.map(|root| root.selected_root().clone())
.collect()
}
#[cfg(test)]
pub(crate) async fn mcp_startup_cancellation_token(&self) -> CancellationToken {
self.services
codex-rs/core/src/session/mcp_runtime.rs
+6 -5
View File
@@ -4,6 +4,7 @@ use std::sync::Arc;
use codex_mcp::McpConfig;
use codex_mcp::McpConnectionManager;
use codex_mcp::McpRuntimeContext;
use codex_protocol::capabilities::SelectedCapabilityRoot;
/// MCP config, plugin availability, exact environment bindings, and manager for one request.
pub struct McpRuntimeSnapshot {
@@ -11,7 +12,7 @@ pub struct McpRuntimeSnapshot {
plugins_available: bool,
manager: Arc<McpConnectionManager>,
runtime_context: McpRuntimeContext,
available_environment_ids: Vec<String>,
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
}
impl McpRuntimeSnapshot {
@@ -20,14 +21,14 @@ impl McpRuntimeSnapshot {
plugins_available: bool,
manager: Arc<McpConnectionManager>,
runtime_context: McpRuntimeContext,
available_environment_ids: Vec<String>,
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
) -> Self {
Self {
config,
plugins_available,
manager,
runtime_context,
available_environment_ids,
ready_selected_capability_roots,
}
}
@@ -51,8 +52,8 @@ impl McpRuntimeSnapshot {
&self.runtime_context
}
pub(crate) fn available_environment_ids(&self) -> &[String] {
&self.available_environment_ids
pub(crate) fn ready_selected_capability_roots(&self) -> &[SelectedCapabilityRoot] {
&self.ready_selected_capability_roots
}
#[cfg(test)]
codex-rs/core/src/session/session.rs
+2 -2
View File
@@ -685,7 +685,7 @@ impl Session {
mcp_thread_init_for_startup,
thread_extension_data_for_mcp,
&mcp_originator,
/*available_environment_ids*/ &[],
/*ready_selected_capability_roots*/ &[],
)
.await;
let mcp_config = &mcp_projection.config;
@@ -1234,7 +1234,7 @@ impl Session {
Arc::new(mcp_projection.config),
mcp_projection.plugins_available,
mcp_runtime_context,
/*available_environment_ids*/ Vec::new(),
/*ready_selected_capability_roots*/ Vec::new(),
mcp_connection_manager,
)
.await?;
codex-rs/core/src/session/tests.rs
+62 -15
View File
@@ -7686,9 +7686,26 @@ async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() {
}
#[tokio::test]
async fn plugin_availability_change_reuses_the_mcp_manager() {
async fn deferred_environment_roots_refresh_plugin_availability() {
struct ReadyPluginContributor;
struct PendingNoiseConnectProvider;
impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider {
fn connect_bundle(
&self,
_: codex_exec_server::NoiseChannelPublicKey,
) -> futures::future::BoxFuture<
'_,
Result<
codex_exec_server::NoiseRendezvousConnectBundle,
codex_exec_server::ExecServerError,
>,
> {
Box::pin(futures::future::pending())
}
}
impl codex_extension_api::McpServerContributor<Config> for ReadyPluginContributor {
fn id(&self) -> &'static str {
"ready_plugin_test"
@@ -7701,8 +7718,8 @@ async fn plugin_availability_change_reuses_the_mcp_manager() {
{
Box::pin(async move {
let available = context
.available_environment_ids()
.is_some_and(|ids| ids.iter().any(|id| id == "executor"));
.ready_selected_capability_roots()
.is_some_and(|roots| roots.iter().any(|root| root.id == "skill-only"));
available
.then(
|| codex_extension_api::McpServerContribution::SelectedPluginPackage {
@@ -7746,23 +7763,53 @@ async fn plugin_availability_change_reuses_the_mcp_manager() {
.expect("selected capability root URI"),
},
};
let environment = Arc::clone(
&turn_context
.environments
.primary()
.expect("ready test environment")
.environment,
let environment_manager = session.services.turn_environments.environment_manager();
let registration = environment_manager
.register_deferred_noise_environment(
"executor".to_string(),
Arc::new(PendingNoiseConnectProvider),
)
.expect("register deferred environment");
let environment = environment_manager
.get_environment("executor")
.expect("deferred environment");
assert!(environment.selected_capability_roots().is_empty());
registration
.complete(Ok(codex_exec_server::EnvironmentReadyInfo {
selected_capability_roots: vec![selected_root.clone()],
}))
.expect("complete deferred environment");
assert_eq!(
environment.selected_capability_roots(),
std::slice::from_ref(&selected_root)
);
let captured_environments = HashMap::from([("executor".to_string(), Some(environment))]);
let local_environment = turn_context
.environments
.primary()
.expect("ready local environment");
let environments = TurnEnvironmentSnapshot {
turn_environments: vec![TurnEnvironment::new(
"executor".to_string(),
environment,
local_environment.cwd().clone(),
local_environment.workspace_roots().to_vec(),
local_environment.shell.clone(),
)],
starting: Vec::new(),
};
let resolved_roots = session
.services
.turn_environments
.environment_manager()
.resolve_selected_capability_roots(&[selected_root], &captured_environments)
.resolve_selected_capability_roots_for_step(&environments)
.await;
assert_eq!(
resolved_roots
.iter()
.map(|root| root.selected_root().clone())
.collect::<Vec<_>>(),
vec![selected_root]
);
let new_runtime = session
.mcp_runtime_for_step(&turn_context, &turn_context.environments, &resolved_roots)
.mcp_runtime_for_step(&turn_context, &environments, &resolved_roots)
.await;
assert!(!old_runtime.plugins_available());
codex-rs/core/src/state/service.rs
+4 -4
View File
@@ -111,14 +111,14 @@ impl SessionServices {
config: Arc<McpConfig>,
plugins_available: bool,
runtime_context: McpRuntimeContext,
available_environment_ids: Vec<String>,
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
manager: McpConnectionManager,
) -> Result<()> {
let runtime = self.publish_mcp_runtime(
config,
plugins_available,
runtime_context,
available_environment_ids,
ready_selected_capability_roots,
manager,
);
runtime.manager().validate_required_servers().await
@@ -129,7 +129,7 @@ impl SessionServices {
config: Arc<McpConfig>,
plugins_available: bool,
runtime_context: McpRuntimeContext,
available_environment_ids: Vec<String>,
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
manager: McpConnectionManager,
) -> Arc<McpRuntimeSnapshot> {
let manager = Arc::new(manager);
@@ -141,7 +141,7 @@ impl SessionServices {
plugins_available,
manager,
runtime_context,
available_environment_ids,
ready_selected_capability_roots,
));
self.mcp_runtime.store(Some(Arc::clone(&runtime)));
runtime
codex-rs/core/src/thread_manager_tests.rs
+2 -2
View File
@@ -668,7 +668,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
&first_session.services.mcp_thread_init,
&first_session.services.thread_extension_data,
&first_originator,
/*available_environment_ids*/ &[],
/*ready_selected_capability_roots*/ &[],
)
.await;
let second_session = &second_thread.thread.session;
@@ -681,7 +681,7 @@ async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors()
&second_session.services.mcp_thread_init,
&second_session.services.thread_extension_data,
&second_originator,
/*available_environment_ids*/ &[],
/*ready_selected_capability_roots*/ &[],
)
.await;
codex-rs/core/tests/suite/remote_env.rs
+2 -1
View File
@@ -7,6 +7,7 @@ use codex_core::compact::SUMMARIZATION_PROMPT;
use codex_core::config::Constrained;
use codex_exec_server::CopyOptions;
use codex_exec_server::CreateDirectoryOptions;
use codex_exec_server::EnvironmentReadyInfo;
use codex_exec_server::ExecServerError;
use codex_exec_server::FileSystemSandboxContext;
use codex_exec_server::LOCAL_ENVIRONMENT_ID;
@@ -585,7 +586,7 @@ async fn deferred_executor_starts_noise_connection_after_registration() -> Resul
wait_for_response_request_count(&response_mock, /*expected_count*/ 1).await;
assert_eq!(response_mock.requests().len(), 1);
assert_eq!(provider.calls.load(Ordering::Relaxed), 0);
registration.complete(Ok(()))?;
registration.complete(Ok(EnvironmentReadyInfo::default()))?;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
codex-rs/exec-server/src/environment.rs
+86 -8
View File
@@ -1,8 +1,12 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::OnceLock;
use std::sync::RwLock;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use futures::FutureExt;
use crate::ExecServerError;
@@ -72,9 +76,23 @@ pub struct EnvironmentManager {
local_runtime_paths: Option<ExecServerRuntimePaths>,
}
/// Information supplied by the environment owner when a deferred environment is ready.
#[derive(Clone, Debug, Default, Eq, PartialEq)]
pub struct EnvironmentReadyInfo {
/// Ordered capability roots selected for this environment.
pub selected_capability_roots: Vec<SelectedCapabilityRoot>,
}
/// The one-shot capability to complete a deferred environment registration.
#[must_use = "the deferred environment cannot connect until registration is completed"]
pub struct DeferredEnvironmentRegistration(oneshot::Sender<Result<(), String>>);
pub struct DeferredEnvironmentRegistration {
completion: oneshot::Sender<Result<(), String>>,
environment_id: String,
ready_info: Arc<OnceLock<EnvironmentReadyInfo>>,
}
/// Maximum capability roots accepted from deferred environment ready information.
pub const MAX_SELECTED_CAPABILITY_ROOTS: usize = 256;
pub const LOCAL_ENVIRONMENT_ID: &str = "local";
pub const REMOTE_ENVIRONMENT_ID: &str = "remote";
@@ -367,15 +385,22 @@ impl EnvironmentManager {
validate_environment_id(&environment_id)?;
let identity = noise_channel_identity()?;
let (completion, readiness) = oneshot::channel();
let environment = Arc::new(Environment::remote_with_transport(
let ready_info = Arc::new(OnceLock::new());
let mut environment = Environment::remote_with_transport(
ExecServerTransportParams::Deferred(Box::new(crate::client_api::Deferred {
readiness: readiness.shared(),
transport: ExecServerTransportParams::NoiseRendezvous { provider, identity },
})),
self.local_runtime_paths.clone(),
));
self.insert_environment(environment_id, environment);
Ok(DeferredEnvironmentRegistration(completion))
);
environment.ready_info = Some(Arc::clone(&ready_info));
let environment = Arc::new(environment);
self.insert_environment(environment_id.clone(), environment);
Ok(DeferredEnvironmentRegistration {
completion,
environment_id,
ready_info,
})
}
/// Adds or replaces a named remote environment that connects through an
@@ -408,9 +433,48 @@ impl EnvironmentManager {
}
impl DeferredEnvironmentRegistration {
/// Completes provisioning with readiness or a terminal error message.
pub fn complete(self, result: Result<(), String>) -> Result<(), ExecServerError> {
self.0.send(result).map_err(|_| {
/// Completes provisioning with ready information or a terminal error message.
pub fn complete(
self,
result: Result<EnvironmentReadyInfo, String>,
) -> Result<(), ExecServerError> {
let result = match result {
Ok(ready_info) => {
if ready_info.selected_capability_roots.len() > MAX_SELECTED_CAPABILITY_ROOTS {
let error = ExecServerError::Protocol(format!(
"environment ready info contains more than {MAX_SELECTED_CAPABILITY_ROOTS} selected capability roots"
));
let _ = self.completion.send(Err(error.to_string()));
return Err(error);
}
let mut root_ids =
HashSet::with_capacity(ready_info.selected_capability_roots.len());
for root in &ready_info.selected_capability_roots {
let CapabilityRootLocation::Environment { environment_id, .. } = &root.location;
if root.id.trim().is_empty()
|| environment_id != &self.environment_id
|| !root_ids.insert(root.id.as_str())
{
let error = ExecServerError::Protocol(format!(
"selected capability roots must have unique non-empty IDs and belong to environment `{}`",
self.environment_id
));
let _ = self.completion.send(Err(error.to_string()));
return Err(error);
}
}
if self.ready_info.set(ready_info).is_err() {
let error = ExecServerError::Protocol(
"deferred environment ready info was already set".to_string(),
);
let _ = self.completion.send(Err(error.to_string()));
return Err(error);
}
Ok(())
}
Err(message) => Err(message),
};
self.completion.send(result).map_err(|_| {
ExecServerError::Disconnected("deferred environment registration is inactive".into())
})
}
@@ -499,6 +563,7 @@ fn optional_environment_value(name: &str) -> Option<String> {
#[derive(Clone)]
pub struct Environment {
remote_client: Option<LazyRemoteExecServerClient>,
ready_info: Option<Arc<OnceLock<EnvironmentReadyInfo>>>,
// Dropping the environment stops unfinished background startup work.
startup_task: Arc<Mutex<Option<AbortOnDropHandle<()>>>>,
exec_backend: Arc<dyn ExecBackend>,
@@ -512,6 +577,7 @@ impl Environment {
pub fn default_for_tests() -> Self {
Self {
remote_client: None,
ready_info: None,
startup_task: Arc::new(Mutex::new(None)),
exec_backend: Arc::new(LocalProcess::default()),
filesystem: Arc::new(LocalFileSystem::unsandboxed()),
@@ -566,6 +632,7 @@ impl Environment {
pub(crate) fn local(local_runtime_paths: ExecServerRuntimePaths) -> Self {
Self {
remote_client: None,
ready_info: None,
startup_task: Arc::new(Mutex::new(None)),
exec_backend: Arc::new(LocalProcess::with_local_runtime_paths(
local_runtime_paths.clone(),
@@ -602,6 +669,7 @@ impl Environment {
Self {
remote_client: Some(client.clone()),
ready_info: None,
startup_task: Arc::new(Mutex::new(None)),
exec_backend,
filesystem,
@@ -614,6 +682,16 @@ impl Environment {
self.remote_client.is_some()
}
/// Returns capability roots supplied with the deferred environment's ready signal.
pub fn selected_capability_roots(&self) -> &[SelectedCapabilityRoot] {
self.ready_info
.as_ref()
.and_then(|ready_info| ready_info.get())
.map_or(&[], |ready_info| {
ready_info.selected_capability_roots.as_slice()
})
}
/// Subscribes to the current connection state for this remote environment.
pub fn subscribe_connection_state(
&self,
codex-rs/exec-server/src/lib.rs
+2
View File
@@ -70,7 +70,9 @@ pub use environment::Environment;
pub use environment::EnvironmentConnectionState;
pub use environment::EnvironmentManager;
pub use environment::EnvironmentObservedStatus;
pub use environment::EnvironmentReadyInfo;
pub use environment::LOCAL_ENVIRONMENT_ID;
pub use environment::MAX_SELECTED_CAPABILITY_ROOTS;
pub use environment::REMOTE_ENVIRONMENT_ID;
pub use environment_provider::DefaultEnvironmentProvider;
pub use environment_provider::EnvironmentProvider;
codex-rs/exec-server/tests/deferred_environment.rs
+59 -3
View File
@@ -3,10 +3,14 @@ use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::EnvironmentReadyInfo;
use codex_exec_server::ExecServerError;
use codex_exec_server::NoiseChannelPublicKey;
use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
use futures::FutureExt;
use futures::future::BoxFuture;
use futures::poll;
@@ -38,6 +42,18 @@ impl NoiseRendezvousConnectProvider for FailingNoiseConnectProvider {
}
}
fn ready_info(root_id: &str, environment_id: &str) -> anyhow::Result<EnvironmentReadyInfo> {
Ok(EnvironmentReadyInfo {
selected_capability_roots: vec![SelectedCapabilityRoot {
id: root_id.to_string(),
location: CapabilityRootLocation::Environment {
environment_id: environment_id.to_string(),
path: PathUri::parse("file:///plugins/root")?,
},
}],
})
}
#[tokio::test]
async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> {
let manager = EnvironmentManager::without_environments();
@@ -52,8 +68,14 @@ async fn deferred_environment_waits_before_connecting() -> anyhow::Result<()> {
assert!(poll!(&mut readiness).is_pending());
assert_eq!(provider.calls(), 0);
assert!(environment.selected_capability_roots().is_empty());
registration.complete(Ok(()))?;
let ready_info = ready_info("selected-root", "tools")?;
registration.complete(Ok(ready_info.clone()))?;
assert_eq!(
environment.selected_capability_roots(),
ready_info.selected_capability_roots
);
let error = readiness.await.unwrap_err();
assert!(error.to_string().contains("test Noise provider called"));
assert_eq!(provider.calls(), 1);
@@ -94,6 +116,29 @@ async fn failure_and_dropped_registration_are_terminal() -> anyhow::Result<()> {
Ok(())
}
#[tokio::test]
async fn invalid_ready_info_is_terminal() -> anyhow::Result<()> {
let manager = EnvironmentManager::without_environments();
let provider = Arc::new(FailingNoiseConnectProvider::default());
let registration =
manager.register_deferred_noise_environment("tools".to_string(), provider.clone())?;
let environment = manager.get_environment("tools").expect("environment");
let error = registration
.complete(Ok(ready_info("selected-root", "other")?))
.unwrap_err();
assert!(matches!(error, ExecServerError::Protocol(_)));
let readiness_error = environment.wait_until_ready().await.unwrap_err();
assert!(
readiness_error
.to_string()
.contains("belong to environment")
);
assert!(environment.selected_capability_roots().is_empty());
assert_eq!(provider.calls(), 0);
Ok(())
}
#[tokio::test]
async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> {
let manager = EnvironmentManager::without_environments();
@@ -106,7 +151,13 @@ async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> {
.register_deferred_noise_environment("tools".to_string(), current_provider.clone())?;
let current = manager.get_environment("tools").expect("current");
old_registration.complete(Ok(()))?;
let old_ready_info = ready_info("old-root", "tools")?;
old_registration.complete(Ok(old_ready_info.clone()))?;
assert_eq!(
old_environment.selected_capability_roots(),
old_ready_info.selected_capability_roots
);
assert!(current.selected_capability_roots().is_empty());
let old_error = old_environment.wait_until_ready().await.unwrap_err();
assert!(old_error.to_string().contains("test Noise provider called"));
assert_eq!(old_provider.calls(), 1);
@@ -114,7 +165,12 @@ async fn late_completion_is_isolated_from_replacement() -> anyhow::Result<()> {
assert!(poll!(&mut current_readiness).is_pending());
assert_eq!(current_provider.calls(), 0);
current_registration.complete(Ok(()))?;
let current_ready_info = ready_info("current-root", "tools")?;
current_registration.complete(Ok(current_ready_info.clone()))?;
assert_eq!(
current.selected_capability_roots(),
current_ready_info.selected_capability_roots
);
let current_error = current_readiness.await.unwrap_err();
assert!(
current_error
codex-rs/exec-server/tests/relay.rs
+21 -1
View File
@@ -17,6 +17,7 @@ use base64::engine::general_purpose::STANDARD;
use codex_api::AuthProvider;
use codex_exec_server::EnvironmentConnectionState;
use codex_exec_server::EnvironmentManager;
use codex_exec_server::EnvironmentReadyInfo;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecResponse;
use codex_exec_server::ExecServerClient;
@@ -30,6 +31,8 @@ use codex_exec_server::NoiseRendezvousConnectBundle;
use codex_exec_server::NoiseRendezvousConnectProvider;
use codex_exec_server::ProcessId;
use codex_exec_server::RemoteEnvironmentConfig;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use codex_utils_path_uri::PathUri;
use futures::SinkExt;
use futures::StreamExt;
@@ -164,7 +167,16 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
.context("remote environment connection state")?;
assert_eq!(provider.calls(), 0);
registration.complete(Ok(()))?;
let selected_capability_roots = vec![SelectedCapabilityRoot {
id: "executor-plugin".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: ENVIRONMENT_ID.to_string(),
path: PathUri::parse("file:///plugins/executor-plugin")?,
},
}];
registration.complete(Ok(EnvironmentReadyInfo {
selected_capability_roots: selected_capability_roots.clone(),
}))?;
let harness_websocket = accept_websocket(&listener, "harness").await?;
let first_relay = tokio::spawn(proxy_relay_frames(
environment_websocket,
@@ -174,6 +186,10 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
let initial_info = timeout(TEST_TIMEOUT, environment.info())
.await
.context("deferred Noise environment should become ready")??;
assert_eq!(
environment.selected_capability_roots(),
selected_capability_roots
);
assert_eq!(provider.calls(), 1);
assert_eq!(
next_connection_state(&mut connection_state).await?,
@@ -198,6 +214,10 @@ async fn deferred_noise_environment_connects_and_reconnects_with_fresh_bundle()
.context("deferred Noise environment should reconnect")??;
assert_eq!(recovered_info, initial_info);
assert_eq!(
environment.selected_capability_roots(),
selected_capability_roots
);
assert_eq!(provider.calls(), 2);
assert_eq!(
next_connection_state(&mut connection_state).await?,
codex-rs/ext/extension-api/src/contributors/mcp.rs
+9 -11
View File
@@ -1,4 +1,5 @@
use codex_config::McpServerConfig;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use crate::ExtensionData;
use crate::ExtensionDataInit;
@@ -17,8 +18,8 @@ pub struct McpServerContributionContext<'a, C> {
thread_init: Option<&'a ExtensionDataInit>,
/// Effective request originator for the active thread, when resolution is thread-scoped.
originator: Option<&'a str>,
/// Environment IDs whose selected roots may contribute to this exact step.
available_environment_ids: Option<&'a [String]>,
/// Selected roots resolved against ready environments for this exact step.
ready_selected_capability_roots: Option<&'a [SelectedCapabilityRoot]>,
}
impl<C> Clone for McpServerContributionContext<'_, C> {
@@ -37,7 +38,7 @@ impl<'a, C> McpServerContributionContext<'a, C> {
thread_store: None,
thread_init: None,
originator: None,
available_environment_ids: None,
ready_selected_capability_roots: None,
}
}
@@ -47,14 +48,14 @@ impl<'a, C> McpServerContributionContext<'a, C> {
thread_init: &'a ExtensionDataInit,
thread_store: &'a ExtensionData,
originator: &'a str,
available_environment_ids: &'a [String],
ready_selected_capability_roots: &'a [SelectedCapabilityRoot],
) -> Self {
Self {
config,
thread_store: Some(thread_store),
thread_init: Some(thread_init),
originator: Some(originator),
available_environment_ids: Some(available_environment_ids),
ready_selected_capability_roots: Some(ready_selected_capability_roots),
}
}
@@ -78,12 +79,9 @@ impl<'a, C> McpServerContributionContext<'a, C> {
self.originator
}
/// Returns the exact environment availability projection for a model step.
///
/// `Some` means contributors must omit selected roots whose environment ID is absent from the
/// slice. Global resolution returns `None` because it has no thread environments.
pub fn available_environment_ids(&self) -> Option<&'a [String]> {
self.available_environment_ids
/// Returns selected roots resolved against the ready environments for this model step.
pub fn ready_selected_capability_roots(&self) -> Option<&'a [SelectedCapabilityRoot]> {
self.ready_selected_capability_roots
}
}
codex-rs/ext/mcp/src/executor_plugin.rs
+1 -17
View File
@@ -6,7 +6,6 @@ use codex_extension_api::ExtensionFuture;
use codex_extension_api::McpServerContribution;
use codex_extension_api::McpServerContributionContext;
use codex_extension_api::McpServerContributor;
use codex_protocol::capabilities::CapabilityRootLocation;
use codex_protocol::capabilities::SelectedCapabilityRoot;
use std::collections::HashMap;
use std::sync::Arc;
@@ -148,31 +147,16 @@ impl McpServerContributor<Config> for SelectedExecutorPluginMcpContributor {
context: McpServerContributionContext<'a, Config>,
) -> ExtensionFuture<'a, Vec<McpServerContribution>> {
Box::pin(async move {
let Some(thread_init) = context.thread_init() else {
return Vec::new();
};
let Some(thread_store) = context.thread_store() else {
return Vec::new();
};
let Some(selected_roots) = thread_init.get::<Vec<SelectedCapabilityRoot>>() else {
let Some(selected_roots) = context.ready_selected_capability_roots() else {
return Vec::new();
};
let state = thread_store.get_or_init(SelectedExecutorPluginMcpState::default);
let mut contributions = Vec::new();
for (selection_order, selected_root) in selected_roots.iter().enumerate() {
let CapabilityRootLocation::Environment { environment_id, .. } =
&selected_root.location;
if context
.available_environment_ids()
.is_some_and(|available| {
!available
.iter()
.any(|available| available == environment_id)
})
{
continue;
}
let Some(plugin) = self.metadata_for_root(&state, selected_root).await else {
continue;
};
codex-rs/ext/mcp/src/lib_tests.rs
+1 -1
View File
@@ -27,7 +27,7 @@ async fn hosted_plugin_runtime_forwards_thread_originator() -> Result<(), Box<dy
&thread_init,
&thread_store,
"codex_work_desktop",
/*available_environment_ids*/ &[],
/*ready_selected_capability_roots*/ &[],
))
.await;
let [McpServerContribution::Set { config: server, .. }] = contributions.as_slice() else {
codex-rs/ext/mcp/tests/executor_plugin_mcp.rs
+4 -5
View File
@@ -186,16 +186,15 @@ async fn raw_selected_plugin_contributions(
Arc::new(EnvironmentManager::default_for_tests()),
);
let registry = builder.build();
let mut thread_init = ExtensionDataInit::new();
thread_init.insert(vec![SelectedCapabilityRoot {
let thread_init = ExtensionDataInit::new();
let selected_capability_roots = vec![SelectedCapabilityRoot {
id: "selected-root".to_string(),
location: CapabilityRootLocation::Environment {
environment_id: LOCAL_ENVIRONMENT_ID.to_string(),
path: PathUri::from_host_native_path(plugin_root)?,
},
}]);
}];
let thread_store = ExtensionData::new_with_init("test-thread", thread_init.clone());
let available_environment_ids = vec![LOCAL_ENVIRONMENT_ID.to_string()];
Ok(registry.mcp_server_contributors()[0]
.contribute(McpServerContributionContext::for_step(
@@ -203,7 +202,7 @@ async fn raw_selected_plugin_contributions(
&thread_init,
&thread_store,
"test_originator",
&available_environment_ids,
&selected_capability_roots,
))
.await)
}