Extract MCP binding clients from the connection manager (#34561)

## What changed

- Add `McpBindingClients` to hold the ready MCP clients captured for a model step.
- Move concurrent, paginated resource and resource-template listing into the new helper.
- Keep server filtering in `McpConnectionManager` while preserving per-client timeouts, duplicate-cursor detection, and warning behavior.

GitOrigin-RevId: 0e3f8f135126b17d73058ca84e98493a27543ca0

codex-rs/codex-mcp/src/binding_clients.rs

@@ -0,0 +1,128 @@
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use anyhow::Result;
+use anyhow::anyhow;
+use rmcp::model::PaginatedRequestParams;
+use rmcp::model::Resource;
+use rmcp::model::ResourceTemplate;
+use tokio::task::JoinSet;
+use tracing::warn;
+
+use crate::rmcp_client::ManagedClient;
+
+/// The ready clients captured for one model step.
+pub(crate) struct McpBindingClients {
+    clients: HashMap<String, Arc<ManagedClient>>,
+}
+
+impl McpBindingClients {
+    pub(crate) fn new(clients: HashMap<String, Arc<ManagedClient>>) -> Self {
+        Self { clients }
+    }
+
+    pub(crate) async fn list_all_resources(
+        &self,
+        include_server: impl Fn(&str) -> bool,
+    ) -> HashMap<String, Vec<Resource>> {
+        let mut join_set = JoinSet::new();
+        for (server_name, managed) in self
+            .clients
+            .iter()
+            .filter(|(server_name, _)| include_server(server_name))
+        {
+            let server_name = server_name.clone();
+            let client = Arc::clone(&managed.client);
+            let timeout = managed.tool_timeout;
+            join_set.spawn(async move {
+                let mut collected = Vec::new();
+                let mut cursor: Option<String> = None;
+                loop {
+                    let params = cursor.as_ref().map(|next| {
+                        PaginatedRequestParams::default().with_cursor(Some(next.clone()))
+                    });
+                    let response = match client.list_resources(params, timeout).await {
+                        Ok(result) => result,
+                        Err(error) => return (server_name, Err(error)),
+                    };
+                    collected.extend(response.resources);
+                    match response.next_cursor {
+                        Some(next) if cursor.as_ref() == Some(&next) => {
+                            return (
+                                server_name,
+                                Err(anyhow!("resources/list returned duplicate cursor")),
+                            );
+                        }
+                        Some(next) => cursor = Some(next),
+                        None => return (server_name, Ok(collected)),
+                    }
+                }
+            });
+        }
+        collect_resource_results(&mut join_set, "resources").await
+    }
+
+    pub(crate) async fn list_all_resource_templates(
+        &self,
+        include_server: impl Fn(&str) -> bool,
+    ) -> HashMap<String, Vec<ResourceTemplate>> {
+        let mut join_set = JoinSet::new();
+        for (server_name, managed) in self
+            .clients
+            .iter()
+            .filter(|(server_name, _)| include_server(server_name))
+        {
+            let server_name = server_name.clone();
+            let client = Arc::clone(&managed.client);
+            let timeout = managed.tool_timeout;
+            join_set.spawn(async move {
+                let mut collected = Vec::new();
+                let mut cursor: Option<String> = None;
+                loop {
+                    let params = cursor.as_ref().map(|next| {
+                        PaginatedRequestParams::default().with_cursor(Some(next.clone()))
+                    });
+                    let response = match client.list_resource_templates(params, timeout).await {
+                        Ok(result) => result,
+                        Err(error) => return (server_name, Err(error)),
+                    };
+                    collected.extend(response.resource_templates);
+                    match response.next_cursor {
+                        Some(next) if cursor.as_ref() == Some(&next) => {
+                            return (
+                                server_name,
+                                Err(anyhow!(
+                                    "resources/templates/list returned duplicate cursor"
+                                )),
+                            );
+                        }
+                        Some(next) => cursor = Some(next),
+                        None => return (server_name, Ok(collected)),
+                    }
+                }
+            });
+        }
+        collect_resource_results(&mut join_set, "resource templates").await
+    }
+}
+
+async fn collect_resource_results<T: Send + 'static>(
+    join_set: &mut JoinSet<(String, Result<Vec<T>>)>,
+    kind: &str,
+) -> HashMap<String, Vec<T>> {
+    let mut resources = HashMap::new();
+    while let Some(result) = join_set.join_next().await {
+        match result {
+            Ok((server, Ok(server_resources))) => {
+                resources.insert(server, server_resources);
+            }
+            Ok((server, Err(error))) => {
+                warn!("Failed to list {kind} for MCP server '{server}': {error:#}");
+            }
+            Err(error) => {
+                warn!("Task panic when listing {kind} for MCP server: {error:#}");
+            }
+        }
+    }
+    resources
+}

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

@@ -17,6 +17,7 @@ use std::sync::Arc;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 
+use crate::binding_clients::McpBindingClients;
 use crate::elicitation::ElicitationRequestManager;
 use crate::elicitation::ElicitationRequestRouter;
 use crate::elicitation::ElicitationReviewerHandle;
