From e497325a6a1743cfadeee41a6b5f05ebf7fd0221 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 23 Jul 2026 11:22:58 +0000 Subject: [PATCH] Centralize thread MCP state in `McpRuntime` (#34930) ## 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 --- .../app-server/src/effective_plugin_change.rs | 7 +- codex-rs/app-server/src/mcp_refresh.rs | 126 ++-- .../request_processors/account_processor.rs | 7 +- .../apps_processor/installed.rs | 56 +- .../request_processors/config_processor.rs | 18 +- .../src/request_processors/mcp_processor.rs | 11 +- codex-rs/codex-mcp/src/binding.rs | 8 + codex-rs/codex-mcp/src/binding_tests.rs | 37 +- codex-rs/codex-mcp/src/connection_manager.rs | 198 ++----- .../src/connection_manager/required.rs | 8 +- .../src/connection_manager/tool_catalog.rs | 62 +- .../codex-mcp/src/connection_manager_tests.rs | 50 +- codex-rs/codex-mcp/src/elicitation.rs | 71 +-- codex-rs/codex-mcp/src/lib.rs | 6 +- codex-rs/codex-mcp/src/mcp/mod.rs | 38 +- codex-rs/codex-mcp/src/mcp/mod_tests.rs | 4 + codex-rs/codex-mcp/src/resource_client.rs | 8 +- codex-rs/codex-mcp/src/rmcp_client.rs | 13 - codex-rs/codex-mcp/src/runtime.rs | 295 +++++++++- codex-rs/core-api/src/lib.rs | 1 - codex-rs/core/src/codex_delegate.rs | 25 +- codex-rs/core/src/codex_thread.rs | 51 +- codex-rs/core/src/compact.rs | 3 +- codex-rs/core/src/compact_tests.rs | 6 +- codex-rs/core/src/config/mod.rs | 4 + codex-rs/core/src/connectors.rs | 87 +-- codex-rs/core/src/mcp.rs | 1 + codex-rs/core/src/mcp_tool_call.rs | 550 +++++++++--------- codex-rs/core/src/mcp_tool_call_tests.rs | 269 ++------- codex-rs/core/src/session/handlers.rs | 23 +- codex-rs/core/src/session/mcp.rs | 460 +++++---------- codex-rs/core/src/session/mcp_runtime.rs | 251 ++++---- codex-rs/core/src/session/mod.rs | 118 ++-- codex-rs/core/src/session/session.rs | 119 ++-- codex-rs/core/src/session/step_context.rs | 6 +- codex-rs/core/src/session/tests.rs | 342 +++++------ codex-rs/core/src/session/turn.rs | 26 +- codex-rs/core/src/session/turn_context.rs | 12 +- codex-rs/core/src/state/service.rs | 62 +- codex-rs/core/src/state/turn.rs | 18 + .../list_mcp_resource_templates.rs | 6 +- .../mcp_resource/list_mcp_resources.rs | 6 +- .../mcp_resource/read_mcp_resource.rs | 4 +- .../tools/handlers/request_plugin_install.rs | 26 +- codex-rs/core/src/tools/router.rs | 3 +- codex-rs/core/src/tools/spec_plan.rs | 9 +- codex-rs/core/src/tools/spec_plan_tests.rs | 5 +- .../core/tests/common/apps_test_server.rs | 19 +- codex-rs/core/tests/suite/mcp_auth_refresh.rs | 64 +- codex-rs/core/tests/suite/mcp_tool_cache.rs | 23 +- .../core/tests/suite/mcp_tool_exposure.rs | 161 +++-- .../tests/suite/request_plugin_install.rs | 19 +- codex-rs/ext/skills/src/extension.rs | 16 +- codex-rs/protocol/src/protocol.rs | 12 +- 54 files changed, 1792 insertions(+), 2038 deletions(-) diff --git a/codex-rs/app-server/src/effective_plugin_change.rs b/codex-rs/app-server/src/effective_plugin_change.rs index c23572a103..fb4b493245 100644 --- a/codex-rs/app-server/src/effective_plugin_change.rs +++ b/codex-rs/app-server/src/effective_plugin_change.rs @@ -32,16 +32,11 @@ pub(crate) fn effective_plugins_changed_callback( thread_manager.skills_service().clear_cache(); let refresh_thread_manager = Arc::clone(&thread_manager); - let refresh_config_manager = config_manager.clone(); tokio::spawn(async move { if refresh_thread_manager.list_thread_ids().await.is_empty() { return; } - crate::mcp_refresh::queue_best_effort_refresh( - &refresh_thread_manager, - &refresh_config_manager, - ) - .await; + crate::mcp_refresh::invalidate_loaded_threads(&refresh_thread_manager).await; }); if change.materialized_remote_plugins.is_empty() { diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index b1ef983082..5a31e8e4dc 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -1,6 +1,7 @@ use crate::config_manager::ConfigManager; use codex_core::CodexThread; use codex_core::ThreadManager; +use codex_core::config::Config; use codex_protocol::ThreadId; use codex_protocol::protocol::McpServerRefreshConfig; use codex_protocol::protocol::Op; @@ -8,7 +9,7 @@ use std::io; use std::sync::Arc; use tracing::warn; -pub(crate) async fn queue_strict_refresh( +pub(crate) async fn reload_mcp_config( thread_manager: &Arc, config_manager: &ConfigManager, ) -> io::Result<()> { @@ -25,15 +26,19 @@ pub(crate) async fn queue_strict_refresh( refreshes.push((thread_id, thread, config)); } for (thread_id, thread, config) in refreshes { - queue_refresh(thread_id, thread, config).await?; + thread + .submit(Op::ReloadMcpConfig { config }) + .await + .map_err(|err| { + io::Error::other(format!( + "failed to queue MCP config reload for thread {thread_id}: {err}" + )) + })?; } Ok(()) } -pub(crate) async fn queue_best_effort_refresh( - thread_manager: &Arc, - config_manager: &ConfigManager, -) { +pub(crate) async fn invalidate_loaded_threads(thread_manager: &Arc) { for thread_id in thread_manager.list_thread_ids().await { let thread = match thread_manager.get_thread(thread_id).await { Ok(thread) => thread, @@ -42,31 +47,29 @@ pub(crate) async fn queue_best_effort_refresh( continue; } }; - let config = match build_refresh_config(thread.as_ref(), config_manager).await { - Ok(config) => config, - Err(err) => { - warn!("failed to build MCP refresh config for thread {thread_id}: {err}"); - continue; - } - }; - if let Err(err) = queue_refresh(thread_id, thread, config).await { + if let Err(err) = queue_invalidation(thread_id, thread).await { warn!("{err}"); } } } +async fn load_refresh_config( + thread: &CodexThread, + config_manager: &ConfigManager, +) -> io::Result { + let thread_config = thread.config().await; + config_manager + .load_latest_config_for_thread(thread_config.as_ref()) + .await +} + async fn build_refresh_config( thread: &CodexThread, config_manager: &ConfigManager, ) -> io::Result { - let thread_config = thread.config().await; - let config = config_manager - .load_latest_config_for_thread(thread_config.as_ref()) - .await?; - let mcp_config = thread.runtime_mcp_config(&config).await; - let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config); + let config = load_refresh_config(thread, config_manager).await?; Ok(McpServerRefreshConfig { - mcp_servers: serde_json::to_value(mcp_servers).map_err(io::Error::other)?, + mcp_servers: serde_json::to_value(config.mcp_servers.get()).map_err(io::Error::other)?, mcp_oauth_credentials_store_mode: serde_json::to_value( config.mcp_oauth_credentials_store_mode, ) @@ -76,13 +79,9 @@ async fn build_refresh_config( }) } -async fn queue_refresh( - thread_id: ThreadId, - thread: Arc, - config: McpServerRefreshConfig, -) -> io::Result<()> { +async fn queue_invalidation(thread_id: ThreadId, thread: Arc) -> io::Result<()> { thread - .submit(Op::RefreshMcpServers { config }) + .submit(Op::RefreshMcpServers) .await .map(|_| ()) .map_err(|err| { @@ -127,35 +126,50 @@ mod tests { #[tokio::test] async fn strict_refresh_reports_thread_planning_failures() -> anyhow::Result<()> { - let (_temp_dir, thread_manager, config_manager, _loader) = refresh_test_state().await?; - - let err = queue_strict_refresh(&thread_manager, &config_manager) - .await - .expect_err("strict refresh should fail"); - - assert_eq!(err.to_string(), "failed to load refresh config"); - Ok(()) - } - - #[tokio::test] - async fn best_effort_refresh_attempts_every_loaded_thread() -> anyhow::Result<()> { - let (_temp_dir, thread_manager, config_manager, loader) = refresh_test_state().await?; - - queue_best_effort_refresh(&thread_manager, &config_manager).await; - - assert_eq!(loader.good_loads.load(Ordering::Relaxed), 1); - assert_eq!(loader.bad_loads.load(Ordering::Relaxed), 1); - Ok(()) - } - - #[tokio::test] - async fn refresh_config_uses_latest_auth_keyring_backend() -> anyhow::Result<()> { let (temp_dir, thread_manager, config_manager, _loader) = refresh_test_state().await?; std::fs::write( temp_dir.path().join(codex_config::CONFIG_TOML_FILE), "[features]\nsecret_auth_storage = true\n", )?; + let err = reload_mcp_config(&thread_manager, &config_manager) + .await + .expect_err("strict refresh should fail"); + + assert_eq!(err.to_string(), "failed to load refresh config"); + for thread_id in thread_manager.list_thread_ids().await { + assert_eq!( + thread_manager + .get_thread(thread_id) + .await? + .config() + .await + .auth_keyring_backend_kind(), + AuthKeyringBackendKind::Direct + ); + } + Ok(()) + } + + #[tokio::test] + async fn invalidation_does_not_reload_thread_config() -> anyhow::Result<()> { + let (_temp_dir, thread_manager, _config_manager, loader) = refresh_test_state().await?; + + invalidate_loaded_threads(&thread_manager).await; + + assert_eq!(loader.good_loads.load(Ordering::Relaxed), 0); + assert_eq!(loader.bad_loads.load(Ordering::Relaxed), 0); + Ok(()) + } + + #[tokio::test] + async fn mcp_config_reload_only_applies_mcp_inputs() -> anyhow::Result<()> { + let (temp_dir, thread_manager, config_manager, _loader) = refresh_test_state().await?; + std::fs::write( + temp_dir.path().join(codex_config::CONFIG_TOML_FILE), + "model = \"unrelated-model-change\"\n[features]\nsecret_auth_storage = true\n", + )?; + let mut good_thread = None; for thread_id in thread_manager.list_thread_ids().await { let thread = thread_manager.get_thread(thread_id).await?; @@ -166,17 +180,15 @@ mod tests { } } let thread = good_thread.expect("good test thread should exist"); + let original_model = thread.config().await.model.clone(); let refresh_config = build_refresh_config(thread.as_ref(), &config_manager).await?; - let backend = serde_json::from_value::( + let keyring_backend_kind = serde_json::from_value::( refresh_config.auth_keyring_backend_kind, )?; - assert_eq!( - thread.config().await.auth_keyring_backend_kind(), - AuthKeyringBackendKind::Direct - ); - assert_eq!(backend, AuthKeyringBackendKind::Secrets); + assert_eq!(keyring_backend_kind, AuthKeyringBackendKind::Secrets); + assert_eq!(thread.config().await.model, original_model); Ok(()) } @@ -253,7 +265,7 @@ enabled = false "#, )?; - queue_strict_refresh(&thread_manager, &config_manager).await?; + reload_mcp_config(&thread_manager, &config_manager).await?; assert_eq!( thread.config().await.mcp_servers.get(), diff --git a/codex-rs/app-server/src/request_processors/account_processor.rs b/codex-rs/app-server/src/request_processors/account_processor.rs index e22cfc4c10..44f519522b 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -258,7 +258,12 @@ impl AccountRequestProcessor { if thread_manager.list_thread_ids().await.is_empty() { return; } - crate::mcp_refresh::queue_best_effort_refresh(&thread_manager, &config_manager).await; + if let Err(err) = + crate::mcp_refresh::reload_mcp_config(&thread_manager, &config_manager).await + { + warn!(%err, "failed to reload MCP configuration after account or plugin change"); + crate::mcp_refresh::invalidate_loaded_threads(&thread_manager).await; + } }); } diff --git a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs index 927abca9c2..3756466c33 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs @@ -6,12 +6,12 @@ use codex_connectors::connector_tool_is_synthetic; use codex_connectors::installed_connector_runtime; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; -use codex_mcp::McpConnectionSet; +use codex_mcp::McpRuntime; +use codex_mcp::McpRuntimeInput; use codex_mcp::ToolInfo; use codex_mcp::effective_mcp_servers; use codex_mcp::host_owned_codex_apps_enabled; use codex_mcp::tool_is_model_visible; -use codex_mcp::tool_plugin_provenance; use codex_protocol::models::PermissionProfile; const CONNECTOR_RUNTIME_REFRESH_TIMEOUT: Duration = Duration::from_secs(30); @@ -52,7 +52,10 @@ impl AppsRequestProcessor { let runtime_enabled = apps_enabled && workspace_enabled; let mcp_manager = self.thread_manager.mcp_manager(); - let mcp_config = mcp_manager.runtime_config(&config).await; + let mut mcp_config = mcp_manager.runtime_config(&config).await; + // Installed-app discovery has no active turn or reviewer. + mcp_config.permission_profile = PermissionProfile::default(); + let mcp_config = Arc::new(mcp_config); let mut mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); mcp_servers.retain(|name, _| name == CODEX_APPS_MCP_SERVER_NAME); let cache_key = connector_runtime_context_key(auth.as_ref()); @@ -78,34 +81,31 @@ impl AppsRequestProcessor { let codex_apps_auth_manager = host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()) .then(|| Arc::clone(&self.auth_manager)); - let connection_manager = McpConnectionSet::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), - &config.permissions.approval_policy, - APPS_INSTALLED_SUBMIT_ID.to_string(), - /*tx_event*/ None, - cancellation_token.clone(), - PermissionProfile::default(), + let runtime = McpRuntime::new(McpRuntimeInput { + config: Arc::clone(&mcp_config), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers, + submit_id: APPS_INSTALLED_SUBMIT_ID.to_string(), + tx_event: None, + startup_cancellation_token: cancellation_token.clone(), runtime_context, - mcp_config.codex_home.clone(), - mcp_manager.codex_apps_tools_cache(), - mcp_manager.tool_catalog_cache(), - cache_key.clone(), - mcp_config.prefix_mcp_tool_names, - mcp_config.client_elicitation_capability.clone(), - /*supports_openai_form_elicitation*/ false, - tool_plugin_provenance(&mcp_config), - auth.as_ref(), + codex_apps_tools_cache: mcp_manager.codex_apps_tools_cache(), + tool_catalog_cache: mcp_manager.tool_catalog_cache(), + codex_apps_tools_cache_key: cache_key.clone(), + supports_openai_form_elicitation: false, + auth: auth.clone(), codex_apps_auth_manager, - /*elicitation_reviewer*/ None, - /*elicitation_lifecycle*/ None, - codex_mcp::ElicitationRequestRouter::default(), - ) + elicitation_reviewer: None, + elicitation_lifecycle: None, + }) .await; - let result = if connection_manager - .wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, startup_timeout) + let result = if runtime + .latest_wait_for_server_ready( + CODEX_APPS_MCP_SERVER_NAME, + startup_timeout, + ) .await { mcp_manager @@ -122,7 +122,7 @@ impl AppsRequestProcessor { )) }; cancellation_token.cancel(); - connection_manager.shutdown().await; + runtime.shutdown().await; result } .await; diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index 9cef41e3dd..4237d29260 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -276,8 +276,8 @@ impl ConfigRequestProcessor { } async fn reload_user_config(&self) { - let next_config = match self.load_latest_config(/*fallback_cwd*/ None).await { - Ok(config) => config, + match self.load_latest_config(/*fallback_cwd*/ None).await { + Ok(_) => {} Err(err) => { tracing::warn!( "failed to rebuild user config for runtime refresh: {}", @@ -291,7 +291,19 @@ impl ConfigRequestProcessor { let Ok(thread) = self.thread_manager.get_thread(thread_id).await else { continue; }; - thread.refresh_runtime_config(next_config.clone()).await; + let current_config = thread.config().await; + let next_config = match self + .config_manager + .load_latest_config_for_thread(current_config.as_ref()) + .await + { + Ok(config) => config, + Err(err) => { + tracing::warn!(%thread_id, %err, "failed to reload thread configuration"); + continue; + } + }; + thread.refresh_runtime_config(next_config).await; } } diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 1cfbc65672..7563a80ebd 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -78,7 +78,7 @@ impl McpRequestProcessor { &self, _params: Option<()>, ) -> Result { - crate::mcp_refresh::queue_strict_refresh(&self.thread_manager, &self.config_manager) + crate::mcp_refresh::reload_mcp_config(&self.thread_manager, &self.config_manager) .await .map_err(|err| internal_error(format!("failed to refresh MCP servers: {err}")))?; Ok(McpServerRefreshResponse {}) @@ -125,8 +125,9 @@ impl McpRequestProcessor { let (mcp_config, runtime_context) = match thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; - let runtime = thread.current_mcp_runtime().await; - (runtime.config().clone(), runtime.runtime_context().clone()) + let (config, runtime_context) = + thread.current_mcp_config_and_runtime_context().await; + ((*config).clone(), runtime_context) } None => { let config = self.load_latest_config(/*fallback_cwd*/ None).await?; @@ -249,8 +250,8 @@ impl McpRequestProcessor { let (mcp_config, runtime_context) = match thread { Some(thread) => { let mcp_config = thread.runtime_mcp_config(&config).await; - let runtime = thread.current_mcp_runtime().await; - (mcp_config, runtime.runtime_context().clone()) + let (_, runtime_context) = thread.current_mcp_config_and_runtime_context().await; + (mcp_config, runtime_context) } None => { let mcp_config = mcp_manager.runtime_config(&config).await; diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index 20c92cb604..d1f1577a70 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -148,6 +148,7 @@ impl fmt::Debug for McpBinding { pub struct PreparedMcpCall { _connections: Arc, client: Arc, + config: Arc, catalog_revision: u64, catalog_revision_source: Arc>, tool_info: ToolInfo, @@ -165,6 +166,7 @@ impl PreparedMcpCall { pub(crate) fn new( connections: Arc, client: Arc, + config: Arc, catalog_revision: u64, catalog_revision_source: Arc>, tool_info: ToolInfo, @@ -176,6 +178,7 @@ impl PreparedMcpCall { Self { _connections: connections, client, + config, catalog_revision, catalog_revision_source, tool_info, @@ -190,6 +193,11 @@ impl PreparedMcpCall { &self.tool_info } + /// Returns the configuration and approval authority captured with this client. + pub fn config(&self) -> &McpConfig { + &self.config + } + pub fn server_name(&self) -> &str { &self.server_name } diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs index 5447dfeab3..c3be31999d 100644 --- a/codex-rs/codex-mcp/src/binding_tests.rs +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -6,6 +6,7 @@ use std::sync::atomic::Ordering; use codex_config::AppToolApproval; use codex_config::Constrained; +use codex_config::types::ApprovalsReviewer; use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; @@ -97,15 +98,20 @@ async fn test_step( SERVER_NAME.to_string(), Arc::clone(&managed_client), )]))); - let connections = Arc::new(McpConnectionSet::new_uninitialized_with_permission_profile( - &Constrained::allow_any(AskForApproval::OnRequest), - &PermissionProfile::default(), - /*prefix_mcp_tool_names*/ true, - )); + let connections = Arc::new(McpConnectionSet::empty(/*prefix_mcp_tool_names*/ true)); let tool_catalog_revision = Arc::new(tokio::sync::RwLock::new(0)); + let mut config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); + if label == "old" { + config.approval_policy = Constrained::allow_any(AskForApproval::Never); + config.permission_profile = PermissionProfile::Disabled; + } else { + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + } + let config = Arc::new(config); let prepared = PreparedMcpCall::new( Arc::clone(&connections), managed_client, + Arc::clone(&config), /*catalog_revision*/ 0, Arc::clone(&tool_catalog_revision), tool.clone(), @@ -128,7 +134,7 @@ async fn test_step( step: Arc::new(McpBinding::new( connections, clients, - Arc::new(crate::mcp::tests::test_mcp_config(std::env::temp_dir())), + config, /*plugins_available*/ false, vec![tool], calls, @@ -206,6 +212,25 @@ async fn prepared_call_keeps_captured_connection_and_authority_after_refresh() - ); assert!(Arc::ptr_eq(&old_call.client.client, &old.client)); assert!(!Arc::ptr_eq(&old.client, &new.client)); + assert_eq!( + ( + old_call.config().approval_policy.value(), + &old_call.config().permission_profile, + old_call.config().approvals_reviewer, + ), + ( + AskForApproval::Never, + &PermissionProfile::Disabled, + ApprovalsReviewer::User, + ) + ); + assert_eq!( + ( + new_call.config().approval_policy.value(), + new_call.config().approvals_reviewer, + ), + (AskForApproval::OnRequest, ApprovalsReviewer::AutoReview) + ); drop(old.step); assert!( diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 215e28bf4e..520de0c9cd 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -1,10 +1,9 @@ //! Aggregates MCP server connections for Codex. //! -//! [`McpConnectionSet`] owns the set of running async RMCP clients keyed by -//! MCP server name. It coordinates startup status events, keeps server origin -//! metadata, aggregates tools/resources/templates across servers, routes tool -//! calls to the right client, and exposes the public manager API used by -//! `codex-core`. +//! [`McpConnectionSet`] is the private connection set behind +//! [`crate::McpRuntime`] and [`crate::McpBinding`]. It coordinates startup status +//! events, keeps server metadata, and aggregates tools and resources across +//! running RMCP clients. #[path = "connection_manager/required.rs"] mod required; @@ -27,6 +26,7 @@ use crate::rmcp_client::AsyncManagedClient; use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::StartupOutcomeError; +use crate::runtime::McpPublicationGate; use crate::runtime::McpRuntimeContext; use crate::server::EffectiveMcpServer; use crate::server::McpServerMetadata; @@ -58,19 +58,14 @@ use codex_protocol::protocol::McpStartupFailure; use codex_protocol::protocol::McpStartupFailureReason; use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpStartupUpdateEvent; -use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::McpAuthState; use codex_rmcp_client::McpLoginRequirement; use codex_rmcp_client::determine_streamable_http_auth_status_from_credentials; use rmcp::model::ElicitationCapability; -use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; -use rmcp::model::RequestId; -use rmcp::model::Resource; -use rmcp::model::ResourceTemplate; use serde_json::Value as JsonValue; use tokio::sync::Mutex; use tokio::sync::RwLock; @@ -107,7 +102,7 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { } /// A thin wrapper around a set of running [`RmcpClient`] instances. -pub struct McpConnectionSet { +pub(crate) struct McpConnectionSet { clients: HashMap, server_metadata: HashMap, required_servers: Vec, @@ -116,7 +111,6 @@ pub struct McpConnectionSet { codex_apps_refresh_lock: Mutex<()>, tool_plugin_provenance: Arc, prefix_mcp_tool_names: bool, - elicitation_requests: ElicitationRequestManager, startup_cancellation_token: CancellationToken, } @@ -147,15 +141,16 @@ impl McpConnectionSet { elicitation_reviewer: Option, elicitation_lifecycle: Option, elicitation_router: ElicitationRequestRouter, + publication_gate: McpPublicationGate, ) -> Self { + let mut clients = HashMap::new(); + let mut server_metadata = HashMap::new(); let mut required_servers = mcp_servers .iter() .filter(|(_, server)| server.enabled() && server.required()) - .map(|(name, _)| name.clone()) + .map(|(server_name, _)| server_name.clone()) .collect::>(); required_servers.sort(); - let mut clients = HashMap::new(); - let mut server_metadata = HashMap::new(); let mut join_set = JoinSet::new(); let elicitation_requests = ElicitationRequestManager::new( approval_policy.value(), @@ -165,7 +160,8 @@ impl McpConnectionSet { elicitation_router, ); let tool_plugin_provenance = Arc::new(tool_plugin_provenance); - let startup_submit_id = submit_id.clone(); + let startup_submit_id = submit_id; + let startup_publication_gate = publication_gate.clone(); let static_chatgpt_auth_provider = auth .filter(|auth| auth.uses_codex_backend()) .map(codex_model_provider::auth_provider_from_auth); @@ -181,17 +177,6 @@ impl McpConnectionSet { { server_metadata.insert(server_name.clone(), McpServerMetadata::from(&server)); let cancel_token = startup_cancellation_token.child_token(); - if let Some(tx_event) = tx_event.as_ref() { - let _ = emit_update( - startup_submit_id.as_str(), - tx_event, - McpStartupUpdateEvent { - server: server_name.clone(), - status: McpStartupStatus::Starting, - }, - ) - .await; - } let configured_config = server.configured_config().cloned(); let resolved_environment = configured_config.as_ref().map_or_else( || Ok(None), @@ -273,7 +258,22 @@ impl McpConnectionSet { clients.insert(server_name.clone(), async_managed_client.clone()); let tx_event = tx_event.clone(); let submit_id = startup_submit_id.clone(); + let publication_gate = publication_gate.clone(); join_set.spawn(async move { + if !publication_gate.wait().await { + return (server_name, Err(StartupOutcomeError::Cancelled)); + } + if let Some(tx_event) = tx_event.as_ref() { + let _ = emit_update( + submit_id.as_str(), + tx_event, + McpStartupUpdateEvent { + server: server_name.clone(), + status: McpStartupStatus::Starting, + }, + ) + .await; + } let mut outcome = async_managed_client.client().await; if cancel_token.is_cancelled() { outcome = Err(StartupOutcomeError::Cancelled); @@ -362,11 +362,13 @@ impl McpConnectionSet { codex_apps_refresh_lock: Mutex::new(()), tool_plugin_provenance, prefix_mcp_tool_names, - elicitation_requests: elicitation_requests.clone(), - startup_cancellation_token: startup_cancellation_token.clone(), + startup_cancellation_token, }; tokio::spawn(async move { let outcomes = join_set.join_all().await; + if !startup_publication_gate.wait().await { + return; + } if let Some(tx_event) = tx_event { let mut summary = McpStartupCompleteEvent::default(); for (server_name, outcome) in outcomes { @@ -392,11 +394,7 @@ impl McpConnectionSet { manager } - pub fn new_uninitialized_with_permission_profile( - approval_policy: &Constrained, - permission_profile: &PermissionProfile, - prefix_mcp_tool_names: bool, - ) -> Self { + pub fn empty(prefix_mcp_tool_names: bool) -> Self { Self { clients: HashMap::new(), server_metadata: HashMap::new(), @@ -406,29 +404,18 @@ impl McpConnectionSet { codex_apps_refresh_lock: Mutex::new(()), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), prefix_mcp_tool_names, - elicitation_requests: ElicitationRequestManager::new( - approval_policy.value(), - permission_profile.clone(), - /*reviewer*/ None, - /*lifecycle*/ None, - ElicitationRequestRouter::default(), - ), startup_cancellation_token: CancellationToken::new(), } } - pub fn empty(prefix_mcp_tool_names: bool) -> Self { - Self::new_uninitialized_with_permission_profile( - &Constrained::allow_any(AskForApproval::OnRequest), - &PermissionProfile::default(), - prefix_mcp_tool_names, - ) - } - pub fn has_servers(&self) -> bool { !self.clients.is_empty() } + pub(crate) fn cancel_startup(&self) { + self.startup_cancellation_token.cancel(); + } + pub(crate) fn contains_server(&self, server_name: &str) -> bool { self.clients.contains_key(server_name) } @@ -448,25 +435,6 @@ impl McpConnectionSet { } } - pub fn server_origin(&self, server_name: &str) -> Option<&str> { - self.server_metadata - .get(server_name) - .and_then(|metadata| metadata.origin.as_ref()) - .map(super::server::McpServerOrigin::as_str) - } - - pub fn server_environment_id(&self, server_name: &str) -> Option<&str> { - self.server_metadata - .get(server_name) - .map(|metadata| metadata.environment_id.as_str()) - } - - pub fn server_pollutes_memory(&self, server_name: &str) -> bool { - self.server_metadata - .get(server_name) - .is_none_or(|metadata| metadata.pollutes_memory) - } - pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option<&str> { self.tool_plugin_provenance .plugin_id_for_mcp_server_name(server_name) @@ -477,56 +445,6 @@ impl McpConnectionSet { .is_selected_plugin_mcp_server(server_name) } - pub fn tool_approval_mode( - &self, - server_name: &str, - tool_name: &str, - ) -> codex_config::AppToolApproval { - self.server_metadata - .get(server_name) - .map(|metadata| metadata.tool_approval_mode(tool_name)) - .unwrap_or_default() - } - - pub fn is_host_owned_codex_apps_server(&self, server_name: &str) -> bool { - server_name == CODEX_APPS_MCP_SERVER_NAME && self.server_metadata.contains_key(server_name) - } - - pub fn set_approval_policy(&self, approval_policy: &Constrained) { - if let Ok(mut policy) = self.elicitation_requests.approval_policy.lock() { - *policy = approval_policy.value(); - } - } - - pub fn set_permission_profile(&self, permission_profile: PermissionProfile) { - if let Ok(mut profile) = self.elicitation_requests.permission_profile.lock() { - *profile = permission_profile; - } - } - - pub fn elicitations_auto_deny(&self) -> bool { - self.elicitation_requests.auto_deny() - } - - pub fn set_elicitations_auto_deny(&self, auto_deny: bool) { - self.elicitation_requests.set_auto_deny(auto_deny); - } - - pub fn elicitation_router(&self) -> ElicitationRequestRouter { - self.elicitation_requests.router() - } - - pub async fn resolve_elicitation( - &self, - server_name: String, - id: RequestId, - response: ElicitationResponse, - ) -> Result<()> { - self.elicitation_requests - .resolve(server_name, id, response) - .await - } - pub async fn wait_for_server_ready(&self, server_name: &str, timeout: Duration) -> bool { let Some(async_managed_client) = self.clients.get(server_name) else { return false; @@ -538,24 +456,20 @@ impl McpConnectionSet { } } - /// Returns resources from servers selected by `include_server`. Each key - /// is the server name and the value is a vector of resources. pub async fn list_all_resources( &self, include_server: impl Fn(&str) -> bool, - ) -> HashMap> { + ) -> HashMap> { self.ready_clients_matching(&include_server) .await .list_all_resources(|_| true) .await } - /// Returns resource templates from servers selected by `include_server`. - /// Each key is the server name and the value is a vector of templates. pub async fn list_all_resource_templates( &self, include_server: impl Fn(&str) -> bool, - ) -> HashMap> { + ) -> HashMap> { self.ready_clients_matching(&include_server) .await .list_all_resource_templates(|_| true) @@ -617,16 +531,6 @@ impl McpConnectionSet { }) } - pub async fn server_supports_sandbox_state_meta_capability( - &self, - server: &str, - ) -> Result { - Ok(self - .client_by_name(server) - .await? - .server_supports_sandbox_state_meta_capability) - } - /// List resources from the specified server. pub async fn list_resources( &self, @@ -643,22 +547,6 @@ impl McpConnectionSet { .with_context(|| format!("resources/list failed for `{server}`")) } - /// List resource templates from the specified server. - pub async fn list_resource_templates( - &self, - server: &str, - params: Option, - ) -> Result { - let managed = self.client_by_name(server).await?; - let client = managed.client.clone(); - let timeout = managed.tool_timeout; - - client - .list_resource_templates(params, timeout) - .await - .with_context(|| format!("resources/templates/list failed for `{server}`")) - } - /// Read a resource from the specified server. pub async fn read_resource( &self, @@ -713,15 +601,11 @@ impl McpConnectionSet { #[cfg(test)] fn new_uninitialized( - approval_policy: &Constrained, - permission_profile: &Constrained, + _approval_policy: &Constrained, + _permission_profile: &Constrained, prefix_mcp_tool_names: bool, ) -> Self { - Self::new_uninitialized_with_permission_profile( - approval_policy, - permission_profile.get(), - prefix_mcp_tool_names, - ) + Self::empty(prefix_mcp_tool_names) } } diff --git a/codex-rs/codex-mcp/src/connection_manager/required.rs b/codex-rs/codex-mcp/src/connection_manager/required.rs index fd4ccda3fa..82b9a7a36d 100644 --- a/codex-rs/codex-mcp/src/connection_manager/required.rs +++ b/codex-rs/codex-mcp/src/connection_manager/required.rs @@ -10,9 +10,9 @@ use crate::rmcp_client::StartupOutcomeError; impl McpConnectionSet { /// Waits for every required server and reports their startup failures together. /// - /// Callers must make the manager reachable to request handlers before awaiting this method, - /// because server initialization may require client elicitation. - pub async fn validate_required_servers(&self) -> Result<()> { + /// The manager must already be reachable through [`crate::McpRuntime`] so + /// startup-time elicitation can resolve while validation waits. + pub(crate) async fn validate_required_servers(&self) -> Result<()> { let failures = async { let mut failures = Vec::new(); for server_name in &self.required_servers { @@ -55,7 +55,7 @@ impl McpConnectionSet { } } -pub(super) fn startup_outcome_error_message(error: StartupOutcomeError) -> String { +fn startup_outcome_error_message(error: StartupOutcomeError) -> String { match error { StartupOutcomeError::Cancelled => "MCP startup cancelled".to_string(), StartupOutcomeError::Failed { error, .. } => error, diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index 98b193fba8..c9e7a25498 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::atomic::Ordering; use std::time::Instant; use anyhow::Context; @@ -35,18 +36,25 @@ impl McpConnectionSet { for (server_name, managed_client) in &self.clients { managed_client.reconnect_failed_startup().await; let has_cached_tools = managed_client.has_cached_tools(); - let startup_complete = managed_client - .startup_complete - .load(std::sync::atomic::Ordering::Acquire); - let Some(server_tools) = managed_client - .listed_tools() - .instrument(trace_span!( - "list_tools_for_server", - server_name = %server_name, - has_cached_tools, - startup_complete - )) - .await + let startup_complete = managed_client.startup_complete.load(Ordering::Acquire); + let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { + self.codex_apps_tools_override.read().await.clone() + } else { + None + }; + let Some(server_tools) = async { + match catalog_override { + Some(tools) => Some(managed_client.prepare_tools(tools)), + None => managed_client.listed_tools().await, + } + } + .instrument(trace_span!( + "list_tools_for_server", + server_name = %server_name, + has_cached_tools, + startup_complete + )) + .await else { unavailable_server_count += 1; trace!( @@ -74,23 +82,11 @@ impl McpConnectionSet { tools } - /// Returns one tool from the current live connection. - pub async fn tool_info(&self, server: &str, tool: &str) -> Option { - let client = self.clients.get(server)?; - let managed_client = client.client().await.ok()?; - let tool = client - .prepare_tools(managed_client.listed_tools()) - .into_iter() - .find(|tool_info| tool_info.tool.name == tool)?; - Some(self.with_server_metadata(tool)) - } - #[expect( clippy::await_holding_invalid_type, reason = "catalog capture must remain serialized with catalog replacement" )] - /// Captures the ready clients, their exact tools, and the supplied runtime metadata. - pub async fn capture_binding_with_metadata( + pub(crate) async fn capture_binding_with_metadata( self: &Arc, config: Arc, plugins_available: bool, @@ -127,10 +123,14 @@ impl McpConnectionSet { let mut tools = Vec::with_capacity(listed_tools.len()); let mut calls = std::collections::HashMap::with_capacity(listed_tools.len()); for tool_info in listed_tools { + if !crate::tool_is_model_visible(&tool_info) { + continue; + } let Some(client) = clients.client(&tool_info.server_name) else { continue; }; - let Some(call) = self.prepare_call(&tool_info, client, *revision) else { + let Some(call) = self.prepare_call(&tool_info, client, Arc::clone(&config), *revision) + else { trace!( server_name = %tool_info.server_name, tool_name = %tool_info.tool.name, @@ -161,12 +161,14 @@ impl McpConnectionSet { self: &Arc, tool_info: &ToolInfo, client: Arc, + config: Arc, tool_catalog_revision: u64, ) -> Option { let server_name = &tool_info.server_name; Some(PreparedMcpCall::new( Arc::clone(self), client, + config, tool_catalog_revision, Arc::clone(&self.tool_catalog_revision), tool_info.clone(), @@ -177,16 +179,12 @@ impl McpConnectionSet { )) } - /// Force-refresh codex apps tools by bypassing the in-process cache. - /// - /// On success, the refreshed tools replace shared cache contents when the - /// cache is enabled and the latest filtered tools are returned directly to - /// the caller. On failure, existing shared cache contents remain unchanged. + /// Force-refresh Codex Apps tools and publish one new exact catalog revision. #[expect( clippy::await_holding_invalid_type, reason = "catalog publication must remain serialized with captured tool calls" )] - pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { + pub(crate) async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { let _refresh = self.codex_apps_refresh_lock.lock().await; let refresh_start = Instant::now(); let managed_client = self diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 9b4bbc8033..15545d266e 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,4 +1,3 @@ -use super::required::startup_outcome_error_message; use super::*; use crate::McpBinding; use crate::elicitation::ElicitationLifecycle; @@ -34,6 +33,7 @@ use codex_protocol::ToolName; use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::GranularApprovalConfig; +use codex_rmcp_client::ElicitationResponse; use codex_rmcp_client::InProcessTransportFactory; use codex_rmcp_client::RmcpClient; use futures::FutureExt; @@ -508,7 +508,7 @@ async fn shared_elicitation_router_targets_the_exact_pending_request() { PermissionProfile::default(), /*reviewer*/ None, Some(lifecycle), - router, + router.clone(), ); let (tx_event, rx_event) = async_channel::bounded(2); let sender_a = manager_a.make_sender("server".to_string(), Some(tx_event.clone())); @@ -552,7 +552,7 @@ async fn shared_elicitation_router_targets_the_exact_pending_request() { content: Some(serde_json::json!({"runtime": "a"})), meta: None, }; - manager_b + router .resolve( "server".to_string(), NumberOrString::String(request_a_id.into()), @@ -565,7 +565,7 @@ async fn shared_elicitation_router_targets_the_exact_pending_request() { content: Some(serde_json::json!({"runtime": "b"})), meta: None, }; - manager_a + router .resolve( "server".to_string(), NumberOrString::String(request_b_id.into()), @@ -1962,42 +1962,6 @@ fn server_metadata_preserves_tool_approval_policy() { ); } -#[test] -fn host_owned_codex_apps_requires_server_metadata() { - let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); - let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let manager = McpConnectionSet::new_uninitialized( - &approval_policy, - &permission_profile, - /*prefix_mcp_tool_names*/ true, - ); - - assert!(!manager.is_host_owned_codex_apps_server(CODEX_APPS_MCP_SERVER_NAME)); -} - -#[test] -fn host_owned_codex_apps_matches_reserved_name_with_server_metadata() { - let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); - let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionSet::new_uninitialized( - &approval_policy, - &permission_profile, - /*prefix_mcp_tool_names*/ true, - ); - let server = EffectiveMcpServer::configured(crate::codex_apps_mcp_server_config( - "https://chatgpt.com", - /*apps_mcp_product_sku*/ None, - /*originator*/ None, - )); - manager.server_metadata.insert( - CODEX_APPS_MCP_SERVER_NAME.to_string(), - McpServerMetadata::from(&server), - ); - - assert!(manager.is_host_owned_codex_apps_server(CODEX_APPS_MCP_SERVER_NAME)); - assert!(!manager.is_host_owned_codex_apps_server("docs")); -} - #[tokio::test] async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -2087,6 +2051,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { /*elicitation_reviewer*/ None, /*elicitation_lifecycle*/ None, ElicitationRequestRouter::default(), + crate::runtime::McpPublicationGate::already_published(), ) .await; @@ -2108,8 +2073,11 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { Ok(_) => panic!("local stdio MCP startup should fail"), Err(error) => error, }; + let StartupOutcomeError::Failed { error, .. } = error else { + panic!("local stdio MCP startup should fail rather than be cancelled"); + }; assert_eq!( - startup_outcome_error_message(error), + error, "local stdio MCP server `stdio` requires a local environment" ); cancel_token.cancel(); diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index de4ce1448c..e2d9ee1eb2 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -7,7 +7,7 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; @@ -85,12 +85,21 @@ struct ActiveElicitation { /// generated by Codex rather than copied from the MCP connection, so separate runtimes may reuse /// the same server request ID without colliding. #[derive(Clone, Default)] -pub struct ElicitationRequestRouter { +pub(crate) struct ElicitationRequestRouter { requests: Arc>, + auto_deny: Arc, } impl ElicitationRequestRouter { - async fn resolve( + pub(crate) fn auto_deny(&self) -> bool { + self.auto_deny.load(Ordering::Relaxed) + } + + pub(crate) fn set_auto_deny(&self, auto_deny: bool) { + self.auto_deny.store(auto_deny, Ordering::Relaxed); + } + + pub(crate) async fn resolve( &self, server_name: String, id: RequestId, @@ -109,9 +118,8 @@ impl ElicitationRequestRouter { #[derive(Clone)] pub(crate) struct ElicitationRequestManager { router: ElicitationRequestRouter, - pub(crate) approval_policy: Arc>, - pub(crate) permission_profile: Arc>, - auto_deny: Arc>, + approval_policy: AskForApproval, + permission_profile: PermissionProfile, reviewer: Option, lifecycle: Option, } @@ -126,66 +134,33 @@ impl ElicitationRequestManager { ) -> Self { Self { router, - approval_policy: Arc::new(StdMutex::new(approval_policy)), - permission_profile: Arc::new(StdMutex::new(permission_profile)), - auto_deny: Arc::new(StdMutex::new(false)), + approval_policy, + permission_profile, reviewer, lifecycle, } } - pub(crate) fn auto_deny(&self) -> bool { - self.auto_deny - .lock() - .map(|auto_deny| *auto_deny) - .unwrap_or(false) - } - - pub(crate) fn set_auto_deny(&self, auto_deny: bool) { - if let Ok(mut current) = self.auto_deny.lock() { - *current = auto_deny; - } - } - - pub(crate) async fn resolve( - &self, - server_name: String, - id: RequestId, - response: ElicitationResponse, - ) -> Result<()> { - self.router.resolve(server_name, id, response).await - } - - pub(crate) fn router(&self) -> ElicitationRequestRouter { - self.router.clone() - } - pub(crate) fn make_sender( &self, server_name: String, tx_event: Option>, ) -> SendElicitation { let router = self.router.clone(); - let approval_policy = self.approval_policy.clone(); + let approval_policy = self.approval_policy; let permission_profile = self.permission_profile.clone(); - let auto_deny = self.auto_deny.clone(); let reviewer = self.reviewer.clone(); let lifecycle = self.lifecycle.clone(); Box::new(move |id, elicitation| { let router = router.clone(); let tx_event = tx_event.clone(); let server_name = server_name.clone(); - let approval_policy = approval_policy.clone(); + let approval_policy = approval_policy; let permission_profile = permission_profile.clone(); - let auto_deny = auto_deny.clone(); let reviewer = reviewer.clone(); let lifecycle = lifecycle.clone(); async move { - let auto_deny = auto_deny - .lock() - .map(|auto_deny| *auto_deny) - .unwrap_or(false); - if auto_deny { + if router.auto_deny() { return Ok(ElicitationResponse { action: ElicitationAction::Decline, content: None, @@ -193,14 +168,6 @@ impl ElicitationRequestManager { }); } - let approval_policy = approval_policy - .lock() - .map(|policy| *policy) - .unwrap_or(AskForApproval::Never); - let permission_profile = permission_profile - .lock() - .map(|profile| profile.clone()) - .unwrap_or_default(); if mcp_permission_prompt_is_auto_approved( approval_policy, &permission_profile, diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index c4d7e2233a..4cc075d35d 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -1,9 +1,7 @@ pub use binding::McpBinding; pub use binding::PreparedMcpCall; -pub use connection_manager::McpConnectionSet; pub use connection_manager::tool_is_model_visible; pub use elicitation::ElicitationLifecycle; -pub use elicitation::ElicitationRequestRouter; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; pub use elicitation::ElicitationReviewerHandle; @@ -14,13 +12,11 @@ pub use resource_client::McpResourceReadResult; pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; pub use runtime::McpRuntime; pub use runtime::McpRuntimeContext; +pub use runtime::McpRuntimeInput; pub use runtime::SandboxState; pub use tool_catalog_cache::McpToolCatalogCache; pub use tools::ToolInfo; -/// Backward-compatible name for the MCP connection set. -pub type McpConnectionManager = McpConnectionSet; - /// Backward-compatible name for the shared Codex Apps tools runtime. pub type CodexAppsToolsCache = codex_connectors::ConnectorRuntimeManager; /// Backward-compatible name for the Codex Apps runtime context key. diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 113dfa5b66..92321ac327 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -19,11 +19,13 @@ use std::env; use std::path::PathBuf; use std::time::Duration; +use codex_config::ConfigLayerStack; use codex_config::Constrained; use codex_config::McpServerAuth; use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; use codex_config::types::AppToolApproval; +use codex_config::types::ApprovalsReviewer; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_connectors::ConnectorRuntimeManager; @@ -38,6 +40,7 @@ use codex_protocol::mcp::Tool; use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::McpAuthStatus; +use codex_utils_path_uri::PathUri; use rmcp::model::ElicitationCapability; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; @@ -105,13 +108,10 @@ pub struct McpPermissionPromptAutoApproveContext { /// MCP runtime settings derived from `codex_core::config::Config`. /// -/// This struct should contain only long-lived configuration values that the -/// `codex-mcp` crate needs to construct server transports, enforce MCP -/// approval/sandbox policy, locate OAuth state, and merge plugin-provided MCP -/// servers. Request-scoped or auth-scoped state should not be stored here; -/// thread those values explicitly into runtime entry points such as -/// [`effective_mcp_servers`] and snapshot collection helpers so config objects -/// do not go stale when auth changes. +/// Each published runtime and prepared call owns one immutable copy of these +/// settings, so its connection, approval policy, and sandbox authority cannot +/// change independently. Auth remains separate and is supplied explicitly to +/// runtime entry points such as [`effective_mcp_servers`]. #[derive(Debug, Clone)] pub struct McpConfig { /// Base URL for ChatGPT-hosted app MCP servers, copied from the root config. @@ -132,6 +132,14 @@ pub struct McpConfig { pub skill_mcp_dependency_install_enabled: bool, /// Approval policy used for MCP tool calls and MCP elicitation requests. pub approval_policy: Constrained, + /// Permission profile captured with the connections and approval policy. + pub permission_profile: PermissionProfile, + /// Configuration layers used to evaluate Apps tool policy and reviewer selection. + pub config_layer_stack: ConfigLayerStack, + /// Default reviewer used when an Apps tool has no reviewer override. + pub approvals_reviewer: ApprovalsReviewer, + /// Working directories for the exact environment handles used by this runtime. + pub environment_cwds: HashMap, /// Optional path to `codex-linux-sandbox` for sandboxed MCP tool execution. pub codex_linux_sandbox_exe: Option, /// Whether to use legacy Landlock behavior in the MCP sandbox state. @@ -334,6 +342,7 @@ pub async fn read_mcp_resource( /*elicitation_reviewer*/ None, /*elicitation_lifecycle*/ None, crate::elicitation::ElicitationRequestRouter::default(), + crate::runtime::McpPublicationGate::already_published(), ) .await; @@ -388,7 +397,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( let server_names = mcp_servers.keys().cloned().collect(); let cancel_token = CancellationToken::new(); - let mcp_connection_manager = McpConnectionSet::new( + let mcp_connection_set = McpConnectionSet::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, config.auth_keyring_backend_kind, @@ -411,11 +420,12 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( /*elicitation_reviewer*/ None, /*elicitation_lifecycle*/ None, crate::elicitation::ElicitationRequestRouter::default(), + crate::runtime::McpPublicationGate::already_published(), ) .await; let snapshot = collect_mcp_server_status_snapshot_from_manager( - &mcp_connection_manager, + &mcp_connection_set, auth_status_entries, server_names, detail, @@ -642,27 +652,27 @@ fn convert_mcp_resource_templates( } async fn collect_mcp_server_status_snapshot_from_manager( - mcp_connection_manager: &McpConnectionSet, + mcp_connection_set: &McpConnectionSet, auth_status_entries: HashMap, server_names: Vec, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let ((server_infos, tools), resources, resource_templates) = tokio::join!( async { - let server_infos = mcp_connection_manager.list_available_server_infos().await; - let tools = mcp_connection_manager.list_all_tools().await; + let server_infos = mcp_connection_set.list_available_server_infos().await; + let tools = mcp_connection_set.list_all_tools().await; (server_infos, tools) }, async { if detail.include_resources() { - mcp_connection_manager.list_all_resources(|_| true).await + mcp_connection_set.list_all_resources(|_| true).await } else { HashMap::new() } }, async { if detail.include_resources() { - mcp_connection_manager + mcp_connection_set .list_all_resource_templates(|_| true) .await } else { diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 6c48a9f75b..994c6ea96f 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -28,6 +28,10 @@ pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig { mcp_oauth_callback_url: None, skill_mcp_dependency_install_enabled: true, approval_policy: Constrained::allow_any(AskForApproval::OnRequest), + permission_profile: PermissionProfile::default(), + config_layer_stack: codex_config::ConfigLayerStack::default(), + approvals_reviewer: codex_config::types::ApprovalsReviewer::default(), + environment_cwds: HashMap::new(), codex_linux_sandbox_exe: None, use_legacy_landlock: false, apps_enabled: false, diff --git a/codex-rs/codex-mcp/src/resource_client.rs b/codex-rs/codex-mcp/src/resource_client.rs index c3145b2de1..eb4fa7bc34 100644 --- a/codex-rs/codex-mcp/src/resource_client.rs +++ b/codex-rs/codex-mcp/src/resource_client.rs @@ -61,14 +61,14 @@ impl McpResourceClient { /// Returns the identity of the connection set used by this client. pub fn cache_key(&self) -> McpResourceClientCacheKey { - McpResourceClientCacheKey(Arc::downgrade(&self.runtime.snapshot())) + McpResourceClientCacheKey(Arc::downgrade(&self.runtime.latest_connections())) } /// Returns whether this client can address the named server. /// /// This does not wait for server startup. pub async fn has_server(&self, server: &str) -> bool { - self.runtime.snapshot().contains_server(server) + self.runtime.latest_connections().contains_server(server) } /// Lists one resource page from the named server. @@ -81,7 +81,7 @@ impl McpResourceClient { cursor.map(|cursor| PaginatedRequestParams::default().with_cursor(Some(cursor))); let result = self .runtime - .snapshot() + .latest_connections() .list_resources(server, params) .await?; let resources = result @@ -100,7 +100,7 @@ impl McpResourceClient { let params = ReadResourceRequestParams::new(uri.to_string()); let result = self .runtime - .snapshot() + .latest_connections() .read_resource(server, params) .await?; let contents = result diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 912a5375b1..ee6161add8 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -478,19 +478,6 @@ impl AsyncManagedClient { ), ) }); - if codex_apps_tools_cache_context - .as_ref() - .is_some_and(ConnectorRuntimeContext::has_current_tools) - || tool_catalog_cache_context - .as_ref() - .is_some_and(McpToolCatalogCacheContext::has_tools) - { - let startup_task = client.clone(); - tokio::spawn(async move { - let _ = startup_task.await; - }); - } - Self { client, is_codex_apps_mcp_server, diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 436e5f601f..009ad279f3 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -5,49 +5,301 @@ //! [`crate::rmcp_client`] and connection-set behavior lives in //! [`crate::connection_manager`]. +use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use arc_swap::ArcSwap; +use async_channel::Sender; +use codex_connectors::ConnectorRuntimeContextKey; +use codex_connectors::ConnectorRuntimeManager; use codex_exec_server::Environment; use codex_exec_server::EnvironmentManager; use codex_exec_server::HttpClient; use codex_exec_server::ReqwestHttpClient; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use codex_protocol::capabilities::SelectedCapabilityRoot; +use codex_protocol::mcp::CallToolResult; use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::Event; +use codex_rmcp_client::ElicitationResponse; use codex_utils_path_uri::PathUri; +use rmcp::model::ReadResourceRequestParams; +use rmcp::model::ReadResourceResult; +use rmcp::model::RequestId; use serde::Deserialize; use serde::Serialize; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; -use crate::McpConnectionSet; +use crate::McpConfig; +use crate::binding::McpBinding; +use crate::connection_manager::McpConnectionSet; +use crate::elicitation::ElicitationLifecycle; +use crate::elicitation::ElicitationRequestRouter; +use crate::elicitation::ElicitationReviewerHandle; +use crate::server::EffectiveMcpServer; +use crate::tool_catalog_cache::McpToolCatalogCache; +use crate::tools::ToolInfo; -/// Owns the currently published MCP connection set for one Codex thread. +/// Everything needed to materialize one exact MCP configuration. +pub struct McpRuntimeInput { + pub config: Arc, + pub plugins_available: bool, + pub ready_selected_capability_roots: Vec, + pub mcp_servers: HashMap, + pub submit_id: String, + pub tx_event: Option>, + pub startup_cancellation_token: CancellationToken, + pub runtime_context: McpRuntimeContext, + pub codex_apps_tools_cache: ConnectorRuntimeManager, + pub tool_catalog_cache: McpToolCatalogCache, + pub codex_apps_tools_cache_key: ConnectorRuntimeContextKey, + pub supports_openai_form_elicitation: bool, + pub auth: Option, + pub codex_apps_auth_manager: Option>, + pub elicitation_reviewer: Option, + pub elicitation_lifecycle: Option, +} + +/// Owns all mutable MCP state for one Codex thread. /// -/// Replacements are published atomically. Callers that already hold a snapshot -/// keep the previous connection set alive until their work completes. +/// Publication replaces the latest state atomically. Existing bindings retain +/// their exact connections and configuration for as long as they are needed. pub struct McpRuntime { - connections: ArcSwap, + current: ArcSwap, + elicitation_router: ElicitationRequestRouter, +} + +struct PublishedMcpRuntime { + connections: Arc, + config: Option>, + plugins_available: bool, + ready_selected_capability_roots: Vec, +} + +#[derive(Clone)] +pub(crate) struct McpPublicationGate { + published: Option>, +} + +impl McpPublicationGate { + fn pending() -> (watch::Sender, Self) { + let (publish, published) = watch::channel(false); + ( + publish, + Self { + published: Some(published), + }, + ) + } + + pub(crate) fn already_published() -> Self { + Self { published: None } + } + + pub(crate) async fn wait(mut self) -> bool { + let Some(published) = self.published.as_mut() else { + return true; + }; + loop { + if *published.borrow() { + return true; + } + if published.changed().await.is_err() { + return false; + } + } + } } impl McpRuntime { - pub fn new(connections: Arc) -> Self { + /// Creates a runtime with no configured servers. + /// + /// This is useful while constructing a thread that must publish a stable + /// runtime handle before its full MCP inputs are available. + pub fn empty(prefix_mcp_tool_names: bool) -> Self { Self { - connections: ArcSwap::from(connections), + current: ArcSwap::from_pointee(PublishedMcpRuntime { + connections: Arc::new(McpConnectionSet::empty(prefix_mcp_tool_names)), + config: None, + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + }), + elicitation_router: ElicitationRequestRouter::default(), } } - pub fn snapshot(&self) -> Arc { - self.connections.load_full() + pub async fn new(input: McpRuntimeInput) -> Self { + let runtime = Self::empty(input.config.prefix_mcp_tool_names); + runtime.replace(input).await; + runtime } - pub fn replace(&self, connections: McpConnectionSet) -> Arc { - let connections = Arc::new(connections); - self.connections.store(Arc::clone(&connections)); - connections + /// Rebuilds configured servers and publishes their immutable runtime snapshot. + pub async fn replace(&self, input: McpRuntimeInput) { + let (publish, publication_gate) = McpPublicationGate::pending(); + let config = Arc::clone(&input.config); + let plugins_available = input.plugins_available; + let ready_selected_capability_roots = input.ready_selected_capability_roots.clone(); + let connections = Arc::new( + Self::materialize(input, self.elicitation_router.clone(), publication_gate).await, + ); + self.current.store(Arc::new(PublishedMcpRuntime { + connections, + config: Some(config), + plugins_available, + ready_selected_capability_roots, + })); + let _ = publish.send(true); + } + + /// Captures the latest published configuration and live client handles. + pub async fn current_binding(&self) -> Option> { + let current = self.current.load_full(); + let config = Arc::clone(current.config.as_ref()?); + Some(Arc::new( + current + .connections + .capture_binding_with_metadata(config, current.plugins_available) + .await, + )) + } + + pub fn current_ready_selected_capability_roots(&self) -> Vec { + self.current.load().ready_selected_capability_roots.clone() + } + + pub fn elicitations_auto_deny(&self) -> bool { + self.elicitation_router.auto_deny() + } + + pub fn set_elicitations_auto_deny(&self, auto_deny: bool) { + self.elicitation_router.set_auto_deny(auto_deny); + } + + pub async fn resolve_elicitation( + &self, + server_name: String, + id: RequestId, + response: ElicitationResponse, + ) -> anyhow::Result<()> { + self.elicitation_router + .resolve(server_name, id, response) + .await + } + + pub async fn latest_hard_refresh_codex_apps_tools_cache( + &self, + ) -> anyhow::Result> { + self.latest_connections() + .hard_refresh_codex_apps_tools_cache() + .await + } + + /// Lists the latest known tools for non-model discovery surfaces. + /// + /// Unlike [`Self::current_binding`], this may return cached tools while their + /// client reconnects because callers only inspect tool metadata. + pub async fn latest_list_all_tools(&self) -> Vec { + self.latest_connections().list_all_tools().await + } + + pub async fn latest_call_tool( + &self, + server: &str, + tool: &str, + arguments: Option, + meta: Option, + ) -> anyhow::Result { + self.latest_connections() + .call_tool(server, tool, arguments, meta) + .await + } + + pub async fn latest_read_resource( + &self, + server: &str, + params: ReadResourceRequestParams, + ) -> anyhow::Result { + self.latest_connections() + .read_resource(server, params) + .await + } + + pub async fn latest_wait_for_server_ready(&self, server: &str, timeout: Duration) -> bool { + self.latest_connections() + .wait_for_server_ready(server, timeout) + .await + } + + pub async fn validate_required_servers(&self) -> anyhow::Result<()> { + self.latest_connections().validate_required_servers().await + } + + pub fn cancel_startup(&self) { + self.current.load().connections.cancel_startup(); + } + + pub(crate) fn latest_connections(&self) -> Arc { + Arc::clone(&self.current.load().connections) } pub async fn shutdown(&self) { - self.snapshot().shutdown().await; + self.latest_connections().shutdown().await; + } + + async fn materialize( + input: McpRuntimeInput, + elicitation_router: ElicitationRequestRouter, + publication_gate: McpPublicationGate, + ) -> McpConnectionSet { + let McpRuntimeInput { + config, + plugins_available: _, + ready_selected_capability_roots: _, + mcp_servers, + submit_id, + tx_event, + startup_cancellation_token, + runtime_context, + codex_apps_tools_cache, + tool_catalog_cache, + codex_apps_tools_cache_key, + supports_openai_form_elicitation, + auth, + codex_apps_auth_manager, + elicitation_reviewer, + elicitation_lifecycle, + } = input; + McpConnectionSet::new( + &mcp_servers, + config.mcp_oauth_credentials_store_mode, + config.auth_keyring_backend_kind, + &config.approval_policy, + submit_id, + tx_event, + startup_cancellation_token, + config.permission_profile.clone(), + runtime_context, + config.codex_home.clone(), + codex_apps_tools_cache, + tool_catalog_cache, + codex_apps_tools_cache_key, + config.prefix_mcp_tool_names, + config.client_elicitation_capability.clone(), + supports_openai_form_elicitation, + crate::mcp::tool_plugin_provenance(&config), + auth.as_ref(), + codex_apps_auth_manager, + elicitation_reviewer, + elicitation_lifecycle, + elicitation_router, + publication_gate, + ) + .await } } @@ -179,6 +431,21 @@ mod tests { } } + #[tokio::test] + async fn publication_gate_opens_only_for_the_winning_candidate() { + let (publish, gate) = McpPublicationGate::pending(); + let wait = tokio::spawn(gate.wait()); + tokio::task::yield_now().await; + assert!(!wait.is_finished()); + + publish.send(true).expect("publish candidate"); + assert!(wait.await.expect("gate task")); + + let (publish, gate) = McpPublicationGate::pending(); + drop(publish); + assert!(!gate.wait().await); + } + fn http_server(environment_id: &str) -> McpServerConfig { McpServerConfig { auth: Default::default(), diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 4f8de06651..865e50fbc1 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -102,7 +102,6 @@ pub use codex_protocol::openai_models::ModelPreset; pub use codex_protocol::protocol::AskForApproval; pub use codex_protocol::protocol::EventMsg; pub use codex_protocol::protocol::InitialHistory; -pub use codex_protocol::protocol::McpServerRefreshConfig; pub use codex_protocol::protocol::Op; pub use codex_protocol::protocol::SessionConfiguredEvent; pub use codex_protocol::protocol::SessionSource; diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index bc16482174..5f463b7d08 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -44,7 +44,6 @@ use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC; use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::mcp_tool_call::build_guardian_mcp_tool_review_request; use crate::mcp_tool_call::is_mcp_tool_approval_question_id; -use crate::mcp_tool_call::lookup_mcp_tool_metadata; use crate::mcp_tool_call::mcp_approvals_reviewer; use crate::session::GitEnrichmentPolicy; use crate::session::SUBMISSION_CHANNEL_CAPACITY; @@ -372,24 +371,12 @@ async fn forward_events( id, msg: EventMsg::McpToolCallBegin(event), } => { - // Runtime refreshes are published before a request step is captured, so - // the child runtime at call begin is the one executing this invocation. - // Cache its metadata now; the later approval event has only a call ID. - let metadata = if let Some(turn_context) = - session.turn_context_for_sub_id(&id).await - { - let mcp = session.services.latest_mcp_runtime(); - lookup_mcp_tool_metadata( - session.as_ref(), - turn_context.as_ref(), - mcp.manager(), - &event.invocation.server, - &event.invocation.tool, - ) - .await - } else { - None - }; + // The later approval event has only a call ID. Retain the exact facts + // captured before this begin event instead of consulting the latest + // runtime after a refresh. + let metadata = session + .mcp_tool_approval_metadata(&id, &event.call_id) + .await; pending_mcp_invocations .lock() .await diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 4885b7b753..12c76d181b 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -606,31 +606,14 @@ impl CodexThread { self.session.runtime_mcp_config(config).await } - /// Returns the exact MCP config, environment bindings, and manager most recently published. - pub async fn current_mcp_runtime(&self) -> Arc { - let turn_context = self.session.new_default_turn().await; - let environments = turn_context.environments.refresh_readiness(); - let selected_capability_roots = self - .session - .resolve_selected_capability_roots_for_step(&environments) - .await; - let ready_selected_capability_roots = - Session::ready_selected_capability_roots(&selected_capability_roots); - let executor_capability_discovery = self - .session - .executor_capability_discovery_for_step( - &turn_context.config, - &ready_selected_capability_roots, - ) - .await; - self.session - .mcp_runtime_for_step( - turn_context.as_ref(), - &environments, - &selected_capability_roots, - executor_capability_discovery.as_deref(), - ) - .await + /// Captures the exact MCP config and environment bindings for the current thread state. + pub async fn current_mcp_config_and_runtime_context( + &self, + ) -> (Arc, codex_mcp::McpRuntimeContext) { + let config = self.session.get_config().await; + let (mcp_config, runtime_context) = + self.session.runtime_mcp_config_and_context(&config).await; + (Arc::new(mcp_config), runtime_context) } pub fn multi_agent_version(&self) -> Option { @@ -662,11 +645,12 @@ impl CodexThread { server: &str, uri: &str, ) -> anyhow::Result { + self.session.refresh_mcp_if_dirty().await; let result = self - .current_mcp_runtime() - .await - .manager_arc() - .read_resource(server, ReadResourceRequestParams::new(uri)) + .session + .services + .mcp_runtime + .latest_read_resource(server, ReadResourceRequestParams::new(uri)) .await?; Ok(serde_json::to_value(result)?) @@ -679,10 +663,11 @@ impl CodexThread { arguments: Option, meta: Option, ) -> anyhow::Result { - self.current_mcp_runtime() - .await - .manager_arc() - .call_tool(server, tool, arguments, meta) + self.session.refresh_mcp_if_dirty().await; + self.session + .services + .mcp_runtime + .latest_call_tool(server, tool, arguments, meta) .await } diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 293a93eaca..783b6bdb05 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -92,10 +92,9 @@ pub(crate) async fn build_compaction_initial_context( step_context, } => { let items = sess - .build_initial_context_with_world_state_and_mcp( + .build_initial_context_with_world_state( step_context.turn.as_ref(), world_state.as_ref(), - step_context.mcp.as_ref(), ) .await; (items, Some(Arc::clone(world_state))) diff --git a/codex-rs/core/src/compact_tests.rs b/codex-rs/core/src/compact_tests.rs index f3f50e021e..6a1a7c40ec 100644 --- a/codex-rs/core/src/compact_tests.rs +++ b/codex-rs/core/src/compact_tests.rs @@ -20,11 +20,7 @@ async fn process_compacted_history_with_test_session( crate::session::step_context::StepContext::for_test(Arc::clone(&turn_context)); let world_state = Arc::new(session.build_world_state_for_step(&step_context).await); let initial_context = session - .build_initial_context_with_world_state_and_mcp( - &turn_context, - world_state.as_ref(), - step_context.mcp.as_ref(), - ) + .build_initial_context_with_world_state(&turn_context, world_state.as_ref()) .await; let initial_context_injection = InitialContextInjection::BeforeLastUserMessage { world_state, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 0d59d618f9..4271f343f5 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1632,6 +1632,10 @@ impl Config { .features .enabled(Feature::SkillMcpDependencyInstall), approval_policy: self.permissions.approval_policy.clone(), + permission_profile: self.permissions.permission_profile().clone(), + config_layer_stack: self.config_layer_stack.clone(), + approvals_reviewer: self.approvals_reviewer, + environment_cwds: HashMap::new(), codex_linux_sandbox_exe: self.codex_linux_sandbox_exe.clone(), use_legacy_landlock: self.features.use_legacy_landlock(), apps_enabled: self.features.enabled(Feature::Apps), diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index a278eef74b..f6126824cb 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -15,7 +15,6 @@ use codex_connectors::apps_config_from_layer_stack; use codex_connectors::connector_runtime_context_key; use codex_exec_server::EnvironmentManager; use codex_exec_server::ExecServerRuntimePaths; -use codex_protocol::models::PermissionProfile; use codex_tools::DiscoverableTool; use tokio_util::sync::CancellationToken; use tracing::instrument; @@ -33,12 +32,14 @@ use codex_login::AuthManager; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; -use codex_mcp::McpConnectionSet; +use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeContext; +use codex_mcp::McpRuntimeInput; use codex_mcp::ToolInfo; use codex_mcp::ToolPluginProvenance; use codex_mcp::effective_mcp_servers; use codex_mcp::tool_plugin_provenance; +use codex_protocol::models::PermissionProfile; const CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS: Duration = Duration::from_secs(30); @@ -211,7 +212,10 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( }); } let cache_key = accessible_connectors_cache_key(config, auth.as_ref()); - let mcp_config = mcp_manager.runtime_config(config).await; + let mut mcp_config = mcp_manager.runtime_config(config).await; + // Discovery has no active turn or reviewer and must never inherit execution authority. + mcp_config.permission_profile = PermissionProfile::default(); + let mcp_config = Arc::new(mcp_config); let tool_plugin_provenance = tool_plugin_provenance(&mcp_config); if !force_refetch && let Some(cached_connectors) = read_cached_accessible_connectors(&cache_key) { @@ -238,37 +242,31 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( let codex_apps_auth_manager = codex_mcp::host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()) .then(|| Arc::clone(&auth_manager)); - let mcp_connection_manager = McpConnectionSet::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), - &config.permissions.approval_policy, - INITIAL_SUBMIT_ID.to_owned(), - /*tx_event*/ None, - cancel_token.clone(), - PermissionProfile::default(), + let mcp_runtime = McpRuntime::new(McpRuntimeInput { + config: Arc::clone(&mcp_config), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers: mcp_servers.clone(), + submit_id: INITIAL_SUBMIT_ID.to_owned(), + tx_event: None, + startup_cancellation_token: cancel_token.clone(), // Connector discovery is threadless. Use an actually configured env if // one exists, but do not reintroduce the old hidden-local fallback. runtime_context, - config.codex_home.to_path_buf(), - mcp_manager.codex_apps_tools_cache(), - mcp_manager.tool_catalog_cache(), - connector_runtime_context_key(auth.as_ref()), - mcp_config.prefix_mcp_tool_names, - mcp_config.client_elicitation_capability, - /*supports_openai_form_elicitation*/ false, - ToolPluginProvenance::default(), - auth.as_ref(), + codex_apps_tools_cache: mcp_manager.codex_apps_tools_cache(), + tool_catalog_cache: mcp_manager.tool_catalog_cache(), + codex_apps_tools_cache_key: connector_runtime_context_key(auth.as_ref()), + supports_openai_form_elicitation: false, + auth: auth.clone(), codex_apps_auth_manager, - /*elicitation_reviewer*/ None, - /*elicitation_lifecycle*/ None, - codex_mcp::ElicitationRequestRouter::default(), - ) + elicitation_reviewer: None, + elicitation_lifecycle: None, + }) .await; let refreshed_tools = if force_refetch { - match mcp_connection_manager - .hard_refresh_codex_apps_tools_cache() + match mcp_runtime + .latest_hard_refresh_codex_apps_tools_cache() .await { Ok(tools) => Some(tools), @@ -287,14 +285,14 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( let mut tools = if let Some(tools) = refreshed_tools { tools } else { - mcp_connection_manager.list_all_tools().await + mcp_runtime.latest_list_all_tools().await }; let mut should_reload_tools = false; let codex_apps_ready = if refreshed_tools_succeeded { true } else if let Some(cfg) = mcp_servers.get(CODEX_APPS_MCP_SERVER_NAME) { - let immediate_ready = mcp_connection_manager - .wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, Duration::ZERO) + let immediate_ready = mcp_runtime + .latest_wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, Duration::ZERO) .await; if immediate_ready { true @@ -303,8 +301,8 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( .configured_config() .and_then(|config| config.startup_timeout_sec) .unwrap_or(CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS); - let ready = mcp_connection_manager - .wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, timeout) + let ready = mcp_runtime + .latest_wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, timeout) .await; should_reload_tools = ready; ready @@ -315,7 +313,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( false }; if should_reload_tools { - tools = mcp_connection_manager.list_all_tools().await; + tools = mcp_runtime.latest_list_all_tools().await; } if codex_apps_ready { cancel_token.cancel(); @@ -327,7 +325,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( } let accessible_connectors = with_app_plugin_sources(accessible_connectors, &tool_plugin_provenance); - mcp_connection_manager.shutdown().await; + mcp_runtime.shutdown().await; Ok(AccessibleConnectorsStatus { connectors: accessible_connectors, codex_apps_ready, @@ -530,9 +528,23 @@ pub(crate) fn mcp_approvals_reviewer( config: &Config, server_name: &str, connector_id: Option<&str>, +) -> ApprovalsReviewer { + mcp_approvals_reviewer_from_layers( + &config.config_layer_stack, + config.approvals_reviewer, + server_name, + connector_id, + ) +} + +pub(crate) fn mcp_approvals_reviewer_from_layers( + config_layer_stack: &codex_config::ConfigLayerStack, + default_reviewer: ApprovalsReviewer, + server_name: &str, + connector_id: Option<&str>, ) -> ApprovalsReviewer { let app_reviewer = if server_name == CODEX_APPS_MCP_SERVER_NAME { - apps_config_from_layer_stack(&config.config_layer_stack).and_then(|apps_config| { + apps_config_from_layer_stack(config_layer_stack).and_then(|apps_config| { connector_id .and_then(|connector_id| apps_config.apps.get(connector_id)) .and_then(|app| app.approvals_reviewer) @@ -547,8 +559,7 @@ pub(crate) fn mcp_approvals_reviewer( }; if let Some(reviewer) = app_reviewer - && config - .config_layer_stack + && config_layer_stack .requirements() .approvals_reviewer .can_set(&reviewer) @@ -557,7 +568,7 @@ pub(crate) fn mcp_approvals_reviewer( return reviewer; } - config.approvals_reviewer + default_reviewer } #[cfg(test)] diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index d5e6fb1efc..5ae3315e70 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -30,6 +30,7 @@ use codex_protocol::capabilities::SelectedCapabilityRoot; const LEGACY_CODEX_APPS_REGISTRATION_ID: &str = "legacy_codex_apps"; /// MCP configuration and capability availability derived from the same inputs. +#[derive(Clone)] pub(crate) struct McpRuntimeProjection { pub(crate) config: McpConfig, pub(crate) plugins_available: bool, diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 6f3689cb2c..e26192ff3f 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -10,7 +10,6 @@ use crate::guardian::GuardianApprovalRequest; use crate::guardian::GuardianMcpAnnotations; use crate::guardian::new_guardian_review_id; use crate::guardian::review_approval_request; -use crate::guardian::routes_approval_to_guardian_with_reviewer; use crate::hook_runtime::run_permission_request_hooks; use crate::mcp_openai_file::rewrite_mcp_tool_arguments_for_openai_files; use crate::mcp_tool_approval_templates::RenderedMcpToolApprovalParam; @@ -34,13 +33,12 @@ use codex_features::Feature; use codex_hooks::PermissionRequestDecision; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; -use codex_mcp::McpConnectionSet; use codex_mcp::McpPermissionPromptAutoApproveContext; +use codex_mcp::PreparedMcpCall; use codex_mcp::SandboxState; use codex_mcp::auth_elicitation_completed_result; use codex_mcp::build_auth_elicitation_plan; use codex_mcp::mcp_permission_prompt_is_auto_approved; -use codex_mcp::tool_is_model_visible; use codex_protocol::approvals::ElicitationRequest; use codex_protocol::items::McpToolCallError; use codex_protocol::items::McpToolCallItem; @@ -119,7 +117,6 @@ pub(crate) async fn handle_mcp_tool_call( arguments: String, ) -> HandledMcpToolCall { let turn_context = &step_context.turn; - let manager = step_context.mcp.manager(); // Parse the `arguments` as JSON. An empty string is OK, but invalid JSON // is not. let arguments_value = if arguments.trim().is_empty() { @@ -143,16 +140,14 @@ pub(crate) async fn handle_mcp_tool_call( arguments: arguments_value.clone(), }; - let metadata = lookup_mcp_tool_metadata( - sess.as_ref(), - turn_context.as_ref(), - manager, - &server, - &tool_name, - ) - .await; - let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, metadata.as_ref()); - if metadata.is_none() { + sess.refresh_mcp_if_dirty().await; + let current_binding = sess.services.mcp_runtime.current_binding().await; + let Some(prepared_call) = current_binding + .as_ref() + .and_then(|binding| binding.prepare_call(&server, &tool_name)) + else { + let item_metadata = + McpToolCallItemMetadata::from_tool_metadata(&server, /*metadata*/ None); let result = notify_mcp_tool_call_skip( sess.as_ref(), turn_context.as_ref(), @@ -168,48 +163,30 @@ pub(crate) async fn handle_mcp_tool_call( tool_input: arguments_value .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())), }; - } + }; + let metadata = mcp_tool_metadata(&prepared_call); + let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, Some(&metadata)); + let runtime_config = prepared_call.config(); let app_tool_policy = if server == CODEX_APPS_MCP_SERVER_NAME { - let annotations = metadata - .as_ref() - .and_then(|metadata| metadata.annotations.as_ref()); - AppToolPolicyEvaluator::new(&turn_context.config.config_layer_stack).policy( - AppToolPolicyInput { - connector_id: metadata - .as_ref() - .and_then(|metadata| metadata.connector_id.as_deref()), - tool_name: &tool_name, - tool_title: metadata - .as_ref() - .and_then(|metadata| metadata.tool_title.as_deref()), - destructive_hint: annotations.and_then(|annotations| annotations.destructive_hint), - open_world_hint: annotations.and_then(|annotations| annotations.open_world_hint), - }, - ) + let annotations = metadata.annotations.as_ref(); + AppToolPolicyEvaluator::new(&runtime_config.config_layer_stack).policy(AppToolPolicyInput { + connector_id: metadata.connector_id.as_deref(), + tool_name: &tool_name, + tool_title: metadata.tool_title.as_deref(), + destructive_hint: annotations.and_then(|annotations| annotations.destructive_hint), + open_world_hint: annotations.and_then(|annotations| annotations.open_world_hint), + }) } else { AppToolPolicy::default() }; let approval_mode = if server == CODEX_APPS_MCP_SERVER_NAME { app_tool_policy.approval - } else if let Some(approval_mode) = { - // Selected-plugin registrations are absent from config.toml and the legacy plugin manager, - // so their resolved catalog entry is the authoritative source for tool approval policy. - manager - .is_selected_plugin_mcp_server(&server) - .then(|| manager.tool_approval_mode(&server, &tool_name)) - } { - approval_mode } else { - custom_mcp_tool_approval_mode(sess.as_ref(), turn_context.as_ref(), &server, &tool_name) - .await + prepared_call.tool_approval_mode() }; - let connector_id = metadata - .as_ref() - .and_then(|metadata| metadata.connector_id.clone()); - let connector_name = metadata - .as_ref() - .and_then(|metadata| metadata.connector_name.clone()); + let connector_id = metadata.connector_id.clone(); + let connector_name = metadata.connector_name.clone(); if server == CODEX_APPS_MCP_SERVER_NAME && !app_tool_policy.enabled { let result = notify_mcp_tool_call_skip( @@ -239,6 +216,8 @@ pub(crate) async fn handle_mcp_tool_call( .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())), }; } + sess.register_mcp_tool_approval_metadata(turn_context, &call_id, metadata.clone()) + .await; notify_mcp_tool_call_started( sess.as_ref(), turn_context.as_ref(), @@ -248,28 +227,39 @@ pub(crate) async fn handle_mcp_tool_call( ) .await; + let approval_policy = if prepared_call.is_selected_plugin_server() { + McpToolApprovalPolicy::for_selected_plugin(approval_mode) + } else { + McpToolApprovalPolicy::for_server(approval_mode) + }; if let Some(decision) = maybe_request_mcp_tool_approval( &sess, step_context, &call_id, &invocation, &hook_tool_name, - metadata.as_ref(), - approval_mode, + &metadata, + prepared_call.config(), + approval_policy, ) .await { let result = match decision { - McpToolApprovalDecision::Accept + decision @ (McpToolApprovalDecision::Accept | McpToolApprovalDecision::AcceptForSession - | McpToolApprovalDecision::AcceptAndRemember => { + | McpToolApprovalDecision::AcceptAndRemember) => { return handle_approved_mcp_tool_call( - sess.as_ref(), + &sess, step_context.as_ref(), &call_id, invocation, - metadata.as_ref(), + prepared_call, + metadata, item_metadata, + McpToolApprovalApplication::Apply { + decision, + policy: approval_policy, + }, ) .await; } @@ -321,12 +311,14 @@ pub(crate) async fn handle_mcp_tool_call( } handle_approved_mcp_tool_call( - sess.as_ref(), + &sess, step_context.as_ref(), &call_id, invocation, - metadata.as_ref(), + prepared_call, + metadata, item_metadata, + McpToolApprovalApplication::NotRequired, ) .await } @@ -372,53 +364,116 @@ impl McpToolCallItemMetadata { } } +#[expect( + clippy::too_many_arguments, + reason = "MCP approval must be applied inside the prepared call's catalog lease" +)] async fn handle_approved_mcp_tool_call( - sess: &Session, + sess: &Arc, step_context: &StepContext, call_id: &str, invocation: McpInvocation, - metadata: Option<&McpToolApprovalMetadata>, + prepared_call: PreparedMcpCall, + metadata: McpToolApprovalMetadata, item_metadata: McpToolCallItemMetadata, + approval_application: McpToolApprovalApplication, ) -> HandledMcpToolCall { let turn_context = step_context.turn.as_ref(); - let manager = step_context.mcp.manager(); let server = invocation.server.clone(); - maybe_mark_thread_memory_mode_polluted(sess, turn_context, manager, &server).await; let tool_name = invocation.tool.clone(); let arguments_value = invocation.arguments.clone(); - let connector_id = metadata.and_then(|metadata| metadata.connector_id.as_deref()); - let connector_name = metadata.and_then(|metadata| metadata.connector_name.as_deref()); - let server_origin = manager.server_origin(&server).map(str::to_string); + let connector_id = metadata.connector_id.as_deref(); + let connector_name = metadata.connector_name.as_deref(); + let server_origin = prepared_call.server_origin().map(str::to_string); let start = Instant::now(); - let rewrite = rewrite_mcp_tool_arguments_for_openai_files( - sess, - turn_context, - arguments_value.clone(), - metadata.and_then(|metadata| metadata.openai_file_input_optional_fields.as_ref()), - ) - .await; - let tool_input = match &rewrite { - Ok(Some(rewritten_arguments)) => rewritten_arguments.clone(), - Ok(None) | Err(_) => arguments_value - .clone() - .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())), - }; + let mut tool_input = arguments_value + .clone() + .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())); let result = async { let result = async { - let rewritten_arguments = rewrite?; - let request_meta = - build_mcp_tool_call_request_meta(turn_context, &server, call_id, metadata); - execute_mcp_tool_call( + let result = prepared_call + .call_with_preparation(|| async { + if let McpToolApprovalApplication::Apply { decision, policy } = + &approval_application + { + let session_approval_key = session_mcp_tool_approval_key( + &invocation, + Some(&metadata), + policy.mode, + ); + let persistent_approval_key = if policy.allow_persistent { + persistent_mcp_tool_approval_key( + &invocation, + Some(&metadata), + policy.mode, + ) + } else { + None + }; + apply_mcp_tool_approval_decision( + sess, + turn_context, + decision, + session_approval_key, + persistent_approval_key, + ) + .await; + } + maybe_mark_thread_memory_mode_polluted(sess, turn_context, &prepared_call) + .await; + let rewritten_arguments = rewrite_mcp_tool_arguments_for_openai_files( + sess, + turn_context, + arguments_value, + metadata.openai_file_input_optional_fields.as_ref(), + ) + .await + .map_err(anyhow::Error::msg)?; + if let Some(rewritten_arguments) = rewritten_arguments.as_ref() { + tool_input = rewritten_arguments.clone(); + } + let request_meta = build_mcp_tool_call_request_meta( + turn_context, + &server, + call_id, + Some(&metadata), + ); + let request_meta = with_mcp_tool_call_thread_id_meta( + request_meta, + &sess.thread_id.to_string(), + ); + let request_meta = augment_mcp_tool_request_meta_with_sandbox_state( + step_context, + &prepared_call, + request_meta, + ) + .await?; + let mcp_call_trace = sess + .services + .rollout_thread_trace + .start_mcp_call_trace(call_id); + Ok(( + rewritten_arguments, + mcp_call_trace.add_request_meta(request_meta), + )) + }) + .await + .map_err(|error| format!("tool call error: {error:?}"))?; + let result = sanitize_mcp_tool_result_for_model( + &turn_context.model_info.input_modalities, + Ok(result), + )?; + Ok(maybe_request_codex_apps_auth_elicitation( sess, - step_context, + turn_context, + prepared_call.config().approval_policy.value(), call_id, - &invocation, - rewritten_arguments, - metadata, - request_meta, + &invocation.server, + Some(&metadata), + result, ) - .await + .await) } .await; record_mcp_result_span_telemetry(&Span::current(), &result); @@ -451,7 +506,7 @@ async fn handle_approved_mcp_tool_call( truncate_mcp_tool_result_for_event(&result), ) .await; - maybe_track_codex_app_used(sess, turn_context, manager, &server, &tool_name).await; + maybe_track_codex_app_used(sess, turn_context, &server, &metadata).await; let outcome = mcp_call_metric_outcome(&result); emit_mcp_call_metrics( @@ -576,64 +631,16 @@ fn truncate_str_to_char_boundary(value: &str, max_chars: usize) -> &str { } } -async fn execute_mcp_tool_call( - sess: &Session, - step_context: &StepContext, - call_id: &str, - invocation: &McpInvocation, - rewritten_arguments: Option, - metadata: Option<&McpToolApprovalMetadata>, - request_meta: Option, -) -> Result { - let turn_context = step_context.turn.as_ref(); - let manager = step_context.mcp.manager(); - let request_meta = with_mcp_tool_call_thread_id_meta(request_meta, &sess.thread_id.to_string()); - let request_meta = augment_mcp_tool_request_meta_with_sandbox_state( - step_context, - manager, - &invocation.server, - request_meta, - ) - .await - .map_err(|e| format!("failed to build MCP tool request metadata: {e:#}"))?; - let mcp_call_trace = sess - .services - .rollout_thread_trace - .start_mcp_call_trace(call_id); - let request_meta = mcp_call_trace.add_request_meta(request_meta); - let result = manager - .call_tool( - &invocation.server, - &invocation.tool, - rewritten_arguments, - request_meta, - ) - .await - .map_err(|e| format!("tool call error: {e:?}"))?; - let result = - sanitize_mcp_tool_result_for_model(&turn_context.model_info.input_modalities, Ok(result))?; - Ok(maybe_request_codex_apps_auth_elicitation( - sess, - turn_context, - manager, - call_id, - &invocation.server, - metadata, - result, - ) - .await) -} - async fn maybe_request_codex_apps_auth_elicitation( - sess: &Session, + sess: &Arc, turn_context: &TurnContext, - manager: &McpConnectionSet, + approval_policy: AskForApproval, call_id: &str, server: &str, metadata: Option<&McpToolApprovalMetadata>, result: CallToolResult, ) -> CallToolResult { - if !manager.is_host_owned_codex_apps_server(server) { + if server != CODEX_APPS_MCP_SERVER_NAME { return result; } @@ -645,7 +652,7 @@ async fn maybe_request_codex_apps_auth_elicitation( return result; } - match turn_context.approval_policy.value() { + match approval_policy { AskForApproval::Never => return result, AskForApproval::Granular(granular_config) if !granular_config.allows_mcp_elicitations() => { return result; @@ -691,16 +698,12 @@ async fn maybe_request_codex_apps_auth_elicitation( return result; } - refresh_codex_apps_after_connector_auth(sess, turn_context, manager).await; + refresh_codex_apps_after_connector_auth(sess, turn_context).await; auth_elicitation_completed_result(&plan.auth_failure, result.meta) } -async fn refresh_codex_apps_after_connector_auth( - sess: &Session, - turn_context: &TurnContext, - manager: &McpConnectionSet, -) { - let mcp_tools_result = manager.hard_refresh_codex_apps_tools_cache().await; +async fn refresh_codex_apps_after_connector_auth(sess: &Arc, turn_context: &TurnContext) { + let mcp_tools_result = sess.hard_refresh_latest_codex_apps_tools().await; match mcp_tools_result { Ok(mcp_tools) => { @@ -719,31 +722,33 @@ async fn refresh_codex_apps_after_connector_auth( async fn augment_mcp_tool_request_meta_with_sandbox_state( step_context: &StepContext, - manager: &McpConnectionSet, - server: &str, + prepared_call: &PreparedMcpCall, mut meta: Option, ) -> anyhow::Result> { - let turn_context = step_context.turn.as_ref(); - let supports_sandbox_state_meta = manager - .server_supports_sandbox_state_meta_capability(server) + let supports_sandbox_state_meta = prepared_call + .server_supports_sandbox_state_meta_capability() .await .unwrap_or(false); if !supports_sandbox_state_meta { return Ok(meta); } - let server_environment_id = manager - .server_environment_id(server) - .unwrap_or(codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID); - let Some(sandbox_cwd) = sandbox_cwd_for_mcp_server(step_context, server_environment_id) else { + let server_environment_id = prepared_call.server_environment_id(); + let Some(sandbox_cwd) = prepared_call + .config() + .environment_cwds + .get(server_environment_id) + .cloned() + .or_else(|| sandbox_cwd_for_mcp_server(step_context, server_environment_id)) + else { return Ok(meta); }; - let permission_profile = turn_context.permission_profile(); + let permission_profile = prepared_call.config().permission_profile.clone(); let sandbox_state = serde_json::to_value(SandboxState { permission_profile, - codex_linux_sandbox_exe: step_context.mcp.config().codex_linux_sandbox_exe.clone(), + codex_linux_sandbox_exe: prepared_call.config().codex_linux_sandbox_exe.clone(), sandbox_cwd, - use_legacy_landlock: step_context.mcp.config().use_legacy_landlock, + use_legacy_landlock: prepared_call.config().use_legacy_landlock, })?; match meta.as_mut() { @@ -787,14 +792,12 @@ fn sandbox_cwd_for_mcp_server(step_context: &StepContext, environment_id: &str) async fn maybe_mark_thread_memory_mode_polluted( sess: &Session, turn_context: &TurnContext, - manager: &McpConnectionSet, - server: &str, + prepared_call: &PreparedMcpCall, ) { if !turn_context.config.memories.disable_on_external_context { return; } - let pollutes_memory = manager.server_pollutes_memory(server); - if !pollutes_memory { + if !prepared_call.server_pollutes_memory() { return; } state_db::mark_thread_memory_mode_polluted( @@ -963,25 +966,17 @@ async fn notify_mcp_tool_call_completed( sess.emit_turn_item_completed(turn_context, item).await; } -struct McpAppUsageMetadata { - connector_id: Option, - app_name: Option, -} - async fn maybe_track_codex_app_used( sess: &Session, turn_context: &TurnContext, - manager: &McpConnectionSet, server: &str, - tool_name: &str, + metadata: &McpToolApprovalMetadata, ) { if server != CODEX_APPS_MCP_SERVER_NAME { return; } - let metadata = lookup_mcp_app_usage_metadata(manager, server, tool_name).await; - let (connector_id, app_name) = metadata - .map(|metadata| (metadata.connector_id, metadata.app_name)) - .unwrap_or((None, None)); + let connector_id = metadata.connector_id.clone(); + let app_name = metadata.connector_name.clone(); let invocation_type = if let Some(connector_id) = connector_id.as_deref() { let mentioned_connector_ids = sess.get_connector_selection().await; if mentioned_connector_ids.contains(connector_id) { @@ -1018,6 +1013,36 @@ enum McpToolApprovalDecision { Cancel, } +#[derive(Clone, Copy)] +struct McpToolApprovalPolicy { + mode: AppToolApproval, + allow_persistent: bool, +} + +enum McpToolApprovalApplication { + NotRequired, + Apply { + decision: McpToolApprovalDecision, + policy: McpToolApprovalPolicy, + }, +} + +impl McpToolApprovalPolicy { + fn for_server(mode: AppToolApproval) -> Self { + Self { + mode, + allow_persistent: true, + } + } + + fn for_selected_plugin(mode: AppToolApproval) -> Self { + Self { + mode, + allow_persistent: false, + } + } +} + #[derive(Clone)] pub(crate) struct McpToolApprovalMetadata { annotations: Option, @@ -1034,6 +1059,40 @@ pub(crate) struct McpToolApprovalMetadata { openai_file_input_optional_fields: Option>>, } +impl Session { + async fn register_mcp_tool_approval_metadata( + &self, + turn_context: &TurnContext, + call_id: &str, + metadata: McpToolApprovalMetadata, + ) { + let Some(turn_state) = self + .input_queue + .turn_state_for_sub_id(&self.active_turn, &turn_context.sub_id) + .await + else { + return; + }; + turn_state + .lock() + .await + .insert_mcp_tool_approval_metadata(call_id.to_string(), metadata); + } + + pub(crate) async fn mcp_tool_approval_metadata( + &self, + sub_id: &str, + call_id: &str, + ) -> Option { + let turn_state = self + .input_queue + .turn_state_for_sub_id(&self.active_turn, sub_id) + .await?; + + turn_state.lock().await.mcp_tool_approval_metadata(call_id) + } +} + const MCP_TOOL_OPENAI_OUTPUT_TEMPLATE_META_KEY: &str = "openai/outputTemplate"; const MCP_TOOL_UI_RESOURCE_URI_META_KEY: &str = "ui/resourceUri"; const MCP_TOOL_LINK_ID_META_KEY: &str = "link_id"; @@ -1042,6 +1101,7 @@ const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId"; const MCP_TOOL_CONNECTED_ACCOUNT_EMAIL_META_KEY: &str = "connected_account_email"; const MCP_TOOL_RESOURCE_URI_META_KEY: &str = "resource_uri"; +#[cfg(test)] async fn custom_mcp_tool_approval_mode( sess: &Session, turn_context: &TurnContext, @@ -1211,38 +1271,45 @@ fn mcp_tool_approval_prompt_options( } } +#[expect(clippy::too_many_arguments)] async fn maybe_request_mcp_tool_approval( sess: &Arc, step_context: &Arc, call_id: &str, invocation: &McpInvocation, hook_tool_name: &HookToolName, - metadata: Option<&McpToolApprovalMetadata>, - approval_mode: AppToolApproval, + metadata: &McpToolApprovalMetadata, + config: &codex_mcp::McpConfig, + policy: McpToolApprovalPolicy, ) -> Option { let turn_context = &step_context.turn; - let manager = step_context.mcp.manager(); - let approvals_reviewer = mcp_approvals_reviewer(turn_context, &invocation.server, metadata); + let approvals_reviewer = connectors::mcp_approvals_reviewer_from_layers( + &config.config_layer_stack, + config.approvals_reviewer, + &invocation.server, + metadata.connector_id.as_deref(), + ); if mcp_permission_prompt_is_auto_approved( - turn_context.approval_policy.value(), - &turn_context.permission_profile(), + config.approval_policy.value(), + &config.permission_profile, McpPermissionPromptAutoApproveContext { - tool_approval_mode: Some(approval_mode), + tool_approval_mode: Some(policy.mode), }, ) { return None; } - let annotations = metadata.and_then(|metadata| metadata.annotations.as_ref()); - if !requires_mcp_tool_approval_for_mode(annotations, approval_mode) { + let annotations = metadata.annotations.as_ref(); + if !requires_mcp_tool_approval_for_mode(annotations, policy.mode) { return None; } - let session_approval_key = session_mcp_tool_approval_key(invocation, metadata, approval_mode); - let persistent_approval_key = if manager.is_selected_plugin_mcp_server(&invocation.server) { - None + let session_approval_key = + session_mcp_tool_approval_key(invocation, Some(metadata), policy.mode); + let persistent_approval_key = if policy.allow_persistent { + persistent_mcp_tool_approval_key(invocation, Some(metadata), policy.mode) } else { - persistent_mcp_tool_approval_key(invocation, metadata, approval_mode) + None }; if let Some(key) = session_approval_key.as_ref() && mcp_tool_approval_is_remembered(sess, key).await @@ -1280,25 +1347,21 @@ async fn maybe_request_mcp_tool_approval( .features .enabled(Feature::ToolCallMcpElicitation); - if routes_approval_to_guardian_with_reviewer(turn_context, approvals_reviewer) { + if matches!( + config.approval_policy.value(), + AskForApproval::OnRequest | AskForApproval::Granular(_) + ) && approvals_reviewer == ApprovalsReviewer::AutoReview + { let review_id = new_guardian_review_id(); let decision = review_approval_request( sess, turn_context, review_id.clone(), - build_guardian_mcp_tool_review_request(call_id, invocation, metadata), + build_guardian_mcp_tool_review_request(call_id, invocation, Some(metadata)), /*retry_reason*/ None, ) .await; let decision = mcp_tool_approval_decision_from_guardian(decision); - apply_mcp_tool_approval_decision( - sess, - turn_context, - &decision, - session_approval_key, - persistent_approval_key, - ) - .await; return Some(decision); } @@ -1310,9 +1373,9 @@ async fn maybe_request_mcp_tool_approval( let question_id = format!("{MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX}_{call_id}"); let rendered_template = render_mcp_tool_approval_template( &invocation.server, - metadata.and_then(|metadata| metadata.connector_id.as_deref()), - metadata.and_then(|metadata| metadata.connector_name.as_deref()), - metadata.and_then(|metadata| metadata.tool_title.as_deref()), + metadata.connector_id.as_deref(), + metadata.connector_name.as_deref(), + metadata.tool_title.as_deref(), invocation.arguments.as_ref(), ); let tool_params_display = rendered_template @@ -1323,7 +1386,7 @@ async fn maybe_request_mcp_tool_approval( question_id.clone(), &invocation.server, &invocation.tool, - metadata.and_then(|metadata| metadata.connector_name.as_deref()), + metadata.connector_name.as_deref(), prompt_options, rendered_template .as_ref() @@ -1336,7 +1399,7 @@ async fn maybe_request_mcp_tool_approval( let request = build_mcp_tool_approval_elicitation_request(McpToolApprovalElicitationRequest { server: &invocation.server, - metadata, + metadata: Some(metadata), tool_params: rendered_template .as_ref() .and_then(|rendered_template| rendered_template.tool_params.as_ref()) @@ -1359,15 +1422,7 @@ async fn maybe_request_mcp_tool_approval( .response, &question_id, ); - let decision = normalize_approval_decision_for_mode(decision, approval_mode); - apply_mcp_tool_approval_decision( - sess, - turn_context, - &decision, - session_approval_key, - persistent_approval_key, - ) - .await; + let decision = normalize_approval_decision_for_mode(decision, policy.mode); return Some(decision); } @@ -1380,16 +1435,8 @@ async fn maybe_request_mcp_tool_approval( .await; let decision = normalize_approval_decision_for_mode( parse_mcp_tool_approval_response(response, &question_id), - approval_mode, + policy.mode, ); - apply_mcp_tool_approval_decision( - sess, - turn_context, - &decision, - session_approval_key, - persistent_approval_key, - ) - .await; Some(decision) } @@ -1478,49 +1525,12 @@ fn mcp_tool_approval_decision_from_guardian(decision: ReviewDecision) -> McpTool } } -pub(crate) async fn lookup_mcp_tool_metadata( - sess: &Session, - turn_context: &TurnContext, - manager: &McpConnectionSet, - server: &str, - tool_name: &str, -) -> Option { - let plugin_id = manager - .plugin_id_for_mcp_server_name(server) - .map(str::to_string); - let tool_info = manager.tool_info(server, tool_name).await?; - if !tool_is_model_visible(&tool_info) { - return None; - } - let connector_description = if server == CODEX_APPS_MCP_SERVER_NAME { - let connectors = match connectors::list_cached_accessible_connectors_from_mcp_tools( - turn_context.config.as_ref(), - ) - .await - { - Some(connectors) => Some(connectors), - None => { - connectors::list_accessible_connectors_from_mcp_tools_with_mcp_manager( - turn_context.config.as_ref(), - /*force_refetch*/ false, - sess.services.turn_environments.environment_manager(), - Arc::clone(&sess.services.mcp_manager), - ) - .await - .ok() - .map(|status| status.connectors) - } - }; - connectors.and_then(|connectors| { - let connector_id = tool_info.connector_id.as_deref()?; - connectors - .into_iter() - .find(|connector| connector.id == connector_id) - .and_then(|connector| connector.description) - }) - } else { - None - }; +fn mcp_tool_metadata(prepared_call: &PreparedMcpCall) -> McpToolApprovalMetadata { + let server = prepared_call.server_name(); + let tool_info = prepared_call.tool_info().clone(); + let connector_description = (server == CODEX_APPS_MCP_SERVER_NAME) + .then(|| tool_info.namespace_description.clone()) + .flatten(); let codex_apps_meta = tool_info .tool @@ -1541,7 +1551,7 @@ pub(crate) async fn lookup_mcp_tool_metadata( None }; - Some(McpToolApprovalMetadata { + McpToolApprovalMetadata { annotations: tool_info.tool.annotations, connector_id: tool_info.connector_id, link_id: tool_info @@ -1554,7 +1564,7 @@ pub(crate) async fn lookup_mcp_tool_metadata( connector_name: tool_info.connector_name, connector_description, connected_account_email, - plugin_id, + plugin_id: prepared_call.plugin_id().map(str::to_string), tool_title: tool_info.tool.title, tool_description: tool_info.tool.description.map(std::borrow::Cow::into_owned), mcp_app_resource_uri: get_mcp_app_resource_uri(tool_info.tool.meta.as_deref()), @@ -1564,7 +1574,7 @@ pub(crate) async fn lookup_mcp_tool_metadata( server, &tool_info.openai_file_input_optional_fields, ), - }) + } } fn openai_file_input_optional_fields_for_server( @@ -1596,18 +1606,6 @@ fn get_mcp_app_resource_uri( }) } -async fn lookup_mcp_app_usage_metadata( - manager: &McpConnectionSet, - server: &str, - tool_name: &str, -) -> Option { - let tool_info = manager.tool_info(server, tool_name).await?; - Some(McpAppUsageMetadata { - connector_id: tool_info.connector_id, - app_name: tool_info.connector_name, - }) -} - fn build_mcp_tool_approval_question( question_id: String, server: &str, diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index bf4ede8461..f5355fc18d 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -5,6 +5,7 @@ use crate::environment_selection::TurnEnvironmentState; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; use crate::session::tests::make_session_and_context_with_rx; +use crate::session::tests::mcp_config_for_test; use crate::session::turn_context::TurnEnvironment; use crate::state::ActiveTurn; use crate::test_support::models_manager_with_provider; @@ -28,12 +29,6 @@ use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::McpInvocation; -use codex_protocol::protocol::SessionSource; -use codex_rollout_trace::ThreadStartedTraceMetadata; -use codex_rollout_trace::ToolDispatchInvocation; -use codex_rollout_trace::ToolDispatchPayload; -use codex_rollout_trace::ToolDispatchRequester; -use codex_rollout_trace::replay_bundle; use codex_utils_path_uri::PathUri; use core_test_support::hooks::trusted_config_layer_stack; use core_test_support::responses::ev_assistant_message; @@ -45,12 +40,8 @@ use core_test_support::responses::start_mock_server; use pretty_assertions::assert_eq; use serde::Deserialize; use std::collections::HashMap; -use std::fs; -use std::path::Path; -use std::path::PathBuf; use std::sync::Arc; use tempfile::tempdir; -use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::Level; use tracing_subscriber::fmt::format::FmtSpan; @@ -93,6 +84,13 @@ fn approval_metadata( } } +fn approval_config(turn_context: &TurnContext) -> codex_mcp::McpConfig { + let mut config = (*mcp_config_for_test(&turn_context.config)).clone(); + config.approval_policy = turn_context.approval_policy.clone(); + config.permission_profile = turn_context.permission_profile.clone(); + config +} + fn mcp_turn_metadata_context(turn_context: &TurnContext) -> McpTurnMetadataContext<'_> { McpTurnMetadataContext { model: turn_context.model_info.slug.as_str(), @@ -134,62 +132,6 @@ fn prompt_options( } } -#[tokio::test] -async fn execute_mcp_tool_call_records_replayable_correlation() -> anyhow::Result<()> { - let temp = tempdir()?; - let (mut session, turn_context) = make_session_and_context().await; - attach_trace_bundle(&mut session, &turn_context, temp.path())?; - - let dispatch_trace = session - .services - .rollout_thread_trace - .start_tool_dispatch_trace(|| { - Some(ToolDispatchInvocation { - thread_id: session.thread_id.to_string(), - codex_turn_id: turn_context.sub_id.clone(), - tool_call_id: "mcp-call".to_string(), - tool_name: "search".to_string(), - tool_namespace: Some("mcp__docs__".to_string()), - requester: ToolDispatchRequester::Model { - model_visible_call_id: "mcp-call".to_string(), - }, - payload: ToolDispatchPayload::Function { - arguments: r#"{"query":"trace"}"#.to_string(), - }, - }) - }); - assert!(dispatch_trace.is_enabled()); - let turn_context = Arc::new(turn_context); - let step_context = StepContext::for_test(Arc::clone(&turn_context)); - - let result = execute_mcp_tool_call( - &session, - step_context.as_ref(), - "mcp-call", - &McpInvocation { - server: "docs".to_string(), - tool: "search".to_string(), - arguments: Some(serde_json::json!({ "query": "trace" })), - }, - /*rewritten_arguments*/ None, - /*metadata*/ None, - /*request_meta*/ None, - ) - .await; - assert!( - result.is_err(), - "the synthetic backend is absent; only trace emission matters", - ); - - let replayed = replay_bundle(single_bundle_dir(temp.path())?)?; - assert!( - replayed.tool_calls["mcp-call"].mcp_call_id.is_some(), - "the real MCP execution path should emit a reducer-visible correlation ID", - ); - - Ok(()) -} - fn install_mcp_permission_request_hook( session: &mut Session, turn_context: &TurnContext, @@ -279,45 +221,6 @@ print({hook_output:?}) log_path.to_path_buf() } -/// Attaches a replayable rollout bundle to one synthetic session under test. -fn attach_trace_bundle( - session: &mut Session, - turn_context: &TurnContext, - root: &Path, -) -> anyhow::Result<()> { - let rollout_thread_trace = - codex_rollout_trace::ThreadTraceContext::start_root_in_root_for_test( - root, - ThreadStartedTraceMetadata { - thread_id: session.thread_id.to_string(), - agent_path: "/root".to_string(), - task_name: None, - nickname: None, - agent_role: None, - session_source: SessionSource::Exec, - cwd: PathBuf::from("/workspace"), - rollout_path: None, - model: "gpt-test".to_string(), - provider_name: "test-provider".to_string(), - approval_policy: "never".to_string(), - sandbox_policy: "danger-full-access".to_string(), - }, - )?; - rollout_thread_trace.record_codex_turn_started(turn_context.sub_id.as_str()); - session.services.rollout_thread_trace = rollout_thread_trace; - Ok(()) -} - -/// Returns the sole bundle emitted under a temporary rollout trace root. -fn single_bundle_dir(root: &Path) -> anyhow::Result { - let mut entries = fs::read_dir(root)? - .map(|entry| entry.map(|entry| entry.path())) - .collect::, _>>()?; - entries.sort(); - assert_eq!(entries.len(), 1); - Ok(entries.remove(0)) -} - #[test] fn mcp_app_resource_uri_reads_known_tool_meta_keys() { let nested = serde_json::json!({ @@ -1442,55 +1345,6 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata { ) } -async fn host_owned_codex_apps_manager( - session: &Session, - turn_context: &TurnContext, -) -> Arc { - let auth = session.services.auth_manager.auth().await; - let startup_cancellation_token = CancellationToken::new(); - startup_cancellation_token.cancel(); - let mcp_servers = HashMap::from([( - CODEX_APPS_MCP_SERVER_NAME.to_string(), - codex_mcp::EffectiveMcpServer::configured(codex_mcp::codex_apps_mcp_server_config( - "https://chatgpt.com", - /*apps_mcp_product_sku*/ None, - Some(&turn_context.originator), - )), - )]); - let manager = codex_mcp::McpConnectionSet::new( - &mcp_servers, - turn_context.config.mcp_oauth_credentials_store_mode, - turn_context.config.auth_keyring_backend_kind(), - &turn_context.approval_policy, - turn_context.sub_id.clone(), - /*tx_event*/ None, - startup_cancellation_token, - turn_context.permission_profile(), - codex_mcp::McpRuntimeContext::new( - session.services.turn_environments.environment_manager(), - { - #[allow(deprecated)] - turn_context.cwd.to_path_buf() - }, - ), - turn_context.config.codex_home.to_path_buf(), - session.services.mcp_manager.codex_apps_tools_cache(), - session.services.mcp_manager.tool_catalog_cache(), - codex_connectors::connector_runtime_context_key(auth.as_ref()), - turn_context.config.prefix_mcp_tool_names(), - rmcp::model::ElicitationCapability::default(), - /*supports_openai_form_elicitation*/ false, - codex_mcp::ToolPluginProvenance::default(), - auth.as_ref(), - /*codex_apps_auth_manager*/ None, - /*elicitation_reviewer*/ None, - /*elicitation_lifecycle*/ None, - codex_mcp::ElicitationRequestRouter::default(), - ) - .await; - Arc::new(manager) -} - #[tokio::test] async fn codex_apps_auth_elicitation_feature_disabled_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; @@ -1498,40 +1352,13 @@ async fn codex_apps_auth_elicitation_feature_disabled_returns_original_result() features.disable(Feature::AuthElicitation); let mutable_turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); Arc::make_mut(&mut mutable_turn_context.config).features = ManagedFeatures::from(features); - let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); let returned = maybe_request_codex_apps_auth_elicitation( &session, &turn_context, - manager.as_ref(), - "call_123", - CODEX_APPS_MCP_SERVER_NAME, - Some(&metadata), - result.clone(), - ) - .await; - - assert_eq!(returned, result); - assert!(rx_event.try_recv().is_err()); -} - -#[tokio::test] -async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_result() { - let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - let mut features = Features::with_defaults(); - features.enable(Feature::AuthElicitation); - let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); - Arc::make_mut(&mut turn_context.config).features = ManagedFeatures::from(features); - let result = codex_apps_auth_failure_result(); - let metadata = codex_apps_auth_failure_metadata(); - let manager = session.services.latest_mcp_runtime().manager_arc(); - - let returned = maybe_request_codex_apps_auth_elicitation( - &session, - turn_context, - manager.as_ref(), + turn_context.approval_policy.value(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1546,22 +1373,17 @@ async fn codex_apps_auth_elicitation_non_host_owned_server_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); - let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); - Arc::make_mut(&mut turn_context.config).features = ManagedFeatures::from(features); - turn_context - .approval_policy - .set(AskForApproval::Never) - .expect("test setup should allow updating approval policy"); + let mutable_turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); + Arc::make_mut(&mut mutable_turn_context.config).features = ManagedFeatures::from(features); let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); let returned = maybe_request_codex_apps_auth_elicitation( &session, - turn_context, - manager.as_ref(), + &turn_context, + AskForApproval::Never, "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1576,12 +1398,11 @@ async fn codex_apps_auth_elicitation_disallowed_by_policy_returns_original_resul #[tokio::test] async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_result() { let (session, mut turn_context, rx_event) = make_session_and_context_with_rx().await; - let manager = host_owned_codex_apps_manager(&session, &turn_context).await; let mut features = Features::with_defaults(); features.enable(Feature::AuthElicitation); - let turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); - Arc::make_mut(&mut turn_context.config).features = ManagedFeatures::from(features); - turn_context + let mutable_turn_context = Arc::get_mut(&mut turn_context).expect("single turn context ref"); + Arc::make_mut(&mut mutable_turn_context.config).features = ManagedFeatures::from(features); + mutable_turn_context .approval_policy .set(AskForApproval::Granular(GranularApprovalConfig { sandbox_approval: true, @@ -1596,8 +1417,8 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu let returned = maybe_request_codex_apps_auth_elicitation( &session, - turn_context, - manager.as_ref(), + &turn_context, + turn_context.approval_policy.value(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -1612,7 +1433,6 @@ async fn codex_apps_auth_elicitation_granular_mcp_disabled_returns_original_resu #[tokio::test] async fn codex_apps_auth_elicitation_enabled_by_default_requests_elicitation() { let (session, turn_context, rx_event) = make_session_and_context_with_rx().await; - let manager = host_owned_codex_apps_manager(&session, &turn_context).await; *session.active_turn.lock().await = Some(ActiveTurn::default()); let result = codex_apps_auth_failure_result(); let metadata = codex_apps_auth_failure_metadata(); @@ -1620,12 +1440,11 @@ async fn codex_apps_auth_elicitation_enabled_by_default_requests_elicitation() { let request_task = tokio::spawn({ let session = Arc::clone(&session); let turn_context = Arc::clone(&turn_context); - let manager = Arc::clone(&manager); async move { maybe_request_codex_apps_auth_elicitation( &session, &turn_context, - manager.as_ref(), + turn_context.approval_policy.value(), "call_123", CODEX_APPS_MCP_SERVER_NAME, Some(&metadata), @@ -2522,8 +2341,9 @@ async fn approve_mode_skips_when_annotations_do_not_require_approval() { "call-1", &invocation, &HookToolName::new("mcp__test__tool"), - Some(&metadata), - AppToolApproval::Approve, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Approve), ) .await; @@ -2598,8 +2418,9 @@ async fn guardian_mode_skips_auto_when_annotations_do_not_require_approval() { "call-guardian", &invocation, &HookToolName::new("mcp__test__tool"), - Some(&metadata), - AppToolApproval::Auto, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -2657,8 +2478,9 @@ async fn permission_request_hook_allows_mcp_tool_call() { "call-mcp-hook", &invocation, &HookToolName::new("mcp__memory__create_entities"), - Some(&metadata), - AppToolApproval::Auto, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -2712,6 +2534,11 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { tool: "create_entities".to_string(), arguments: Some(serde_json::json!({ "entities": [] })), }; + let metadata = approval_metadata( + /*connector_id*/ None, /*connector_name*/ None, + /*connector_description*/ None, /*tool_title*/ None, + /*tool_description*/ None, + ); let decision = maybe_request_mcp_tool_approval( &session, @@ -2719,8 +2546,9 @@ async fn permission_request_hook_uses_hook_tool_name_without_metadata() { "call-mcp-hook-no-metadata", &invocation, &HookToolName::new("mcp__memory__create_entities"), - /*metadata*/ None, - AppToolApproval::Auto, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -2801,8 +2629,9 @@ async fn permission_request_hook_runs_after_remembered_mcp_approval() { "call-mcp-remembered", &invocation, &HookToolName::new("mcp__memory__create_entities"), - Some(&metadata), - AppToolApproval::Auto, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -2884,8 +2713,9 @@ async fn guardian_mode_mcp_denial_returns_rationale_message() { "call-guardian-deny", &invocation, &HookToolName::new("mcp__test__tool"), - Some(&metadata), - AppToolApproval::Auto, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Auto), ) .await; @@ -2944,8 +2774,9 @@ async fn prompt_mode_waits_for_approval_when_annotations_do_not_require_approval "call-prompt", &invocation, &HookToolName::new("mcp__test__tool"), - Some(&metadata), - AppToolApproval::Prompt, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Prompt), ) .await }) @@ -3002,8 +2833,9 @@ async fn full_access_mode_skips_mcp_tool_approval_for_all_approval_modes() { "call-2", &invocation, &HookToolName::new("mcp__test__tool"), - Some(&metadata), - approval_mode, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(approval_mode), ) .await; @@ -3091,8 +2923,9 @@ async fn approve_mode_skips_guardian_in_every_permission_mode() { "call-3", &invocation, &HookToolName::new("mcp__test__tool"), - Some(&metadata), - AppToolApproval::Approve, + &metadata, + &approval_config(&turn_context), + McpToolApprovalPolicy::for_server(AppToolApproval::Approve), ) .await; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 1594b73584..ee2a0d3624 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -32,7 +32,6 @@ use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::GuardianAssessmentEvent; use codex_protocol::protocol::GuardianAssessmentStatus; use codex_protocol::protocol::InterAgentCommunication; -use codex_protocol::protocol::McpServerRefreshConfig; use codex_protocol::protocol::Op; use codex_protocol::protocol::RealtimeConversationListVoicesResponseEvent; use codex_protocol::protocol::RealtimeVoicesList; @@ -231,11 +230,6 @@ pub(super) async fn user_input_or_turn_inner( .set_responsesapi_client_metadata(responsesapi_client_metadata); } current_context.session_telemetry.user_prompt(&items); - sess.refresh_mcp_servers_if_requested( - ¤t_context, - Some(sess.mcp_elicitation_reviewer()), - ) - .await; let additional_context_input = { let mut state = sess.state.lock().await; state.additional_context.merge(additional_context) @@ -427,9 +421,8 @@ pub async fn dynamic_tool_response(sess: &Arc, id: String, response: Dy sess.notify_dynamic_tool_response(&id, response).await; } -pub async fn refresh_mcp_servers(sess: &Arc, refresh_config: McpServerRefreshConfig) { - let mut guard = sess.pending_mcp_server_refresh_config.lock().await; - *guard = Some(refresh_config); +pub fn refresh_mcp_servers(sess: &Session) { + sess.mark_mcp_runtime_dirty(); } pub async fn reload_user_config(sess: &Arc) { @@ -592,6 +585,8 @@ async fn shutdown_session_runtime(sess: &Arc) { if let Err(err) = sess.services.code_mode_service.shutdown().await { warn!("failed to shutdown code mode session: {err}"); } + let _refresh = sess.mcp_refresh_lock.acquire().await; + sess.mcp_refresh_lock.close(); sess.services.mcp_runtime.shutdown().await; sess.guardian_review_session.shutdown().await; @@ -665,8 +660,6 @@ pub async fn review( let turn_context = sess.new_default_turn_with_sub_id(sub_id.clone()).await; sess.maybe_emit_model_warnings_for_turn(turn_context.as_ref()) .await; - sess.refresh_mcp_servers_if_requested(&turn_context, Some(sess.mcp_elicitation_reviewer())) - .await; #[allow(deprecated)] match resolve_review_request(review_request, &turn_context.cwd) { Ok(resolved) => { @@ -784,8 +777,12 @@ pub(super) async fn submission_loop( dynamic_tool_response(&sess, id, response).await; false } - Op::RefreshMcpServers { config } => { - refresh_mcp_servers(&sess, config).await; + Op::RefreshMcpServers => { + refresh_mcp_servers(&sess); + false + } + Op::ReloadMcpConfig { config } => { + sess.refresh_mcp_config(config).await; false } Op::ReloadUserConfig => { diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index d11ec2ad46..aaad252f30 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -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, } +/// 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, 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) { + 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, + ) -> anyhow::Result> { + 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, turn_context: &TurnContext, - environments: &TurnEnvironmentSnapshot, selected_capability_roots: &[ResolvedSelectedCapabilityRoot], - executor_capability_discovery: Option<&ExecutorCapabilityDiscoverySnapshot>, - ) -> Arc { + ) -> Arc { 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, - ) -> Arc { - 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, - ) { - 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::>(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::( - 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::(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, ) { - 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 { - 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 { @@ -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(); } } diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index ffa553d834..03f57463ba 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -1,111 +1,166 @@ -use std::fmt; -use std::sync::Arc; +//! Thread MCP runtime projection and publication. +//! +//! This module owns the small correctness boundary between immutable session +//! inputs and the mutable [`codex_mcp::McpRuntime`]. Background scheduling +//! belongs elsewhere. -use codex_mcp::McpConfig; -use codex_mcp::McpConnectionSet; -use codex_mcp::McpRuntimeContext; +use super::session::SessionConfiguration; +use super::*; +use crate::mcp::McpRuntimeProjection; +use codex_mcp::ElicitationReviewerHandle; use codex_protocol::capabilities::SelectedCapabilityRoot; -/// MCP config, plugin availability, exact environment bindings, and manager for one request. -pub struct McpRuntimeSnapshot { - config: Arc, - plugins_available: bool, - manager: Arc, - runtime_context: McpRuntimeContext, - ready_selected_capability_roots: Vec, +pub(super) struct McpDesiredState { + pub(super) config: Arc, + pub(super) submit_id: String, + pub(super) originator: String, + pub(super) environments: TurnEnvironmentSnapshot, } -impl McpRuntimeSnapshot { - pub(crate) fn new( - config: Arc, - plugins_available: bool, - manager: Arc, - runtime_context: McpRuntimeContext, - ready_selected_capability_roots: Vec, - ) -> Self { - Self { - config, - plugins_available, - manager, - runtime_context, - ready_selected_capability_roots, +impl McpDesiredState { + pub(super) fn local_stdio_fallback_cwd(&self) -> PathBuf { + self.environments + .primary() + .and_then(|environment| environment.cwd().to_abs_path().ok()) + .map(|cwd| cwd.to_path_buf()) + .unwrap_or_else(|| self.config.cwd.to_path_buf()) + } +} + +impl Session { + pub(super) async fn latest_mcp_desired_state(&self) -> McpDesiredState { + let session_configuration = { + let state = self.state.lock().await; + state.session_configuration.clone() + }; + let environments = self.services.turn_environments.snapshot().await; + let cwd = environments + .primary() + .and_then(|environment| environment.cwd().to_abs_path().ok()) + .unwrap_or_else(|| session_configuration.cwd().clone()); + let mut config = Self::build_per_turn_config(&session_configuration, cwd); + config.permissions.approval_policy = session_configuration.approval_policy.clone(); + + McpDesiredState { + config: Arc::new(config), + submit_id: self.next_internal_sub_id(), + originator: session_configuration.originator.clone(), + environments, } } - pub fn config(&self) -> &McpConfig { - self.config.as_ref() - } - - pub(crate) fn plugins_available(&self) -> bool { - self.plugins_available - } - - pub fn manager(&self) -> &McpConnectionSet { - self.manager.as_ref() - } - - pub(crate) fn manager_arc(&self) -> Arc { - Arc::clone(&self.manager) - } - - pub fn runtime_context(&self) -> &McpRuntimeContext { - &self.runtime_context - } - - pub(crate) fn ready_selected_capability_roots(&self) -> &[SelectedCapabilityRoot] { - &self.ready_selected_capability_roots - } - - #[cfg(test)] - pub(crate) fn new_uninitialized_for_test(config: &crate::config::Config) -> Arc { - use codex_exec_server::EnvironmentManager; - use codex_features::Feature; - use codex_mcp::ResolvedMcpCatalog; - use rmcp::model::ElicitationCapability; - - let mcp_config = McpConfig { - chatgpt_base_url: config.chatgpt_base_url.clone(), - apps_mcp_product_sku: config.apps_mcp_product_sku.clone(), - codex_home: config.codex_home.to_path_buf(), - mcp_oauth_credentials_store_mode: config.mcp_oauth_credentials_store_mode, - auth_keyring_backend_kind: config.auth_keyring_backend_kind(), - mcp_oauth_callback_port: config.mcp_oauth_callback_port, - mcp_oauth_callback_url: config.mcp_oauth_callback_url.clone(), - skill_mcp_dependency_install_enabled: config - .features - .enabled(Feature::SkillMcpDependencyInstall), - approval_policy: config.permissions.approval_policy.clone(), - codex_linux_sandbox_exe: config.codex_linux_sandbox_exe.clone(), - use_legacy_landlock: config.features.use_legacy_landlock(), - apps_enabled: config.features.enabled(Feature::Apps), - prefix_mcp_tool_names: config.prefix_mcp_tool_names(), - client_elicitation_capability: ElicitationCapability::default(), - mcp_server_catalog: ResolvedMcpCatalog::default(), - connector_snapshot: codex_connectors::ConnectorSnapshot::default(), + pub(super) async fn install_initial_mcp_runtime( + self: &Arc, + session_configuration: &SessionConfiguration, + mcp_projection: McpRuntimeProjection, + resolved_environments: &TurnEnvironmentSnapshot, + local_stdio_fallback_cwd: PathBuf, + ) -> anyhow::Result<()> { + let cwd = AbsolutePathBuf::from_absolute_path(local_stdio_fallback_cwd) + .unwrap_or_else(|_| session_configuration.cwd().clone()); + let mut config = Self::build_per_turn_config(session_configuration, cwd); + config.permissions.approval_policy = session_configuration.approval_policy.clone(); + let desired = McpDesiredState { + config: Arc::new(config), + submit_id: INITIAL_SUBMIT_ID.to_owned(), + originator: session_configuration.originator.clone(), + environments: resolved_environments.clone(), }; - let manager = McpConnectionSet::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ); - let runtime_context = McpRuntimeContext::new( - Arc::new(EnvironmentManager::default_for_tests()), - config.cwd.to_path_buf(), - ); - Arc::new(Self::new( - Arc::new(mcp_config), - /*plugins_available*/ false, - Arc::new(manager), - runtime_context, - Vec::new(), + self.publish_mcp_runtime( + &desired, + mcp_projection, + /*ready_selected_capability_roots*/ &[], + Some(self.mcp_elicitation_reviewer()), + ) + .instrument(info_span!( + "session_init.mcp_manager_init", + otel.name = "session_init.mcp_manager_init", )) - } -} + .await; -impl fmt::Debug for McpRuntimeSnapshot { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("McpRuntimeSnapshot") - .finish_non_exhaustive() + self.services.mcp_runtime.validate_required_servers().await + } + + #[tracing::instrument(name = "mcp.runtime.refresh", skip_all)] + pub(super) async fn publish_mcp_runtime( + &self, + desired: &McpDesiredState, + mcp_projection: McpRuntimeProjection, + ready_selected_capability_roots: &[SelectedCapabilityRoot], + elicitation_reviewer: Option, + ) { + let input = self + .build_mcp_runtime_input( + desired, + mcp_projection, + ready_selected_capability_roots, + elicitation_reviewer, + ) + .await; + self.services.mcp_runtime.replace(input).await; + } + + async fn build_mcp_runtime_input( + &self, + desired: &McpDesiredState, + mcp_projection: McpRuntimeProjection, + ready_selected_capability_roots: &[SelectedCapabilityRoot], + elicitation_reviewer: Option, + ) -> McpRuntimeInput { + let auth = self.services.auth_manager.auth().await; + let supports_openai_form_elicitation = self + .services + .supports_openai_form_elicitation + .load(std::sync::atomic::Ordering::Acquire); + let McpRuntimeProjection { + mut config, + plugins_available, + } = mcp_projection; + config.approval_policy = desired.config.permissions.approval_policy.clone(); + config.permission_profile = desired.config.permissions.effective_permission_profile(); + config.approvals_reviewer = desired.config.approvals_reviewer; + config.environment_cwds = desired + .environments + .turn_environments() + .map(|environment| { + ( + environment.environment_id.clone(), + environment.cwd().clone(), + ) + }) + .collect(); + config + .environment_cwds + .entry(codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string()) + .or_insert_with(|| PathUri::from_abs_path(&desired.config.cwd)); + let mcp_config = Arc::new(config); + let mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); + let local_stdio_fallback_cwd = desired.local_stdio_fallback_cwd(); + let runtime_context = McpRuntimeContext::new( + self.services.turn_environments.environment_manager(), + local_stdio_fallback_cwd, + ); + let codex_apps_auth_manager = + codex_mcp::host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()) + .then(|| Arc::clone(&self.services.auth_manager)); + + McpRuntimeInput { + config: mcp_config, + plugins_available, + ready_selected_capability_roots: ready_selected_capability_roots.to_vec(), + mcp_servers, + submit_id: desired.submit_id.clone(), + tx_event: Some(self.get_tx_event()), + startup_cancellation_token: CancellationToken::new(), + runtime_context, + codex_apps_tools_cache: self.services.mcp_manager.codex_apps_tools_cache(), + tool_catalog_cache: self.services.mcp_manager.tool_catalog_cache(), + codex_apps_tools_cache_key: connector_runtime_context_key(auth.as_ref()), + supports_openai_form_elicitation, + auth, + codex_apps_auth_manager, + elicitation_reviewer, + elicitation_lifecycle: Some(self.mcp_elicitation_lifecycle()), + } } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c8df79e21c..392246dade 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -53,8 +53,7 @@ use chrono::Utc; use codex_analytics::AnalyticsEventsClient; use codex_analytics::SubAgentThreadStartedInput; use codex_analytics::TurnCodexErrorFact; -use codex_config::types::AuthKeyringBackendKind; -use codex_config::types::OAuthCredentialsStoreMode; +use codex_async_utils::OrCancelExt; use codex_connectors::connector_runtime_context_key; use codex_core_skills::injection::HostSkillsCatalogInWorldState; use codex_exec_server::Environment; @@ -73,10 +72,10 @@ use codex_hooks::HooksConfig; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::auth_env_telemetry::collect_auth_env_telemetry; -use codex_mcp::McpConnectionSet; use codex_mcp::McpResourceClient; use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeContext; +use codex_mcp::McpRuntimeInput; use codex_models_manager::manager::RefreshStrategy; use codex_models_manager::manager::SharedModelsManager; use codex_network_proxy::NetworkProxy; @@ -195,7 +194,9 @@ use crate::thread_rollout_truncation::initial_history_has_prior_user_turns; use codex_config::CONFIG_TOML_FILE; use codex_config::ConfigLayerSource; use codex_config::ConfigLayerStackOrdering; +use codex_config::types::AuthKeyringBackendKind; use codex_config::types::McpServerConfig; +use codex_config::types::OAuthCredentialsStoreMode; use codex_model_provider_info::ModelProviderInfo; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; @@ -231,7 +232,6 @@ use self::handlers::submission_loop; pub(crate) use self::input_queue::InputQueueActivity; pub(crate) use self::input_queue::TurnInput; pub(crate) use self::input_queue::TurnInputQueue; -pub use self::mcp_runtime::McpRuntimeSnapshot; use self::review::spawn_review_thread; use self::session::AppServerClientMetadata; use self::session::Session; @@ -1494,11 +1494,15 @@ impl Session { let updated_permission_profile = updated.permission_profile(); let permission_profile_changed = previous_permission_profile != updated_permission_profile; + let mcp_inputs_changed = state.session_configuration.mcp_inputs_differ(&updated); if updates.environments.is_some() { self.services .turn_environments .update_selections(updated.environment_selections()); } + if mcp_inputs_changed { + self.mark_mcp_runtime_dirty(); + } state.session_configuration = updated; (previous_config, new_config, permission_profile_changed) }; @@ -1507,7 +1511,6 @@ impl Session { self.refresh_managed_network_proxy_for_current_permission_profile() .await; } - Ok(()) } @@ -1540,8 +1543,7 @@ impl Session { }) .await?; self.services - .latest_mcp_runtime() - .manager() + .mcp_runtime .set_elicitations_auto_deny(mcp_elicitations_auto_deny); Ok(()) } @@ -1608,8 +1610,17 @@ impl Session { .with_user_layer_from(&next_config.config_layer_stack); config.tool_suggest = resolve_tool_suggest_config_from_layer_stack(&config.config_layer_stack); + config.mcp_servers = next_config.mcp_servers.clone(); + config.mcp_oauth_credentials_store_mode = next_config.mcp_oauth_credentials_store_mode; + if let Err(err) = config.features.set_enabled( + Feature::SecretAuthStorage, + next_config.features.enabled(Feature::SecretAuthStorage), + ) { + warn!("failed to refresh MCP auth storage config: {err}"); + } let config = Arc::new(config); state.session_configuration.original_config_do_not_use = Arc::clone(&config); + self.mark_mcp_runtime_dirty(); let new_config = notify_config_contributors .then(|| Self::build_effective_session_config(&state.session_configuration)); (previous_config, new_config, config) @@ -1636,6 +1647,54 @@ impl Session { } } + pub(crate) async fn refresh_mcp_config(&self, next_config: McpServerRefreshConfig) { + let McpServerRefreshConfig { + mcp_servers, + mcp_oauth_credentials_store_mode, + auth_keyring_backend_kind, + } = next_config; + let mcp_servers = + match serde_json::from_value::>(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::( + 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::(auth_keyring_backend_kind) { + Ok(kind) => kind, + Err(err) => { + warn!("failed to parse MCP auth keyring backend refresh config: {err}"); + return; + } + }; + let mut state = self.state.lock().await; + let mut config = (*state.session_configuration.original_config_do_not_use).clone(); + if let Err(err) = config.mcp_servers.set(mcp_servers) { + warn!("failed to apply MCP server refresh config: {err}"); + return; + } + config.mcp_oauth_credentials_store_mode = store_mode; + if let Err(err) = config.features.set_enabled( + Feature::SecretAuthStorage, + matches!(keyring_backend_kind, AuthKeyringBackendKind::Secrets), + ) { + warn!("failed to refresh MCP auth storage config: {err}"); + } + state.session_configuration.original_config_do_not_use = Arc::new(config); + self.mark_mcp_runtime_dirty(); + } + fn emit_config_changed_contributors( &self, previous_config: Option<&Config>, @@ -2941,20 +3000,16 @@ impl Session { ) .await; let mcp = self - .mcp_runtime_for_step( - turn_context.as_ref(), - &environments, - &selected_capability_roots, - executor_capability_discovery.as_deref(), - ) - .await; + .mcp_runtime_for_step(turn_context.as_ref(), &selected_capability_roots) + .or_cancel(cancellation_token) + .await?; let (mcp_tools, tool_router) = turn::built_tools( self.as_ref(), turn_context.as_ref(), &environments, mcp.as_ref(), - cancellation_token, ) + .or_cancel(cancellation_token) .await?; Ok(Arc::new(StepContext { turn: turn_context, @@ -3227,22 +3282,10 @@ impl Session { items } - #[cfg(test)] pub(crate) async fn build_initial_context_with_world_state( &self, turn_context: &TurnContext, world_state: &WorldState, - ) -> Vec { - let mcp = self.services.latest_mcp_runtime(); - self.build_initial_context_with_world_state_and_mcp(turn_context, world_state, &mcp) - .await - } - - pub(crate) async fn build_initial_context_with_world_state_and_mcp( - &self, - turn_context: &TurnContext, - world_state: &WorldState, - mcp: &McpRuntimeSnapshot, ) -> Vec { let mut developer_sections = Vec::::with_capacity(8); let mut contextual_user_sections = Vec::::with_capacity(2); @@ -3392,9 +3435,10 @@ impl Session { if turn_context.config.features.enabled(Feature::TokenBudget) && turn_context.model_context_window().is_some() { - let mcp_result = mcp - .manager() - .call_tool( + let mcp_result = self + .services + .mcp_runtime + .latest_call_tool( "notes", "thread_hint", /*arguments*/ None, @@ -3549,11 +3593,7 @@ impl Session { }; let (window_number, window_ids) = window; let context_items = self - .build_initial_context_with_world_state_and_mcp( - turn_context, - world_state.as_ref(), - step_context.mcp.as_ref(), - ) + .build_initial_context_with_world_state(turn_context, world_state.as_ref()) .await; let turn_context_item = turn_context.to_turn_context_item(); self.replace_compacted_history( @@ -3606,11 +3646,7 @@ impl Session { // Full initial context resets the baseline; later turns persist only its changes. let (mut context_items, world_state_item) = if should_inject_full_context { let context_items = self - .build_initial_context_with_world_state_and_mcp( - turn_context, - world_state.as_ref(), - step_context.mcp.as_ref(), - ) + .build_initial_context_with_world_state(turn_context, world_state.as_ref()) .await; let snapshot = world_state.snapshot(); self.state @@ -3966,7 +4002,7 @@ impl Session { let had_active_turn = self.active_turn.lock().await.is_some(); self.abort_all_tasks(TurnAbortReason::Interrupted).await; if !had_active_turn { - self.cancel_mcp_startup().await; + self.cancel_mcp_startup(); } } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 910ade583b..33c28db6f1 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -38,7 +38,9 @@ pub(crate) struct Session { /// session. pub(super) features: ManagedFeatures, pub(super) multi_agent_version: OnceLock, - pub(super) pending_mcp_server_refresh_config: Mutex>, + pub(super) mcp_refresh_pending: std::sync::atomic::AtomicBool, + /// Serializes runtime refreshes without blocking calls that own a snapshot. + pub(super) mcp_refresh_lock: Semaphore, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, pub(crate) input_queue: InputQueue, @@ -143,6 +145,13 @@ impl SessionConfiguration { self.permission_profile_state.permission_profile().clone() } + pub(super) fn mcp_inputs_differ(&self, next: &Self) -> bool { + self.environments != next.environments + || self.approval_policy.value() != next.approval_policy.value() + || self.approvals_reviewer != next.approvals_reviewer + || self.permission_profile() != next.permission_profile() + } + fn materialized_permission_profile(&self) -> PermissionProfile { self.permission_profile() .materialize_project_roots_with_workspace_roots(&self.primary_workspace_roots()) @@ -702,8 +711,6 @@ impl Session { .and_then(|environment| environment.cwd.to_abs_path().ok()) .map(|cwd| cwd.to_path_buf()) .unwrap_or_else(|| session_configuration.cwd().to_path_buf()); - let mcp_runtime_context = - McpRuntimeContext::new(Arc::clone(&environment_manager), mcp_runtime_cwd); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; let mcp_projection = mcp_manager_for_mcp @@ -716,10 +723,7 @@ impl Session { /*executor_capability_discovery*/ None, ) .await; - let mcp_config = &mcp_projection.config; - let mcp_servers = codex_mcp::effective_mcp_servers(mcp_config, auth.as_ref()); - let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(mcp_config); - (auth, mcp_projection, mcp_servers, tool_plugin_provenance) + (auth, mcp_projection) } .instrument(info_span!( "session_init.auth_mcp", @@ -727,11 +731,8 @@ impl Session { )); // Join all independent futures. - let ( - thread_persistence_result, - state_db_ctx, - (auth, mcp_projection, mcp_servers, tool_plugin_provenance), - ) = tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); + let (thread_persistence_result, state_db_ctx, (auth, mcp_projection)) = + tokio::join!(thread_persistence_fut, state_db_fut, auth_and_mcp_fut); let mut live_thread_init = LiveThreadInitGuard::new(thread_persistence_result.map_err(|e| { @@ -808,10 +809,12 @@ impl Session { ) { post_session_configured_events.push(event); } - let auth = auth.as_ref(); - let auth_mode = auth.map(CodexAuth::auth_mode).map(TelemetryAuthMode::from); - let account_id = auth.and_then(CodexAuth::get_account_id); - let account_email = auth.and_then(CodexAuth::get_account_email); + let telemetry_auth = auth.as_ref(); + let auth_mode = telemetry_auth + .map(CodexAuth::auth_mode) + .map(TelemetryAuthMode::from); + let account_id = telemetry_auth.and_then(CodexAuth::get_account_id); + let account_email = telemetry_auth.and_then(CodexAuth::get_account_email); let originator = session_configuration.originator.clone(); let terminal_type = user_agent(); let session_model = session_configuration.collaboration_mode.model().to_string(); @@ -860,6 +863,13 @@ impl Session { )], ); + let mcp_server_names = + codex_mcp::effective_mcp_servers( + &mcp_projection.config, + auth.as_ref(), + ) + .into_keys() + .collect::>(); session_telemetry.conversation_starts( config.model_provider.name.as_str(), session_configuration.collaboration_mode.reasoning_effort(), @@ -872,7 +882,7 @@ impl Session { config .permissions .legacy_sandbox_policy(session_configuration.cwd().as_path()), - mcp_servers.keys().map(String::as_str).collect(), + mcp_server_names.iter().map(String::as_str).collect(), ); let use_zsh_fork_shell = config.features.enabled(Feature::ShellZshFork); @@ -1032,13 +1042,10 @@ impl Session { config.analytics_enabled, ) }); - let mcp_runtime = Arc::new(McpRuntime::new(Arc::new( - McpConnectionSet::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), - ))); + // Extensions need a stable thread-owned resource client before the Session exists. + let mcp_runtime = Arc::new(McpRuntime::empty( + mcp_projection.config.prefix_mcp_tool_names, + )); let session_extension_data = codex_extension_api::ExtensionData::new(session_id.to_string()); let mcp_resource_client = Arc::new(McpResourceClient::new(Arc::clone(&mcp_runtime))); @@ -1058,9 +1065,6 @@ impl Session { // Start with an empty connection set. The initialized set is // published after SessionConfigured so MCP events follow it. mcp_runtime, - mcp_runtime_snapshot: arc_swap::ArcSwapOption::empty(), - mcp_projection_lock: Mutex::new(()), - mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -1145,7 +1149,8 @@ impl Session { managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), features: config.features.clone(), multi_agent_version, - pending_mcp_server_refresh_config: Mutex::new(None), + mcp_refresh_pending: std::sync::atomic::AtomicBool::new(false), + mcp_refresh_lock: Semaphore::new(/*permits*/ 1), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: InputQueue::new(), @@ -1196,59 +1201,13 @@ impl Session { } turn_environments.start_connection_event_forwarding(tx_event.clone()); - let mcp_startup_cancellation_token = { - let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; - cancel_guard.cancel(); - let cancel_token = CancellationToken::new(); - *cancel_guard = cancel_token.clone(); - cancel_token - }; - let codex_apps_auth_manager = - codex_mcp::host_owned_codex_apps_enabled(&mcp_projection.config, auth) - .then(|| Arc::clone(&sess.services.auth_manager)); - let mcp_connection_manager = McpConnectionSet::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind(), - &session_configuration.approval_policy, - INITIAL_SUBMIT_ID.to_owned(), - Some(tx_event.clone()), - mcp_startup_cancellation_token, - session_configuration.permission_profile(), - mcp_runtime_context.clone(), - config.codex_home.to_path_buf(), - sess.services.mcp_manager.codex_apps_tools_cache(), - sess.services.mcp_manager.tool_catalog_cache(), - connector_runtime_context_key(auth), - config.prefix_mcp_tool_names(), - mcp_projection - .config - .client_elicitation_capability - .clone(), - sess.services - .supports_openai_form_elicitation - .load(std::sync::atomic::Ordering::Relaxed), - tool_plugin_provenance, - auth, - codex_apps_auth_manager, - Some(sess.mcp_elicitation_reviewer()), - Some(sess.mcp_elicitation_lifecycle()), - codex_mcp::ElicitationRequestRouter::default(), + sess.install_initial_mcp_runtime( + &session_configuration, + mcp_projection, + &resolved_environments, + mcp_runtime_cwd, ) - .instrument(info_span!( - "session_init.mcp_manager_init", - otel.name = "session_init.mcp_manager_init", - )) - .await; - sess.services - .install_mcp_runtime( - Arc::new(mcp_projection.config), - mcp_projection.plugins_available, - mcp_runtime_context, - /*ready_selected_capability_roots*/ Vec::new(), - mcp_connection_manager, - ) - .await?; + .await?; sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; let session_start_source = match &initial_history { diff --git a/codex-rs/core/src/session/step_context.rs b/codex-rs/core/src/session/step_context.rs index c9a341a220..13dde96fde 100644 --- a/codex-rs/core/src/session/step_context.rs +++ b/codex-rs/core/src/session/step_context.rs @@ -2,11 +2,11 @@ use std::sync::Arc; use crate::agents_md::LoadedAgentsMd; use crate::environment_selection::TurnEnvironmentSnapshot; -use crate::session::McpRuntimeSnapshot; use crate::session::turn_context::TurnContext; use crate::tools::router::ToolRouter; use codex_exec_server::ExecutorCapabilityDiscoverySnapshot; use codex_exec_server::ResolvedSelectedCapabilityRoot; +use codex_mcp::McpBinding; use codex_mcp::ToolInfo; /// Request-scoped state that may change between model sampling requests. @@ -17,8 +17,8 @@ pub(crate) struct StepContext { pub(crate) selected_capability_roots: Vec, /// Executor-materialized capability files shared by MCP and skills in this exact step. pub(crate) executor_capability_discovery: Option>, - /// The exact MCP config and manager used to advertise and execute tools for this step. - pub(crate) mcp: Arc, + /// The exact MCP connections, configuration, and catalog captured for this step. + pub(crate) mcp: Arc, /// The fixed MCP tool list used for this exact sampling request. pub(crate) mcp_tools: Vec, /// The finalized tool plan advertised and executed for this exact sampling request. diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0b2ea4652b..918d9b502d 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -195,6 +195,13 @@ use std::sync::Arc; use std::sync::OnceLock; use std::time::Duration as StdDuration; +pub(crate) fn mcp_config_for_test(config: &crate::config::Config) -> Arc { + Arc::new(config.to_mcp_config_with_loaded_plugins( + &codex_core_plugins::PluginLoadOutcome::default(), + std::iter::empty(), + )) +} + impl StepContext { pub(crate) fn for_test(turn: Arc) -> Arc { let environments = turn.environments.clone(); @@ -203,7 +210,9 @@ impl StepContext { environments, selected_capability_roots: Vec::new(), executor_capability_discovery: None, - mcp: crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config), + mcp: Arc::new(codex_mcp::McpBinding::empty(mcp_config_for_test( + &turn.config, + ))), mcp_tools: Vec::new(), tool_router: Arc::new(ToolRouter::from_parts( ToolRegistry::empty_for_test(), @@ -414,8 +423,7 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled() let (session, turn_context, rx) = make_session_and_context_with_rx().await; session .services - .latest_mcp_runtime() - .manager() + .mcp_runtime .set_elicitations_auto_deny(/*auto_deny*/ true); let response = session @@ -5407,16 +5415,9 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let mcp_runtime_snapshot = - crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); - let mcp_runtime = Arc::new(codex_mcp::McpRuntime::new( - mcp_runtime_snapshot.manager_arc(), - )); + let mcp_runtime = Arc::new(codex_mcp::McpRuntime::empty(config.prefix_mcp_tool_names())); let services = SessionServices { mcp_runtime, - mcp_runtime_snapshot: arc_swap::ArcSwapOption::from(Some(mcp_runtime_snapshot)), - mcp_projection_lock: Mutex::new(()), - mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -5539,7 +5540,8 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), features: config.features.clone(), multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), - pending_mcp_server_refresh_config: Mutex::new(None), + mcp_refresh_pending: std::sync::atomic::AtomicBool::new(true), + mcp_refresh_lock: Semaphore::new(/*permits*/ 1), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), @@ -6634,7 +6636,6 @@ fn op_kind_for_input_and_context_ops() { async fn user_turn_updates_approvals_reviewer() { let (session, turn_context, _rx) = make_session_and_context_with_rx().await; let config = session.get_config().await; - handlers::user_input_or_turn( &session, "sub-1".to_string(), @@ -6673,6 +6674,12 @@ async fn user_turn_updates_approvals_reviewer() { state.session_configuration.approvals_reviewer, codex_config::types::ApprovalsReviewer::AutoReview ); + assert!( + session + .mcp_refresh_pending + .load(std::sync::atomic::Ordering::Acquire), + "server elicitation authority changes must refresh MCP state" + ); } #[tokio::test] @@ -7569,16 +7576,9 @@ where /*bundled_skills_enabled*/ true, )); let network_approval = Arc::new(NetworkApprovalService::default()); - let mcp_runtime_snapshot = - crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(config.as_ref()); - let mcp_runtime = Arc::new(codex_mcp::McpRuntime::new( - mcp_runtime_snapshot.manager_arc(), - )); + let mcp_runtime = Arc::new(codex_mcp::McpRuntime::empty(config.prefix_mcp_tool_names())); let services = SessionServices { mcp_runtime, - mcp_runtime_snapshot: arc_swap::ArcSwapOption::from(Some(mcp_runtime_snapshot)), - mcp_projection_lock: Mutex::new(()), - mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -7701,7 +7701,8 @@ where managed_network_proxy_refresh_lock: Semaphore::new(/*permits*/ 1), features: config.features.clone(), multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), - pending_mcp_server_refresh_config: Mutex::new(None), + mcp_refresh_pending: std::sync::atomic::AtomicBool::new(true), + mcp_refresh_lock: Semaphore::new(/*permits*/ 1), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), @@ -7740,23 +7741,15 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( } #[tokio::test] -async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { +async fn refresh_mcp_servers_uses_latest_state_for_existing_turns() { let (session, turn_context) = make_session_and_context().await; let session = Arc::new(session); let turn_context = Arc::new(turn_context); - let old_runtime = session.services.latest_mcp_runtime(); - let step_context = session + let old_step = session .capture_step_context(Arc::clone(&turn_context), &CancellationToken::new()) .await .expect("a fresh cancellation token cannot be cancelled"); - assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); - let old_token = session.mcp_startup_cancellation_token().await; - assert!(!old_token.is_cancelled()); - let mcp_oauth_credentials_store_mode = - serde_json::to_value(OAuthCredentialsStoreMode::Auto).expect("serialize store mode"); - let auth_keyring_backend_kind = - serde_json::to_value(AuthKeyringBackendKind::Secrets).expect("serialize keyring backend"); let refreshed_mcp_servers = serde_json::from_value::>(json!({ "refreshed": { "url": "https://refreshed.example/mcp", @@ -7764,46 +7757,129 @@ async fn refresh_mcp_servers_keeps_the_previous_runtime_alive() { } })) .expect("parse refreshed MCP servers"); - let refresh_config = McpServerRefreshConfig { - mcp_servers: serde_json::to_value(&refreshed_mcp_servers) - .expect("serialize refreshed MCP servers"), - mcp_oauth_credentials_store_mode, - auth_keyring_backend_kind, - }; { - let mut guard = session.pending_mcp_server_refresh_config.lock().await; - *guard = Some(refresh_config); + let mut state = session.state.lock().await; + let mut config = (*state.session_configuration.original_config_do_not_use).clone(); + config + .mcp_servers + .set(refreshed_mcp_servers.clone()) + .expect("set refreshed MCP servers"); + config.mcp_oauth_credentials_store_mode = + codex_config::types::OAuthCredentialsStoreMode::Auto; + config + .features + .set_enabled(Feature::SecretAuthStorage, /*enabled*/ true) + .expect("enable secret auth storage"); + state.session_configuration.original_config_do_not_use = Arc::new(config); } + session.mark_mcp_runtime_dirty(); - assert!(!old_token.is_cancelled()); - assert!( - session - .pending_mcp_server_refresh_config - .lock() - .await - .is_some() - ); - - session - .refresh_mcp_servers_if_requested(&turn_context, /*elicitation_reviewer*/ None) + let next_turn = session.new_default_turn().await; + let new_step = session + .capture_step_context(next_turn, &CancellationToken::new()) + .await + .expect("a fresh cancellation token cannot be cancelled"); + let rematerialized_old = session + .mcp_runtime_for_step(&turn_context, /*selected_capability_roots*/ &[]) .await; - assert!(!old_token.is_cancelled()); + let configured_servers = codex_mcp::configured_mcp_servers(new_step.mcp.config()); + assert_eq!( + configured_servers.get("refreshed"), + refreshed_mcp_servers.get("refreshed") + ); + assert!( + !codex_mcp::configured_mcp_servers(old_step.mcp.config()).contains_key("refreshed"), + "an already-bound step must keep its captured config" + ); + assert!( + codex_mcp::configured_mcp_servers(rematerialized_old.config()).contains_key("refreshed"), + "an older turn should resolve the latest MCP state" + ); + let current = session + .services + .mcp_runtime + .current_binding() + .await + .expect("current MCP binding"); + assert!( + codex_mcp::configured_mcp_servers(current.config()).contains_key("refreshed"), + "the refreshed state should remain globally current" + ); +} + +#[tokio::test] +async fn refreshed_mcp_binding_captures_current_approval_authority() { + let (session, old_turn) = make_session_and_context().await; + let session = Arc::new(session); + let previous_policy = old_turn.approval_policy.value(); + + session + .update_settings(SessionSettingsUpdate { + approval_policy: Some(AskForApproval::Never), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + permission_profile: Some(PermissionProfile::Disabled), + ..Default::default() + }) + .await + .expect("approval settings should update"); + session.refresh_mcp_if_dirty().await; + + let binding = session + .services + .mcp_runtime + .current_binding() + .await + .expect("refreshed runtime should be available"); + let config = binding.config(); + assert_eq!( + ( + config.approval_policy.value(), + &config.permission_profile, + config.approvals_reviewer, + ), + ( + AskForApproval::Never, + &PermissionProfile::Disabled, + ApprovalsReviewer::AutoReview, + ) + ); + assert_eq!(old_turn.approval_policy.value(), previous_policy); +} + +#[tokio::test] +async fn cancelled_mcp_refresh_remains_pending() { + let (session, _turn_context) = make_session_and_context().await; + let session = Arc::new(session); + + { + let _state = session.state.lock().await; + { + let mut refresh = Box::pin(session.refresh_mcp_if_dirty()); + let mut context = std::task::Context::from_waker(futures::task::noop_waker_ref()); + assert!(std::future::Future::poll(refresh.as_mut(), &mut context).is_pending()); + assert!( + !session + .mcp_refresh_pending + .load(std::sync::atomic::Ordering::Acquire), + "the refresh should have claimed its pending invalidation" + ); + } + } + assert!( session - .pending_mcp_server_refresh_config - .lock() - .await - .is_none() + .mcp_refresh_pending + .load(std::sync::atomic::Ordering::Acquire), + "a cancelled refresh must leave the runtime dirty" ); - let new_token = session.mcp_startup_cancellation_token().await; - assert!(!new_token.is_cancelled()); - let new_runtime = session.services.latest_mcp_runtime(); - assert!(!Arc::ptr_eq(&old_runtime, &new_runtime)); - assert!(Arc::ptr_eq(&step_context.mcp, &old_runtime)); - assert_eq!( - codex_mcp::configured_mcp_servers(new_runtime.config()), - refreshed_mcp_servers + + session.refresh_mcp_if_dirty().await; + assert!( + !session + .mcp_refresh_pending + .load(std::sync::atomic::Ordering::Acquire), + "the next refresh should publish the pending runtime" ); } @@ -7821,126 +7897,6 @@ impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectPr } } -#[tokio::test] -async fn deferred_environment_roots_refresh_plugin_availability() { - struct ReadyPluginContributor; - - impl codex_extension_api::McpServerContributor for ReadyPluginContributor { - fn id(&self) -> &'static str { - "ready_plugin_test" - } - - fn contribute<'a>( - &'a self, - context: codex_extension_api::McpServerContributionContext<'a, Config>, - ) -> codex_extension_api::ExtensionFuture<'a, Vec> - { - Box::pin(async move { - let available = context - .ready_selected_capability_roots() - .is_some_and(|roots| roots.iter().any(|root| root.id == "skill-only")); - available - .then( - || codex_extension_api::McpServerContribution::SelectedPluginPackage { - plugin_id: "skill-only".to_string(), - plugin_display_name: "Skill Only".to_string(), - connector_ids: Vec::new(), - }, - ) - .into_iter() - .collect() - }) - } - } - - let (mut session, turn_context) = make_session_and_context().await; - let mut registry = codex_extension_api::ExtensionRegistryBuilder::new(); - registry.mcp_server_contributor(Arc::new(ReadyPluginContributor)); - let registry = Arc::new(registry.build()); - session.services.extensions = Arc::clone(®istry); - session.services.mcp_manager = Arc::new(McpManager::new_with_extensions( - Arc::clone(&session.services.plugins_manager), - registry, - crate::CodexAppsToolsCache::default(), - )); - let session = Arc::new(session); - session - .refresh_mcp_servers_now( - &turn_context, - &turn_context.config, - /*elicitation_reviewer*/ None, - ) - .await; - let old_runtime = session.services.latest_mcp_runtime(); - let old_manager = old_runtime.manager_arc(); - - let selected_root = codex_protocol::capabilities::SelectedCapabilityRoot { - id: "skill-only".to_string(), - location: codex_protocol::capabilities::CapabilityRootLocation::Environment { - environment_id: "executor".to_string(), - path: PathUri::from_host_native_path(turn_context.config.cwd.as_path()) - .expect("selected capability root URI"), - }, - }; - 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 local_environment = turn_context - .environments - .primary() - .expect("ready local environment"); - let environments = TurnEnvironmentSnapshot { - environments: vec![TurnEnvironmentState::Ready(TurnEnvironment::new( - "executor".to_string(), - environment, - local_environment.cwd().clone(), - local_environment.workspace_roots().to_vec(), - local_environment.shell.clone(), - ))], - }; - let resolved_roots = session - .resolve_selected_capability_roots_for_step(&environments) - .await; - assert_eq!( - resolved_roots - .iter() - .map(|root| root.selected_root().clone()) - .collect::>(), - vec![selected_root] - ); - - let new_runtime = session - .mcp_runtime_for_step( - &turn_context, - &environments, - &resolved_roots, - /*executor_capability_discovery*/ None, - ) - .await; - - assert!(!old_runtime.plugins_available()); - assert!(new_runtime.plugins_available()); - assert!(!Arc::ptr_eq(&old_runtime, &new_runtime)); - assert!(Arc::ptr_eq(&old_manager, &new_runtime.manager_arc())); -} - #[tokio::test] #[tracing_test::traced_test] async fn conflicting_ready_environment_root_ids_keep_first_location() { @@ -8021,7 +7977,7 @@ async fn conflicting_ready_environment_root_ids_keep_first_location() { } #[tokio::test] -async fn built_tools_uses_the_step_mcp_runtime() -> anyhow::Result<()> { +async fn step_context_keeps_its_mcp_runtime_for_tools() -> anyhow::Result<()> { let (session, turn_context) = make_session_and_context().await; let session = Arc::new(session); let turn_context = Arc::new(turn_context); @@ -8065,6 +8021,22 @@ async fn built_tools_uses_the_step_mcp_runtime() -> anyhow::Result<()> { ) .await; + let next_step = session + .capture_step_context(Arc::clone(&step_context.turn), &CancellationToken::new()) + .await + .expect("a fresh cancellation token cannot be cancelled"); + assert!(codex_mcp::configured_mcp_servers(next_step.mcp.config()).contains_key("newer")); + + session.mark_mcp_runtime_dirty(); + session.refresh_mcp_if_dirty().await; + let current = session + .services + .mcp_runtime + .current_binding() + .await + .expect("refreshed runtime should be available"); + assert!(codex_mcp::configured_mcp_servers(current.config()).contains_key("newer")); + let router = &step_context.tool_router; assert!( !router diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 4a0b2ea6d6..3bd3bcc2d1 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -39,7 +39,6 @@ use crate::responses_metadata::CodexResponsesMetadata; use crate::responses_metadata::CodexResponsesRequestKind; use crate::responses_retry::ResponsesStreamRequest; use crate::responses_retry::handle_retryable_response_stream_error; -use crate::session::McpRuntimeSnapshot; use crate::session::PreviousTurnSettings; use crate::session::TurnInput; use crate::session::session::Session; @@ -602,17 +601,7 @@ async fn build_skills_and_plugins( // Plugin mentions need raw MCP/app inventory even when app tools // are normally hidden so we can describe the plugin's currently // usable capabilities for this turn. - match step_context - .mcp - .manager_arc() - .list_all_tools() - .or_cancel(cancellation_token) - .await - { - Ok(mcp_tools) => mcp_tools, - Err(_) if turn_context.apps_enabled() => return None, - Err(_) => Vec::new(), - } + step_context.mcp.tools().to_vec() } else { Vec::new() }; @@ -1281,14 +1270,9 @@ pub(crate) async fn built_tools( sess: &Session, turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, - mcp: &McpRuntimeSnapshot, - cancellation_token: &CancellationToken, -) -> CodexResult<(Vec, Arc)> { - let all_mcp_tools = mcp - .manager() - .list_all_tools() - .or_cancel(cancellation_token) - .await?; + mcp: &codex_mcp::McpBinding, +) -> (Vec, Arc) { + let all_mcp_tools = mcp.tools().to_vec(); let loaded_plugins = sess .services .plugins_manager @@ -1409,7 +1393,7 @@ pub(crate) async fn built_tools( }, &sess.services.tool_search_handler_cache, )); - Ok((all_mcp_tools, tool_router)) + (all_mcp_tools, tool_router) } #[derive(Debug)] diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index 7e364e7b13..e8a077470b 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -603,6 +603,7 @@ impl Session { let mut state = self.state.lock().await; match state.session_configuration.clone().apply(&updates) { Ok(next) => { + let mcp_inputs_changed = state.session_configuration.mcp_inputs_differ(&next); let previous_permission_profile = state.session_configuration.permission_profile(); let next_permission_profile = next.permission_profile(); @@ -618,6 +619,9 @@ impl Session { .turn_environments .update_selections(next.environment_selections()); } + if mcp_inputs_changed { + self.mark_mcp_runtime_dirty(); + } state.session_configuration = next.clone(); Ok(( next, @@ -652,7 +656,6 @@ impl Session { self.refresh_managed_network_proxy_for_current_permission_profile() .await; } - Ok(self .new_turn_from_configuration( sub_id, @@ -711,13 +714,6 @@ impl Session { .and_then(|turn_environment| turn_environment.cwd().to_abs_path().ok()) .unwrap_or_else(|| session_configuration.cwd().clone()); let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); - { - let mcp_runtime = self.services.latest_mcp_runtime(); - let mcp_connection_manager = mcp_runtime.manager(); - mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy); - mcp_connection_manager - .set_permission_profile(session_configuration.permission_profile()); - } let model_info = self .services diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index cb72075be0..a01b428115 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -14,13 +15,11 @@ use crate::environment_selection::ThreadEnvironments; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; -use crate::session::McpRuntimeSnapshot; use crate::tools::code_mode::CodeModeService; use crate::tools::handlers::ToolSearchHandlerCache; use crate::tools::network_approval::NetworkApprovalService; use crate::tools::sandboxing::ApprovalStore; use crate::unified_exec::UnifiedExecProcessManager; -use anyhow::Result; use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; @@ -30,10 +29,7 @@ use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_hooks::Hooks; use codex_login::AuthManager; -use codex_mcp::McpConfig; -use codex_mcp::McpConnectionSet; use codex_mcp::McpRuntime; -use codex_mcp::McpRuntimeContext; use codex_models_manager::manager::SharedModelsManager; use codex_otel::SessionTelemetry; use codex_protocol::capabilities::SelectedCapabilityRoot; @@ -41,19 +37,12 @@ use codex_rollout::state_db::StateDbHandle; use codex_rollout_trace::ThreadTraceContext; use codex_thread_store::LiveThread; use codex_thread_store::ThreadStore; -use std::path::PathBuf; use tokio::runtime::Handle; use tokio::sync::Mutex; -use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { /// The single owner of live MCP connections for this thread. pub(crate) mcp_runtime: Arc, - /// The latest atomically published MCP config and connection snapshot. - pub(crate) mcp_runtime_snapshot: ArcSwapOption, - /// Serializes environment-driven runtime rebuilds. - pub(crate) mcp_projection_lock: Mutex<()>, - pub(crate) mcp_startup_cancellation_token: Mutex, pub(crate) unified_exec_manager: UnifiedExecProcessManager, pub(crate) elicitations: ElicitationService, #[cfg_attr(not(unix), allow(dead_code))] @@ -100,52 +89,3 @@ pub(crate) struct SessionServices { pub(crate) tool_search_handler_cache: ToolSearchHandlerCache, pub(crate) turn_environments: Arc, } - -impl SessionServices { - /// Publishes the initial connections before validating required servers so startup-time - /// elicitation can resolve through the thread runtime while validation waits. - pub(crate) async fn install_mcp_runtime( - &self, - config: Arc, - plugins_available: bool, - runtime_context: McpRuntimeContext, - ready_selected_capability_roots: Vec, - connections: McpConnectionSet, - ) -> Result<()> { - let runtime = self.publish_mcp_runtime( - config, - plugins_available, - runtime_context, - ready_selected_capability_roots, - connections, - ); - runtime.manager().validate_required_servers().await - } - - pub(crate) fn publish_mcp_runtime( - &self, - config: Arc, - plugins_available: bool, - runtime_context: McpRuntimeContext, - ready_selected_capability_roots: Vec, - connections: McpConnectionSet, - ) -> Arc { - let connections = self.mcp_runtime.replace(connections); - let runtime = Arc::new(McpRuntimeSnapshot::new( - config, - plugins_available, - connections, - runtime_context, - ready_selected_capability_roots, - )); - self.mcp_runtime_snapshot.store(Some(Arc::clone(&runtime))); - runtime - } - - pub(crate) fn latest_mcp_runtime(&self) -> Arc { - let Some(runtime) = self.mcp_runtime_snapshot.load_full() else { - unreachable!("MCP runtime must be installed before handling requests"); - }; - runtime - } -} diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index d1e31b7994..b74005f208 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -19,6 +19,7 @@ use rmcp::model::RequestId; use tokio::sync::oneshot; use crate::agent::control::AgentExecutionGuard; +use crate::mcp_tool_call::McpToolApprovalMetadata; use crate::session::TurnInputQueue; use crate::session::turn_context::TurnContext; use crate::tasks::AnySessionTask; @@ -89,6 +90,7 @@ pub(crate) struct TurnState { pending_request_permissions: HashMap, pending_user_input: HashMap>, pending_elicitations: HashMap<(String, RequestId), oneshot::Sender>, + mcp_tool_approval_metadata: HashMap, pending_dynamic_tools: HashMap>, pub(crate) pending_input: TurnInputQueue, mailbox_delivery_phase: MailboxDeliveryPhase, @@ -126,6 +128,7 @@ impl TurnState { self.pending_request_permissions.clear(); self.pending_user_input.clear(); self.pending_elicitations.clear(); + self.mcp_tool_approval_metadata.clear(); self.pending_dynamic_tools.clear(); } @@ -179,6 +182,21 @@ impl TurnState { .remove(&(server_name.to_string(), request_id.clone())) } + pub(crate) fn insert_mcp_tool_approval_metadata( + &mut self, + call_id: String, + metadata: McpToolApprovalMetadata, + ) { + self.mcp_tool_approval_metadata.insert(call_id, metadata); + } + + pub(crate) fn mcp_tool_approval_metadata( + &self, + call_id: &str, + ) -> Option { + self.mcp_tool_approval_metadata.get(call_id).cloned() + } + pub(crate) fn insert_pending_dynamic_tool( &mut self, key: String, diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index 8c7f21fd17..eb28ce3e5a 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -59,7 +59,7 @@ impl ListMcpResourceTemplatesHandler { .. } = invocation; let turn = std::sync::Arc::clone(&step_context.turn); - let manager = step_context.mcp.manager(); + let mcp = &step_context.mcp; let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -91,7 +91,7 @@ impl ListMcpResourceTemplatesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = manager + let result = mcp .list_resource_templates(&server_name, params) .await .map_err(|err| { @@ -110,7 +110,7 @@ impl ListMcpResourceTemplatesHandler { )); } - let templates = manager + let templates = mcp .list_all_resource_templates(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index f5e9bae7e3..464837bc0a 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -59,7 +59,7 @@ impl ListMcpResourcesHandler { .. } = invocation; let turn = std::sync::Arc::clone(&step_context.turn); - let manager = step_context.mcp.manager(); + let mcp = &step_context.mcp; let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -91,7 +91,7 @@ impl ListMcpResourcesHandler { let params = cursor .clone() .map(|value| PaginatedRequestParams::default().with_cursor(Some(value))); - let result = manager + let result = mcp .list_resources(&server_name, params) .await .map_err(|err| { @@ -108,7 +108,7 @@ impl ListMcpResourcesHandler { )); } - let resources = manager + let resources = mcp .list_all_resources(|server_name| { model_can_access_mcp_server(turn.as_ref(), server_name) }) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs index 5e9b85241c..2a83e5961f 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/read_mcp_resource.rs @@ -58,7 +58,7 @@ impl ReadMcpResourceHandler { .. } = invocation; let turn = std::sync::Arc::clone(&step_context.turn); - let manager = step_context.mcp.manager(); + let mcp = &step_context.mcp; let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -86,7 +86,7 @@ impl ReadMcpResourceHandler { let payload_result: Result = async { ensure_model_can_access_mcp_server(turn.as_ref(), &server)?; - let result = manager + let result = mcp .read_resource(&server, ReadResourceRequestParams::new(uri.clone())) .await .map_err(|err| { diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 3365b4683d..5afb527fcd 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -100,7 +100,7 @@ impl RequestPluginInstallHandler { .. } = invocation; let turn = Arc::clone(&step_context.turn); - let manager = step_context.mcp.manager(); + let mcp = &step_context.mcp; let arguments = match payload { ToolPayload::Function { arguments } => arguments, @@ -230,7 +230,7 @@ impl RequestPluginInstallHandler { let auth = session.services.auth_manager.auth().await; let completed = if user_confirmed { - verify_request_plugin_install_completed(&session, &turn, manager, &tool, auth.as_ref()) + verify_request_plugin_install_completed(&session, &turn, mcp, &tool, auth.as_ref()) .await } else { false @@ -347,16 +347,17 @@ fn disabled_install_request(tool: &DiscoverableTool) -> ToolSuggestDisabledTool } async fn verify_request_plugin_install_completed( - session: &crate::session::session::Session, + session: &Arc, turn: &crate::session::turn_context::TurnContext, - manager: &codex_mcp::McpConnectionSet, + mcp: &codex_mcp::McpBinding, tool: &DiscoverableTool, auth: Option<&codex_login::CodexAuth>, ) -> bool { match tool { DiscoverableTool::Connector(connector) => refresh_missing_requested_connectors( + session, turn, - manager, + mcp, auth, std::slice::from_ref(&connector.id), connector.id.as_str(), @@ -375,8 +376,9 @@ async fn verify_request_plugin_install_completed( plugin.id.as_str(), ), refresh_missing_requested_connectors( + session, turn, - manager, + mcp, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -398,8 +400,9 @@ async fn verify_request_plugin_install_completed( session.services.plugins_manager.as_ref(), ); let _ = refresh_missing_requested_connectors( + session, turn, - manager, + mcp, auth, &plugin.app_connector_ids, plugin.id.as_str(), @@ -440,8 +443,9 @@ fn is_remote_plugin_install_suggestion(plugin_id: &str) -> bool { } async fn refresh_missing_requested_connectors( + session: &Arc, turn: &crate::session::turn_context::TurnContext, - manager: &codex_mcp::McpConnectionSet, + mcp: &codex_mcp::McpBinding, auth: Option<&codex_login::CodexAuth>, expected_connector_ids: &[String], tool_id: &str, @@ -450,16 +454,16 @@ async fn refresh_missing_requested_connectors( return Some(Vec::new()); } - let mcp_tools = manager.list_all_tools().await; + let mcp_tools = mcp.tools(); let accessible_connectors = connectors::with_app_enabled_state( - connectors::accessible_connectors_from_mcp_tools(&mcp_tools), + connectors::accessible_connectors_from_mcp_tools(mcp_tools), &turn.config, ); if all_requested_connectors_picked_up(expected_connector_ids, &accessible_connectors) { return Some(accessible_connectors); } - match manager.hard_refresh_codex_apps_tools_cache().await { + match session.hard_refresh_latest_codex_apps_tools().await { Ok(mcp_tools) => { let accessible_connectors = connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(&mcp_tools), diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index dc0b118c1d..d943bbd256 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -1,6 +1,5 @@ use crate::environment_selection::TurnEnvironmentSnapshot; use crate::function_tool::FunctionCallError; -use crate::session::McpRuntimeSnapshot; use crate::session::session::Session; use crate::session::step_context::StepContext; use crate::session::turn_context::TurnContext; @@ -63,7 +62,7 @@ impl ToolRouter { pub(crate) fn from_context( turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, - mcp: &McpRuntimeSnapshot, + mcp: &codex_mcp::McpBinding, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> Self { diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index a0bae39a92..7bf871b561 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -1,7 +1,6 @@ use crate::agent::exceeds_thread_spawn_depth_limit; use crate::agent::next_thread_spawn_depth; use crate::environment_selection::TurnEnvironmentSnapshot; -use crate::session::McpRuntimeSnapshot; use crate::session::turn_context::TurnContext; use crate::tools::code_mode::default_exec_yield_time_override_ms; use crate::tools::code_mode::execute_spec::create_code_mode_tool; @@ -143,7 +142,7 @@ impl PlannedTools { struct CoreToolPlanContext<'a> { turn_context: &'a TurnContext, environments: &'a TurnEnvironmentSnapshot, - mcp: &'a McpRuntimeSnapshot, + mcp: &'a codex_mcp::McpBinding, tool_runtimes: &'a [PlannedRuntime], tool_suggest_candidates: Option<&'a crate::tools::router::ToolSuggestCandidates>, extension_tool_executors: &'a [Arc>], @@ -157,7 +156,7 @@ struct CoreToolPlanContext<'a> { pub(crate) fn build_tool_router( turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, - mcp: &McpRuntimeSnapshot, + mcp: &codex_mcp::McpBinding, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> ToolRouter { @@ -175,7 +174,7 @@ pub(crate) fn build_tool_router( fn build_tool_specs_and_registry( turn_context: &TurnContext, environments: &TurnEnvironmentSnapshot, - mcp: &McpRuntimeSnapshot, + mcp: &codex_mcp::McpBinding, params: ToolRouterParams<'_>, tool_search_handler_cache: &ToolSearchHandlerCache, ) -> (Vec, ToolRegistry) { @@ -701,7 +700,7 @@ fn unified_exec_should_include_shell_parameter( #[instrument(level = "trace", skip_all)] fn add_mcp_resource_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut PlannedTools) { - if context.mcp.manager().has_servers() { + if context.mcp.has_servers() { planned_tools.add(ListMcpResourcesHandler); planned_tools.add(ListMcpResourceTemplatesHandler); planned_tools.add(ReadMcpResourceHandler); diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 94cc92e9a8..556979d443 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/codex-rs/core/src/tools/spec_plan_tests.rs @@ -32,6 +32,7 @@ use crate::config::CurrentTimeReminderConfig; use crate::environment_selection::TurnEnvironmentState; use crate::session::step_context::StepContext; use crate::session::tests::make_session_and_context; +use crate::session::tests::mcp_config_for_test; use crate::session::turn_context::TurnContext; use crate::tools::handlers::McpHandler; use crate::tools::handlers::ToolSearchHandlerCache; @@ -681,7 +682,9 @@ async fn environment_tools_follow_the_step_context() { let environments = turn.environments.clone(); turn.environments.environments.clear(); let turn = Arc::new(turn); - let mcp = crate::session::McpRuntimeSnapshot::new_uninitialized_for_test(&turn.config); + let mcp = Arc::new(codex_mcp::McpBinding::empty(mcp_config_for_test( + &turn.config, + ))); let plan = ToolPlanProbe::from_router(ToolRouter::from_context( turn.as_ref(), diff --git a/codex-rs/core/tests/common/apps_test_server.rs b/codex-rs/core/tests/common/apps_test_server.rs index ba619dd25a..99063164da 100644 --- a/codex-rs/core/tests/common/apps_test_server.rs +++ b/codex-rs/core/tests/common/apps_test_server.rs @@ -8,6 +8,7 @@ use codex_models_manager::bundled_models_response; use serde_json::Value; use serde_json::json; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use wiremock::Mock; @@ -81,10 +82,10 @@ pub enum AppsTestToolLoading { Searchable, } -#[derive(Clone, Copy)] +#[derive(Clone)] enum AppsTestToolsListBehavior { AlwaysAvailable, - AvailableAfterInitialList, + AvailableWhen(Arc), AlwaysUnavailable, } @@ -197,12 +198,13 @@ impl AppsTestServer { )) } - pub async fn mount_with_tools_available_after_initial_list( + pub async fn mount_with_tools_available_when( server: &MockServer, + tools_available: Arc, ) -> Result { Self::mount_with_tools_list_behavior( server, - AppsTestToolsListBehavior::AvailableAfterInitialList, + AppsTestToolsListBehavior::AvailableWhen(tools_available), ) .await } @@ -435,7 +437,6 @@ async fn mount_streamable_http_json_rpc_with_startup_control( searchable, include_app_only_tool, tools_list_behavior, - tools_list_calls: AtomicUsize::new(0), initialize_attempts, remaining_initialize_failures, }) @@ -449,7 +450,6 @@ struct CodexAppsJsonRpcResponder { searchable: bool, include_app_only_tool: bool, tools_list_behavior: AppsTestToolsListBehavior, - tools_list_calls: AtomicUsize, initialize_attempts: Option>, remaining_initialize_failures: Option>, } @@ -515,10 +515,11 @@ impl Respond for CodexAppsJsonRpcResponder { } "notifications/initialized" => ResponseTemplate::new(202), "tools/list" => { - let list_index = self.tools_list_calls.fetch_add(1, Ordering::SeqCst); - let tools_available = match self.tools_list_behavior { + let tools_available = match &self.tools_list_behavior { AppsTestToolsListBehavior::AlwaysAvailable => true, - AppsTestToolsListBehavior::AvailableAfterInitialList => list_index > 0, + AppsTestToolsListBehavior::AvailableWhen(tools_available) => { + tools_available.load(Ordering::SeqCst) + } AppsTestToolsListBehavior::AlwaysUnavailable => false, }; let id = body.get("id").cloned().unwrap_or(Value::Null); diff --git a/codex-rs/core/tests/suite/mcp_auth_refresh.rs b/codex-rs/core/tests/suite/mcp_auth_refresh.rs index 8e17b31ac6..602654619b 100644 --- a/codex-rs/core/tests/suite/mcp_auth_refresh.rs +++ b/codex-rs/core/tests/suite/mcp_auth_refresh.rs @@ -2,9 +2,8 @@ use anyhow::Result; use codex_config::McpServerTransportConfig; -use codex_config::types::OAuthCredentialsStoreMode; +use codex_core::config::ConfigBuilder; use codex_core::config::Constrained; -use codex_login::AuthKeyringBackendKind; use codex_login::AuthManager; use codex_login::CodexAuth; use codex_login::ExternalAuth; @@ -13,18 +12,15 @@ use codex_login::ExternalAuthRefreshContext; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::CodexAppsToolsCache; use codex_mcp::EffectiveMcpServer; -use codex_mcp::ElicitationRequestRouter; -use codex_mcp::McpConnectionSet; +use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeContext; +use codex_mcp::McpRuntimeInput; use codex_mcp::McpToolCatalogCache; -use codex_mcp::ToolPluginProvenance; -use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use pretty_assertions::assert_eq; -use rmcp::model::ElicitationCapability; use serde_json::Value; use serde_json::json; use std::collections::HashMap; @@ -83,34 +79,34 @@ async fn hosted_plugin_runtime_ps_mcp_tool_calls_use_current_auth_manager_token( CODEX_APPS_MCP_SERVER_NAME.to_string(), EffectiveMcpServer::configured(hosted_plugin_runtime_config), )]); - let approval_policy = Constrained::allow_any(AskForApproval::Never); - let manager = McpConnectionSet::new( - &mcp_servers, - OAuthCredentialsStoreMode::default(), - AuthKeyringBackendKind::default(), - &approval_policy, - "test".to_string(), - /*tx_event*/ None, - CancellationToken::new(), - PermissionProfile::default(), - McpRuntimeContext::new( + let mut config = ConfigBuilder::default() + .codex_home(home.path().to_path_buf()) + .build() + .await?; + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::Never); + let plugins_manager = codex_core_plugins::PluginsManager::new(home.path().to_path_buf()); + let mcp_config = Arc::new(config.to_mcp_config(&plugins_manager).await); + let runtime = McpRuntime::new(McpRuntimeInput { + config: mcp_config, + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers, + submit_id: "test".to_string(), + tx_event: None, + startup_cancellation_token: CancellationToken::new(), + runtime_context: McpRuntimeContext::new( Arc::new(codex_exec_server::EnvironmentManager::without_environments()), home.path().to_path_buf(), ), - home.path().to_path_buf(), - CodexAppsToolsCache::default(), - McpToolCatalogCache::default(), - codex_mcp::codex_apps_tools_cache_key(Some(&expected_auth)), - /*prefix_mcp_tool_names*/ true, - ElicitationCapability::default(), - /*supports_openai_form_elicitation*/ false, - ToolPluginProvenance::default(), - Some(&expected_auth), - Some(Arc::clone(&auth_manager)), - /*elicitation_reviewer*/ None, - /*elicitation_lifecycle*/ None, - ElicitationRequestRouter::default(), - ) + codex_apps_tools_cache: CodexAppsToolsCache::default(), + tool_catalog_cache: McpToolCatalogCache::default(), + codex_apps_tools_cache_key: codex_mcp::codex_apps_tools_cache_key(Some(&expected_auth)), + supports_openai_form_elicitation: false, + auth: Some(expected_auth.clone()), + codex_apps_auth_manager: Some(Arc::clone(&auth_manager)), + elicitation_reviewer: None, + elicitation_lifecycle: None, + }) .await; // The model-provider test covers AuthManager reload behavior. Keep this // regression focused on core MCP wiring by updating the same shared @@ -128,8 +124,8 @@ async fn hosted_plugin_runtime_ps_mcp_tool_calls_use_current_auth_manager_token( // The manager and its static fallback were created before the auth update, // so this tool call only sees the new token if the Codex Apps provider // reads the shared AuthManager at request time. - let tool_result = manager - .call_tool( + let tool_result = runtime + .latest_call_tool( CODEX_APPS_MCP_SERVER_NAME, "calendar_create_event", Some(json!({ diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index 5a5766fe12..508b371b53 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Duration; use anyhow::Context; +use codex_config::Constrained; use codex_core::NewThread; use codex_core::StartThreadOptions; use codex_exec_server::ExecutorFileSystem; @@ -107,6 +108,11 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: let fixture = test_codex() .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) .with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::Never); + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("test config should allow disabled permissions"); let app_only_cwd_marker_file = config.cwd.join("cwd-app-only"); let barrier_file = config.cwd.join("allow-initialize"); let pid_file = config.cwd.join("mcp.pid"); @@ -242,31 +248,30 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: .await; anyhow::Ok(called_process) }); + fixture.codex.shutdown_and_wait().await?; + fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) + .await?; tokio::time::timeout(Duration::from_secs(2), async { while cached_response.requests().is_empty() { tokio::time::sleep(Duration::from_millis(10)).await; } }) .await - .context("cached MCP definitions should reach inference before initialization")?; + .context("live MCP definitions should reach inference after initialization")?; assert_definition( &cached_response, - &format!("Tools in the {NAMESPACE} namespace."), - &format!("Echo from {first_process}."), + &format!("Use the tools from {second_process}."), + &format!("Echo from {second_process}."), ); - fixture.codex.shutdown_and_wait().await?; - fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) - .await?; - let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model"); assert_eq!(cached_turn.await??, second_process); let output = cached_done_response .single_request() .function_call_output_text(app_only_call_id) .expect("app-only tool error should be returned to the model"); assert!( - output.contains(&expected_error), - "model-visible tool output should contain the live visibility error: {output}" + output.contains("is not available to the model") || output.contains("unsupported call"), + "app-only tools must be rejected before reaching the MCP server: {output}" ); let output = cached_done_response .single_request() diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index b9788907f7..44841b2f35 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -1,7 +1,5 @@ use anyhow::Result; -use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; -use codex_config::types::McpServerConfig; -use codex_config::types::McpServerTransportConfig; +use codex_config::Constrained; use codex_core::config::Config; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; @@ -10,8 +8,9 @@ use codex_extension_api::ThreadStartInput; use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::McpResourceClient; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::McpServerRefreshConfig; use codex_protocol::protocol::Op; use codex_protocol::request_user_input::RequestUserInputAnswer; use codex_protocol::request_user_input::RequestUserInputResponse; @@ -31,8 +30,8 @@ use core_test_support::responses::namespace_child_tool; use core_test_support::responses::sse; use core_test_support::skip_if_no_network; use core_test_support::wait_for_event; -use core_test_support::wait_for_event_match; use core_test_support::wait_for_mcp_server; +use serde::Deserialize; use serde_json::Value; use serde_json::json; use std::collections::HashMap; @@ -63,17 +62,8 @@ impl ThreadLifecycleContributor for McpResourceClientCapture { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn session_resource_client_follows_published_mcp_runtime() -> Result<()> { +async fn out_of_band_resource_read_reconciles_the_published_mcp_runtime() -> Result<()> { let server = responses::start_mock_server().await; - let response = responses::mount_sse_once( - &server, - sse(vec![ - ev_response_created("resp-1"), - ev_assistant_message("msg-1", "done"), - ev_completed("resp-1"), - ]), - ) - .await; let captured_client = Arc::new(Mutex::new(None)); let mut extensions = ExtensionRegistryBuilder::::new(); @@ -91,50 +81,37 @@ async fn session_resource_client_follows_published_mcp_runtime() -> Result<()> { .expect("thread start should capture the MCP resource client"); assert!(!resource_client.has_server("refreshed").await); - let refreshed_server = McpServerConfig { - transport: McpServerTransportConfig::StreamableHttp { - url: format!("{}/mcp", server.uri()), - bearer_token_env_var: None, - http_headers: None, - env_http_headers: None, - }, - auth: Default::default(), - environment_id: DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), - enabled: true, - required: false, - supports_parallel_tool_calls: false, - disabled_reason: None, - startup_timeout_sec: Some(Duration::from_millis(100)), - tool_timeout_sec: None, - default_tools_approval_mode: None, - enabled_tools: None, - disabled_tools: None, - scopes: None, - oauth: None, - oauth_resource: None, - tools: HashMap::new(), - }; - test.codex - .submit(Op::RefreshMcpServers { - config: McpServerRefreshConfig { - mcp_servers: serde_json::to_value(HashMap::from([( - "refreshed".to_string(), - refreshed_server, - )]))?, - mcp_oauth_credentials_store_mode: serde_json::to_value( - test.config.mcp_oauth_credentials_store_mode, - )?, - auth_keyring_backend_kind: serde_json::to_value( - test.config.auth_keyring_backend_kind(), - )?, - }, - }) - .await?; - test.submit_turn("observe the refreshed MCP runtime") - .await?; + let mut refresh_config = test.config.clone(); + let user_config_path = refresh_config.codex_home.join("config.toml"); + let user_config: toml::Value = toml::from_str(&format!( + r#" +[mcp_servers.refreshed] +url = "{}/mcp" +startup_timeout_sec = 0.1 +"#, + server.uri() + ))?; + let refreshed_servers = user_config + .get("mcp_servers") + .cloned() + .map(HashMap::::deserialize) + .transpose()? + .expect("test config should define MCP servers"); + refresh_config + .mcp_servers + .set(refreshed_servers) + .expect("test config should allow MCP servers"); + refresh_config.config_layer_stack = refresh_config + .config_layer_stack + .with_user_config(&user_config_path, user_config)?; + test.codex.refresh_runtime_config(refresh_config).await; + test.codex.submit(Op::RefreshMcpServers).await?; + let _ = test + .codex + .read_mcp_resource("refreshed", "test://resource") + .await; assert!(resource_client.has_server("refreshed").await); - response.single_request(); Ok(()) } @@ -206,14 +183,14 @@ async fn code_mode_only_exposes_direct_model_only_mcp_namespaces() -> Result<()> } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn apps_guidance_appears_after_background_recovery_within_a_turn() -> Result<()> { +async fn apps_guidance_appears_after_recovery_between_sampling_requests() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; let (apps_server, startup_control) = AppsTestServer::mount_with_startup_control(&server).await?; - // Initial context is rendered twice before the first request, so keep both reads unavailable. - startup_control.fail_next_initialize_attempts(/*attempts*/ 2); + // The exact step is captured once before the first request; its background retry can recover. + startup_control.fail_next_initialize_attempts(/*attempts*/ 1); let call_id = "pause-for-apps"; let response = mount_sse_sequence( &server, @@ -257,7 +234,6 @@ async fn apps_guidance_appears_after_background_recovery_within_a_turn() -> Resu .expect("test config should allow feature update"); }); let test = builder.build(&server).await?; - let mcp_runtime = test.codex.current_mcp_runtime().await; test.codex .submit(Op::UserInput { @@ -271,11 +247,33 @@ async fn apps_guidance_appears_after_background_recovery_within_a_turn() -> Resu thread_settings: Default::default(), }) .await?; - let request = wait_for_event_match(&test.codex, |event| match event { - EventMsg::RequestUserInput(request) => Some(request.clone()), - _ => None, + let request = tokio::time::timeout(Duration::from_secs(3), async { + let mut request = None; + let mut apps_ready = false; + while request.is_none() || !apps_ready { + let event = test + .codex + .next_event() + .await + .expect("event stream should stay open"); + match event.msg { + EventMsg::RequestUserInput(next_request) => request = Some(next_request), + EventMsg::McpStartupUpdate(update) + if update.server == CODEX_APPS_MCP_SERVER_NAME + && matches!( + update.status, + codex_protocol::protocol::McpStartupStatus::Ready + ) => + { + apps_ready = true; + } + _ => {} + } + } + request.expect("request user input event") }) - .await; + .await + .expect("Apps should recover before the second sampling request"); let initial_requests = response.requests(); assert_eq!(initial_requests.len(), 1); @@ -289,16 +287,6 @@ async fn apps_guidance_appears_after_background_recovery_within_a_turn() -> Resu 0 ); - tokio::time::timeout(Duration::from_secs(3), async { - loop { - if !mcp_runtime.manager().list_all_tools().await.is_empty() { - break; - } - tokio::time::sleep(Duration::from_millis(10)).await; - } - }) - .await - .expect("Apps MCP should recover while the turn is paused"); test.codex .submit(Op::UserInputAnswer { id: request.turn_id, @@ -363,6 +351,11 @@ async fn later_follow_up_uses_background_recovered_apps_after_mid_thread_startup let mut builder = search_capable_apps_builder(apps_server.chatgpt_base_url.clone()) .with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::Never); + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("test config should allow disabled permissions"); config .features .enable(Feature::CodeModeOnly) @@ -388,24 +381,10 @@ async fn later_follow_up_uses_background_recovered_apps_after_mid_thread_startup tokio::fs::remove_dir_all(test.codex_home_path().join("cache/codex_apps_tools")).await?; startup_control.fail_next_initialize_attempts(/*attempts*/ 1); - let runtime_mcp_config = test.codex.runtime_mcp_config(&test.config).await; - let refresh_config = McpServerRefreshConfig { - mcp_servers: serde_json::to_value(codex_mcp::configured_mcp_servers(&runtime_mcp_config))?, - mcp_oauth_credentials_store_mode: serde_json::to_value( - runtime_mcp_config.mcp_oauth_credentials_store_mode, - )?, - auth_keyring_backend_kind: serde_json::to_value( - runtime_mcp_config.auth_keyring_backend_kind, - )?, - }; - test.codex - .submit(Op::RefreshMcpServers { - config: refresh_config, - }) - .await?; + test.codex.submit(Op::RefreshMcpServers).await?; test.submit_turn("use Calendar after transient Apps startup failures") .await?; - tokio::time::timeout(Duration::from_secs(1), async { + tokio::time::timeout(Duration::from_secs(5), async { while startup_control.initialize_attempts() < 3 { tokio::time::sleep(Duration::from_millis(1)).await; } diff --git a/codex-rs/core/tests/suite/request_plugin_install.rs b/codex-rs/core/tests/suite/request_plugin_install.rs index da10cd1149..c8ca9308f6 100644 --- a/codex-rs/core/tests/suite/request_plugin_install.rs +++ b/codex-rs/core/tests/suite/request_plugin_install.rs @@ -38,6 +38,9 @@ use core_test_support::wait_for_event; use core_test_support::wait_for_event_match; use serde_json::Value; use serde_json::json; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use wiremock::Mock; @@ -137,7 +140,13 @@ async fn build_test( .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) .with_config({ let apps_base_url = apps_server.chatgpt_base_url.clone(); - move |config| configure_apps_without_search_tool(config, apps_base_url.as_str()) + move |config| { + config + .permissions + .set_permission_profile(PermissionProfile::Disabled) + .expect("test config should allow disabled permissions"); + configure_apps_without_search_tool(config, apps_base_url.as_str()); + } }); builder.build(server).await } @@ -550,9 +559,11 @@ async fn remote_plugin_install_refreshes_plugin_and_apps_tool_caches() -> Result async fn run_remote_plugin_install_refresh_case(refreshed_tools: RefreshedAppsTools) -> Result<()> { let server = start_mock_server().await; + let tools_available = Arc::new(AtomicBool::new(false)); let apps_server = match refreshed_tools { RefreshedAppsTools::Available => { - AppsTestServer::mount_with_tools_available_after_initial_list(&server).await? + AppsTestServer::mount_with_tools_available_when(&server, Arc::clone(&tools_available)) + .await? } RefreshedAppsTools::Missing => AppsTestServer::mount_without_tools(&server).await?, }; @@ -589,6 +600,10 @@ async fn run_remote_plugin_install_refresh_case(refreshed_tools: RefreshedAppsTo let elicitation = start_install_turn(&test, "use Calendar").await?; mount_remote_calendar_installed_plugins(&server).await; drop(initial_remote_installed_plugins); + tools_available.store( + matches!(refreshed_tools, RefreshedAppsTools::Available), + Ordering::SeqCst, + ); resolve_install_elicitation(&test, elicitation, ElicitationAction::Accept).await?; let requests = mock.requests(); diff --git a/codex-rs/ext/skills/src/extension.rs b/codex-rs/ext/skills/src/extension.rs index e70b85fa57..ad75177d5f 100644 --- a/codex-rs/ext/skills/src/extension.rs +++ b/codex-rs/ext/skills/src/extension.rs @@ -24,6 +24,7 @@ use codex_extension_api::TurnInputContext; use codex_extension_api::TurnInputContributor; use codex_extension_api::WorldStateContributionInput; use codex_extension_api::WorldStateSectionContribution; +use codex_mcp::McpResourceClient; use codex_otel::MetricsClient; use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::Event; @@ -306,7 +307,7 @@ where include_host_skills: !host_catalog_in_world_state, include_bundled_skills: config.bundled_skills_enabled, include_orchestrator_skills: thread_state.orchestrator_skills_enabled(), - mcp_resources, + mcp_resources: mcp_resources.clone(), executor_capability_discovery: None, }; let host_query = query.clone(); @@ -366,7 +367,12 @@ where let mut injected_host_skill_prompts = InjectedHostSkillPrompts::default(); for entry in &selected_entries { match self - .read_main_prompt(entry, host_snapshot.clone(), session_store, &thread_state) + .read_main_prompt( + entry, + host_snapshot.clone(), + mcp_resources.clone(), + &thread_state, + ) .await { Ok(read_result) => { @@ -466,7 +472,7 @@ impl SkillsExtension { &self, entry: &SkillCatalogEntry, host_snapshot: Option>, - session_store: &ExtensionData, + mcp_resources: Option>, thread_state: &SkillsThreadState, ) -> Result { thread_state @@ -477,9 +483,7 @@ impl SkillsExtension { package: entry.id.clone(), resource: entry.main_prompt.clone(), host_snapshot, - mcp_resources: session_store - .get::() - .and_then(|state| state.mcp_resources.clone()), + mcp_resources, }, ) .await diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 19ddc61df5..fdfb3a4906 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -192,12 +192,10 @@ pub struct W3cTraceContext { pub tracestate: Option, } -/// Config payload for refreshing MCP servers. +/// Resolved MCP inputs to apply through a thread's submission queue. #[derive(Debug, Clone, PartialEq)] pub struct McpServerRefreshConfig { - /// Complete runtime server map after source and thread-scoped resolution. pub mcp_servers: Value, - /// OAuth credential store mode to use with this server snapshot. pub mcp_oauth_credentials_store_mode: Value, pub auth_keyring_backend_kind: Value, } @@ -641,7 +639,10 @@ pub enum Op { }, /// Request MCP servers to reinitialize and refresh cached tool lists. - RefreshMcpServers { config: McpServerRefreshConfig }, + RefreshMcpServers, + + /// Replace the thread's resolved MCP configuration before its next turn. + ReloadMcpConfig { config: McpServerRefreshConfig }, /// Reload user config layer overrides for the active session. /// @@ -882,7 +883,8 @@ impl Op { Self::UserInputAnswer { .. } => "user_input_answer", Self::RequestPermissionsResponse { .. } => "request_permissions_response", Self::DynamicToolResponse { .. } => "dynamic_tool_response", - Self::RefreshMcpServers { .. } => "refresh_mcp_servers", + Self::RefreshMcpServers => "refresh_mcp_servers", + Self::ReloadMcpConfig { .. } => "reload_mcp_config", Self::ReloadUserConfig => "reload_user_config", Self::Compact => "compact", Self::SetThreadMemoryMode { .. } => "set_thread_memory_mode",