diff --git a/codex-rs/codex-mcp/src/binding_clients.rs b/codex-rs/codex-mcp/src/binding_clients.rs new file mode 100644 index 0000000000..08e5b1c465 --- /dev/null +++ b/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>, +} + +impl McpBindingClients { + pub(crate) fn new(clients: HashMap>) -> Self { + Self { clients } + } + + pub(crate) async fn list_all_resources( + &self, + include_server: impl Fn(&str) -> bool, + ) -> HashMap> { + 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 = 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> { + 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 = 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( + join_set: &mut JoinSet<(String, Result>)>, + kind: &str, +) -> HashMap> { + 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 +} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 19ad27a1b3..1cc3f9fe9b 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/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> { - 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 = Vec::new(); - let mut cursor: Option = 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> = 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> { - 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 = Vec::new(); - let mut cursor: Option = 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> = 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. diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 499bcfc0b7..499fe43cf1 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/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;