mirror of
https://github.com/openai/codex.git
synced 2026-11-13 14:46:14 +00:00
Refresh MCP runtimes across thread startup (#35204)
## Why MCP configuration can change while a thread is still starting, before it appears in the thread manager's loaded-thread list. In that window, invalidation could miss the new thread and leave it with a stale MCP runtime. ## What changed - Centralize MCP runtime invalidation in `ThreadManager` so it refreshes loaded threads and marks in-progress thread startups for refresh after publication. - Use the centralized invalidation path for account and plugin changes. - Invalidate MCP runtimes after successful server and plugin OAuth logins. ## Testing - Add a regression test that pauses thread startup, invalidates MCP runtimes, and verifies that the newly published thread refreshes its MCP projection. GitOrigin-RevId: 8944f5a21816a1405e038a0d4c08b66362ea9368
This commit is contained in:
1 parent
fe8500c0a0
commit
3645a4397c
8 files changed
+154
-39
No files matched your search
@@ -33,10 +33,7 @@ pub(crate) fn effective_plugins_changed_callback(
|
|||||||
|
|
||||||
let refresh_thread_manager = Arc::clone(&thread_manager);
|
let refresh_thread_manager = Arc::clone(&thread_manager);
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
if refresh_thread_manager.list_thread_ids().await.is_empty() {
|
refresh_thread_manager.invalidate_mcp_runtimes().await;
|
||||||
return;
|
|
||||||
}
|
|
||||||
crate::mcp_refresh::invalidate_loaded_threads(&refresh_thread_manager).await;
|
|
||||||
});
|
});
|
||||||
|
|
||||||
if change.materialized_remote_plugins.is_empty() {
|
if change.materialized_remote_plugins.is_empty() {
|
||||||
|
|||||||
@@ -2,12 +2,10 @@ use crate::config_manager::ConfigManager;
|
|||||||
use codex_core::CodexThread;
|
use codex_core::CodexThread;
|
||||||
use codex_core::ThreadManager;
|
use codex_core::ThreadManager;
|
||||||
use codex_core::config::Config;
|
use codex_core::config::Config;
|
||||||
use codex_protocol::ThreadId;
|
|
||||||
use codex_protocol::protocol::McpServerRefreshConfig;
|
use codex_protocol::protocol::McpServerRefreshConfig;
|
||||||
use codex_protocol::protocol::Op;
|
use codex_protocol::protocol::Op;
|
||||||
use std::io;
|
use std::io;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tracing::warn;
|
|
||||||
|
|
||||||
pub(crate) async fn reload_mcp_config(
|
pub(crate) async fn reload_mcp_config(
|
||||||
thread_manager: &Arc<ThreadManager>,
|
thread_manager: &Arc<ThreadManager>,
|
||||||
@@ -38,21 +36,6 @@ pub(crate) async fn reload_mcp_config(
|
|||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn invalidate_loaded_threads(thread_manager: &Arc<ThreadManager>) {
|
|
||||||
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(
|
async fn load_refresh_config(
|
||||||
thread: &CodexThread,
|
thread: &CodexThread,
|
||||||
config_manager: &ConfigManager,
|
config_manager: &ConfigManager,
|
||||||
@@ -79,18 +62,6 @@ async fn build_refresh_config(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn queue_invalidation(thread_id: ThreadId, thread: Arc<CodexThread>) -> 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)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
@@ -155,7 +126,7 @@ mod tests {
|
|||||||
async fn invalidation_does_not_reload_thread_config() -> anyhow::Result<()> {
|
async fn invalidation_does_not_reload_thread_config() -> anyhow::Result<()> {
|
||||||
let (_temp_dir, thread_manager, _config_manager, loader) = refresh_test_state().await?;
|
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.good_loads.load(Ordering::Relaxed), 0);
|
||||||
assert_eq!(loader.bad_loads.load(Ordering::Relaxed), 0);
|
assert_eq!(loader.bad_loads.load(Ordering::Relaxed), 0);
|
||||||
|
|||||||
@@ -255,15 +255,12 @@ impl AccountRequestProcessor {
|
|||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
thread_manager.plugins_manager().clear_cache();
|
thread_manager.plugins_manager().clear_cache();
|
||||||
thread_manager.skills_service().clear_cache();
|
thread_manager.skills_service().clear_cache();
|
||||||
if thread_manager.list_thread_ids().await.is_empty() {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
if let Err(err) =
|
if let Err(err) =
|
||||||
crate::mcp_refresh::reload_mcp_config(&thread_manager, &config_manager).await
|
crate::mcp_refresh::reload_mcp_config(&thread_manager, &config_manager).await
|
||||||
{
|
{
|
||||||
warn!(%err, "failed to reload MCP configuration after account or plugin change");
|
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;
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -201,12 +201,16 @@ impl McpRequestProcessor {
|
|||||||
let notification_name = name.clone();
|
let notification_name = name.clone();
|
||||||
let notification_thread_id = thread_id;
|
let notification_thread_id = thread_id;
|
||||||
let outgoing = Arc::clone(&self.outgoing);
|
let outgoing = Arc::clone(&self.outgoing);
|
||||||
|
let thread_manager = Arc::clone(&self.thread_manager);
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let (success, error) = match handle.wait().await {
|
let (success, error) = match handle.wait().await {
|
||||||
Ok(()) => (true, None),
|
Ok(()) => (true, None),
|
||||||
Err(err) => (false, Some(err.to_string())),
|
Err(err) => (false, Some(err.to_string())),
|
||||||
};
|
};
|
||||||
|
if success {
|
||||||
|
thread_manager.invalidate_mcp_runtimes().await;
|
||||||
|
}
|
||||||
|
|
||||||
let notification = ServerNotification::McpServerOauthLoginCompleted(
|
let notification = ServerNotification::McpServerOauthLoginCompleted(
|
||||||
McpServerOauthLoginCompletedNotification {
|
McpServerOauthLoginCompletedNotification {
|
||||||
|
|||||||
@@ -1866,6 +1866,7 @@ impl PluginRequestProcessor {
|
|||||||
let callback_url = config.mcp_oauth_callback_url.clone();
|
let callback_url = config.mcp_oauth_callback_url.clone();
|
||||||
let outgoing = Arc::clone(&self.outgoing);
|
let outgoing = Arc::clone(&self.outgoing);
|
||||||
let notification_name = name.clone();
|
let notification_name = name.clone();
|
||||||
|
let thread_manager = Arc::clone(&self.thread_manager);
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let oauth_client_id = server.oauth_client_id();
|
let oauth_client_id = server.oauth_client_id();
|
||||||
@@ -1908,6 +1909,9 @@ impl PluginRequestProcessor {
|
|||||||
Ok(()) => (true, None),
|
Ok(()) => (true, None),
|
||||||
Err(err) => (false, Some(err.to_string())),
|
Err(err) => (false, Some(err.to_string())),
|
||||||
};
|
};
|
||||||
|
if success {
|
||||||
|
thread_manager.invalidate_mcp_runtimes().await;
|
||||||
|
}
|
||||||
|
|
||||||
let notification = ServerNotification::McpServerOauthLoginCompleted(
|
let notification = ServerNotification::McpServerOauthLoginCompleted(
|
||||||
McpServerOauthLoginCompletedNotification {
|
McpServerOauthLoginCompletedNotification {
|
||||||
|
|||||||
@@ -6,7 +6,7 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
impl Session {
|
impl Session {
|
||||||
pub(super) fn request_mcp_runtime_refresh(&self) {
|
pub(crate) fn request_mcp_runtime_refresh(&self) {
|
||||||
self.mark_mcp_runtime_dirty();
|
self.mark_mcp_runtime_dirty();
|
||||||
self.schedule_mcp_prewarm();
|
self.schedule_mcp_prewarm();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -271,6 +271,7 @@ pub(crate) struct ThreadManagerState {
|
|||||||
auth_manager: Arc<AuthManager>,
|
auth_manager: Arc<AuthManager>,
|
||||||
models_manager: SharedModelsManager,
|
models_manager: SharedModelsManager,
|
||||||
environment_manager: Arc<EnvironmentManager>,
|
environment_manager: Arc<EnvironmentManager>,
|
||||||
|
starting_mcp_runtimes: std::sync::Mutex<Vec<std::sync::Weak<AtomicBool>>>,
|
||||||
skills_service: Arc<SkillsService>,
|
skills_service: Arc<SkillsService>,
|
||||||
plugins_manager: Arc<PluginsManager>,
|
plugins_manager: Arc<PluginsManager>,
|
||||||
mcp_manager: Arc<McpManager>,
|
mcp_manager: Arc<McpManager>,
|
||||||
@@ -371,6 +372,7 @@ impl ThreadManager {
|
|||||||
thread_created_tx,
|
thread_created_tx,
|
||||||
models_manager,
|
models_manager,
|
||||||
environment_manager,
|
environment_manager,
|
||||||
|
starting_mcp_runtimes: std::sync::Mutex::new(Vec::new()),
|
||||||
skills_service,
|
skills_service,
|
||||||
plugins_manager,
|
plugins_manager,
|
||||||
mcp_manager,
|
mcp_manager,
|
||||||
@@ -503,6 +505,7 @@ impl ThreadManager {
|
|||||||
models_manager: create_model_provider(provider, Some(auth_manager.clone()))
|
models_manager: create_model_provider(provider, Some(auth_manager.clone()))
|
||||||
.models_manager(codex_home, /*config_model_catalog*/ None),
|
.models_manager(codex_home, /*config_model_catalog*/ None),
|
||||||
environment_manager,
|
environment_manager,
|
||||||
|
starting_mcp_runtimes: std::sync::Mutex::new(Vec::new()),
|
||||||
skills_service,
|
skills_service,
|
||||||
plugins_manager,
|
plugins_manager,
|
||||||
mcp_manager,
|
mcp_manager,
|
||||||
@@ -550,6 +553,37 @@ impl ThreadManager {
|
|||||||
self.state.environment_manager.clone()
|
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::<Vec<_>>();
|
||||||
|
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(
|
pub fn default_environment_selections(
|
||||||
&self,
|
&self,
|
||||||
cwd: &AbsolutePathBuf,
|
cwd: &AbsolutePathBuf,
|
||||||
@@ -1605,6 +1639,15 @@ impl ThreadManagerState {
|
|||||||
supports_openai_form_elicitation: bool,
|
supports_openai_form_elicitation: bool,
|
||||||
user_shell_override: Option<crate::shell::Shell>,
|
user_shell_override: Option<crate::shell::Shell>,
|
||||||
) -> CodexResult<NewThread> {
|
) -> CodexResult<NewThread> {
|
||||||
|
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(_));
|
let is_resumed_thread = matches!(&initial_history, InitialHistory::Resumed(_));
|
||||||
if let InitialHistory::Resumed(resumed) = &initial_history {
|
if let InitialHistory::Resumed(resumed) = &initial_history {
|
||||||
let mut threads = self.threads.write().await;
|
let mut threads = self.threads.write().await;
|
||||||
@@ -1695,6 +1738,9 @@ impl ThreadManagerState {
|
|||||||
let new_thread = self
|
let new_thread = self
|
||||||
.finalize_thread_spawn(session, io, tracked_session_source)
|
.finalize_thread_spawn(session, io, tracked_session_source)
|
||||||
.await?;
|
.await?;
|
||||||
|
if source_changed_during_startup.load(Ordering::Acquire) {
|
||||||
|
new_thread.thread.session.request_mcp_runtime_refresh();
|
||||||
|
}
|
||||||
if is_resumed_thread {
|
if is_resumed_thread {
|
||||||
new_thread.thread.emit_thread_resume_lifecycle().await;
|
new_thread.thread.emit_thread_resume_lifecycle().await;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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<Config> 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<Config> 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<codex_extension_api::McpServerContribution>>
|
||||||
|
{
|
||||||
|
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]
|
#[tokio::test]
|
||||||
async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() {
|
async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() {
|
||||||
let temp_dir = tempdir().expect("tempdir");
|
let temp_dir = tempdir().expect("tempdir");
|
||||||
|
|||||||
Reference in new issue
Block a user