@@ -524,69 +525,10 @@ impl McpConnectionManager {
         &self,
         include_server: impl Fn(&str) -> bool,
     ) -> HashMap<String, Vec<Resource>> {
-        let mut join_set = JoinSet::new();
-
-        let clients_snapshot = &self.clients;
-
-        for (server_name, async_managed_client) in clients_snapshot
-            .iter()
-            .filter(|(server_name, _)| include_server(server_name))
-        {
-            let server_name = server_name.clone();
-            let Ok(managed_client) = async_managed_client.client().await else {
-                continue;
-            };
-            let timeout = managed_client.tool_timeout;
-            let client = managed_client.client.clone();
-
-            join_set.spawn(async move {
-                let mut collected: Vec<Resource> = Vec::new();
-                let mut cursor: Option<String> = None;
-
-                loop {
-                    let params = cursor.as_ref().map(|next| {
-                        PaginatedRequestParams::default().with_cursor(Some(next.clone()))
-                    });
-                    let response = match client.list_resources(params, timeout).await {
-                        Ok(result) => result,
-                        Err(err) => return (server_name, Err(err)),
-                    };
-
-                    collected.extend(response.resources);
-
-                    match response.next_cursor {
-                        Some(next) => {
-                            if cursor.as_ref() == Some(&next) {
-                                return (
-                                    server_name,
-                                    Err(anyhow!("resources/list returned duplicate cursor")),
-                                );
-                            }
-                            cursor = Some(next);
-                        }
-                        None => return (server_name, Ok(collected)),
-                    }
-                }
-            });
-        }
-
-        let mut aggregated: HashMap<String, Vec<Resource>> = HashMap::new();
-
-        while let Some(join_res) = join_set.join_next().await {
-            match join_res {
-                Ok((server_name, Ok(resources))) => {
-                    aggregated.insert(server_name, resources);
-                }
-                Ok((server_name, Err(err))) => {
-                    warn!("Failed to list resources for MCP server '{server_name}': {err:#}");
-                }
-                Err(err) => {
-                    warn!("Task panic when listing resources for MCP server: {err:#}");
-                }
-            }
-        }
-
-        aggregated
+        self.ready_clients_matching(&include_server)
+            .await
+            .list_all_resources(|_| true)
+            .await
     }
 
     /// Returns resource templates from servers selected by `include_server`.
@@ -595,73 +537,27 @@ impl McpConnectionManager {
         &self,
         include_server: impl Fn(&str) -> bool,
     ) -> HashMap<String, Vec<ResourceTemplate>> {
-        let mut join_set = JoinSet::new();
+        self.ready_clients_matching(&include_server)
+            .await
+            .list_all_resource_templates(|_| true)
+            .await
+    }
 
-        let clients_snapshot = &self.clients;
-
-        for (server_name, async_managed_client) in clients_snapshot
+    async fn ready_clients_matching(
+        &self,
+        include_server: &impl Fn(&str) -> bool,
+    ) -> McpBindingClients {
+        let mut clients = HashMap::new();
+        for (server, client) in self
+            .clients
             .iter()
-            .filter(|(server_name, _)| include_server(server_name))
+            .filter(|(server, _)| include_server(server))
         {
-            let server_name_cloned = server_name.clone();
-            let Ok(managed_client) = async_managed_client.client().await else {
-                continue;
-            };
-            let client = managed_client.client.clone();
-            let timeout = managed_client.tool_timeout;
-
-            join_set.spawn(async move {
-                let mut collected: Vec<ResourceTemplate> = Vec::new();
-                let mut cursor: Option<String> = None;
-
-                loop {
-                    let params = cursor.as_ref().map(|next| {
-                        PaginatedRequestParams::default().with_cursor(Some(next.clone()))
-                    });
-                    let response = match client.list_resource_templates(params, timeout).await {
-                        Ok(result) => result,
-                        Err(err) => return (server_name_cloned, Err(err)),
-                    };
-
-                    collected.extend(response.resource_templates);
-
-                    match response.next_cursor {
-                        Some(next) => {
-                            if cursor.as_ref() == Some(&next) {
-                                return (
-                                    server_name_cloned,
-                                    Err(anyhow!(
-                                        "resources/templates/list returned duplicate cursor"
-                                    )),
-                                );
-                            }
-                            cursor = Some(next);
-                        }
-                        None => return (server_name_cloned, Ok(collected)),
-                    }
-                }
-            });
-        }
-
-        let mut aggregated: HashMap<String, Vec<ResourceTemplate>> = HashMap::new();
-
-        while let Some(join_res) = join_set.join_next().await {
-            match join_res {
-                Ok((server_name, Ok(templates))) => {
-                    aggregated.insert(server_name, templates);
-                }
-                Ok((server_name, Err(err))) => {
-                    warn!(
-                        "Failed to list resource templates for MCP server '{server_name}': {err:#}"
-                    );
-                }
-                Err(err) => {
-                    warn!("Task panic when listing resource templates for MCP server: {err:#}");
-                }
+            if let Ok(client) = client.client().await {
+                clients.insert(server.clone(), Arc::new(client));
             }
         }
-
-        aggregated
+        McpBindingClients::new(clients)
     }
 
     /// Invoke the tool indicated by the (server, tool) pair.

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

@@ -82,6 +82,7 @@ pub use mcp::mcp_permission_prompt_is_auto_approved;
 pub use mcp::qualified_mcp_tool_name_prefix;
 
 pub(crate) mod auth_elicitation;
+pub(crate) mod binding_clients;
 mod catalog;
 pub(crate) mod codex_apps;
 pub(crate) mod connection_manager;