codex-mcp: serialize connector runtime refreshes

Author: mzeng-openai

codex-rs/codex-mcp/src/connection_manager.rs

@@ -17,13 +17,15 @@ use crate::McpAuthStatusEntry;
 use crate::connector_runtime::CodexAppsToolsCache;
 use crate::connector_runtime::CodexAppsToolsCacheKey;
 use crate::connector_runtime::CodexAppsToolsFetchSource;
+use crate::connector_runtime::ConnectorRuntimeSnapshot;
 use crate::elicitation::ElicitationRequestManager;
 use crate::elicitation::ElicitationRequestRouter;
 use crate::elicitation::ElicitationReviewerHandle;
 use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
 use crate::mcp::ToolPluginProvenance;
 use crate::rmcp_client::AsyncManagedClient;
 use crate::rmcp_client::CODEX_APPS_REFRESH_DURATION_METRIC;
+use crate::rmcp_client::CodexAppsStartupMode;
 use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT;
 use crate::rmcp_client::MCP_TOOLS_LIST_DURATION_METRIC;
 use crate::rmcp_client::ManagedClient;
@@ -144,6 +146,7 @@ impl McpConnectionManager {
         elicitation_reviewer: Option<ElicitationReviewerHandle>,
         elicitation_lifecycle: Option<crate::ElicitationLifecycle>,
         elicitation_router: ElicitationRequestRouter,
+        codex_apps_startup_mode: CodexAppsStartupMode,
     ) -> Self {
         let mut required_servers = mcp_servers
             .iter()
@@ -223,6 +226,7 @@ impl McpConnectionManager {
                 runtime_auth_provider,
                 client_elicitation_capability.clone(),
                 supports_openai_form_elicitation,
+                codex_apps_startup_mode,
             );
             clients.insert(server_name.clone(), async_managed_client.clone());
             let tx_event = tx_event.clone();
@@ -528,13 +532,98 @@ impl McpConnectionManager {
         normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names)
     }
 
+    /// Force-refreshes the connector runtime exactly once under the active
+    /// context's shared lock and returns the exact committed snapshot.
+    pub async fn hard_refresh_codex_apps_runtime(&self) -> Result<Arc<ConnectorRuntimeSnapshot>> {
+        let refresh_start = Instant::now();
+        let async_client = self
+            .clients
+            .get(CODEX_APPS_MCP_SERVER_NAME)
+            .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?;
+        let result = async {
+            let cache_context = async_client
+                .codex_apps_tools_cache_context
+                .as_ref()
+                .ok_or_else(|| anyhow!("connector runtime manager is unavailable"))?;
+            let _refresh_guard = cache_context
+                .lock_explicit_refresh()
+                .await
+                .context("connector runtime context changed before refresh")?;
+            let managed_client = self.client_by_name(CODEX_APPS_MCP_SERVER_NAME).await?;
+            let list_start = Instant::now();
+            let fetch_ticket = cache_context.begin_fetch(CodexAppsToolsFetchSource::HardRefresh);
+            let tools = list_tools_for_client_uncached(
+                CODEX_APPS_MCP_SERVER_NAME,
+                /*is_codex_apps_mcp_server*/ true,
+                /*codex_apps_refresh_trigger*/ "explicit",
+                &managed_client.client,
+                managed_client.tool_timeout,
+                managed_client.server_instructions.as_deref(),
+            )
+            .await
+            .with_context(|| {
+                format!("failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'")
+            })?;
+            let snapshot = cache_context
+                .publish_runtime_if_newest_accepted(
+                    fetch_ticket,
+                    &managed_client.server_info,
+                    tools,
+                )
+                .context("connector runtime context changed while publishing refresh")?;
+            emit_duration(
+                MCP_TOOLS_LIST_DURATION_METRIC,
+                list_start.elapsed(),
+                &[("cache", "miss")],
+            );
+            Ok(snapshot)
+        }
+        .await;
+        let outcome = if result.is_ok() { "success" } else { "error" };
+        let retained_previous = if result.is_err()
+            && async_client
+                .codex_apps_tools_cache_context
+                .as_ref()
+                .and_then(super::connector_runtime::ConnectorRuntimeContext::current_snapshot)
+                .is_some()
+        {
+            "true"
+        } else {
+            "false"
+        };
+        emit_duration(
+            CODEX_APPS_REFRESH_DURATION_METRIC,
+            refresh_start.elapsed(),
+            &[
+                ("path", "new"),
+                ("trigger", "explicit"),
+                ("outcome", outcome),
+                ("retained_previous_snapshot", retained_previous),
+            ],
+        );
+        result
+    }
+
     /// Force-refresh codex apps tools by bypassing the in-process cache.
     ///
     /// On success, the refreshed tools replace shared cache contents when the
     /// cache is enabled and the latest filtered tools are returned directly to
     /// the caller. On failure, existing shared cache contents remain unchanged.
     pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result<Vec<ToolInfo>> {
         let refresh_start = Instant::now();
+        let async_client = self
+            .clients
+            .get(CODEX_APPS_MCP_SERVER_NAME)
+            .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?;
+        let _refresh_guard = match async_client.codex_apps_tools_cache_context.as_ref() {
+            Some(cache_context) => Some(
+                cache_context
+                    .lock_explicit_refresh()
+                    .await
+                    .context("connector runtime context changed before refresh")?,
+            ),
+            None => None,
+        };
         let managed_client = self.client_by_name(CODEX_APPS_MCP_SERVER_NAME).await?;
 
         let list_start = Instant::now();
@@ -555,16 +644,21 @@ impl McpConnectionManager {
             format!("failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'")
         })?;
 
