Encapsulate MCP refresh coordination (#35164)

## What changed

- Add `McpRefresh` to own the pending invalidation state and serialize MCP runtime refreshes.
- Keep cancellation recovery alongside the refresh state so a claimed refresh is marked pending again if publication does not complete.
- Update session initialization, shutdown, refresh paths, and tests to use the new abstraction.

GitOrigin-RevId: 915b28fa8b98a4f4c87ce8f2aa0d0ff99b57577d

codex-rs/core/src/session/handlers.rs

@@ -588,8 +588,8 @@ async fn shutdown_session_runtime(sess: &Arc<Session>) {
     }
     sess.stop_mcp_prewarm_worker().await;
     {
-        let _refresh = sess.mcp_refresh_lock.acquire().await;
-        sess.mcp_refresh_lock.close();
+        let _refresh = sess.mcp_refresh.acquire().await;
+        sess.mcp_refresh.close();
         sess.services.mcp_runtime.shutdown().await;
     }
     sess.guardian_review_session.shutdown().await;

codex-rs/core/src/session/mcp.rs

@@ -1,3 +1,4 @@
+use super::mcp_refresh::McpRefreshInvalidationGuard;
 use super::*;
 use codex_exec_server::ExecutorCapabilityDiscoveryCache;
 use codex_exec_server::ExecutorCapabilityDiscoverySnapshot;
@@ -44,21 +45,6 @@ struct GuardianMcpElicitationReviewer {
     session: std::sync::Weak<Session>,
 }
 
-/// Restores a claimed refresh when its task is cancelled before publication.
-struct McpRefreshInvalidationGuard<'a> {
-    pending: &'a std::sync::atomic::AtomicBool,
-    published: bool,
-}
-
-impl Drop for McpRefreshInvalidationGuard<'_> {
-    fn drop(&mut self) {
-        if !self.published {
-            self.pending
-                .store(true, std::sync::atomic::Ordering::Release);
-        }
-    }
-}
-
 pub(crate) struct McpServerElicitationOutcome {
     pub(crate) response: Option<ElicitationResponse>,
     pub(crate) sent: bool,
@@ -147,7 +133,7 @@ impl Session {
 
     /// Publishes changed MCP state, waiting for any refresh already in progress.
     pub(crate) async fn refresh_mcp_if_dirty(self: &Arc<Self>) {
-        let Ok(_refresh) = self.mcp_refresh_lock.acquire().await else {
+        let Ok(_refresh) = self.mcp_refresh.acquire().await else {
             error!("MCP runtime refresh semaphore closed");
             return;
         };
@@ -165,14 +151,11 @@ impl Session {
                 self.mark_mcp_runtime_dirty();
             }
 
-            if !self
-                .mcp_refresh_pending
-                .swap(false, std::sync::atomic::Ordering::AcqRel)
-            {
+            if !self.mcp_refresh.claim() {
                 return;
             }
             let mut refresh_invalidation = McpRefreshInvalidationGuard {
-                pending: &self.mcp_refresh_pending,
+                refresh: &self.mcp_refresh,
                 published: false,
             };
             let auth = self.services.auth_manager.auth().await;
@@ -211,10 +194,7 @@ impl Session {
             )
             .await;
             refresh_invalidation.published = true;
-            if !self
-                .mcp_refresh_pending
-                .load(std::sync::atomic::Ordering::Acquire)
-            {
+            if !self.mcp_refresh.is_pending() {
                 return;
             }
         }
@@ -226,7 +206,7 @@ impl Session {
     ) -> anyhow::Result<Vec<codex_mcp::ToolInfo>> {
         self.refresh_mcp_if_dirty().await;
         let _refresh = self
-            .mcp_refresh_lock
+            .mcp_refresh
             .acquire()
             .await
             .map_err(|_| anyhow::anyhow!("MCP runtime refresh semaphore closed"))?;
@@ -272,8 +252,7 @@ impl Session {
     }
 
     pub(super) fn mark_mcp_runtime_dirty(&self) {
-        self.mcp_refresh_pending
-            .store(true, std::sync::atomic::Ordering::Release);
+        self.mcp_refresh.invalidate();
     }
 
     #[tracing::instrument(name = "mcp.runtime.resolve_for_step", skip_all)]
@@ -531,7 +510,7 @@ impl Session {
         refresh_config: &Config,
         elicitation_reviewer: Option<ElicitationReviewerHandle>,
     ) {
-        let Ok(_refresh) = self.mcp_refresh_lock.acquire().await else {
+        let Ok(_refresh) = self.mcp_refresh.acquire().await else {
             error!("MCP runtime refresh semaphore closed");
             return;
         };

codex-rs/core/src/session/mcp_refresh.rs

@@ -0,0 +1,54 @@
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use tokio::sync::AcquireError;
+use tokio::sync::Semaphore;
+use tokio::sync::SemaphorePermit;
+
+/// Owns MCP invalidation and the single gate used to publish runtime updates.
+pub(super) struct McpRefresh {
+    pending: AtomicBool,
+    gate: Semaphore,
+}
+
+impl McpRefresh {
+    pub(super) fn new() -> Self {
+        Self {
+            pending: AtomicBool::new(false),
+            gate: Semaphore::new(/*permits*/ 1),
+        }
+    }
+
+    pub(super) fn invalidate(&self) {
+        self.pending.store(true, Ordering::Release);
+    }
+
+    pub(super) fn is_pending(&self) -> bool {
+        self.pending.load(Ordering::Acquire)
+    }
+
+    pub(super) fn claim(&self) -> bool {
+        self.pending.swap(false, Ordering::AcqRel)
+    }
+
+    pub(super) async fn acquire(&self) -> Result<SemaphorePermit<'_>, AcquireError> {
+        self.gate.acquire().await
+    }
+
+    pub(super) fn close(&self) {
+        self.gate.close();
+    }
+}
+
+/// Restores a claimed refresh when its task is cancelled before publication.
+pub(super) struct McpRefreshInvalidationGuard<'a> {
+    pub(super) refresh: &'a McpRefresh,
+    pub(super) published: bool,
+}
+
+impl Drop for McpRefreshInvalidationGuard<'_> {
+    fn drop(&mut self) {
+        if !self.published {
+            self.refresh.invalidate();
+        }
+    }
+}

codex-rs/core/src/session/mod.rs

@@ -212,6 +212,7 @@ mod inject;
 mod input_queue;
 mod mcp;
 mod mcp_prewarm;
+mod mcp_refresh;
 mod mcp_runtime;
 pub(crate) mod multi_agents;
 mod review;

codex-rs/core/src/session/session.rs

@@ -1,4 +1,5 @@
 use super::input_queue::InputQueue;
+use super::mcp_refresh::McpRefresh;
 use super::*;
 use crate::agents_md_manager::AgentsMdManager;
 use crate::config::ConstraintError;
@@ -40,9 +41,8 @@ pub(crate) struct Session {
     pub(crate) windows_sandbox_proxy_settings_mode:
         codex_sandboxing::WindowsSandboxProxySettingsMode,
     pub(super) multi_agent_version: OnceLock<MultiAgentVersion>,
-    pub(super) mcp_refresh_pending: std::sync::atomic::AtomicBool,
-    /// Serializes runtime refreshes without blocking calls that own a snapshot.
-    pub(super) mcp_refresh_lock: Semaphore,
+    /// Owns invalidation and serializes refreshes without blocking captured calls.
+    pub(super) mcp_refresh: McpRefresh,
     pub(super) mcp_elicitation_reviewer_handle: OnceLock<codex_mcp::ElicitationReviewerHandle>,
     pub(super) mcp_elicitation_lifecycle_handle: OnceLock<codex_mcp::ElicitationLifecycle>,
     pub(super) mcp_prewarm_tx: async_channel::Sender<()>,
@@ -1160,8 +1160,7 @@ impl Session {
                 features: config.features.clone(),
                 windows_sandbox_proxy_settings_mode,
                 multi_agent_version,
-                mcp_refresh_pending: std::sync::atomic::AtomicBool::new(false),
-                mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
+                mcp_refresh: McpRefresh::new(),
                 mcp_elicitation_reviewer_handle: OnceLock::new(),
                 mcp_elicitation_lifecycle_handle: OnceLock::new(),
                 mcp_prewarm_tx,

codex-rs/core/src/session/tests.rs

@@ -1,3 +1,4 @@
+use super::mcp_refresh::McpRefresh;
 use super::turn_context::TurnEnvironment;
 use super::*;
 use crate::agents_md_manager::AgentsMdManager;
@@ -5544,8 +5545,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         windows_sandbox_proxy_settings_mode:
             codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile,
         multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()),
-        mcp_refresh_pending: std::sync::atomic::AtomicBool::new(true),
-        mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
+        mcp_refresh: McpRefresh::new(),
         mcp_elicitation_reviewer_handle: OnceLock::new(),
         mcp_elicitation_lifecycle_handle: OnceLock::new(),
         mcp_prewarm_tx: async_channel::bounded(1).0,
@@ -5560,6 +5560,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         next_internal_sub_id: AtomicU64::new(0),
     };
 
+    session.mark_mcp_runtime_dirty();
     (session, turn_context)
 }
 
@@ -6686,9 +6687,7 @@ async fn user_turn_updates_approvals_reviewer() {
         codex_config::types::ApprovalsReviewer::AutoReview
     );
     assert!(
-        session
-            .mcp_refresh_pending
-            .load(std::sync::atomic::Ordering::Acquire),
+        session.mcp_refresh.is_pending(),
         "server elicitation authority changes must refresh MCP state"
     );
 }
@@ -7714,8 +7713,7 @@ where
         windows_sandbox_proxy_settings_mode:
             codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile,
         multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()),
-        mcp_refresh_pending: std::sync::atomic::AtomicBool::new(true),
-        mcp_refresh_lock: Semaphore::new(/*permits*/ 1),
+        mcp_refresh: McpRefresh::new(),
         mcp_elicitation_reviewer_handle: OnceLock::new(),
         mcp_elicitation_lifecycle_handle: OnceLock::new(),
         mcp_prewarm_tx: async_channel::bounded(1).0,
@@ -7730,6 +7728,7 @@ where
         next_internal_sub_id: AtomicU64::new(0),
     });
 
