diff --git a/codex-rs/app-server/src/mcp_refresh.rs b/codex-rs/app-server/src/mcp_refresh.rs index 51cb6dded..f8a1857fd 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -1,7 +1,6 @@ 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; @@ -22,8 +21,7 @@ pub(crate) async fn queue_strict_refresh( .get_thread(thread_id) .await .map_err(|err| io::Error::other(format!("failed to load thread {thread_id}: {err}")))?; - let config = - build_refresh_config(thread_manager, config_manager, thread.config().await).await?; + let config = build_refresh_config(thread.as_ref(), config_manager).await?; refreshes.push((thread_id, thread, config)); } for (thread_id, thread, config) in refreshes { @@ -44,15 +42,13 @@ pub(crate) async fn queue_best_effort_refresh( continue; } }; - let config = - match build_refresh_config(thread_manager, config_manager, thread.config().await).await - { - Ok(config) => config, - Err(err) => { - warn!("failed to build MCP refresh config for thread {thread_id}: {err}"); - 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 { warn!("{err}"); } @@ -60,14 +56,15 @@ pub(crate) async fn queue_best_effort_refresh( } async fn build_refresh_config( - thread_manager: &ThreadManager, + thread: &CodexThread, config_manager: &ConfigManager, - thread_config: Arc, ) -> io::Result { + let thread_config = thread.config().await; let config = config_manager .load_latest_config_for_thread(thread_config.as_ref()) .await?; - let mcp_servers = thread_manager.mcp_manager().runtime_servers(&config).await; + let mcp_config = thread.runtime_mcp_config(&config).await; + let mcp_servers = codex_mcp::configured_mcp_servers(&mcp_config); Ok(McpServerRefreshConfig { mcp_servers: serde_json::to_value(mcp_servers).map_err(io::Error::other)?, mcp_oauth_credentials_store_mode: serde_json::to_value( 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 facd4facc..855236df3 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -203,22 +203,28 @@ impl McpRequestProcessor { let request = request_id.clone(); let outgoing = Arc::clone(&self.outgoing); - let config = match params.thread_id.as_deref() { + let (config, thread) = match params.thread_id.as_deref() { Some(thread_id) => { let (_, thread) = self.load_thread(thread_id).await?; let thread_config = thread.config().await; - self.config_manager + let config = self + .config_manager .load_latest_config_for_thread(thread_config.as_ref()) .await - .map_err(|err| internal_error(format!("failed to reload config: {err}")))? + .map_err(|err| internal_error(format!("failed to reload config: {err}")))?; + (config, Some(thread)) + } + None => (self.load_latest_config(/*fallback_cwd*/ None).await?, None), + }; + let mcp_config = match thread { + Some(thread) => thread.runtime_mcp_config(&config).await, + None => { + self.thread_manager + .mcp_manager() + .runtime_config(&config) + .await } - None => self.load_latest_config(/*fallback_cwd*/ None).await?, }; - let mcp_config = self - .thread_manager - .mcp_manager() - .runtime_config(&config) - .await; let auth = self.auth_manager.auth().await; let environment_manager = self.thread_manager.environment_manager(); // This status path has no turn-selected environment. Use config cwd diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index c4cb1e8e1..3b8d972d7 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -550,6 +550,11 @@ impl CodexThread { self.codex.session.get_config().await } + /// Resolves the MCP runtime configuration using this thread's extension data. + pub async fn runtime_mcp_config(&self, config: &crate::config::Config) -> codex_mcp::McpConfig { + self.codex.session.runtime_mcp_config(config).await + } + pub fn multi_agent_version(&self) -> Option { self.codex.session.multi_agent_version() } diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 5af247b69..4ff56fb62 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -4,8 +4,10 @@ use std::sync::Arc; use crate::config::Config; use codex_config::McpServerConfig; use codex_core_plugins::PluginsManager; +use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_extension_api::McpServerContribution; +use codex_extension_api::McpServerContributionContext; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::EffectiveMcpServer; @@ -45,6 +47,24 @@ impl McpManager { /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { + self.runtime_config_with_context(config, /*thread_init*/ None) + .await + } + + pub(crate) async fn runtime_config_for_thread( + &self, + config: &Config, + thread_init: &ExtensionDataInit, + ) -> McpConfig { + self.runtime_config_with_context(config, Some(thread_init)) + .await + } + + async fn runtime_config_with_context( + &self, + config: &Config, + thread_init: Option<&ExtensionDataInit>, + ) -> McpConfig { let mut mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; let mut catalog = mcp_config.mcp_server_catalog.to_builder(); if mcp_config.apps_enabled { @@ -63,9 +83,13 @@ impl McpManager { ); } + let context = match thread_init { + Some(thread_init) => McpServerContributionContext::for_thread(config, thread_init), + None => McpServerContributionContext::global(config), + }; let mut contribution_order = 0; for contributor in self.extensions.mcp_server_contributors() { - for contribution in contributor.contribute(config).await { + for contribution in contributor.contribute(context).await { match contribution { McpServerContribution::Set { name, diff --git a/codex-rs/core/src/mcp_skill_dependencies.rs b/codex-rs/core/src/mcp_skill_dependencies.rs index 60915ccea..2c5507a01 100644 --- a/codex-rs/core/src/mcp_skill_dependencies.rs +++ b/codex-rs/core/src/mcp_skill_dependencies.rs @@ -53,11 +53,7 @@ pub(crate) async fn maybe_prompt_and_install_mcp_dependencies( return; } - let installed = sess - .services - .mcp_manager - .runtime_servers(config.as_ref()) - .await; + let installed = sess.runtime_mcp_servers(config.as_ref()).await; let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed); if missing.is_empty() { return; @@ -98,7 +94,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( } let codex_home = config.codex_home.clone(); - let installed = sess.services.mcp_manager.runtime_servers(config).await; + let installed = sess.runtime_mcp_servers(config).await; let missing = collect_missing_mcp_dependencies(mentioned_skills, &installed); if missing.is_empty() { return; @@ -201,11 +197,7 @@ pub(crate) async fn maybe_install_mcp_dependencies( warn!("failed to refresh MCP dependencies for mentioned skills: {err}"); return; } - let refresh_servers = sess - .services - .mcp_manager - .runtime_servers(&refresh_config) - .await; + let refresh_servers = sess.runtime_mcp_servers(&refresh_config).await; sess.refresh_mcp_servers_now( turn_context, refresh_servers, diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index 7f52da219..5e52e6d18 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -74,6 +74,20 @@ impl ElicitationReviewer for GuardianMcpElicitationReviewer { } impl Session { + pub(crate) async fn runtime_mcp_config(&self, config: &Config) -> McpConfig { + self.services + .mcp_manager + .runtime_config_for_thread(config, &self.services.mcp_thread_init) + .await + } + + pub(crate) async fn runtime_mcp_servers( + &self, + config: &Config, + ) -> HashMap { + codex_mcp::configured_mcp_servers(&self.runtime_mcp_config(config).await) + } + pub(crate) fn mcp_elicitation_reviewer(self: &Arc) -> ElicitationReviewerHandle { Arc::new(GuardianMcpElicitationReviewer::new(self)) } @@ -289,11 +303,7 @@ impl Session { ) { let auth = self.services.auth_manager.auth().await; let config = self.get_config().await; - let mcp_config = self - .services - .mcp_manager - .runtime_config(config.as_ref()) - .await; + let mcp_config = self.runtime_mcp_config(config.as_ref()).await; let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); let mcp_servers = effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref()); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index d82bac104..b6241408b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -323,6 +323,7 @@ use crate::unified_exec::UnifiedExecProcessManager; use crate::windows_sandbox::WindowsSandboxLevelExt; use codex_core_plugins::PluginsManager; use codex_git_utils::get_git_repo_root; +use codex_mcp::McpConfig; use codex_mcp::compute_auth_statuses; use codex_mcp::effective_mcp_servers_from_configured; use codex_mcp::host_owned_codex_apps_enabled; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index df60870a6..dc118e850 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -516,6 +516,11 @@ impl Session { } InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, }; + let mcp_thread_init = thread_extension_init.clone(); + let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( + thread_id.to_string(), + thread_extension_init, + ); // Kick off independent async setup tasks in parallel to reduce startup latency. // // - initialize thread persistence with new or resumed session info @@ -598,9 +603,12 @@ impl Session { let auth_manager_clone = Arc::clone(&auth_manager); let config_for_mcp = Arc::clone(&config); let mcp_manager_for_mcp = Arc::clone(&mcp_manager); + let mcp_thread_init_for_startup = &mcp_thread_init; let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; - let mcp_config = mcp_manager_for_mcp.runtime_config(&config_for_mcp).await; + let mcp_config = mcp_manager_for_mcp + .runtime_config_for_thread(&config_for_mcp, mcp_thread_init_for_startup) + .await; let mcp_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); let tool_plugin_provenance = codex_mcp::tool_plugin_provenance(&mcp_config); let auth_statuses = compute_auth_statuses( @@ -955,10 +963,6 @@ impl Session { session_extension_data.insert(McpResourceClient::new(Arc::clone( &mcp_connection_manager, ))); - let thread_extension_data = codex_extension_api::ExtensionData::new_with_init( - thread_id.to_string(), - thread_extension_init, - ); for contributor in extensions.thread_lifecycle_contributors() { contributor.on_thread_start(codex_extension_api::ThreadStartInput { config: config.as_ref(), @@ -1005,6 +1009,7 @@ impl Session { // TODO(jif): extract session to share between sub-agents session_extension_data, thread_extension_data, + mcp_thread_init, agent_control, network_proxy: arc_swap::ArcSwapOption::from(network_proxy.map(Arc::new)), network_proxy_audit_metadata, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 214489a18..9ebbdbe0a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4978,6 +4978,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { agent_control.session_id().to_string(), ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), + mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(), @@ -6977,6 +6978,7 @@ where agent_control.session_id().to_string(), ), thread_extension_data: codex_extension_api::ExtensionData::new(thread_id.to_string()), + mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index d8d91e85a..937e84899 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -22,6 +22,7 @@ use codex_analytics::AnalyticsEventsClient; use codex_core_plugins::PluginsManager; use codex_exec_server::EnvironmentManager; use codex_extension_api::ExtensionData; +use codex_extension_api::ExtensionDataInit; use codex_extension_api::ExtensionRegistry; use codex_hooks::Hooks; use codex_login::AuthManager; @@ -67,6 +68,7 @@ pub(crate) struct SessionServices { pub(crate) extensions: Arc>, pub(crate) session_extension_data: ExtensionData, pub(crate) thread_extension_data: ExtensionData, + pub(crate) mcp_thread_init: ExtensionDataInit, pub(crate) agent_control: AgentControl, pub(crate) network_proxy: ArcSwapOption, pub(crate) network_proxy_audit_metadata: NetworkProxyAuditMetadata, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index e76b43f16..fefa5b8ae 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -9,6 +9,8 @@ use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; use codex_extension_api::empty_extension_registry; use codex_models_manager::manager::RefreshStrategy; +use codex_protocol::capabilities::CapabilityRootLocation; +use codex_protocol::capabilities::SelectedCapabilityRoot; use codex_protocol::models::ContentItem; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::models::ResponseItem; @@ -387,11 +389,10 @@ async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { } #[tokio::test] -async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() { - struct InitialMarker(&'static str); - +async fn start_thread_seeds_extension_data_for_mcp_and_lifecycle_contributors() { struct InitialDataRecorder { - observed: Arc>>, + lifecycle_observed: Arc>>, + mcp_observed: Arc>>, } impl codex_extension_api::ThreadLifecycleContributor for InitialDataRecorder { @@ -400,17 +401,56 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() { input: codex_extension_api::ThreadStartInput<'a, Config>, ) -> codex_extension_api::ExtensionFuture<'a, ()> { Box::pin(async move { - let marker = input + let selected_root = input .thread_store - .get::() - .expect("initial extension data should be available"); - *self - .observed + .get::>() + .and_then(|roots| roots.first().cloned()) + .expect("selected root should be available"); + self.lifecycle_observed .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(( - input.thread_store.level_id().to_string(), - marker.0.to_string(), - )); + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push((input.thread_store.level_id().to_string(), selected_root.id)); + input + .thread_store + .insert(Vec::::new()); + }) + } + } + + impl codex_extension_api::McpServerContributor for InitialDataRecorder { + fn id(&self) -> &'static str { + "selected_root_test" + } + + fn contribute<'a>( + &'a self, + context: codex_extension_api::McpServerContributionContext<'a, Config>, + ) -> codex_extension_api::ExtensionFuture<'a, Vec> + { + Box::pin(async move { + let thread_init = context + .thread_init() + .expect("initial MCP resolution should be thread-scoped"); + let selected_root = thread_init + .get::>() + .and_then(|roots| roots.first().cloned()) + .expect("selected root should be available"); + self.mcp_observed + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(selected_root.id.clone()); + let mut server = codex_mcp::codex_apps_mcp_server_config( + "https://selected.invalid", + /*apps_mcp_product_sku*/ None, + ); + let CapabilityRootLocation::Environment { environment_id, .. } = + &selected_root.location; + server.environment_id = environment_id.clone(); + server.enabled = false; + vec![codex_extension_api::McpServerContribution::Set { + name: selected_root.id, + config: Box::new(server), + }] }) } } @@ -421,11 +461,15 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() { config.cwd = config.codex_home.abs(); std::fs::create_dir_all(&config.codex_home).expect("create codex home"); - let observed = Arc::new(std::sync::Mutex::new(None)); + let lifecycle_observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let mcp_observed = Arc::new(std::sync::Mutex::new(Vec::new())); + let recorder = Arc::new(InitialDataRecorder { + lifecycle_observed: Arc::clone(&lifecycle_observed), + mcp_observed: Arc::clone(&mcp_observed), + }); let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); - extensions.thread_lifecycle_contributor(Arc::new(InitialDataRecorder { - observed: Arc::clone(&observed), - })); + extensions.thread_lifecycle_contributor(recorder.clone()); + extensions.mcp_server_contributor(recorder); let manager = ThreadManager::new( &config, AuthManager::from_auth_for_testing(CodexAuth::create_dummy_chatgpt_auth_for_testing()), @@ -439,12 +483,21 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() { TEST_INSTALLATION_ID.to_string(), /*attestation_provider*/ None, ); - let mut thread_extension_init = codex_extension_api::ExtensionDataInit::new(); - thread_extension_init.insert(InitialMarker("seeded")); + let selected_root_init = |id: &str, environment_id: &str| { + let mut init = codex_extension_api::ExtensionDataInit::new(); + init.insert(vec![SelectedCapabilityRoot { + id: id.to_string(), + location: CapabilityRootLocation::Environment { + environment_id: environment_id.to_string(), + path: format!("/plugins/{id}"), + }, + }]); + init + }; - let thread = manager + let first_thread = manager .start_thread_with_options(StartThreadOptions { - config, + config: config.clone(), initial_history: InitialHistory::New, session_source: None, thread_source: None, @@ -452,17 +505,64 @@ async fn start_thread_seeds_extension_data_before_lifecycle_contributors_run() { metrics_service_name: None, parent_trace: None, environments: Vec::new(), - thread_extension_init, + thread_extension_init: selected_root_init("selected-a", "env-a"), }) .await - .expect("start thread"); + .expect("start first thread"); + let second_thread = manager + .start_thread_with_options(StartThreadOptions { + config: config.clone(), + initial_history: InitialHistory::New, + session_source: None, + thread_source: None, + dynamic_tools: Vec::new(), + metrics_service_name: None, + parent_trace: None, + environments: Vec::new(), + thread_extension_init: selected_root_init("selected-b", "env-b"), + }) + .await + .expect("start second thread"); + let first_resolved = first_thread.thread.runtime_mcp_config(&config).await; + let second_resolved = second_thread.thread.runtime_mcp_config(&config).await; assert_eq!( - observed + *lifecycle_observed .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .clone(), - Some((thread.thread_id.to_string(), "seeded".to_string())) + .unwrap_or_else(std::sync::PoisonError::into_inner), + vec![ + (first_thread.thread_id.to_string(), "selected-a".to_string()), + ( + second_thread.thread_id.to_string(), + "selected-b".to_string() + ), + ] + ); + assert_eq!( + *mcp_observed + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + vec![ + "selected-a".to_string(), + "selected-b".to_string(), + "selected-a".to_string(), + "selected-b".to_string(), + ] + ); + let selected_servers = |config: &codex_mcp::McpConfig| { + codex_mcp::configured_mcp_servers(config) + .into_iter() + .filter(|(name, _)| name.starts_with("selected-")) + .map(|(name, server)| (name, server.environment_id)) + .collect::>() + }; + assert_eq!( + selected_servers(&first_resolved), + std::collections::BTreeMap::from([("selected-a".to_string(), "env-a".to_string())]) + ); + assert_eq!( + selected_servers(&second_resolved), + std::collections::BTreeMap::from([("selected-b".to_string(), "env-b".to_string())]) ); } diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index efc22f5fd..6edaf1c02 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -19,6 +19,7 @@ mod turn_input; mod turn_lifecycle; pub use mcp::McpServerContribution; +pub use mcp::McpServerContributionContext; pub use prompt::PromptFragment; pub use prompt::PromptSlot; pub use thread_lifecycle::ThreadIdleInput; @@ -45,13 +46,19 @@ pub type ExtensionFuture<'a, T> = Pin + Send + 'a>>; /// Contributors run in registration order. Later contributions for the same /// name replace earlier ones. Implementations must contribute only names they /// own and must apply any source-specific policy before returning a server. +/// Thread-scoped resolution exposes the host-seeded thread inputs; global +/// resolution exposes none and must not imply a local fallback. Thread inputs +/// are frozen for the runtime and do not include lifecycle-contributor state. /// Plugin-owned servers and their provenance continue to be resolved by the /// plugin manager until that ownership moves into an extension explicitly. pub trait McpServerContributor: Send + Sync { /// Stable identity used for registration provenance and conflict diagnostics. fn id(&self) -> &'static str; - fn contribute<'a>(&'a self, config: &'a C) -> ExtensionFuture<'a, Vec>; + fn contribute<'a>( + &'a self, + context: McpServerContributionContext<'a, C>, + ) -> ExtensionFuture<'a, Vec>; } /// Extension contribution that adds prompt fragments during prompt assembly. @@ -69,8 +76,7 @@ pub trait ContextContributor: Send + Sync { /// extension-private thread state. Heavy dependencies belong on the extension /// value created by the host, not in these inputs. pub trait ThreadLifecycleContributor: Send + Sync { - /// Called after thread-scoped extension stores are created, before later - /// contributors can read from them. + /// Called after host startup has initialized the thread-scoped store. fn on_thread_start<'a>(&'a self, input: ThreadStartInput<'a, C>) -> ExtensionFuture<'a, ()> { Box::pin(async move { let _self = self; diff --git a/codex-rs/ext/extension-api/src/contributors/mcp.rs b/codex-rs/ext/extension-api/src/contributors/mcp.rs index 7fdedcdda..399cc0009 100644 --- a/codex-rs/ext/extension-api/src/contributors/mcp.rs +++ b/codex-rs/ext/extension-api/src/contributors/mcp.rs @@ -1,5 +1,55 @@ use codex_config::McpServerConfig; +use crate::ExtensionDataInit; + +/// Input supplied while resolving MCP server contributions. +/// +/// Thread-scoped implementations can read the immutable host-seeded inputs +/// through [`Self::thread_init`]. Implementations should not retain borrowed +/// context after contribution completes. +pub struct McpServerContributionContext<'a, C> { + /// Host configuration visible during MCP resolution. + config: &'a C, + /// Initial inputs for the active thread, when resolution is thread-scoped. + thread_init: Option<&'a ExtensionDataInit>, +} + +impl Clone for McpServerContributionContext<'_, C> { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for McpServerContributionContext<'_, C> {} + +impl<'a, C> McpServerContributionContext<'a, C> { + /// Creates context for resolution that is not associated with a running thread. + pub fn global(config: &'a C) -> Self { + Self { + config, + thread_init: None, + } + } + + /// Creates context for one active thread runtime. + pub fn for_thread(config: &'a C, thread_init: &'a ExtensionDataInit) -> Self { + Self { + config, + thread_init: Some(thread_init), + } + } + + /// Returns the host configuration visible during resolution. + pub fn config(&self) -> &'a C { + self.config + } + + /// Returns the frozen initial inputs when resolving for a running thread. + pub fn thread_init(&self) -> Option<&'a ExtensionDataInit> { + self.thread_init + } +} + /// One extension-owned overlay for the runtime MCP server configuration. #[derive(Clone, Debug)] pub enum McpServerContribution { diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 829ef295b..e646a53a5 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -36,6 +36,7 @@ pub use contributors::ConfigContributor; pub use contributors::ContextContributor; pub use contributors::ExtensionFuture; pub use contributors::McpServerContribution; +pub use contributors::McpServerContributionContext; pub use contributors::McpServerContributor; pub use contributors::PromptFragment; pub use contributors::PromptSlot; diff --git a/codex-rs/ext/extension-api/src/state.rs b/codex-rs/ext/extension-api/src/state.rs index 97152df5a..d2a55534c 100644 --- a/codex-rs/ext/extension-api/src/state.rs +++ b/codex-rs/ext/extension-api/src/state.rs @@ -9,9 +9,11 @@ type ErasedData = Arc; /// Typed values supplied before an [`ExtensionData`] scope is created. /// -/// Hosts consume this value once to seed a scope before lifecycle contributors -/// run. It does not install extensions or provide persistence. -#[derive(Debug, Default)] +/// Hosts may retain a clone when later operations must use the same initial +/// inputs. Cloning freezes the attachment map and shares each value by `Arc`; +/// values with interior mutability remain shared. This type does not install +/// extensions or provide persistence. +#[derive(Clone, Debug, Default)] pub struct ExtensionDataInit { entries: HashMap, } @@ -31,6 +33,15 @@ impl ExtensionDataInit { .insert(TypeId::of::(), Arc::new(value)) .map(downcast_data) } + + /// Returns a host-supplied initial attachment without creating a mutable scope. + pub fn get(&self) -> Option> + where + T: Any + Send + Sync, + { + let value = self.entries.get(&TypeId::of::())?.clone(); + Some(downcast_data(value)) + } } /// Typed extension-owned data attached to one host object. diff --git a/codex-rs/ext/mcp/src/lib.rs b/codex-rs/ext/mcp/src/lib.rs index 3843e0684..e1a1007d4 100644 --- a/codex-rs/ext/mcp/src/lib.rs +++ b/codex-rs/ext/mcp/src/lib.rs @@ -2,6 +2,7 @@ use codex_core::config::Config; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; +use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::hosted_plugin_runtime_mcp_server_config; @@ -15,9 +16,10 @@ impl McpServerContributor for HostedPluginRuntimeExtension { fn contribute<'a>( &'a self, - config: &'a Config, + context: McpServerContributionContext<'a, Config>, ) -> ExtensionFuture<'a, Vec> { Box::pin(async move { + let config = context.config(); let name = CODEX_APPS_MCP_SERVER_NAME.to_string(); if !config.features.enabled(codex_features::Feature::Apps) { return vec![McpServerContribution::Remove { name }]; diff --git a/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs index 4a3902630..3836234b8 100644 --- a/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs +++ b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs @@ -7,6 +7,7 @@ use codex_core::config::ConfigBuilder; use codex_core_plugins::PluginsManager; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; +use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; @@ -191,7 +192,7 @@ impl McpServerContributor for RemoveCodexApps { fn contribute<'a>( &'a self, - _config: &'a Config, + _context: McpServerContributionContext<'a, Config>, ) -> codex_extension_api::ExtensionFuture<'a, Vec> { Box::pin(async move { vec![McpServerContribution::Remove { diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 059c5692c..845cb730d 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -177,7 +177,9 @@ pub struct W3cTraceContext { /// Config payload for refreshing MCP servers. #[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, }