Move connector invalidation into runtime

2026-11-12 21:20:34 +00:00 · 2026-07-09 21:05:23 -07:00
parent 519c6e99f3
commit d15ebdd08d
5 changed files with 136 additions and 122 deletions
@@ -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<ManagedClient> {
-        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)]
@@ -139,6 +139,7 @@ async fn create_ready_async_managed_client(tools: Vec<ToolInfo>) -> 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]
@@ -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<ConnectorRuntimeManagerInner>,
-}
-
-#[derive(Default)]
-struct ConnectorRuntimeManagerInner {
-    active: Mutex<Option<ActiveConnectorRuntime>>,
-    next_activation: AtomicU64,
-}
-
-struct ActiveConnectorRuntime {
-    identity: ConnectorRuntimeIdentity,
-    activation: u64,
-    entry: Arc<ConnectorRuntimeEntry>,
-}
-
-/// Handle to the active account/workspace connector runtime.
-#[derive(Clone)]
-pub(crate) struct ConnectorRuntimeContext {
-    manager: Arc<ConnectorRuntimeManagerInner>,
-    activation: u64,
-    pub(crate) entry: Arc<ConnectorRuntimeEntry>,
+    active: Arc<Mutex<Option<Arc<ConnectorRuntimeEntry>>>>,
 }

 /// 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<Arc<ConnectorRuntimeSnapshot>> {
        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<Mutex<Option<Arc<ConnectorRuntimeEntry>>>>,
+    pub(crate) entry: Arc<ConnectorRuntimeEntry>,
+}
+
+/// Compatibility alias for existing cache call sites.
+pub(crate) type CodexAppsToolsCacheContext = ConnectorRuntimeContext;
+
 impl ConnectorRuntimeContext {
    pub fn current_snapshot(&self) -> Option<Arc<ConnectorRuntimeSnapshot>> {
-        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<Arc<ConnectorRuntimeSnapshot>, 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<ToolInfo>) {
-        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<ConnectorRuntimeEntry>>) -> 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<ConnectorRuntimeSnapshot>,
+    routing_cancellation_token: CancellationToken,
    next_fetch_generation: AtomicU64,
    last_accepted_generation: Mutex<u64>,
 }
@@ -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),
        }
@@ -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!(
@@ -381,6 +381,7 @@ pub(crate) struct AsyncManagedClient {
    pub(crate) startup_reconnect: Option<Arc<CodexAppsStartupReconnect>>,
    pub(crate) tool_plugin_provenance: Arc<ToolPluginProvenance>,
    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<Event>,
        elicitation_requests: ElicitationRequestManager,
        codex_apps_tools_cache_context: Option<CodexAppsToolsCacheContext>,
@@ -466,10 +468,22 @@ impl AsyncManagedClient {
            startup_reconnect,
            tool_plugin_provenance,
            cancel_token,
+            routing_cancellation_token,
        }
    }

    pub(crate) async fn client(&self) -> Result<ManagedClient, StartupOutcomeError> {
+        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<ManagedClient, StartupOutcomeError> {
        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<Vec<ToolInfo>> {
        self.codex_apps_tools_cache_context
            .as_ref()
@@ -530,9 +533,6 @@ impl AsyncManagedClient {
    }

    pub(crate) async fn listed_tools(&self) -> Option<Vec<ToolInfo>> {
-        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 {