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
@@ -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]
|
||||
async fn start_thread_keeps_internal_threads_hidden_from_normal_lookups() {
|
||||
let temp_dir = tempdir().expect("tempdir");
|
||||
|
||||
Reference in new issue
Block a user