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, ClientInfo, 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, manager::ToolListChangedRequest};
16
17pub struct McpClient {
18    client_info: ClientInfo,
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: ClientInfo, 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) -> ClientInfo {
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        Ok(self.dispatch_elicitation(request).await)
102    }
103
104    async fn on_custom_notification(
105        &self,
106        notification: CustomNotification,
107        _context: NotificationContext<RoleClient>,
108    ) {
109        if notification.method != "notifications/elicitation/complete" {
110            return;
111        }
112        let params: Option<ElicitationCompleteParams> =
113            notification.params.and_then(|params| serde_json::from_value(params).ok());
114        let Some(params) = params else {
115            tracing::warn!("Ignoring malformed MCP elicitation completion notification");
116            return;
117        };
118        let _ = self
119            .event_sender
120            .send(McpClientEvent::ElicitationComplete {
121                server_name: self.server_name.clone(),
122                elicitation_id: params.elicitation_id,
123            })
124            .await;
125    }
126
127    async fn on_tool_list_changed(&self, context: NotificationContext<RoleClient>) {
128        let Some(sender) = &self.tool_refresh_sender else {
129            return;
130        };
131        let request = ToolListChangedRequest::new(self.server_name.clone(), self.connection_generation, context.peer);
132        if sender.send(request).await.is_err() {
133            tracing::debug!(server = %self.server_name, "MCP tool refresh receiver closed");
134        }
135    }
136}
137
138#[derive(serde::Deserialize)]
139#[serde(rename_all = "camelCase")]
140struct ElicitationCompleteParams {
141    elicitation_id: String,
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147    use rmcp::model::{ElicitationSchema, Implementation};
148    use std::collections::BTreeMap;
149
150    fn test_client_info() -> ClientInfo {
151        ClientInfo::new(client_capabilities(), Implementation::new("test", "0.1.0"))
152    }
153
154    fn make_client(event_sender: mpsc::Sender<McpClientEvent>) -> McpClient {
155        McpClient::new(test_client_info(), "test-server".to_string(), event_sender)
156    }
157
158    fn unwrap_elicitation(event: McpClientEvent) -> ElicitationRequest {
159        match event {
160            McpClientEvent::Elicitation(req) => *req,
161            other => panic!("expected Elicitation, got {other:?}"),
162        }
163    }
164
165    #[tokio::test]
166    async fn dispatch_elicitation_dropped_sender_returns_cancel() {
167        let (event_tx, _) = mpsc::channel(1);
168        let client = make_client(event_tx);
169
170        let request = ElicitRequestParams::FormElicitationParams {
171            meta: None,
172            message: "test".to_string(),
173            requested_schema: ElicitationSchema::new(BTreeMap::new()),
174        };
175
176        let result = client.dispatch_elicitation(request).await;
177        assert_eq!(result.action, ElicitationAction::Cancel, "dropped sender should return Cancel, not Decline");
178        assert!(result.content.is_none());
179    }
180
181    #[tokio::test]
182    async fn dispatch_elicitation_dropped_receiver_returns_cancel() {
183        let (event_tx, mut event_rx) = mpsc::channel(1);
184        let client = make_client(event_tx);
185
186        let request = ElicitRequestParams::FormElicitationParams {
187            meta: None,
188            message: "test".to_string(),
189            requested_schema: ElicitationSchema::new(BTreeMap::new()),
190        };
191
192        let handle = tokio::spawn(async move {
193            let event = event_rx.recv().await.unwrap();
194            let elicitation = unwrap_elicitation(event);
195            drop(elicitation.response_sender);
196        });
197
198        let result = client.dispatch_elicitation(request).await;
199        handle.await.unwrap();
200
201        assert_eq!(result.action, ElicitationAction::Cancel, "dropped receiver should return Cancel, not Decline");
202        assert!(result.content.is_none());
203    }
204
205    #[tokio::test]
206    async fn dispatch_elicitation_forwards_request_with_server_name() {
207        let (event_tx, mut event_rx) = mpsc::channel(1);
208        let client = make_client(event_tx);
209
210        let request = ElicitRequestParams::UrlElicitationParams {
211            meta: None,
212            message: "Auth".to_string(),
213            url: "https://example.com/auth".to_string(),
214            elicitation_id: "el-123".to_string(),
215        };
216
217        let handle = tokio::spawn(async move {
218            let event = event_rx.recv().await.unwrap();
219            let elicitation = unwrap_elicitation(event);
220            assert_eq!(elicitation.server_name, "test-server");
221            let _ = elicitation.response_sender.send(ElicitResult::new(ElicitationAction::Accept));
222        });
223
224        let result = client.dispatch_elicitation(request).await;
225        handle.await.unwrap();
226        assert_eq!(result.action, ElicitationAction::Accept);
227    }
228
229    #[test]
230    fn capabilities_include_form_url_and_tasks() {
231        let info = test_client_info();
232        let caps = &info.capabilities;
233        let elicitation = caps.elicitation.as_ref().expect("elicitation capability should be set");
234        assert!(elicitation.form.is_some(), "form capability should be advertised");
235        assert!(elicitation.url.is_some(), "url capability should be advertised");
236        assert!(
237            caps.extensions.as_ref().is_some_and(|extensions| extensions.contains_key("io.modelcontextprotocol/tasks"))
238        );
239    }
240}
mcp_utils/client/mcp_client.rs

mcp_utils/client/
mcp_client.rs
