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, ConstString, CustomNotification, ElicitRequestParams, ElicitResult, ElicitationAction,
7        ElicitationResponseNotificationMethod, 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: ElicitRequestParams) -> ElicitResult {
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    pub async fn forward_url_elicitation_complete(&self, elicitation_id: String) {
49        let event = McpClientEvent::UrlElicitationComplete(super::UrlElicitationCompleteParams {
50            server_name: self.server_name.clone(),
51            elicitation_id,
52        });
53        if self.event_sender.send(event).await.is_err() {
54            tracing::warn!("Failed to forward URL elicitation completion: receiver dropped");
55        }
56    }
57}
58
59pub fn cancel_result() -> ElicitResult {
60    ElicitResult::new(ElicitationAction::Cancel)
61}
62
63impl ClientHandler for McpClient {
64    fn get_info(&self) -> ClientInfo {
65        self.client_info.clone()
66    }
67
68    async fn on_progress(&self, params: ProgressNotificationParam, _context: NotificationContext<RoleClient>) -> () {
69        self.progress_dispatcher.handle_notification(params).await;
70    }
71
72    async fn create_elicitation(
73        &self,
74        request: ElicitRequestParams,
75        _context: RequestContext<RoleClient>,
76    ) -> Result<ElicitResult, ErrorData> {
77        Ok(self.dispatch_elicitation(request).await)
78    }
79
80    async fn on_custom_notification(
81        &self,
82        notification: CustomNotification,
83        _context: NotificationContext<RoleClient>,
84    ) {
85        if notification.method != ElicitationResponseNotificationMethod::VALUE {
86            return;
87        }
88
89        let Some(elicitation_id) = notification
90            .params
91            .as_ref()
92            .and_then(|params| params.get("elicitationId"))
93            .and_then(serde_json::Value::as_str)
94        else {
95            tracing::warn!("URL elicitation completion notification is missing elicitationId");
96            return;
97        };
98
99        self.forward_url_elicitation_complete(elicitation_id.to_string()).await;
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use rmcp::model::{
107        ClientCapabilities, ElicitationSchema, FormElicitationCapability, Implementation, UrlElicitationCapability,
108    };
109    use std::collections::BTreeMap;
110
111    fn test_client_info() -> ClientInfo {
112        let mut capabilities = ClientCapabilities::builder().enable_elicitation().build();
113        if let Some(elicitation) = capabilities.elicitation.as_mut() {
114            elicitation.form = Some(FormElicitationCapability::default());
115            elicitation.url = Some(UrlElicitationCapability::default());
116        }
117        ClientInfo::new(capabilities, Implementation::new("test", "0.1.0"))
118    }
119
120    fn make_client(event_sender: mpsc::Sender<McpClientEvent>) -> McpClient {
121        McpClient::new(test_client_info(), "test-server".to_string(), event_sender)
122    }
123
124    fn unwrap_elicitation(event: McpClientEvent) -> ElicitationRequest {
125        match event {
126            McpClientEvent::Elicitation(req) => req,
127            other => panic!("expected Elicitation, got {other:?}"),
128        }
129    }
130
131    #[tokio::test]
132    async fn dispatch_elicitation_dropped_sender_returns_cancel() {
133        let (event_tx, _) = mpsc::channel(1);
134        let client = make_client(event_tx);
135
136        let request = ElicitRequestParams::FormElicitationParams {
137            meta: None,
138            message: "test".to_string(),
139            requested_schema: ElicitationSchema::new(BTreeMap::new()),
140        };
141
142        let result = client.dispatch_elicitation(request).await;
143        assert_eq!(result.action, ElicitationAction::Cancel, "dropped sender should return Cancel, not Decline");
144        assert!(result.content.is_none());
145    }
146
147    #[tokio::test]
148    async fn dispatch_elicitation_dropped_receiver_returns_cancel() {
149        let (event_tx, mut event_rx) = mpsc::channel(1);
150        let client = make_client(event_tx);
151
152        let request = ElicitRequestParams::FormElicitationParams {
153            meta: None,
154            message: "test".to_string(),
155            requested_schema: ElicitationSchema::new(BTreeMap::new()),
156        };
157
158        let handle = tokio::spawn(async move {
159            let event = event_rx.recv().await.unwrap();
160            let elicitation = unwrap_elicitation(event);
161            drop(elicitation.response_sender);
162        });
163
164        let result = client.dispatch_elicitation(request).await;
165        handle.await.unwrap();
166
167        assert_eq!(result.action, ElicitationAction::Cancel, "dropped receiver should return Cancel, not Decline");
168        assert!(result.content.is_none());
169    }
170
171    #[tokio::test]
172    async fn dispatch_elicitation_forwards_request_with_server_name() {
173        let (event_tx, mut event_rx) = mpsc::channel(1);
174        let client = make_client(event_tx);
175
176        let request = ElicitRequestParams::UrlElicitationParams {
177            meta: None,
178            message: "Auth".to_string(),
179            url: "https://example.com/auth".to_string(),
180            elicitation_id: "el-123".to_string(),
181        };
182
183        let handle = tokio::spawn(async move {
184            let event = event_rx.recv().await.unwrap();
185            let elicitation = unwrap_elicitation(event);
186            assert_eq!(elicitation.server_name, "test-server");
187            let _ = elicitation.response_sender.send(ElicitResult::new(ElicitationAction::Accept));
188        });
189
190        let result = client.dispatch_elicitation(request).await;
191        handle.await.unwrap();
192        assert_eq!(result.action, ElicitationAction::Accept);
193    }
194
195    #[tokio::test]
196    async fn forward_url_elicitation_complete_uses_server_name_and_id() {
197        let (event_tx, mut event_rx) = mpsc::channel(1);
198        let client = make_client(event_tx);
199
200        client.forward_url_elicitation_complete("el-456".to_string()).await;
201
202        let event = event_rx.recv().await.unwrap();
203        match event {
204            McpClientEvent::UrlElicitationComplete(params) => {
205                assert_eq!(params.server_name, "test-server");
206                assert_eq!(params.elicitation_id, "el-456");
207            }
208            other => panic!("expected UrlElicitationComplete, got {other:?}"),
209        }
210    }
211
212    #[tokio::test]
213    async fn forward_url_elicitation_complete_swallows_dropped_receiver() {
214        let (event_tx, event_rx) = mpsc::channel(1);
215        drop(event_rx);
216        let client = make_client(event_tx);
217
218        // Should not panic even though the receiver is dropped.
219        client.forward_url_elicitation_complete("el-gone".to_string()).await;
220    }
221
222    #[test]
223    fn capabilities_include_form_and_url() {
224        let info = test_client_info();
225        let caps = &info.capabilities;
226        let elicitation = caps.elicitation.as_ref().expect("elicitation capability should be set");
227        assert!(elicitation.form.is_some(), "form capability should be advertised");
228        assert!(elicitation.url.is_some(), "url capability should be advertised");
229    }
230}
mcp_utils/client/mcp_client.rs

mcp_utils/client/
mcp_client.rs
