mcp_utils/client/
mcp_client.rs

1use super::error::ClientError;
2use super::transport::ClientService;
3use super::{ClientOptions, ToolCall, ToolCallOptions, Transport, handler::Handler};
4use futures::future::{BoxFuture, FutureExt, Shared};
5use futures::stream::{self, Stream};
6use rmcp::{
7    RoleClient,
8    model::{CallToolRequestParams, GetPromptRequestParams, GetPromptResult, Prompt, Tool},
9    service::{Peer, QuitReason},
10};
11use serde_json::{Map, Value};
12use std::fmt;
13use std::pin::pin;
14use std::sync::Arc;
15use std::time::Duration;
16use tokio::task::JoinError;
17use tokio::time::timeout;
18use tokio_util::sync::{CancellationToken, DropGuard};
19
20#[derive(Clone)]
21pub struct McpClient {
22    pub(super) peer: Peer<RoleClient>,
23    pub(super) handler: Arc<Handler>,
24    shutdown: Arc<Shutdown>,
25}
26
27impl McpClient {
28    pub async fn connect(
29        name: impl Into<String>,
30        transport: Transport,
31        options: &ClientOptions,
32    ) -> Result<Self, ClientError> {
33        transport.connect(Handler::new(name.into(), options), options).await.map(Self::from_service)
34    }
35
36    pub async fn authorize(
37        name: impl Into<String>,
38        transport: Transport,
39        options: &ClientOptions,
40        challenge: Option<String>,
41    ) -> Result<Self, ClientError> {
42        transport.authorize(Handler::new(name.into(), options), options, challenge).await.map(Self::from_service)
43    }
44
45    pub fn name(&self) -> &str {
46        self.handler.server()
47    }
48
49    pub fn description(&self) -> Option<String> {
50        let info = self.peer.peer_info()?;
51        non_empty(info.server_info.as_ref()?.description.as_ref())
52    }
53
54    pub fn instructions(&self) -> Option<String> {
55        non_empty(self.peer.peer_info()?.instructions.as_ref())
56    }
57
58    pub async fn list_tools(&self) -> Result<Vec<Tool>, ClientError> {
59        self.peer.list_all_tools().await.map_err(|source| ClientError::request(self.name(), source))
60    }
61
62    pub async fn list_prompts(&self) -> Result<Vec<Prompt>, ClientError> {
63        if !self.supports_prompts() {
64            return Ok(Vec::new());
65        }
66        self.peer.list_all_prompts().await.map_err(|source| ClientError::request(self.name(), source))
67    }
68
69    pub async fn get_prompt(&self, name: &str, args: Map<String, Value>) -> Result<GetPromptResult, ClientError> {
70        let mut request = GetPromptRequestParams::new(name);
71        if !args.is_empty() {
72            request = request.with_arguments(args);
73        }
74        self.peer.get_prompt(request).await.map_err(|source| ClientError::request(self.name(), source))
75    }
76
77    pub fn tool_lists(&self) -> impl Stream<Item = Result<Vec<Tool>, ClientError>> + Send + use<> {
78        let mut changes = self.handler.tool_list_changes();
79        changes.mark_changed();
80        stream::unfold((self.clone(), changes), |(client, mut changes)| async move {
81            tokio::select! {
82                biased;
83                () = client.shutdown.closed.clone() => return None,
84                notified = changes.changed() => notified.ok()?,
85            }
86            let tools = client.list_tools().await;
87            Some((tools, (client, changes)))
88        })
89    }
90
91    pub fn call_tool(&self, name: &str, args: Map<String, Value>, options: ToolCallOptions) -> ToolCall {
92        let params = CallToolRequestParams::new(name.to_string()).with_arguments(args);
93        ToolCall::new(self.clone(), params, options)
94    }
95
96    pub async fn close(&self) {
97        self.shutdown.cancel.cancel();
98        self.shutdown.closed.clone().await;
99    }
100
101    pub(crate) fn close_on_drop(&self) -> DropGuard {
102        self.shutdown.cancel.clone().drop_guard()
103    }
104
105    #[cfg(any(test, feature = "testing"))]
106    pub fn peer(&self) -> &Peer<RoleClient> {
107        &self.peer
108    }
109
110    fn supports_prompts(&self) -> bool {
111        self.peer.peer_info().is_some_and(|info| info.capabilities.prompts.is_some())
112    }
113
114    fn from_service(service: ClientService) -> Self {
115        let peer = service.peer().clone();
116        let handler = Arc::clone(service.service());
117        let cancel = CancellationToken::new();
118        let closed = tokio::spawn(run_until_closed(service, cancel.clone())).map(|_| ()).boxed().shared();
119        let shutdown = Shutdown { _cancel_on_drop: cancel.clone().drop_guard(), cancel, closed };
120        Self { peer, handler, shutdown: Arc::new(shutdown) }
121    }
122}
123
124impl fmt::Debug for McpClient {
125    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
126        formatter.debug_struct("McpClient").field("server", &self.name()).finish_non_exhaustive()
127    }
128}
129
130struct Shutdown {
131    cancel: CancellationToken,
132    _cancel_on_drop: DropGuard,
133    closed: Shared<BoxFuture<'static, ()>>,
134}
135
136const CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
137
138async fn run_until_closed(service: ClientService, cancel: CancellationToken) {
139    let server = service.service().server().to_string();
140    let stop = service.cancellation_token();
141    let mut running = pin!(service.waiting());
142    let closed = tokio::select! {
143        closed = &mut running => closed.map(Some),
144        () = cancel.cancelled() => {
145            stop.cancel();
146            match timeout(CLOSE_TIMEOUT, running).await {
147                Ok(closed) => closed.map(Some),
148                Err(_) => Ok(None),
149            }
150        }
151    };
152    log_close(&server, closed);
153}
154
155fn non_empty(text: Option<&String>) -> Option<String> {
156    text.filter(|text| !text.is_empty()).cloned()
157}
158
159fn log_close(server: &str, result: Result<Option<QuitReason>, JoinError>) {
160    match result {
161        Ok(Some(reason)) => tracing::debug!(server, ?reason, "MCP server connection closed"),
162        Ok(None) => tracing::warn!(server, "MCP server connection did not close within {CLOSE_TIMEOUT:?}"),
163        Err(error) => tracing::warn!(server, %error, "MCP server connection task failed while closing"),
164    }
165}
mcp_utils/client/mcp_client.rs

mcp_utils/client/
mcp_client.rs
