From d15ebdd08d0dce1e94bdb66ca1f3e44ab27e1d26 Mon Sep 17 00:00:00 2001 From: Matthew Zeng Date: Thu, 9 Jul 2026 21:05:23 -0700 Subject: [PATCH] Move connector invalidation into runtime --- codex-rs/codex-mcp/src/connection_manager.rs | 54 +++----- .../codex-mcp/src/connection_manager_tests.rs | 40 +++++- codex-rs/codex-mcp/src/connector_runtime.rs | 118 ++++++++---------- .../codex-mcp/src/connector_runtime_tests.rs | 12 ++ codex-rs/codex-mcp/src/rmcp_client.rs | 34 ++--- 5 files changed, 136 insertions(+), 122 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 481354fe4d..94dd5f4ab0 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -200,6 +200,11 @@ impl McpConnectionManager { codex_apps_tools_cache .context(codex_home.clone(), codex_apps_tools_cache_key.clone()) }); + let routing_cancellation_token = codex_apps_tools_cache_context + .as_ref() + .map_or_else(CancellationToken::new, |context| { + context.routing_cancellation_token() + }); // If Codex Apps has an env bearer token, that is its auth path. Do // not also attach the ambient CodexAuth provider. let runtime_auth_provider = @@ -215,6 +220,7 @@ impl McpConnectionManager { store_mode, keyring_backend_kind, cancel_token.clone(), + routing_cancellation_token, tx_event.clone(), elicitation_requests.clone(), codex_apps_tools_cache_context, @@ -472,18 +478,10 @@ impl McpConnectionManager { } 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; - }; - if !async_managed_client.connector_runtime_context_is_active() { - return false; - } - - let ready = match tokio::time::timeout(timeout, async_managed_client.client()).await { + match tokio::time::timeout(timeout, self.client_by_name(server_name)).await { Ok(Ok(_)) => true, Ok(Err(_)) | Err(_) => false, - }; - ready && async_managed_client.connector_runtime_context_is_active() + } } /// Returns all tools with model-visible names normalized. @@ -595,20 +593,14 @@ impl McpConnectionManager { let clients_snapshot = &self.clients; - for (server_name, async_managed_client) in clients_snapshot + for (server_name, _) in clients_snapshot .iter() .filter(|(server_name, _)| include_server(server_name)) { let server_name = server_name.clone(); - if !async_managed_client.connector_runtime_context_is_active() { - continue; - } - let Ok(managed_client) = async_managed_client.client().await else { + let Ok(managed_client) = self.client_by_name(&server_name).await else { continue; }; - if !async_managed_client.connector_runtime_context_is_active() { - continue; - } let timeout = managed_client.tool_timeout; let client = managed_client.client.clone(); @@ -672,20 +664,14 @@ impl McpConnectionManager { let clients_snapshot = &self.clients; - for (server_name, async_managed_client) in clients_snapshot + for (server_name, _) in clients_snapshot .iter() .filter(|(server_name, _)| include_server(server_name)) { let server_name_cloned = server_name.clone(); - if !async_managed_client.connector_runtime_context_is_active() { - continue; - } - let Ok(managed_client) = async_managed_client.client().await else { + let Ok(managed_client) = self.client_by_name(server_name).await else { continue; }; - if !async_managed_client.connector_runtime_context_is_active() { - continue; - } let client = managed_client.client.clone(); let timeout = managed_client.tool_timeout; @@ -881,18 +867,12 @@ impl McpConnectionManager { } async fn client_by_name(&self, name: &str) -> Result { - let client = self - .clients + self.clients .get(name) - .ok_or_else(|| anyhow!("unknown MCP server '{name}'"))?; - if !client.connector_runtime_context_is_active() { - return Err(anyhow!("connector runtime context was discarded")); - } - let managed = client.client().await.context("failed to get client")?; - if !client.connector_runtime_context_is_active() { - return Err(anyhow!("connector runtime context was discarded")); - } - Ok(managed) + .ok_or_else(|| anyhow!("unknown MCP server '{name}'"))? + .client() + .await + .context("failed to get client") } #[cfg(test)] diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 2eef231c43..4a87021d5d 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -139,6 +139,7 @@ async fn create_ready_async_managed_client(tools: Vec) -> AsyncManaged startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), } } @@ -157,6 +158,7 @@ async fn create_blocked_codex_apps_client( let (started_tx, started_rx) = tokio::sync::oneshot::channel(); let release = Arc::new(tokio::sync::Notify::new()); let release_for_client = Arc::clone(&release); + let routing_cancellation_token = cache_context.routing_cancellation_token(); let client = async move { let _ = started_tx.send(()); release_for_client.notified().await; @@ -175,6 +177,7 @@ async fn create_blocked_codex_apps_client( startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token, }, started_rx, release, @@ -217,6 +220,7 @@ fn create_test_manager_with_failed_apps_startup( startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), }, ); manager @@ -834,6 +838,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), }, ); @@ -874,6 +879,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), }, ); @@ -975,6 +981,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), }, ); @@ -983,6 +990,22 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { assert!(timeout_result.is_err()); } +#[tokio::test] +async fn cancelling_startup_does_not_disable_a_ready_client() { + let client = create_ready_async_managed_client(vec![create_test_tool("ready", "search")]).await; + + client.cancel_token.cancel(); + + let managed = client + .client() + .await + .expect("startup cancellation should not disable a ready client"); + assert_eq!( + model_tool_names(&managed.tools), + HashSet::from([ToolName::namespaced("ready", "search")]) + ); +} + #[tokio::test] async fn shutdown_cancels_pending_tool_listing() { let cancel_token = CancellationToken::new(); @@ -1014,6 +1037,7 @@ async fn shutdown_cancels_pending_tool_listing() { startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token, + routing_cancellation_token: CancellationToken::new(), }, ); let manager = Arc::new(manager); @@ -1061,6 +1085,7 @@ async fn shutdown_continues_after_caller_is_aborted() { startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), }, ); let manager = Arc::new(manager); @@ -1114,6 +1139,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), }, ); @@ -1164,6 +1190,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token: CancellationToken::new(), }, ); @@ -1222,6 +1249,7 @@ async fn context_discard_while_checking_failed_startup_does_not_reconnect() { .boxed() .shared() }); + let routing_cancellation_token = context.routing_cancellation_token(); let client = AsyncManagedClient { client, is_codex_apps_mcp_server: true, @@ -1232,6 +1260,7 @@ async fn context_discard_while_checking_failed_startup_does_not_reconnect() { startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), + routing_cancellation_token, }; let reconnect_task = tokio::spawn({ let client = client.clone(); @@ -1549,6 +1578,7 @@ async fn discarded_codex_apps_context_hides_tools_and_rejects_calls() { )]) .await; client.is_codex_apps_mcp_server = true; + client.routing_cancellation_token = context.routing_cancellation_token(); client.codex_apps_tools_cache_context = Some(context); let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -1581,7 +1611,10 @@ async fn discarded_codex_apps_context_hides_tools_and_rejects_calls() { ) .await .expect_err("discarded context should reject tool calls"); - assert_eq!(error.to_string(), "connector runtime context was discarded"); + assert_eq!( + format!("{error:#}"), + "failed to get client: MCP startup cancelled" + ); } #[tokio::test] @@ -1757,7 +1790,10 @@ async fn context_discard_during_startup_rejects_codex_apps_calls() { .expect("tool call should finish") .expect("tool call task should not panic") .expect_err("discarded context should reject tool calls"); - assert_eq!(error.to_string(), "connector runtime context was discarded"); + assert_eq!( + format!("{error:#}"), + "failed to get client: MCP startup cancelled" + ); } #[test] diff --git a/codex-rs/codex-mcp/src/connector_runtime.rs b/codex-rs/codex-mcp/src/connector_runtime.rs index 688fc995a2..65a347f9c6 100644 --- a/codex-rs/codex-mcp/src/connector_runtime.rs +++ b/codex-rs/codex-mcp/src/connector_runtime.rs @@ -19,6 +19,7 @@ use codex_login::CodexAuth; use codex_protocol::mcp::McpServerInfo; use serde::Deserialize; use serde::Serialize; +use tokio_util::sync::CancellationToken; use crate::connector_runtime_persistence::load_cached_codex_apps_server_info; use crate::connector_runtime_persistence::load_cached_connector_runtime_for_identity; @@ -91,67 +92,13 @@ pub fn codex_apps_tools_cache_key(auth: Option<&CodexAuth>) -> CodexAppsToolsCac /// prevents account A state from bleeding into account B. #[derive(Clone, Default)] pub struct ConnectorRuntimeManager { - inner: Arc, -} - -#[derive(Default)] -struct ConnectorRuntimeManagerInner { - active: Mutex>, - next_activation: AtomicU64, -} - -struct ActiveConnectorRuntime { - identity: ConnectorRuntimeIdentity, - activation: u64, - entry: Arc, -} - -/// Handle to the active account/workspace connector runtime. -#[derive(Clone)] -pub(crate) struct ConnectorRuntimeContext { - manager: Arc, - activation: u64, - pub(crate) entry: Arc, + active: Arc>>>, } /// Compatibility alias for existing cache call sites. pub type CodexAppsToolsCache = ConnectorRuntimeManager; -/// Compatibility alias for existing cache call sites. -pub(crate) type CodexAppsToolsCacheContext = ConnectorRuntimeContext; - impl ConnectorRuntimeManager { - pub(crate) fn context( - &self, - codex_home: PathBuf, - key: ConnectorRuntimeContextKey, - ) -> ConnectorRuntimeContext { - let identity = ConnectorRuntimeIdentity { codex_home, key }; - let mut active = lock_unpoisoned(&self.inner.active); - if let Some(active) = active.as_ref() - && active.identity == identity - { - return ConnectorRuntimeContext { - manager: Arc::clone(&self.inner), - activation: active.activation, - entry: Arc::clone(&active.entry), - }; - } - - let activation = self.inner.next_activation.fetch_add(1, Ordering::Relaxed) + 1; - let entry = Arc::new(ConnectorRuntimeEntry::new(identity.clone())); - *active = Some(ActiveConnectorRuntime { - identity, - activation, - entry: Arc::clone(&entry), - }); - ConnectorRuntimeContext { - manager: Arc::clone(&self.inner), - activation, - entry, - } - } - pub fn current_snapshot( &self, codex_home: PathBuf, @@ -159,11 +106,48 @@ impl ConnectorRuntimeManager { ) -> Option> { self.context(codex_home, key).current_snapshot() } + + pub(crate) fn context( + &self, + codex_home: PathBuf, + key: ConnectorRuntimeContextKey, + ) -> ConnectorRuntimeContext { + let identity = ConnectorRuntimeIdentity { codex_home, key }; + let mut active = lock_unpoisoned(&self.active); + if let Some(active) = active.as_ref() + && active.identity == identity + { + return ConnectorRuntimeContext { + active: Arc::clone(&self.active), + entry: Arc::clone(active), + }; + } + + let entry = Arc::new(ConnectorRuntimeEntry::new(identity)); + if let Some(previous) = active.as_ref() { + previous.routing_cancellation_token.cancel(); + } + *active = Some(Arc::clone(&entry)); + ConnectorRuntimeContext { + active: Arc::clone(&self.active), + entry, + } + } } +/// Handle to the active account/workspace connector runtime. +#[derive(Clone)] +pub(crate) struct ConnectorRuntimeContext { + active: Arc>>>, + pub(crate) entry: Arc, +} + +/// Compatibility alias for existing cache call sites. +pub(crate) type CodexAppsToolsCacheContext = ConnectorRuntimeContext; + impl ConnectorRuntimeContext { pub fn current_snapshot(&self) -> Option> { - let active = lock_unpoisoned(&self.manager.active); + let active = lock_unpoisoned(&self.active); self.matches_active(active.as_ref()) .then(|| self.entry.current_snapshot.load_full()) .flatten() @@ -186,6 +170,11 @@ impl ConnectorRuntimeContext { self.current_snapshot().is_some() } + /// Shared by MCP clients attached to this context and cancelled before an identity switch. + pub(crate) fn routing_cancellation_token(&self) -> CancellationToken { + self.entry.routing_cancellation_token.clone() + } + pub(crate) fn begin_fetch( &self, source: CodexAppsToolsFetchSource, @@ -222,7 +211,7 @@ impl ConnectorRuntimeContext { persist: impl FnOnce(&ConnectorRuntimeContext, &McpServerInfo, &ConnectorRuntimeSnapshot), ) -> Result, ConnectorRuntimeContextDiscarded> { let publish_start = Instant::now(); - let active = lock_unpoisoned(&self.manager.active); + let active = lock_unpoisoned(&self.active); if !self.matches_active(active.as_ref()) { drop(active); emit_duration( @@ -281,7 +270,7 @@ impl ConnectorRuntimeContext { #[cfg(test)] pub(crate) fn store_current_tools_for_test(&self, tools: Vec) { - if !self.is_active() { + if !self.matches_active(lock_unpoisoned(&self.active).as_ref()) { return; } let snapshot = ConnectorRuntimeSnapshot { @@ -291,14 +280,8 @@ impl ConnectorRuntimeContext { self.entry.current_snapshot.store(Some(Arc::new(snapshot))); } - pub(crate) fn is_active(&self) -> bool { - self.matches_active(lock_unpoisoned(&self.manager.active).as_ref()) - } - - fn matches_active(&self, active: Option<&ActiveConnectorRuntime>) -> bool { - active.is_some_and(|active| { - active.activation == self.activation && Arc::ptr_eq(&active.entry, &self.entry) - }) + fn matches_active(&self, active: Option<&Arc>) -> bool { + active.is_some_and(|active| Arc::ptr_eq(active, &self.entry)) } } @@ -326,9 +309,11 @@ pub(crate) struct CodexAppsToolsFetchTicket { #[error("connector runtime context was discarded")] pub(crate) struct ConnectorRuntimeContextDiscarded; +/// All live state owned by one activated connector identity. pub(crate) struct ConnectorRuntimeEntry { pub(crate) identity: ConnectorRuntimeIdentity, pub(crate) current_snapshot: ArcSwapOption, + routing_cancellation_token: CancellationToken, next_fetch_generation: AtomicU64, last_accepted_generation: Mutex, } @@ -339,6 +324,7 @@ impl ConnectorRuntimeEntry { Self { identity, current_snapshot: ArcSwapOption::from(current_snapshot), + routing_cancellation_token: CancellationToken::new(), next_fetch_generation: AtomicU64::new(0), last_accepted_generation: Mutex::new(0), } diff --git a/codex-rs/codex-mcp/src/connector_runtime_tests.rs b/codex-rs/codex-mcp/src/connector_runtime_tests.rs index 39c6e39dbe..f179b5d1e4 100644 --- a/codex-rs/codex-mcp/src/connector_runtime_tests.rs +++ b/codex-rs/codex-mcp/src/connector_runtime_tests.rs @@ -492,6 +492,17 @@ fn activating_new_context_discards_old_context_without_state_bleed() { tools_a.clone(), ) .expect("publish context A"); + let client_a = context_a.routing_cancellation_token(); + + let _same_context_a = manager.context( + codex_home.path().to_path_buf(), + ConnectorRuntimeContextKey { + account_id: Some("account-a".to_string()), + chatgpt_user_id: Some("user-a".to_string()), + is_workspace_account: false, + }, + ); + assert!(!client_a.is_cancelled()); let context_b = manager.context( codex_home.path().to_path_buf(), @@ -502,6 +513,7 @@ fn activating_new_context_discards_old_context_without_state_bleed() { }, ); + assert!(client_a.is_cancelled()); assert!(context_a.current_snapshot().is_none()); assert!(context_b.current_snapshot().is_none()); assert!( diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 690644d927..41d2b1b072 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -381,6 +381,7 @@ pub(crate) struct AsyncManagedClient { pub(crate) startup_reconnect: Option>, pub(crate) tool_plugin_provenance: Arc, pub(crate) cancel_token: CancellationToken, + pub(crate) routing_cancellation_token: CancellationToken, } impl AsyncManagedClient { @@ -395,6 +396,7 @@ impl AsyncManagedClient { store_mode: OAuthCredentialsStoreMode, keyring_backend_kind: AuthKeyringBackendKind, cancel_token: CancellationToken, + routing_cancellation_token: CancellationToken, tx_event: Sender, elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option, @@ -466,10 +468,22 @@ impl AsyncManagedClient { startup_reconnect, tool_plugin_provenance, cancel_token, + routing_cancellation_token, } } pub(crate) async fn client(&self) -> Result { + if self.routing_cancellation_token.is_cancelled() { + return Err(StartupOutcomeError::Cancelled); + } + let outcome = self.client_unchecked().await; + if self.routing_cancellation_token.is_cancelled() { + return Err(StartupOutcomeError::Cancelled); + } + outcome + } + + async fn client_unchecked(&self) -> Result { if let Some(client) = self .startup_reconnect .as_ref() @@ -481,9 +495,6 @@ impl AsyncManagedClient { } pub(crate) async fn reconnect_failed_startup(&self) { - if !self.connector_runtime_context_is_active() { - return; - } let Some(startup_reconnect) = self.startup_reconnect.as_ref() else { return; }; @@ -491,7 +502,7 @@ impl AsyncManagedClient { return; } if matches!(self.client().await, Err(StartupOutcomeError::Failed { .. })) - && self.connector_runtime_context_is_active() + && !self.routing_cancellation_token.is_cancelled() { startup_reconnect.reconnect_in_background(); } @@ -499,7 +510,7 @@ impl AsyncManagedClient { pub(crate) async fn shutdown(&self) { self.cancel_token.cancel(); - match self.client().await { + match self.client_unchecked().await { Ok(client) => client.client.shutdown().await, Err(StartupOutcomeError::Cancelled) => {} Err(error) => { @@ -514,14 +525,6 @@ impl AsyncManagedClient { .is_some_and(CodexAppsToolsCacheContext::has_current_tools) } - pub(crate) fn connector_runtime_context_is_active(&self) -> bool { - !self.is_codex_apps_mcp_server - || self - .codex_apps_tools_cache_context - .as_ref() - .is_none_or(CodexAppsToolsCacheContext::is_active) - } - fn cached_tools(&self) -> Option> { self.codex_apps_tools_cache_context .as_ref() @@ -530,9 +533,6 @@ impl AsyncManagedClient { } pub(crate) async fn listed_tools(&self) -> Option> { - if !self.connector_runtime_context_is_active() { - return None; - } // Keep cache payloads raw; plugin provenance is resolved per-session at read time. let tools = if !self.startup_complete.load(Ordering::Acquire) && let Some(startup_tools) = self.cached_tools() @@ -544,7 +544,7 @@ impl AsyncManagedClient { Err(_) => self.cached_tools(), } }?; - if !self.connector_runtime_context_is_active() { + if self.routing_cancellation_token.is_cancelled() { return None; } Some(if self.is_codex_apps_mcp_server {