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        ClientCapabilities, ClientConfig, CustomNotification, ElicitRequestParams, ElicitResult, ElicitationAction,
7        ElicitationCapability, ErrorData, FormElicitationCapability, ProgressNotificationParam,
8        UrlElicitationCapability,
9    },
10    service::{NotificationContext, RequestContext},
11};
12use std::result::Result;
13use tokio::sync::{mpsc, oneshot};
14
15use crate::client::{ElicitationRequest, McpClientEvent, elicitation::with_meta, manager::ToolListChangedRequest};
16
17pub struct McpClient {
18    client_info: ClientConfig,
19    server_name: String,
20    pub(crate) progress_dispatcher: ProgressDispatcher,
21    event_sender: mpsc::Sender<McpClientEvent>,
22    tool_refresh_sender: Option<mpsc::Sender<ToolListChangedRequest>>,
23    connection_generation: u64,
24}
25
26impl McpClient {
27    pub fn new(client_info: ClientConfig, server_name: String, event_sender: mpsc::Sender<McpClientEvent>) -> Self {
28        Self {
29            client_info,
30            server_name,
31            progress_dispatcher: ProgressDispatcher::new(),
32            event_sender,
33            tool_refresh_sender: None,
34            connection_generation: 0,
35        }
36    }
37
38    pub(super) fn with_tool_refresh(
39        mut self,
40        sender: mpsc::Sender<ToolListChangedRequest>,
41        connection_generation: u64,
42    ) -> Self {
43        self.tool_refresh_sender = Some(sender);
44        self.connection_generation = connection_generation;
45        self
46    }
47
48    pub fn server_name(&self) -> &str {
49        &self.server_name
50    }
51
52    /// Dispatch an elicitation request through the shared event channel.
53    ///
54    /// Used by both the `create_elicitation` handler and the MRTR round loop
55    /// in `call_tool_mrtr` to ensure the same user-facing flow.
56    pub async fn dispatch_elicitation(&self, request: ElicitRequestParams) -> ElicitResult {
57        let (response_tx, response_rx) = oneshot::channel();
58        let elicitation_request =
59            ElicitationRequest { server_name: self.server_name.clone(), request, response_sender: response_tx };
60
61        if self.event_sender.send(McpClientEvent::Elicitation(Box::new(elicitation_request))).await.is_err() {
62            return cancel_result();
63        }
64        response_rx.await.unwrap_or_else(|_| cancel_result())
65    }
66}
67
68pub fn cancel_result() -> ElicitResult {
69    ElicitResult::new(ElicitationAction::Cancel)
70}
71
72pub fn client_capabilities() -> ClientCapabilities {
73    client_capabilities_for(true, true)
74}
75
76pub fn client_capabilities_for(form: bool, url: bool) -> ClientCapabilities {
77    let mut capabilities = ClientCapabilities::builder().enable_tasks().build();
78    if form || url {
79        let mut elicitation = ElicitationCapability::new();
80        elicitation.form = form.then(FormElicitationCapability::default);
81        elicitation.url = url.then(UrlElicitationCapability::default);
82        capabilities.elicitation = Some(elicitation);
83    }
84    capabilities
85}
86
87impl ClientHandler for McpClient {
88    fn get_info(&self) -> ClientConfig {
89        self.client_info.clone()
90    }
91
92    async fn on_progress(&self, params: ProgressNotificationParam, _context: NotificationContext<RoleClient>) -> () {
93        self.progress_dispatcher.handle_notification(params).await;
94    }
95
96    async fn create_elicitation(
97        &self,
98        request: ElicitRequestParams,
99        context: RequestContext<RoleClient>,
100    ) -> Result<ElicitResult, ErrorData> {
101        let meta = (!context.meta.is_empty()).then_some(context.meta);
102        Ok(self.dispatch_elicitation(with_meta(request, meta)).await)
103    }
104
105    async fn on_custom_notification(
106        &self,
107        notification: CustomNotification,
108        _context: NotificationContext<RoleClient>,
109    ) {
110        if notification.method != "notifications/elicitation/complete" {
111            return;
112        }
113        let params: Option<ElicitationCompleteParams> =
114            notification.params.and_then(|params| serde_json::from_value(params).ok());
115        let Some(params) = params else {
116            tracing::warn!("Ignoring malformed MCP elicitation completion notification");
117            return;
118        };
119        let _ = self
120            .event_sender
121            .send(McpClientEvent::ElicitationComplete {
122                server_name: self.server_name.clone(),
123                elicitation_id: params.elicitation_id,
124            })
125            .await;
126    }
127
128    async fn on_tool_list_changed(&self, context: NotificationContext<RoleClient>) {
129        let Some(sender) = &self.tool_refresh_sender else {
130            return;
131        };
132        let request = ToolListChangedRequest::new(self.server_name.clone(), self.connection_generation, context.peer);
133        if sender.send(request).await.is_err() {
134            tracing::debug!(server = %self.server_name, "MCP tool refresh receiver closed");
135        }
136    }
137}
138
139#[derive(serde::Deserialize)]
140#[serde(rename_all = "camelCase")]
141struct ElicitationCompleteParams {
142    elicitation_id: String,
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use rmcp::model::{ElicitationSchema, Implementation};
149    use std::collections::BTreeMap;
150
151    fn test_client_info() -> ClientConfig {
152        ClientConfig::new(client_capabilities(), Implementation::new("test", "0.1.0"))
153    }
154
155    fn make_client(event_sender: mpsc::Sender<McpClientEvent>) -> McpClient {
156        McpClient::new(test_client_info(), "test-server".to_string(), event_sender)
157    }
158
159    fn unwrap_elicitation(event: McpClientEvent) -> ElicitationRequest {
160        match event {
161            McpClientEvent::Elicitation(req) => *req,
162            other => panic!("expected Elicitation, got {other:?}"),
163        }
164    }
165
166    #[tokio::test]
167    async fn dispatch_elicitation_dropped_sender_returns_cancel() {
168        let (event_tx, _) = mpsc::channel(1);
169        let client = make_client(event_tx);
170
171        let request = ElicitRequestParams::FormElicitationParams {
172            meta: None,
173            message: "test".to_string(),
174            requested_schema: ElicitationSchema::new(BTreeMap::new()),
175        };
176
177        let result = client.dispatch_elicitation(request).await;
178        assert_eq!(result.action, ElicitationAction::Cancel, "dropped sender should return Cancel, not Decline");
179        assert!(result.content.is_none());
180    }
181
182    #[tokio::test]
183    async fn dispatch_elicitation_dropped_receiver_returns_cancel() {
184        let (event_tx, mut event_rx) = mpsc::channel(1);
185        let client = make_client(event_tx);
186
187        let request = ElicitRequestParams::FormElicitationParams {
188            meta: None,
189            message: "test".to_string(),
190            requested_schema: ElicitationSchema::new(BTreeMap::new()),
191        };
192
193        let handle = tokio::spawn(async move {
194            let event = event_rx.recv().await.unwrap();
195            let elicitation = unwrap_elicitation(event);
196            drop(elicitation.response_sender);
197        });
198
199        let result = client.dispatch_elicitation(request).await;
200        handle.await.unwrap();
201
202        assert_eq!(result.action, ElicitationAction::Cancel, "dropped receiver should return Cancel, not Decline");
203        assert!(result.content.is_none());
204    }
205
206    #[tokio::test]
207    async fn dispatch_elicitation_forwards_request_with_server_name() {
208        let (event_tx, mut event_rx) = mpsc::channel(1);
209        let client = make_client(event_tx);
210
211        let request = ElicitRequestParams::UrlElicitationParams {
212            meta: None,
213            message: "Auth".to_string(),
214            url: "https://example.com/auth".to_string(),
215            elicitation_id: "el-123".to_string(),
216        };
217
218        let handle = tokio::spawn(async move {
219            let event = event_rx.recv().await.unwrap();
220            let elicitation = unwrap_elicitation(event);
221            assert_eq!(elicitation.server_name, "test-server");
222            let _ = elicitation.response_sender.send(ElicitResult::new(ElicitationAction::Accept));
223        });
224
225        let result = client.dispatch_elicitation(request).await;
226        handle.await.unwrap();
227        assert_eq!(result.action, ElicitationAction::Accept);
228    }
229
230    #[test]
231    fn capabilities_include_form_url_and_tasks() {
232        let info = test_client_info();
233        let caps = &info.capabilities;
234        let elicitation = caps.elicitation.as_ref().expect("elicitation capability should be set");
235        assert!(elicitation.form.is_some(), "form capability should be advertised");
236        assert!(elicitation.url.is_some(), "url capability should be advertised");
237        assert!(
238            caps.extensions.as_ref().is_some_and(|extensions| extensions.contains_key("io.modelcontextprotocol/tasks"))
239        );
240    }
241}
mcp_utils/client/mcp_client.rs

mcp_utils/client/
mcp_client.rs