+    session.mark_mcp_runtime_dirty();
     (session, turn_context, rx_event)
 }
 
@@ -7877,26 +7876,20 @@ async fn cancelled_mcp_refresh_remains_pending() {
             let mut context = std::task::Context::from_waker(futures::task::noop_waker_ref());
             assert!(std::future::Future::poll(refresh.as_mut(), &mut context).is_pending());
             assert!(
-                !session
-                    .mcp_refresh_pending
-                    .load(std::sync::atomic::Ordering::Acquire),
+                !session.mcp_refresh.is_pending(),
                 "the refresh should have claimed its pending invalidation"
             );
         }
     }
 
     assert!(
-        session
-            .mcp_refresh_pending
-            .load(std::sync::atomic::Ordering::Acquire),
+        session.mcp_refresh.is_pending(),
         "a cancelled refresh must leave the runtime dirty"
     );
 
     session.refresh_mcp_if_dirty().await;
     assert!(
-        !session
-            .mcp_refresh_pending
-            .load(std::sync::atomic::Ordering::Acquire),
+        !session.mcp_refresh.is_pending(),
         "the next refresh should publish the pending runtime"
     );
 }
@@ -7929,11 +7922,7 @@ async fn mcp_policy_changes_schedule_runtime_refresh() {
         .await
         .expect("approval policy update should succeed");
 
-    assert!(
-        session
-            .mcp_refresh_pending
-            .load(std::sync::atomic::Ordering::Acquire)
-    );
+    assert!(session.mcp_refresh.is_pending());
 }
 
 #[tokio::test]
@@ -7948,9 +7937,7 @@ async fn mcp_refresh_updates_plugin_auth_mode_before_checking_pending_state() {
     let auth_mode = session.services.auth_manager.get_api_auth_mode();
 
     assert_ne!(session.services.plugins_manager.auth_mode(), auth_mode);
-    session
-        .mcp_refresh_pending
-        .store(false, std::sync::atomic::Ordering::Release);
+    session.mcp_refresh.claim();
 
     session.refresh_mcp_if_dirty().await;
 
@@ -7981,9 +7968,7 @@ async fn mcp_refresh_updates_plugin_auth_mode_before_checking_pending_state() {
         Some("new-api-key".to_string())
     );
     assert_eq!(session.services.plugins_manager.auth_mode(), auth_mode);
-    session
-        .mcp_refresh_pending
-        .store(false, std::sync::atomic::Ordering::Release);
+    session.mcp_refresh.claim();
 
     session.refresh_mcp_if_dirty().await;