-        let tools =
-            match (
-                managed_client.codex_apps_tools_cache_context.as_ref(),
-                fetch_ticket,
-            ) {
-                (Some(cache_context), Some(fetch_ticket)) => cache_context
-                    .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools)?,
-                (None, None) => tools,
-                _ => unreachable!("Codex Apps fetch ticket requires cache context"),
-            };
+        let tools = match (
+            managed_client.codex_apps_tools_cache_context.as_ref(),
+            fetch_ticket,
+        ) {
+            (Some(cache_context), Some(fetch_ticket)) => cache_context
+                .publish_runtime_if_newest_accepted(
+                    fetch_ticket,
+                    &managed_client.server_info,
+                    tools,
+                )
+                .map(|snapshot| snapshot.tools().to_vec())
+                .context("connector runtime context changed while publishing refresh")?,
+            (None, None) => tools,
+            _ => unreachable!("Codex Apps fetch ticket requires cache context"),
+        };
         emit_duration(
             MCP_TOOLS_LIST_DURATION_METRIC,
             list_start.elapsed(),

codex-rs/codex-mcp/src/connection_manager_tests.rs

@@ -8,6 +8,7 @@ use crate::elicitation::ElicitationRequestRouter;
 use crate::elicitation::elicitation_is_rejected_by_policy;
 use crate::rmcp_client::AsyncManagedClient;
 use crate::rmcp_client::CODEX_APPS_RECONNECT_INITIAL_BACKOFF;
+use crate::rmcp_client::CodexAppsStartupMode;
 use crate::rmcp_client::CodexAppsStartupReconnect;
 use crate::rmcp_client::ManagedClient;
 use crate::rmcp_client::ManagedClientFuture;
@@ -1911,6 +1912,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() {
         /*elicitation_reviewer*/ None,
         /*elicitation_lifecycle*/ None,
         ElicitationRequestRouter::default(),
+        CodexAppsStartupMode::ListTools,
     )
     .await;
 

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::sync::OwnedMutexGuard;
 
 use crate::connector_runtime_persistence::load_cached_codex_apps_server_info;
 use crate::connector_runtime_persistence::load_cached_connector_runtime_for_identity;
@@ -200,6 +201,21 @@ impl ConnectorRuntimeContext {
         }
     }
 
