github/codex
4
0 Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-14 06:48:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Centralize thread MCP state in McpRuntime (#34930)

Browse Source 
## What changed

- Make `McpRuntime` own the published MCP configuration, connections, elicitation routing, and selected capability roots for a thread.
- Capture immutable MCP bindings for model steps and tool calls so in-flight work keeps a consistent connection set and approval authority while refreshed state is published atomically.
- Mark MCP state dirty when relevant configuration, plugins, environments, authentication, or elicitation capabilities change, then rebuild it before the next sampling request or out-of-band MCP operation.
- Separate config reloads from server invalidation: `ReloadMcpConfig` applies resolved MCP inputs, while `RefreshMcpServers` requests reinitialization from the thread's latest state.

## Testing

- Cover refreshed state visibility for existing turns, stable step bindings, current approval authority, cancelled refresh retries, resource-client reconciliation, and Apps recovery between sampling requests.

GitOrigin-RevId: 59eabb1aa8dc083426bd18ef4d3630508f376401
This commit is contained in:
jif authored and copyberry committed 2026-07-23 11:27:34 +00:00
1 parent 39a2438d16
commit e497325a6a
54 files changed
+1792 -2038
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

codex-rs/core/src/session/mcp.rs
+149 -311
View File
@@ -1,5 +1,4 @@
use super::*;
use crate::mcp::McpRuntimeProjection;
use codex_exec_server::ExecutorCapabilityDiscoveryCache;
use codex_exec_server::ExecutorCapabilityDiscoverySnapshot;
use codex_exec_server::MAX_SELECTED_CAPABILITY_ROOTS;
@@ -7,7 +6,6 @@ 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;
@@ -45,6 +43,21 @@ struct GuardianMcpElicitationReviewer {
session: std::sync::Weak<Session>,
}
/// Restores a claimed refresh when its task is cancelled before publication.
struct McpRefreshInvalidationGuard<'a> {
pending: &'a std::sync::atomic::AtomicBool,
published: bool,
}
impl Drop for McpRefreshInvalidationGuard<'_> {
fn drop(&mut self) {
if !self.published {
self.pending
.store(true, std::sync::atomic::Ordering::Release);
}
}
}
pub(crate) struct McpServerElicitationOutcome {
pub(crate) response: Option<ElicitationResponse>,
pub(crate) sent: bool,
@@ -82,6 +95,13 @@ impl ElicitationReviewer for GuardianMcpElicitationReviewer {
impl Session {
pub(crate) async fn runtime_mcp_config(&self, config: &Config) -> McpConfig {
self.runtime_mcp_config_and_context(config).await.0
}
pub(crate) async fn runtime_mcp_config_and_context(
&self,
config: &Config,
) -> (McpConfig, McpRuntimeContext) {
let originator = self.originator().await;
let environments = self.services.turn_environments.snapshot().await;
let selected_capability_roots = self
@@ -92,7 +112,8 @@ impl Session {
let executor_capability_discovery = self
.executor_capability_discovery_for_step(config, &ready_selected_capability_roots)
.await;
self.services
let mcp_config = self
.services
.mcp_manager
.runtime_config_for_step(
config,
@@ -103,7 +124,17 @@ impl Session {
executor_capability_discovery.as_deref(),
)
.await
.config
.config;
let local_stdio_fallback_cwd = environments
.primary()
.and_then(|environment| environment.cwd().to_abs_path().ok())
.map(|cwd| cwd.to_path_buf())
.unwrap_or_else(|| config.cwd.to_path_buf());
let runtime_context = McpRuntimeContext::new(
self.services.turn_environments.environment_manager(),
local_stdio_fallback_cwd,
);
(mcp_config, runtime_context)
}
pub(crate) async fn runtime_mcp_servers(
@@ -113,91 +144,107 @@ impl Session {
codex_mcp::configured_mcp_servers(&self.runtime_mcp_config(config).await)
}
#[expect(
clippy::await_holding_invalid_type,
reason = "MCP runtime comparison and publication must remain serialized"
)]
/// Publishes changed MCP state, waiting for any refresh already in progress.
pub(crate) async fn refresh_mcp_if_dirty(self: &Arc<Self>) {
let Ok(_refresh) = self.mcp_refresh_lock.acquire().await else {
error!("MCP runtime refresh semaphore closed");
return;
};
loop {
if !self
.mcp_refresh_pending
.swap(false, std::sync::atomic::Ordering::AcqRel)
{
return;
}
let mut refresh_invalidation = McpRefreshInvalidationGuard {
pending: &self.mcp_refresh_pending,
published: false,
};
let desired = self.latest_mcp_desired_state().await;
let selected_capability_roots = self
.resolve_selected_capability_roots_for_step(&desired.environments)
.await;
let ready_selected_capability_roots =
Self::ready_selected_capability_roots(&selected_capability_roots);
let executor_capability_discovery = self
.executor_capability_discovery_for_step(
&desired.config,
&ready_selected_capability_roots,
)
.await;
let mcp_projection = self
.services
.mcp_manager
.runtime_config_for_step(
&desired.config,
&self.services.mcp_thread_init,
&self.services.thread_extension_data,
&desired.originator,
&ready_selected_capability_roots,
executor_capability_discovery.as_deref(),
)
.await;
self.publish_mcp_runtime(
&desired,
mcp_projection,
&ready_selected_capability_roots,
Some(self.mcp_elicitation_reviewer()),
)
.await;
refresh_invalidation.published = true;
if !self
.mcp_refresh_pending
.load(std::sync::atomic::Ordering::Acquire)
{
return;
}
}
}
/// Refreshes the future Apps catalog without making an exact step a refresh owner.
pub(crate) async fn hard_refresh_latest_codex_apps_tools(
self: &Arc<Self>,
) -> anyhow::Result<Vec<codex_mcp::ToolInfo>> {
self.refresh_mcp_if_dirty().await;
let _refresh = self
.mcp_refresh_lock
.acquire()
.await
.map_err(|_| anyhow::anyhow!("MCP runtime refresh semaphore closed"))?;
self.services
.mcp_runtime
.latest_hard_refresh_codex_apps_tools_cache()
.await
}
pub(super) fn mark_mcp_runtime_dirty(&self) {
self.mcp_refresh_pending
.store(true, std::sync::atomic::Ordering::Release);
}
#[tracing::instrument(name = "mcp.runtime.resolve_for_step", skip_all)]
pub(crate) async fn mcp_runtime_for_step(
self: &Arc<Self>,
turn_context: &TurnContext,
environments: &TurnEnvironmentSnapshot,
selected_capability_roots: &[ResolvedSelectedCapabilityRoot],
executor_capability_discovery: Option<&ExecutorCapabilityDiscoverySnapshot>,
) -> Arc<McpRuntimeSnapshot> {
) -> Arc<codex_mcp::McpBinding> {
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.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.ready_selected_capability_roots() == ready_selected_capability_roots {
return current;
}
let mcp_projection = self
if self
.services
.mcp_manager
.runtime_config_for_step(
&turn_context.config,
&self.services.mcp_thread_init,
&self.services.thread_extension_data,
&turn_context.originator,
&ready_selected_capability_roots,
executor_capability_discovery,
)
.await;
let mcp_config = &mcp_projection.config;
let changed_environment_is_used_by_mcp = mcp_config
.mcp_server_catalog
.configured_servers()
.values()
.any(|server| {
let was_available = current
.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
});
if !changed_environment_is_used_by_mcp
&& current
.config()
.mcp_server_catalog
.has_same_servers(&mcp_config.mcp_server_catalog)
&& current.config().connector_snapshot == mcp_config.connector_snapshot
.mcp_runtime
.current_ready_selected_capability_roots()
!= ready_selected_capability_roots
{
// 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(),
ready_selected_capability_roots,
));
self.services
.mcp_runtime_snapshot
.store(Some(Arc::clone(&runtime)));
return runtime;
self.mark_mcp_runtime_dirty();
}
self.refresh_mcp_servers_inner(
turn_context,
mcp_projection,
environments,
&ready_selected_capability_roots,
Some(self.mcp_elicitation_reviewer()),
)
.await
self.refresh_mcp_if_dirty().await;
if let Some(binding) = self.services.mcp_runtime.current_binding().await {
return binding;
}
let config = Arc::new(self.runtime_mcp_config(&turn_context.config).await);
Arc::new(codex_mcp::McpBinding::empty(config))
}
#[tracing::instrument(
@@ -299,12 +346,7 @@ impl Session {
request_id: RequestId,
request: ElicitationRequest,
) -> McpServerElicitationOutcome {
if self
.services
.latest_mcp_runtime()
.manager()
.elicitations_auto_deny()
{
if self.services.mcp_runtime.elicitations_auto_deny() {
return McpServerElicitationOutcome {
response: Some(ElicitationResponse {
action: codex_rmcp_client::ElicitationAction::Accept,
@@ -398,187 +440,11 @@ impl Session {
}
self.services
.latest_mcp_runtime()
.manager_arc()
.mcp_runtime
.resolve_elicitation(server_name, id, response)
.await
}
#[tracing::instrument(name = "mcp.runtime.refresh", skip_all)]
async fn refresh_mcp_servers_inner(
&self,
turn_context: &TurnContext,
mcp_projection: McpRuntimeProjection,
environments: &TurnEnvironmentSnapshot,
ready_selected_capability_roots: &[SelectedCapabilityRoot],
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) -> Arc<McpRuntimeSnapshot> {
let auth = self.services.auth_manager.auth().await;
let McpRuntimeProjection {
config: mcp_config,
plugins_available,
} = mcp_projection;
let mcp_config = Arc::new(mcp_config);
let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config);
let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref());
let environment_manager = self.services.turn_environments.environment_manager();
// TODO(anp): Migrate MCP runtime cwd plumbing to PathUri so foreign environment cwd
// values can be used without falling back to the legacy host cwd.
let cwd = environments
.primary()
.and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok())
.map(|cwd| cwd.to_path_buf())
.unwrap_or_else(|| {
#[allow(deprecated)]
turn_context.cwd.to_path_buf()
});
let mcp_runtime_context = McpRuntimeContext::new(environment_manager, cwd);
let mcp_startup_cancellation_token = {
let mut guard = self.services.mcp_startup_cancellation_token.lock().await;
// The previous runtime owns the old token and may still be serving an in-flight step.
// Its manager cancels that token when the last runtime handle is dropped.
let cancellation_token = CancellationToken::new();
*guard = cancellation_token.clone();
cancellation_token
};
let current_runtime = self.services.latest_mcp_runtime();
let codex_apps_auth_manager =
codex_mcp::host_owned_codex_apps_enabled(&mcp_config, auth.as_ref())
.then(|| Arc::clone(&self.services.auth_manager));
let refreshed_manager = McpConnectionSet::new(
&mcp_servers,
mcp_config.mcp_oauth_credentials_store_mode,
mcp_config.auth_keyring_backend_kind,
&turn_context.approval_policy,
turn_context.sub_id.clone(),
Some(self.get_tx_event()),
mcp_startup_cancellation_token,
turn_context.permission_profile(),
mcp_runtime_context.clone(),
mcp_config.codex_home.clone(),
self.services.mcp_manager.codex_apps_tools_cache(),
self.services.mcp_manager.tool_catalog_cache(),
connector_runtime_context_key(auth.as_ref()),
mcp_config.prefix_mcp_tool_names,
mcp_config.client_elicitation_capability.clone(),
self.services
.supports_openai_form_elicitation
.load(std::sync::atomic::Ordering::Relaxed),
tool_plugin_provenance,
auth.as_ref(),
codex_apps_auth_manager,
elicitation_reviewer,
Some(self.mcp_elicitation_lifecycle()),
current_runtime.manager().elicitation_router(),
)
.await;
refreshed_manager
.set_elicitations_auto_deny(current_runtime.manager().elicitations_auto_deny());
self.services.publish_mcp_runtime(
mcp_config,
plugins_available,
mcp_runtime_context,
ready_selected_capability_roots.to_vec(),
refreshed_manager,
)
}
#[expect(
clippy::await_holding_invalid_type,
reason = "MCP runtime refresh and publication must remain serialized"
)]
pub(crate) async fn refresh_mcp_servers_if_requested(
&self,
turn_context: &TurnContext,
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) {
let refresh_config = { self.pending_mcp_server_refresh_config.lock().await.take() };
let Some(refresh_config) = refresh_config else {
return;
};
let McpServerRefreshConfig {
mcp_servers,
mcp_oauth_credentials_store_mode,
auth_keyring_backend_kind,
} = refresh_config;
let mcp_servers =
match serde_json::from_value::<HashMap<String, McpServerConfig>>(mcp_servers) {
Ok(servers) => servers,
Err(err) => {
warn!("failed to parse MCP server refresh config: {err}");
return;
}
};
let store_mode = match serde_json::from_value::<OAuthCredentialsStoreMode>(
mcp_oauth_credentials_store_mode,
) {
Ok(mode) => mode,
Err(err) => {
warn!("failed to parse MCP OAuth refresh config: {err}");
return;
}
};
let keyring_backend_kind =
match serde_json::from_value::<AuthKeyringBackendKind>(auth_keyring_backend_kind) {
Ok(kind) => kind,
Err(err) => {
warn!("failed to parse MCP auth keyring backend refresh config: {err}");
return;
}
};
let mut refresh_config = self.get_config().await.as_ref().clone();
refresh_config.mcp_oauth_credentials_store_mode = store_mode;
let secret_auth_storage_enabled = match keyring_backend_kind {
AuthKeyringBackendKind::Direct => false,
AuthKeyringBackendKind::Secrets => true,
};
if let Err(err) = refresh_config
.features
.set_enabled(Feature::SecretAuthStorage, secret_auth_storage_enabled)
{
warn!("failed to apply MCP auth keyring backend refresh config: {err}");
return;
}
let _guard = self.services.mcp_projection_lock.lock().await;
let current_runtime = self.services.latest_mcp_runtime();
let ready_selected_capability_roots =
current_runtime.ready_selected_capability_roots().to_vec();
let executor_capability_discovery = self
.executor_capability_discovery_for_step(
&refresh_config,
&ready_selected_capability_roots,
)
.await;
let mut mcp_projection = self
.services
.mcp_manager
.runtime_config_for_step(
&refresh_config,
&self.services.mcp_thread_init,
&self.services.thread_extension_data,
&turn_context.originator,
&ready_selected_capability_roots,
executor_capability_discovery.as_deref(),
)
.await;
mcp_projection.config.mcp_server_catalog = mcp_projection
.config
.mcp_server_catalog
.with_materialized_servers(mcp_servers);
self.refresh_mcp_servers_inner(
turn_context,
mcp_projection,
&turn_context.environments,
&ready_selected_capability_roots,
elicitation_reviewer,
)
.await;
}
pub(crate) async fn set_openai_form_elicitation_support(
&self,
supported: bool,
@@ -592,35 +458,33 @@ impl Session {
return Ok(());
}
let config = self.get_config().await;
let refresh_config = McpServerRefreshConfig {
mcp_servers: serde_json::to_value(config.mcp_servers.get())?,
mcp_oauth_credentials_store_mode: serde_json::to_value(
config.mcp_oauth_credentials_store_mode,
)?,
auth_keyring_backend_kind: serde_json::to_value(config.auth_keyring_backend_kind())?,
};
self.services
.supports_openai_form_elicitation
.store(supported, std::sync::atomic::Ordering::Relaxed);
*self.pending_mcp_server_refresh_config.lock().await = Some(refresh_config);
self.mark_mcp_runtime_dirty();
Ok(())
}
#[expect(
clippy::await_holding_invalid_type,
reason = "MCP runtime refresh and publication must remain serialized"
)]
pub(crate) async fn refresh_mcp_servers_now(
&self,
turn_context: &TurnContext,
refresh_config: &Config,
elicitation_reviewer: Option<ElicitationReviewerHandle>,
) {
let _guard = self.services.mcp_projection_lock.lock().await;
let current_runtime = self.services.latest_mcp_runtime();
let ready_selected_capability_roots =
current_runtime.ready_selected_capability_roots().to_vec();
let Ok(_refresh) = self.mcp_refresh_lock.acquire().await else {
error!("MCP runtime refresh semaphore closed");
return;
};
{
let mut state = self.state.lock().await;
let mut config = (*state.session_configuration.original_config_do_not_use).clone();
config.mcp_servers = refresh_config.mcp_servers.clone();
state.session_configuration.original_config_do_not_use = Arc::new(config);
}
let ready_selected_capability_roots = self
.services
.mcp_runtime
.current_ready_selected_capability_roots();
let executor_capability_discovery = self
.executor_capability_discovery_for_step(
refresh_config,
@@ -639,30 +503,17 @@ impl Session {
executor_capability_discovery.as_deref(),
)
.await;
self.refresh_mcp_servers_inner(
turn_context,
let mut desired = self.latest_mcp_desired_state().await;
desired.config = Arc::new(refresh_config.clone());
self.publish_mcp_runtime(
&desired,
mcp_projection,
&turn_context.environments,
&ready_selected_capability_roots,
elicitation_reviewer,
)
.await;
}
fn available_selected_environment_ids(
selected_capability_roots: &[ResolvedSelectedCapabilityRoot],
) -> Vec<String> {
let mut available = Vec::new();
for root in selected_capability_roots {
let CapabilityRootLocation::Environment { environment_id, .. } =
&root.selected_root().location;
if !available.contains(environment_id) {
available.push(environment_id.clone());
}
}
available
}
pub(crate) fn ready_selected_capability_roots(
selected_capability_roots: &[ResolvedSelectedCapabilityRoot],
) -> Vec<SelectedCapabilityRoot> {
@@ -672,21 +523,8 @@ impl Session {
.collect()
}
#[cfg(test)]
pub(crate) async fn mcp_startup_cancellation_token(&self) -> CancellationToken {
self.services
.mcp_startup_cancellation_token
.lock()
.await
.clone()
}
pub(crate) async fn cancel_mcp_startup(&self) {
self.services
.mcp_startup_cancellation_token
.lock()
.await
.cancel();
pub(crate) fn cancel_mcp_startup(&self) {
self.services.mcp_runtime.cancel_startup();
}
}