diff --git a/codex-rs/app-server/src/effective_plugin_change.rs b/codex-rs/app-server/src/effective_plugin_change.rs index fb4b493245..d35474169f 100644 --- a/codex-rs/app-server/src/effective_plugin_change.rs +++ b/codex-rs/app-server/src/effective_plugin_change.rs @@ -33,10 +33,7 @@ pub(crate) fn effective_plugins_changed_callback( let refresh_thread_manager = Arc::clone(&thread_manager); tokio::spawn(async move { - if refresh_thread_manager.list_thread_ids().await.is_empty() { - return; - } - crate::mcp_refresh::invalidate_loaded_threads(&refresh_thread_manager).await; + refresh_thread_manager.invalidate_mcp_runtimes().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 5a31e8e4dc..cdd4670092 100644 --- a/codex-rs/app-server/src/mcp_refresh.rs +++ b/codex-rs/app-server/src/mcp_refresh.rs @@ -2,12 +2,10 @@ 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; use std::io; use std::sync::Arc; -use tracing::warn; pub(crate) async fn reload_mcp_config( thread_manager: &Arc, @@ -38,21 +36,6 @@ pub(crate) async fn reload_mcp_config( Ok(()) } -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, - Err(err) => { - warn!("failed to load thread {thread_id} for MCP refresh: {err}"); - continue; - } - }; - if let Err(err) = queue_invalidation(thread_id, thread).await { - warn!("{err}"); - } - } -} - async fn load_refresh_config( thread: &CodexThread, config_manager: &ConfigManager, @@ -79,18 +62,6 @@ async fn build_refresh_config( }) } -async fn queue_invalidation(thread_id: ThreadId, thread: Arc) -> io::Result<()> { - thread - .submit(Op::RefreshMcpServers) - .await - .map(|_| ()) - .map_err(|err| { - io::Error::other(format!( - "failed to queue MCP refresh for thread {thread_id}: {err}" - )) - }) -} - #[cfg(test)] mod tests { use super::*; @@ -155,7 +126,7 @@ mod tests { 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; + thread_manager.invalidate_mcp_runtimes().await; assert_eq!(loader.good_loads.load(Ordering::Relaxed), 0); assert_eq!(loader.bad_loads.load(Ordering::Relaxed), 0); 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 44f519522b..94ffa0a5d9 100644 --- a/codex-rs/app-server/src/request_processors/account_processor.rs +++ b/codex-rs/app-server/src/request_processors/account_processor.rs @@ -255,15 +255,12 @@ impl AccountRequestProcessor { tokio::spawn(async move { thread_manager.plugins_manager().clear_cache(); thread_manager.skills_service().clear_cache(); - if thread_manager.list_thread_ids().await.is_empty() { - return; - } 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; } + thread_manager.invalidate_mcp_runtimes().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 6b154fb572..881995bb7f 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -201,12 +201,16 @@ impl McpRequestProcessor { let notification_name = name.clone(); let notification_thread_id = thread_id; let outgoing = Arc::clone(&self.outgoing); + let thread_manager = Arc::clone(&self.thread_manager); tokio::spawn(async move { let (success, error) = match handle.wait().await { Ok(()) => (true, None), Err(err) => (false, Some(err.to_string())), }; + if success { + thread_manager.invalidate_mcp_runtimes().await; + } let notification = ServerNotification::McpServerOauthLoginCompleted( McpServerOauthLoginCompletedNotification { diff --git a/codex-rs/app-server/src/request_processors/plugins.rs b/codex-rs/app-server/src/request_processors/plugins.rs index 61e0ab7a78..f318e1d005 100644 --- a/codex-rs/app-server/src/request_processors/plugins.rs +++ b/codex-rs/app-server/src/request_processors/plugins.rs @@ -1866,6 +1866,7 @@ impl PluginRequestProcessor { let callback_url = config.mcp_oauth_callback_url.clone(); let outgoing = Arc::clone(&self.outgoing); let notification_name = name.clone(); + let thread_manager = Arc::clone(&self.thread_manager); tokio::spawn(async move { let oauth_client_id = server.oauth_client_id(); @@ -1908,6 +1909,9 @@ impl PluginRequestProcessor { Ok(()) => (true, None), Err(err) => (false, Some(err.to_string())), }; + if success { + thread_manager.invalidate_mcp_runtimes().await; + } let notification = ServerNotification::McpServerOauthLoginCompleted( McpServerOauthLoginCompletedNotification { diff --git a/codex-rs/core/src/session/mcp_prewarm.rs b/codex-rs/core/src/session/mcp_prewarm.rs index e923b6b86f..f4b45b6e21 100644 --- a/codex-rs/core/src/session/mcp_prewarm.rs +++ b/codex-rs/core/src/session/mcp_prewarm.rs @@ -6,7 +6,7 @@ use super::*; impl Session { - pub(super) fn request_mcp_runtime_refresh(&self) { + pub(crate) fn request_mcp_runtime_refresh(&self) { self.mark_mcp_runtime_dirty(); self.schedule_mcp_prewarm(); } diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 105505c1a4..2ea29c869c 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -271,6 +271,7 @@ pub(crate) struct ThreadManagerState { auth_manager: Arc, models_manager: SharedModelsManager, environment_manager: Arc, + starting_mcp_runtimes: std::sync::Mutex>>, skills_service: Arc, plugins_manager: Arc, mcp_manager: Arc, @@ -371,6 +372,7 @@ impl ThreadManager { thread_created_tx, models_manager, environment_manager, + starting_mcp_runtimes: std::sync::Mutex::new(Vec::new()), skills_service, plugins_manager, mcp_manager, @@ -503,6 +505,7 @@ impl ThreadManager { models_manager: create_model_provider(provider, Some(auth_manager.clone())) .models_manager(codex_home, /*config_model_catalog*/ None), environment_manager, + starting_mcp_runtimes: std::sync::Mutex::new(Vec::new()), skills_service, plugins_manager, mcp_manager, @@ -550,6 +553,37 @@ impl ThreadManager { self.state.environment_manager.clone() } + /// Refreshes every loaded thread and marks threads that are still being created. + pub async fn invalidate_mcp_runtimes(&self) { + self.invalidate_starting_mcp_runtimes(); + let threads = self + .state + .threads + .read() + .await + .values() + .cloned() + .collect::>(); + for thread in threads { + thread.session.request_mcp_runtime_refresh(); + } + } + + fn invalidate_starting_mcp_runtimes(&self) { + let mut starting = self + .state + .starting_mcp_runtimes + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + starting.retain(|runtime| { + let Some(runtime) = runtime.upgrade() else { + return false; + }; + runtime.store(true, Ordering::Release); + true + }); + } + pub fn default_environment_selections( &self, cwd: &AbsolutePathBuf, @@ -1605,6 +1639,15 @@ impl ThreadManagerState { supports_openai_form_elicitation: bool, user_shell_override: Option, ) -> CodexResult { + let source_changed_during_startup = Arc::new(AtomicBool::new(false)); + { + let mut starting = self + .starting_mcp_runtimes + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + starting.retain(|runtime| runtime.strong_count() != 0); + starting.push(Arc::downgrade(&source_changed_during_startup)); + } let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_)); if let InitialHistory::Resumed(resumed) = &initial_history { let mut threads = self.threads.write().await; @@ -1695,6 +1738,9 @@ impl ThreadManagerState { let new_thread = self .finalize_thread_spawn(session, io, tracked_session_source) .await?; + if source_changed_during_startup.load(Ordering::Acquire) { + new_thread.thread.session.request_mcp_runtime_refresh(); + } if is_resumed_thread { new_thread.thread.emit_thread_resume_lifecycle().await; } diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 0d29bcbe08..21a459ea43 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -479,6 +479,102 @@ async fn code_mode_session_provider_is_shared_across_threads() { ); } +#[tokio::test] +async fn mcp_invalidation_refreshes_threads_that_are_still_starting() { + struct BlockingThreadStartup { + entered: tokio::sync::Notify, + release: tokio::sync::Notify, + refreshed: tokio::sync::Notify, + projections: std::sync::atomic::AtomicUsize, + } + + impl codex_extension_api::ThreadLifecycleContributor for BlockingThreadStartup { + fn on_thread_start<'a>( + &'a self, + _input: codex_extension_api::ThreadStartInput<'a, Config>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + self.entered.notify_one(); + self.release.notified().await; + }) + } + } + + impl codex_extension_api::McpServerContributor for BlockingThreadStartup { + fn id(&self) -> &'static str { + "starting_mcp_runtime_refresh_test" + } + + fn contribute<'a>( + &'a self, + _context: codex_extension_api::McpServerContributionContext<'a, Config>, + ) -> codex_extension_api::ExtensionFuture<'a, Vec> + { + Box::pin(async move { + if self.projections.fetch_add(1, Ordering::AcqRel) != 0 { + self.refreshed.notify_one(); + } + Vec::new() + }) + } + } + + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let observer = Arc::new(BlockingThreadStartup { + entered: tokio::sync::Notify::new(), + release: tokio::sync::Notify::new(), + refreshed: tokio::sync::Notify::new(), + projections: std::sync::atomic::AtomicUsize::new(0), + }); + let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(observer.clone()); + extensions.mcp_server_contributor(observer.clone()); + let auth_manager = AuthManager::from_auth_for_testing(CodexAuth::from_api_key("dummy")); + let manager = Arc::new(ThreadManager::new( + &config, + Arc::clone(&auth_manager), + build_models_manager(&config, auth_manager), + crate::CodexAppsToolsCache::default(), + SessionSource::Exec, + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + Arc::new(extensions.build()), + Arc::new(crate::test_support::EmptyUserInstructionsProvider), + /*analytics_events_client*/ None, + thread_store_from_config(&config, /*state_db*/ None), + /*agent_graph_store*/ None, + TEST_INSTALLATION_ID.to_string(), + /*attestation_provider*/ None, + /*external_time_provider*/ None, + )); + let starting = tokio::spawn({ + let manager = Arc::clone(&manager); + async move { manager.start_thread(StartThreadOptions::new(config)).await } + }); + + tokio::time::timeout(Duration::from_secs(5), observer.entered.notified()) + .await + .expect("thread should enter its startup lifecycle"); + assert!(manager.list_thread_ids().await.is_empty()); + manager.invalidate_mcp_runtimes().await; + observer.release.notify_one(); + starting + .await + .expect("thread startup task should finish") + .expect("thread should start"); + tokio::time::timeout(Duration::from_secs(5), observer.refreshed.notified()) + .await + .expect("invalidation during startup should refresh the newly published thread"); + let shutdown = manager + .shutdown_all_threads_bounded(Duration::from_secs(5)) + .await; + assert!(shutdown.timed_out.is_empty()); +} + #[tokio::test] async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() { let temp_dir = tempdir().expect("tempdir");