+    pub(crate) async fn lock_explicit_refresh(
+        &self,
+    ) -> Result<OwnedMutexGuard<()>, ConnectorRuntimeContextDiscarded> {
+        if !self.is_active() {
+            return Err(ConnectorRuntimeContextDiscarded);
+        }
+        let guard = Arc::clone(&self.entry.explicit_refresh_lock)
+            .lock_owned()
+            .await;
+        if !self.is_active() {
+            return Err(ConnectorRuntimeContextDiscarded);
+        }
+        Ok(guard)
+    }
+
     pub(crate) fn publish_runtime_if_newest_accepted(
         &self,
         ticket: CodexAppsToolsFetchTicket,
@@ -269,6 +285,7 @@ impl ConnectorRuntimeContext {
         Ok(snapshot)
     }
 
+    #[cfg(test)]
     pub(crate) fn publish_if_newest_accepted(
         &self,
         ticket: CodexAppsToolsFetchTicket,
@@ -331,6 +348,7 @@ pub(crate) struct ConnectorRuntimeEntry {
     pub(crate) current_snapshot: ArcSwapOption<ConnectorRuntimeSnapshot>,
     next_fetch_generation: AtomicU64,
     last_accepted_generation: Mutex<u64>,
+    explicit_refresh_lock: Arc<tokio::sync::Mutex<()>>,
 }
 
 impl ConnectorRuntimeEntry {
@@ -341,6 +359,7 @@ impl ConnectorRuntimeEntry {
             current_snapshot: ArcSwapOption::from(current_snapshot),
             next_fetch_generation: AtomicU64::new(0),
             last_accepted_generation: Mutex::new(0),
+            explicit_refresh_lock: Arc::new(tokio::sync::Mutex::new(())),
         }
     }
 }

codex-rs/codex-mcp/src/connector_runtime_tests.rs

@@ -16,6 +16,7 @@ use std::collections::HashSet;
 use std::os::unix::ffi::OsStringExt;
 use std::path::PathBuf;
 use std::sync::Arc;
+use std::time::Duration;
 use tempfile::tempdir;
 
 fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo {
@@ -751,3 +752,65 @@ fn live_publish_sets_timestamp_and_stale_publish_preserves_it() {
     assert!(Arc::ptr_eq(&current, &stale));
     assert_eq!(stale.refreshed_at(), current.refreshed_at());
 }
+
+#[tokio::test]
+async fn explicit_refresh_lock_serializes_concurrent_refreshes() {
+    let codex_home = tempdir().expect("tempdir");
+    let context = create_codex_apps_tools_cache_context(
+        codex_home.path().to_path_buf(),
+        Some("account-one"),
+        Some("user-one"),
+    );
+    let first_guard = context
+        .lock_explicit_refresh()
+        .await
+        .expect("acquire first refresh lock");
+    let second_context = context.clone();
+    let (attempting_tx, attempting_rx) = tokio::sync::oneshot::channel();
+    let (acquired_tx, mut acquired_rx) = tokio::sync::oneshot::channel();
+    let second_refresh = tokio::spawn(async move {
+        attempting_tx.send(()).expect("signal lock attempt");
+        let _guard = second_context
+            .lock_explicit_refresh()
+            .await
+            .expect("acquire second refresh lock");
+        acquired_tx.send(()).expect("signal lock acquisition");
+    });
+    attempting_rx.await.expect("second refresh attempted lock");
+
+    assert!(
+        tokio::time::timeout(Duration::from_millis(20), &mut acquired_rx)
+            .await
+            .is_err()
+    );
+    drop(first_guard);
+    tokio::time::timeout(Duration::from_secs(1), acquired_rx)
+        .await
+        .expect("second refresh should acquire after first releases")
+        .expect("second refresh acquisition signal");
+    second_refresh.await.expect("second refresh task");
+}
+
+#[tokio::test]
+async fn discarded_context_cannot_acquire_explicit_refresh_lock() {
+    let codex_home = tempdir().expect("tempdir");
+    let manager = ConnectorRuntimeManager::default();
+    let 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,
+        },
+    );
+    let _context_b = manager.context(
+        codex_home.path().to_path_buf(),
+        ConnectorRuntimeContextKey {
+            account_id: Some("account-b".to_string()),
+            chatgpt_user_id: Some("user-b".to_string()),
+            is_workspace_account: false,
+        },
+    );
+
+    assert!(context_a.lock_explicit_refresh().await.is_err());
+}

codex-rs/codex-mcp/src/lib.rs

@@ -9,6 +9,7 @@ pub use resource_client::McpResourceClient;
 pub use resource_client::McpResourceClientCacheKey;
 pub use resource_client::McpResourcePage;
 pub use resource_client::McpResourceReadResult;
+pub use rmcp_client::CodexAppsStartupMode;
 pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY;
 pub use runtime::McpRuntimeContext;
 pub use runtime::SandboxState;

codex-rs/codex-mcp/src/mcp/mod.rs

@@ -47,6 +47,7 @@ use crate::ResolvedMcpCatalog;
 use crate::connection_manager::McpConnectionManager;
 use crate::connector_runtime::CodexAppsToolsCache;
 use crate::connector_runtime::codex_apps_tools_cache_key;
+use crate::rmcp_client::CodexAppsStartupMode;
 use crate::runtime::McpRuntimeContext;
 use crate::server::EffectiveMcpServer;
 
@@ -341,6 +342,7 @@ pub async fn read_mcp_resource(
         /*elicitation_reviewer*/ None,
         /*elicitation_lifecycle*/ None,
         crate::elicitation::ElicitationRequestRouter::default(),
+        CodexAppsStartupMode::ListTools,
     )
     .await;
 
@@ -419,6 +421,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
         /*elicitation_reviewer*/ None,
         /*elicitation_lifecycle*/ None,
         crate::elicitation::ElicitationRequestRouter::default(),
+        CodexAppsStartupMode::ListTools,
     )
     .await;