mcp_utils/client/
mcp_client.rs

1// Don't use custom Result type here as we need to return rmcp::ErrorData
2use rmcp::{
3    ClientHandler, RoleClient,
4    handler::client::progress::ProgressDispatcher,
5    model::{
6        ClientInfo, CreateElicitationRequestParams, CreateElicitationResult, ElicitationAction,
7        ElicitationResponseNotificationParam, ErrorData, ProgressNotificationParam,
8    },
9    service::{NotificationContext, RequestContext},
10};
11use std::result::Result;
12use tokio::sync::{mpsc, oneshot};
13
14use crate::client::{ElicitationRequest, McpClientEvent};
15
16pub struct McpClient {
17    client_info: ClientInfo,
18    server_name: String,
19    pub progress_dispatcher: ProgressDispatcher,
20    event_sender: mpsc::Sender<McpClientEvent>,
21}
22
23impl McpClient {
24    pub fn new(client_info: ClientInfo, server_name: String, event_sender: mpsc::Sender<McpClientEvent>) -> Self {
25        Self { client_info, server_name, progress_dispatcher: ProgressDispatcher::new(), event_sender }
26    }
27
28    pub fn server_name(&self) -> &str {
29        &self.server_name
30    }
31
32    /// Dispatch an elicitation request through the shared event channel.
33    ///
34    /// Used by both the `create_elicitation` handler and the `-32042`
35    /// `URL_ELICITATION_REQUIRED` error path to ensure the same user-facing flow.
36    pub async fn dispatch_elicitation(&self, request: CreateElicitationRequestParams) -> CreateElicitationResult {
37        let (response_tx, response_rx) = oneshot::channel();
38        let elicitation_request =
39            ElicitationRequest { server_name: self.server_name.clone(), request, response_sender: response_tx };
40
41        if self.event_sender.send(McpClientEvent::Elicitation(elicitation_request)).await.is_err() {
42            return cancel_result();
43        }
44        response_rx.await.unwrap_or_else(|_| cancel_result())
45    }
46
47    /// Forward a URL elicitation completion through the shared event channel.
48    ///
49    /// Split out from `on_url_elicitation_notification_complete` so it can be
50    /// tested without constructing a `NotificationContext`.
51    pub async fn forward_url_elicitation_complete(&self, elicitation_id: String) {
52        let event = McpClientEvent::UrlElicitationComplete(super::UrlElicitationCompleteParams {
53            server_name: self.server_name.clone(),
54            elicitation_id,
55        });
56        if self.event_sender.send(event).await.is_err() {
57            tracing::warn!("Failed to forward URL elicitation completion: receiver dropped");
58        }
59    }
60}
61
62pub fn cancel_result() -> CreateElicitationResult {
63    CreateElicitationResult { action: ElicitationAction::Cancel, content: None, meta: Option::default() }
64}
65
66impl ClientHandler for McpClient {
67    fn get_info(&self) -> ClientInfo {
68        self.client_info.clone()
69    }
70
71    async fn on_progress(&self, params: ProgressNotificationParam, _context: NotificationContext<RoleClient>) -> () {
72        self.progress_dispatcher.handle_notification(params).await;
73    }
74
75    async fn create_elicitation(
76        &self,
77        request: CreateElicitationRequestParams,
78        _context: RequestContext<RoleClient>,
79    ) -> Result<CreateElicitationResult, ErrorData> {
80        Ok(self.dispatch_elicitation(request).await)
81    }
82
83    async fn on_url_elicitation_notification_complete(
84        &self,
85        params: ElicitationResponseNotificationParam,
86        _context: NotificationContext<RoleClient>,
87    ) {
88        self.forward_url_elicitation_complete(params.elicitation_id).await;
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use rmcp::model::{
96        ClientCapabilities, ElicitationSchema, FormElicitationCapability, Implementation, UrlElicitationCapability,
97    };
98    use std::collections::BTreeMap;
99
100    fn test_client_info() -> ClientInfo {
101        let mut capabilities = ClientCapabilities::builder().enable_elicitation().build();
102        if let Some(elicitation) = capabilities.elicitation.as_mut() {
103            elicitation.form = Some(FormElicitationCapability::default());
104            elicitation.url = Some(UrlElicitationCapability::default());
105        }
106        ClientInfo::new(capabilities, Implementation::new("test", "0.1.0"))
107    }
108
109    fn make_client(event_sender: mpsc::Sender<McpClientEvent>) -> McpClient {
110        McpClient::new(test_client_info(), "test-server".to_string(), event_sender)
111    }
112
113    fn unwrap_elicitation(event: McpClientEvent) -> ElicitationRequest {
114        match event {
115            McpClientEvent::Elicitation(req) => req,
116            other => panic!("expected Elicitation, got {other:?}"),
117        }
118    }
119
120    #[tokio::test]
121    async fn dispatch_elicitation_dropped_sender_returns_cancel() {
122        let (event_tx, _) = mpsc::channel(1);
123        let client = make_client(event_tx);
124
125        let request = CreateElicitationRequestParams::FormElicitationParams {
126            meta: None,
127            message: "test".to_string(),
128            requested_schema: ElicitationSchema::new(BTreeMap::new()),
129        };
130
131        let result = client.dispatch_elicitation(request).await;
132        assert_eq!(result.action, ElicitationAction::Cancel, "dropped sender should return Cancel, not Decline");
133        assert!(result.content.is_none());
134    }
135
136    #[tokio::test]
137    async fn dispatch_elicitation_dropped_receiver_returns_cancel() {
138        let (event_tx, mut event_rx) = mpsc::channel(1);
139        let client = make_client(event_tx);
140
141        let request = CreateElicitationRequestParams::FormElicitationParams {
142            meta: None,
143            message: "test".to_string(),
144            requested_schema: ElicitationSchema::new(BTreeMap::new()),
145        };
146
147        let handle = tokio::spawn(async move {
148            let event = event_rx.recv().await.unwrap();
149            let elicitation = unwrap_elicitation(event);
150            drop(elicitation.response_sender);
151        });
152
153        let result = client.dispatch_elicitation(request).await;
154        handle.await.unwrap();
155
156        assert_eq!(result.action, ElicitationAction::Cancel, "dropped receiver should return Cancel, not Decline");
157        assert!(result.content.is_none());
158    }
159
160    #[tokio::test]
161    async fn dispatch_elicitation_forwards_request_with_server_name() {
162        let (event_tx, mut event_rx) = mpsc::channel(1);
163        let client = make_client(event_tx);
164
165        let request = CreateElicitationRequestParams::UrlElicitationParams {
166            meta: None,
167            message: "Auth".to_string(),
168            url: "https://example.com/auth".to_string(),
169            elicitation_id: "el-123".to_string(),
170        };
171
172        let handle = tokio::spawn(async move {
173            let event = event_rx.recv().await.unwrap();
174            let elicitation = unwrap_elicitation(event);
175            assert_eq!(elicitation.server_name, "test-server");
176            let _ = elicitation.response_sender.send(CreateElicitationResult {
177                action: ElicitationAction::Accept,
178                content: None,
179                meta: Option::default(),
180            });
181        });
182
183        let result = client.dispatch_elicitation(request).await;
184        handle.await.unwrap();
185        assert_eq!(result.action, ElicitationAction::Accept);
186    }
187
188    #[tokio::test]
189    async fn forward_url_elicitation_complete_uses_server_name_and_id() {
190        let (event_tx, mut event_rx) = mpsc::channel(1);
191        let client = make_client(event_tx);
192
193        client.forward_url_elicitation_complete("el-456".to_string()).await;
194
195        let event = event_rx.recv().await.unwrap();
196        match event {
197            McpClientEvent::UrlElicitationComplete(params) => {
198                assert_eq!(params.server_name, "test-server");
199                assert_eq!(params.elicitation_id, "el-456");
200            }
201            other => panic!("expected UrlElicitationComplete, got {other:?}"),
202        }
203    }
204
205    #[tokio::test]
206    async fn forward_url_elicitation_complete_swallows_dropped_receiver() {
207        let (event_tx, event_rx) = mpsc::channel(1);
208        drop(event_rx);
209        let client = make_client(event_tx);
210
211        // Should not panic even though the receiver is dropped.
212        client.forward_url_elicitation_complete("el-gone".to_string()).await;
213    }
214
215    #[test]
216    fn capabilities_include_form_and_url() {
217        let info = test_client_info();
218        let caps = &info.capabilities;
219        let elicitation = caps.elicitation.as_ref().expect("elicitation capability should be set");
220        assert!(elicitation.form.is_some(), "form capability should be advertised");
221        assert!(elicitation.url.is_some(), "url capability should be advertised");
222    }
223}
mcp_utils/client/mcp_client.rs

mcp_utils/client/
mcp_client.rs
