rmcp/service/
client.rs

1// Sampling/Roots/Logging are SEP-2577-deprecated; internal references are expected.
2#![expect(deprecated)]
3use std::borrow::Cow;
4
5use thiserror::Error;
6
7use super::*;
8use crate::{
9    model::{
10        ArgumentInfo, CallToolRequest, CallToolRequestParams, CallToolResult,
11        CancelledNotification, CancelledNotificationParam, ClientInfo, ClientJsonRpcMessage,
12        ClientNotification, ClientRequest, ClientResult, CompleteRequest, CompleteRequestParams,
13        CompleteResult, CompletionContext, CompletionInfo, ErrorData, GetPromptRequest,
14        GetPromptRequestParams, GetPromptResult, InitializeRequest, InitializedNotification,
15        JsonRpcResponse, ListPromptsRequest, ListPromptsResult, ListResourceTemplatesRequest,
16        ListResourceTemplatesResult, ListResourcesRequest, ListResourcesResult, ListToolsRequest,
17        ListToolsResult, PaginatedRequestParams, ProgressNotification, ProgressNotificationParam,
18        ReadResourceRequest, ReadResourceRequestParams, ReadResourceResult, Reference, RequestId,
19        RootsListChangedNotification, ServerInfo, ServerJsonRpcMessage, ServerNotification,
20        ServerRequest, ServerResult, SetLevelRequest, SetLevelRequestParams, SubscribeRequest,
21        SubscribeRequestParams, UnsubscribeRequest, UnsubscribeRequestParams,
22    },
23    transport::DynamicTransportError,
24};
25
26/// It represents the error that may occur when serving the client.
27///
28/// if you want to handle the error, you can use `serve_client_with_ct` or `serve_client` with `Result<RunningService<RoleClient, S>, ClientError>`
29#[derive(Error, Debug)]
30#[non_exhaustive]
31pub enum ClientInitializeError {
32    #[error("expect initialized response, but received: {0:?}")]
33    ExpectedInitResponse(Option<ServerJsonRpcMessage>),
34
35    #[error("expect initialized result, but received: {0:?}")]
36    ExpectedInitResult(Option<ServerResult>),
37
38    #[error("conflict initialized response id: expected {0}, got {1}")]
39    ConflictInitResponseId(RequestId, RequestId),
40
41    #[error("connection closed: {0}")]
42    ConnectionClosed(String),
43
44    #[error("Send message error {error}, when {context}")]
45    TransportError {
46        error: DynamicTransportError,
47        context: Cow<'static, str>,
48    },
49
50    #[error("JSON-RPC error: {0}")]
51    JsonRpcError(ErrorData),
52
53    #[error("Cancelled")]
54    Cancelled,
55}
56
57impl ClientInitializeError {
58    pub fn transport<T: Transport<RoleClient> + 'static>(
59        error: T::Error,
60        context: impl Into<Cow<'static, str>>,
61    ) -> Self {
62        Self::TransportError {
63            error: DynamicTransportError::new::<T, _>(error),
64            context: context.into(),
65        }
66    }
67}
68
69/// Helper function to get the next message from the stream
70async fn expect_next_message<T>(
71    transport: &mut T,
72    context: &str,
73) -> Result<ServerJsonRpcMessage, ClientInitializeError>
74where
75    T: Transport<RoleClient>,
76{
77    transport
78        .receive()
79        .await
80        .ok_or_else(|| ClientInitializeError::ConnectionClosed(context.to_string()))
81}
82
83/// Helper function to expect a response from the stream
84async fn expect_response<T, S>(
85    transport: &mut T,
86    context: &str,
87    service: &S,
88    peer: Peer<RoleClient>,
89) -> Result<(ServerResult, RequestId), ClientInitializeError>
90where
91    T: Transport<RoleClient>,
92    S: Service<RoleClient>,
93{
94    loop {
95        let message = expect_next_message(transport, context).await?;
96        match message {
97            // Expected message to complete the initialization
98            ServerJsonRpcMessage::Response(JsonRpcResponse { id, result, .. }) => {
99                break Ok((result, id));
100            }
101            // Handle JSON-RPC error responses
102            ServerJsonRpcMessage::Error(error) => {
103                break Err(ClientInitializeError::JsonRpcError(error.error));
104            }
105            // Server could send logging messages before handshake
106            ServerJsonRpcMessage::Notification(mut notification) => {
107                let ServerNotification::LoggingMessageNotification(logging) =
108                    &mut notification.notification
109                else {
110                    tracing::warn!(?notification, "Received unexpected message");
111                    continue;
112                };
113
114                let mut context = NotificationContext {
115                    peer: peer.clone(),
116                    meta: Meta::default(),
117                    extensions: Extensions::default(),
118                };
119
120                if let Some(meta) = logging.extensions.get_mut::<Meta>() {
121                    std::mem::swap(&mut context.meta, meta);
122                }
123                std::mem::swap(&mut context.extensions, &mut logging.extensions);
124
125                if let Err(error) = service
126                    .handle_notification(notification.notification, context)
127                    .await
128                {
129                    tracing::warn!(?error, "Handle logging before handshake failed.");
130                }
131            }
132            // Server could send pings before handshake
133            ServerJsonRpcMessage::Request(ref request)
134                if matches!(request.request, ServerRequest::PingRequest(_)) =>
135            {
136                tracing::trace!("Received ping request. Ignored.")
137            }
138            // Server SHOULD NOT send any other messages before handshake. We ignore them anyway
139            _ => tracing::warn!(?message, "Received unexpected message"),
140        }
141    }
142}
143
144#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
145#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
146pub struct RoleClient;
147
148impl ServiceRole for RoleClient {
149    type Req = ClientRequest;
150    type Resp = ClientResult;
151    type Not = ClientNotification;
152    type PeerReq = ServerRequest;
153    type PeerResp = ServerResult;
154    type PeerNot = ServerNotification;
155    type Info = ClientInfo;
156    type PeerInfo = ServerInfo;
157    type InitializeError = ClientInitializeError;
158    const IS_CLIENT: bool = true;
159}
160
161pub type ServerSink = Peer<RoleClient>;
162
163impl<S: Service<RoleClient>> ServiceExt<RoleClient> for S {
164    fn serve_with_ct<T, E, A>(
165        self,
166        transport: T,
167        ct: CancellationToken,
168    ) -> impl Future<Output = Result<RunningService<RoleClient, Self>, ClientInitializeError>>
169    + MaybeSendFuture
170    where
171        T: IntoTransport<RoleClient, E, A>,
172        E: std::error::Error + Send + Sync + 'static,
173        Self: Sized,
174    {
175        serve_client_with_ct(self, transport, ct)
176    }
177}
178
179pub async fn serve_client<S, T, E, A>(
180    service: S,
181    transport: T,
182) -> Result<RunningService<RoleClient, S>, ClientInitializeError>
183where
184    S: Service<RoleClient>,
185    T: IntoTransport<RoleClient, E, A>,
186    E: std::error::Error + Send + Sync + 'static,
187{
188    serve_client_with_ct(service, transport, Default::default()).await
189}
190
191pub async fn serve_client_with_ct<S, T, E, A>(
192    service: S,
193    transport: T,
194    ct: CancellationToken,
195) -> Result<RunningService<RoleClient, S>, ClientInitializeError>
196where
197    S: Service<RoleClient>,
198    T: IntoTransport<RoleClient, E, A>,
199    E: std::error::Error + Send + Sync + 'static,
200{
201    tokio::select! {
202        result = serve_client_with_ct_inner(service, transport.into_transport(), ct.clone()) => { result }
203        _ = ct.cancelled() => {
204            Err(ClientInitializeError::Cancelled)
205        }
206    }
207}
208
209async fn serve_client_with_ct_inner<S, T>(
210    service: S,
211    transport: T,
212    ct: CancellationToken,
213) -> Result<RunningService<RoleClient, S>, ClientInitializeError>
214where
215    S: Service<RoleClient>,
216    T: Transport<RoleClient> + 'static,
217{
218    let mut transport = transport.into_transport();
219    let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
220
221    // service
222    let id = id_provider.next_request_id();
223    let init_request = InitializeRequest {
224        method: Default::default(),
225        params: service.get_info(),
226        extensions: Default::default(),
227    };
228    transport
229        .send(ClientJsonRpcMessage::request(
230            ClientRequest::InitializeRequest(init_request),
231            id.clone(),
232        ))
233        .await
234        .map_err(|error| ClientInitializeError::TransportError {
235            error: DynamicTransportError::new::<T, _>(error),
236            context: "send initialize request".into(),
237        })?;
238
239    let (peer, peer_rx) = Peer::new(id_provider, None);
240
241    let (response, response_id) = expect_response(
242        &mut transport,
243        "initialize response",
244        &service,
245        peer.clone(),
246    )
247    .await?;
248
249    if id != response_id {
250        return Err(ClientInitializeError::ConflictInitResponseId(
251            id,
252            response_id,
253        ));
254    }
255
256    let ServerResult::InitializeResult(initialize_result) = response else {
257        return Err(ClientInitializeError::ExpectedInitResult(Some(response)));
258    };
259    peer.set_peer_info(initialize_result);
260
261    // send notification
262    let notification = ClientJsonRpcMessage::notification(
263        ClientNotification::InitializedNotification(InitializedNotification {
264            method: Default::default(),
265            extensions: Default::default(),
266        }),
267    );
268    transport.send(notification).await.map_err(|error| {
269        ClientInitializeError::transport::<T>(error, "send initialized notification")
270    })?;
271    Ok(serve_inner(service, transport, peer, peer_rx, ct))
272}
273
274macro_rules! method {
275    ($(#[$meta:meta])* peer_req $method:ident $Req:ident() => $Resp: ident ) => {
276        $(#[$meta])*
277        pub async fn $method(&self) -> Result<$Resp, ServiceError> {
278            let result = self
279                .send_request(ClientRequest::$Req($Req {
280                    method: Default::default(),
281                }))
282                .await?;
283            match result {
284                ServerResult::$Resp(result) => Ok(result),
285                _ => Err(ServiceError::UnexpectedResponse),
286            }
287        }
288    };
289    ($(#[$meta:meta])* peer_req $method:ident $Req:ident($Param: ident) => $Resp: ident ) => {
290        $(#[$meta])*
291        pub async fn $method(&self, params: $Param) -> Result<$Resp, ServiceError> {
292            let result = self
293                .send_request(ClientRequest::$Req($Req {
294                    method: Default::default(),
295                    params,
296                    extensions: Default::default(),
297                }))
298                .await?;
299            match result {
300                ServerResult::$Resp(result) => Ok(result),
301                _ => Err(ServiceError::UnexpectedResponse),
302            }
303        }
304    };
305    ($(#[$meta:meta])* peer_req $method:ident $Req:ident($Param: ident)? => $Resp: ident ) => {
306        $(#[$meta])*
307        pub async fn $method(&self, params: Option<$Param>) -> Result<$Resp, ServiceError> {
308            let result = self
309                .send_request(ClientRequest::$Req($Req {
310                    method: Default::default(),
311                    params,
312                    extensions: Default::default(),
313                }))
314                .await?;
315            match result {
316                ServerResult::$Resp(result) => Ok(result),
317                _ => Err(ServiceError::UnexpectedResponse),
318            }
319        }
320    };
321    ($(#[$meta:meta])* peer_req $method:ident $Req:ident($Param: ident)) => {
322        $(#[$meta])*
323        pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
324            let result = self
325                .send_request(ClientRequest::$Req($Req {
326                    method: Default::default(),
327                    params,
328                    extensions: Default::default(),
329                }))
330                .await?;
331            match result {
332                ServerResult::EmptyResult(_) => Ok(()),
333                _ => Err(ServiceError::UnexpectedResponse),
334            }
335        }
336    };
337
338    ($(#[$meta:meta])* peer_not $method:ident $Not:ident($Param: ident)) => {
339        $(#[$meta])*
340        pub async fn $method(&self, params: $Param) -> Result<(), ServiceError> {
341            self.send_notification(ClientNotification::$Not($Not {
342                method: Default::default(),
343                params,
344                extensions: Default::default(),
345            }))
346            .await?;
347            Ok(())
348        }
349    };
350    ($(#[$meta:meta])* peer_not $method:ident $Not:ident) => {
351        $(#[$meta])*
352        pub async fn $method(&self) -> Result<(), ServiceError> {
353            self.send_notification(ClientNotification::$Not($Not {
354                method: Default::default(),
355                extensions: Default::default(),
356            }))
357            .await?;
358            Ok(())
359        }
360    };
361}
362
363impl Peer<RoleClient> {
364    method!(peer_req complete CompleteRequest(CompleteRequestParams) => CompleteResult);
365    method!(
366        #[deprecated(
367            since = "1.8.0",
368            note = "Logging is deprecated by SEP-2577 and will be removed in a future release. See https://github.com/modelcontextprotocol/modelcontextprotocol/pull/2577"
369        )]
370        peer_req set_level SetLevelRequest(SetLevelRequestParams)
371    );
372    method!(peer_req get_prompt GetPromptRequest(GetPromptRequestParams) => GetPromptResult);
373    method!(peer_req list_prompts ListPromptsRequest(PaginatedRequestParams)? => ListPromptsResult);
374    method!(peer_req list_resources ListResourcesRequest(PaginatedRequestParams)? => ListResourcesResult);
375    method!(peer_req list_resource_templates ListResourceTemplatesRequest(PaginatedRequestParams)? => ListResourceTemplatesResult);
376    method!(peer_req read_resource ReadResourceRequest(ReadResourceRequestParams) => ReadResourceResult);
377    method!(peer_req subscribe SubscribeRequest(SubscribeRequestParams) );
378    method!(peer_req unsubscribe UnsubscribeRequest(UnsubscribeRequestParams));
379    method!(peer_req call_tool CallToolRequest(CallToolRequestParams) => CallToolResult);
380    method!(peer_req list_tools ListToolsRequest(PaginatedRequestParams)? => ListToolsResult);
381
382    method!(peer_not notify_cancelled CancelledNotification(CancelledNotificationParam));
383    method!(peer_not notify_progress ProgressNotification(ProgressNotificationParam));
384    method!(peer_not notify_initialized InitializedNotification);
385    method!(peer_not notify_roots_list_changed RootsListChangedNotification);
386}
387
388impl Peer<RoleClient> {
389    /// A wrapper method for [`Peer<RoleClient>::list_tools`].
390    ///
391    /// This function will call [`Peer<RoleClient>::list_tools`] multiple times until all tools are listed.
392    pub async fn list_all_tools(&self) -> Result<Vec<crate::model::Tool>, ServiceError> {
393        let mut tools = Vec::new();
394        let mut cursor = None;
395        loop {
396            let result = self
397                .list_tools(Some(PaginatedRequestParams { meta: None, cursor }))
398                .await?;
399            tools.extend(result.tools);
400            cursor = result.next_cursor;
401            if cursor.is_none() {
402                break;
403            }
404        }
405        Ok(tools)
406    }
407
408    /// A wrapper method for [`Peer<RoleClient>::list_prompts`].
409    ///
410    /// This function will call [`Peer<RoleClient>::list_prompts`] multiple times until all prompts are listed.
411    pub async fn list_all_prompts(&self) -> Result<Vec<crate::model::Prompt>, ServiceError> {
412        let mut prompts = Vec::new();
413        let mut cursor = None;
414        loop {
415            let result = self
416                .list_prompts(Some(PaginatedRequestParams { meta: None, cursor }))
417                .await?;
418            prompts.extend(result.prompts);
419            cursor = result.next_cursor;
420            if cursor.is_none() {
421                break;
422            }
423        }
424        Ok(prompts)
425    }
426
427    /// A wrapper method for [`Peer<RoleClient>::list_resources`].
428    ///
429    /// This function will call [`Peer<RoleClient>::list_resources`] multiple times until all resources are listed.
430    pub async fn list_all_resources(&self) -> Result<Vec<crate::model::Resource>, ServiceError> {
431        let mut resources = Vec::new();
432        let mut cursor = None;
433        loop {
434            let result = self
435                .list_resources(Some(PaginatedRequestParams { meta: None, cursor }))
436                .await?;
437            resources.extend(result.resources);
438            cursor = result.next_cursor;
439            if cursor.is_none() {
440                break;
441            }
442        }
443        Ok(resources)
444    }
445
446    /// A wrapper method for [`Peer<RoleClient>::list_resource_templates`].
447    ///
448    /// This function will call [`Peer<RoleClient>::list_resource_templates`] multiple times until all resource templates are listed.
449    pub async fn list_all_resource_templates(
450        &self,
451    ) -> Result<Vec<crate::model::ResourceTemplate>, ServiceError> {
452        let mut resource_templates = Vec::new();
453        let mut cursor = None;
454        loop {
455            let result = self
456                .list_resource_templates(Some(PaginatedRequestParams { meta: None, cursor }))
457                .await?;
458            resource_templates.extend(result.resource_templates);
459            cursor = result.next_cursor;
460            if cursor.is_none() {
461                break;
462            }
463        }
464        Ok(resource_templates)
465    }
466
467    /// Convenient method to get completion suggestions for a prompt argument
468    ///
469    /// # Arguments
470    /// * `prompt_name` - Name of the prompt being completed
471    /// * `argument_name` - Name of the argument being completed  
472    /// * `current_value` - Current partial value of the argument
473    /// * `context` - Optional context with previously resolved arguments
474    ///
475    /// # Returns
476    /// CompletionInfo with suggestions for the specified prompt argument
477    pub async fn complete_prompt_argument(
478        &self,
479        prompt_name: impl Into<String>,
480        argument_name: impl Into<String>,
481        current_value: impl Into<String>,
482        context: Option<CompletionContext>,
483    ) -> Result<CompletionInfo, ServiceError> {
484        let request = CompleteRequestParams {
485            meta: None,
486            r#ref: Reference::for_prompt(prompt_name),
487            argument: ArgumentInfo {
488                name: argument_name.into(),
489                value: current_value.into(),
490            },
491            context,
492        };
493
494        let result = self.complete(request).await?;
495        Ok(result.completion)
496    }
497
498    /// Convenient method to get completion suggestions for a resource URI argument
499    ///
500    /// # Arguments
501    /// * `uri_template` - URI template pattern being completed
502    /// * `argument_name` - Name of the URI parameter being completed
503    /// * `current_value` - Current partial value of the parameter
504    /// * `context` - Optional context with previously resolved arguments
505    ///
506    /// # Returns
507    /// CompletionInfo with suggestions for the specified resource URI argument
508    pub async fn complete_resource_argument(
509        &self,
510        uri_template: impl Into<String>,
511        argument_name: impl Into<String>,
512        current_value: impl Into<String>,
513        context: Option<CompletionContext>,
514    ) -> Result<CompletionInfo, ServiceError> {
515        let request = CompleteRequestParams {
516            meta: None,
517            r#ref: Reference::for_resource(uri_template),
518            argument: ArgumentInfo {
519                name: argument_name.into(),
520                value: current_value.into(),
521            },
522            context,
523        };
524
525        let result = self.complete(request).await?;
526        Ok(result.completion)
527    }
528
529    /// Simple completion for a prompt argument without context
530    ///
531    /// This is a convenience wrapper around `complete_prompt_argument` for
532    /// simple completion scenarios that don't require context awareness.
533    pub async fn complete_prompt_simple(
534        &self,
535        prompt_name: impl Into<String>,
536        argument_name: impl Into<String>,
537        current_value: impl Into<String>,
538    ) -> Result<Vec<String>, ServiceError> {
539        let completion = self
540            .complete_prompt_argument(prompt_name, argument_name, current_value, None)
541            .await?;
542        Ok(completion.values)
543    }
544
545    /// Simple completion for a resource URI argument without context
546    ///
547    /// This is a convenience wrapper around `complete_resource_argument` for
548    /// simple completion scenarios that don't require context awareness.
549    pub async fn complete_resource_simple(
550        &self,
551        uri_template: impl Into<String>,
552        argument_name: impl Into<String>,
553        current_value: impl Into<String>,
554    ) -> Result<Vec<String>, ServiceError> {
555        let completion = self
556            .complete_resource_argument(uri_template, argument_name, current_value, None)
557            .await?;
558        Ok(completion.values)
559    }
560}
rmcp/service/client.rs

rmcp/service/
client.rs
