aether_core/mcp/
run_mcp_task.rs

1use crate::events::TraceContext;
2use mcp_utils::client::{
3    McpClient, McpConnectAttempt, McpConnectionAttemptManager, McpError, McpManager, McpServer, McpServerStatusEntry,
4};
5use mcp_utils::display_meta::ToolResultMeta;
6
7use futures::future::Either;
8use futures::stream::{self, StreamExt};
9use llm::{ToolCallError, ToolCallRequest, ToolCallResult};
10use rmcp::RoleClient;
11use rmcp::model::{
12    CallToolRequestParams, ElicitRequestParams, ErrorCode, GetPromptResult, ProgressNotificationParam, Prompt,
13    RequestMetaObject,
14};
15use rmcp::service::RunningService;
16use std::collections::HashSet;
17use std::sync::Arc;
18use std::time::Duration;
19use tokio::select;
20use tokio::sync::mpsc;
21use tokio::sync::oneshot;
22
23/// Events emitted during tool execution lifecycle
24#[derive(Debug)]
25pub enum ToolExecutionEvent {
26    Progress { tool_id: String, progress: ProgressNotificationParam },
27    Complete { tool_id: String, result: Result<ToolCallResult, ToolCallError>, result_meta: Option<ToolResultMeta> },
28}
29
30const MCP_AUTH_TIMEOUT: Duration = Duration::from_mins(3);
31const URL_ELICITATION_REQUIRED: ErrorCode = ErrorCode(-32042);
32
33/// Commands that can be sent to the MCP manager task
34#[derive(Debug)]
35pub enum McpCommand {
36    ExecuteTool {
37        request: ToolCallRequest,
38        trace_context: Option<TraceContext>,
39        timeout: Duration,
40        tx: mpsc::Sender<ToolExecutionEvent>,
41    },
42    ListPrompts {
43        tx: oneshot::Sender<Result<Vec<Prompt>, String>>,
44    },
45    GetPrompt {
46        name: String,
47        arguments: Option<serde_json::Map<String, serde_json::Value>>,
48        tx: oneshot::Sender<Result<GetPromptResult, String>>,
49    },
50    GetServerStatuses {
51        tx: oneshot::Sender<Vec<McpServerStatusEntry>>,
52    },
53    AuthenticateServer {
54        name: String,
55    },
56}
57
58pub async fn run_mcp_task(
59    mut mcp: McpManager,
60    mut command_rx: mpsc::Receiver<McpCommand>,
61    pending_servers: Vec<McpServer>,
62) {
63    let mut mcp_connection_attempts = McpConnectionAttemptManager::default();
64    let mut pending_connections: HashSet<String> = pending_servers.iter().map(|server| server.name.clone()).collect();
65    for server in pending_servers {
66        let name = server.name.clone();
67        let task = mcp.connect_pending_task(server);
68        mcp_connection_attempts.spawn(name, task);
69    }
70    if pending_connections.is_empty() {
71        mcp.emit_connection_ready().await;
72    }
73
74    loop {
75        select! {
76            command = command_rx.recv() => {
77                let Some(command) = command else { break; };
78                on_command(command, &mut mcp, &mut mcp_connection_attempts).await;
79            }
80
81            Some(joined) = mcp_connection_attempts.join_next(), if !mcp_connection_attempts.is_empty() => {
82                match joined {
83                    Ok(attempt) => {
84                        let was_bootstrap = pending_connections.remove(&attempt.name);
85                        mcp.apply_connection_attempt(attempt).await;
86                        if was_bootstrap && pending_connections.is_empty() {
87                            mcp.emit_connection_ready().await;
88                        }
89                    }
90                    Err(e) => tracing::error!("MCP auth task did not complete normally: {e:?}"),
91                }
92            }
93        }
94    }
95
96    mcp_connection_attempts.shutdown().await;
97    mcp.shutdown().await;
98    tracing::debug!("MCP manager task ended");
99}
100
101async fn on_command(command: McpCommand, mcp: &mut McpManager, auth_tasks: &mut McpConnectionAttemptManager) {
102    match command {
103        McpCommand::ExecuteTool { request, trace_context, timeout, tx } => {
104            let tool_id = request.id.clone();
105
106            match mcp.get_client_for_tool(&request.name, &request.arguments) {
107                Ok((client, params)) => {
108                    let trace_meta = trace_context.as_ref().map(TraceContext::to_meta);
109                    tokio::spawn(async move {
110                        let outcome = execute_mcp_call(
111                            client,
112                            &request,
113                            params,
114                            trace_meta,
115                            timeout,
116                            tool_id.clone(),
117                            tx.clone(),
118                        )
119                        .await;
120                        let (result, result_meta) = match outcome {
121                            Ok((r, m)) => (Ok(r), m),
122                            Err(e) => (Err(e), None),
123                        };
124                        let _ = tx.send(ToolExecutionEvent::Complete { tool_id, result, result_meta }).await;
125                    });
126                }
127                Err(e) => {
128                    tracing::error!("Failed to get client for tool {}: {e}", request.name);
129                    let error = ToolCallError::from_request(&request, format!("Failed to get client: {e}"));
130                    let _ =
131                        tx.send(ToolExecutionEvent::Complete { tool_id, result: Err(error), result_meta: None }).await;
132                }
133            }
134        }
135
136        McpCommand::ListPrompts { tx } => {
137            let result = mcp.list_prompts().await.map_err(|e| format!("Failed to list prompts: {e}"));
138            let _ = tx.send(result);
139        }
140
141        McpCommand::GetPrompt { name: namespaced_name, arguments, tx } => {
142            let result =
143                mcp.get_prompt(&namespaced_name, arguments).await.map_err(|e| format!("Failed to get prompt: {e}"));
144            let _ = tx.send(result);
145        }
146
147        McpCommand::GetServerStatuses { tx } => {
148            let _ = tx.send(mcp.server_statuses());
149        }
150
151        McpCommand::AuthenticateServer { name } => match mcp.authenticate_server_task(&name).await {
152            Ok(task) => {
153                let server_name = name.clone();
154                auth_tasks.spawn(name, async move {
155                    match tokio::time::timeout(MCP_AUTH_TIMEOUT, task).await {
156                        Ok(attempt) => attempt,
157                        Err(_) => McpConnectAttempt::failed(
158                            server_name,
159                            McpError::ConnectionFailed("authentication timed out after 3 minutes".to_string()),
160                            false,
161                        ),
162                    }
163                });
164            }
165            Err(e) => tracing::warn!("Authentication failed for '{name}': {e}"),
166        },
167    }
168}
169
170/// Shared logic for sending an MCP tool call, streaming progress events,
171/// and collecting the result.
172async fn execute_mcp_call(
173    client: Arc<RunningService<RoleClient, McpClient>>,
174    request: &ToolCallRequest,
175    params: CallToolRequestParams,
176    trace_meta: Option<RequestMetaObject>,
177    timeout: Duration,
178    tool_call_id: String,
179    event_tx: mpsc::Sender<ToolExecutionEvent>,
180) -> Result<(ToolCallResult, Option<ToolResultMeta>), ToolCallError> {
181    use super::tool_bridge::mcp_result_to_tool_call_result;
182    use rmcp::model::{ClientRequest::CallToolRequest, Request, ServerResult};
183    use rmcp::service::PeerRequestOptions;
184
185    let handle = client
186        .send_cancellable_request(CallToolRequest(Request::new(params)), {
187            let mut opts = PeerRequestOptions::default();
188            opts.timeout = Some(timeout);
189            opts.meta = trace_meta;
190            opts
191        })
192        .await
193        .map_err(|e| ToolCallError::from_request(request, format!("Failed to send tool request: {e}")))?;
194
195    let progress_subscriber = client.service().progress_dispatcher.subscribe(handle.progress_token.clone()).await;
196
197    let progress_stream = progress_subscriber
198        .map(move |progress| Either::Left(ToolExecutionEvent::Progress { tool_id: tool_call_id.clone(), progress }));
199
200    let result_stream = stream::once(handle.await_response()).map(Either::Right);
201    let combined_stream = stream::select(progress_stream, result_stream);
202    tokio::pin!(combined_stream);
203
204    let server_result = loop {
205        match combined_stream.next().await {
206            Some(Either::Left(progress_event)) => {
207                let _ = event_tx.send(progress_event).await;
208            }
209            Some(Either::Right(result)) => {
210                break match result {
211                    Ok(server_result) => server_result,
212                    Err(e) => {
213                        if let rmcp::service::ServiceError::McpError(ref error_data) = e
214                            && error_data.code == URL_ELICITATION_REQUIRED
215                        {
216                            return Err(handle_url_elicitation_required(&client, request, error_data).await);
217                        }
218                        return Err(ToolCallError::from_request(request, format!("Tool execution failed: {e}")));
219                    }
220                };
221            }
222            None => {
223                return Err(ToolCallError::from_request(request, "Stream ended without result"));
224            }
225        }
226    };
227
228    let ServerResult::CallToolResult(mcp_result) = server_result else {
229        return Err(ToolCallError::from_request(request, "Unexpected response type from MCP server"));
230    };
231
232    mcp_result_to_tool_call_result(request, mcp_result)
233}
234
235#[derive(serde::Deserialize)]
236struct UrlElicitationRequiredData {
237    elicitations: Vec<ElicitRequestParams>,
238}
239
240#[derive(Debug)]
241enum UrlElicitationRequiredParseError {
242    MissingData,
243    InvalidData(serde_json::Error),
244    NoUrlRequests,
245}
246
247impl std::fmt::Display for UrlElicitationRequiredParseError {
248    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
249        match self {
250            Self::MissingData => write!(f, "missing error data"),
251            Self::InvalidData(error) => write!(f, "malformed error data: {error}"),
252            Self::NoUrlRequests => write!(f, "provided no URL elicitation requests"),
253        }
254    }
255}
256
257fn parse_required_url_elicitations(
258    error_data: &rmcp::model::ErrorData,
259) -> Result<Vec<ElicitRequestParams>, UrlElicitationRequiredParseError> {
260    let data = error_data.data.as_ref().ok_or(UrlElicitationRequiredParseError::MissingData)?;
261    let parsed: UrlElicitationRequiredData =
262        serde_json::from_value(data.clone()).map_err(UrlElicitationRequiredParseError::InvalidData)?;
263
264    let url_elicitations = parsed
265        .elicitations
266        .into_iter()
267        .filter(|elicitation| matches!(elicitation, ElicitRequestParams::UrlElicitationParams { .. }))
268        .collect::<Vec<_>>();
269
270    if url_elicitations.is_empty() {
271        return Err(UrlElicitationRequiredParseError::NoUrlRequests);
272    }
273
274    Ok(url_elicitations)
275}
276
277/// Handle a `URL_ELICITATION_REQUIRED` (-32042) error by dispatching each
278/// URL elicitation through the same consent channel used by normal
279/// `create_elicitation` requests.
280async fn handle_url_elicitation_required(
281    client: &Arc<RunningService<RoleClient, McpClient>>,
282    request: &ToolCallRequest,
283    error_data: &rmcp::model::ErrorData,
284) -> ToolCallError {
285    let server_name = client.service().server_name().to_string();
286    let url_elicitations = match parse_required_url_elicitations(error_data) {
287        Ok(url_elicitations) => url_elicitations,
288        Err(UrlElicitationRequiredParseError::NoUrlRequests) => {
289            return ToolCallError::from_request(
290                request,
291                format!("Server '{server_name}' requires URL elicitation but provided no URL elicitation requests"),
292            );
293        }
294        Err(parse_error) => {
295            return ToolCallError::from_request(
296                request,
297                format!("Server '{server_name}' sent an invalid URL elicitation response: {parse_error}"),
298            );
299        }
300    };
301
302    tracing::info!("Server '{server_name}' requires {} URL elicitation(s)", url_elicitations.len());
303
304    for elicitation in url_elicitations {
305        let result = client.service().dispatch_elicitation(elicitation).await;
306        match result.action {
307            rmcp::model::ElicitationAction::Decline => {
308                return ToolCallError::from_request(
309                    request,
310                    format!("Required browser interaction for server '{server_name}' was declined"),
311                );
312            }
313            rmcp::model::ElicitationAction::Cancel => {
314                return ToolCallError::from_request(
315                    request,
316                    format!("Required browser interaction for server '{server_name}' was cancelled"),
317                );
318            }
319            rmcp::model::ElicitationAction::Accept => {
320                tracing::info!("User accepted URL elicitation for server '{server_name}'");
321            }
322            _ => {
323                return ToolCallError::from_request(
324                    request,
325                    format!("Required browser interaction for server '{server_name}' returned an unsupported response"),
326                );
327            }
328        }
329    }
330
331    ToolCallError::from_request(
332        request,
333        format!(
334            "Server '{server_name}' requires a browser flow. The URL has been opened for your approval. Retry the previous request after completing the browser flow."
335        ),
336    )
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn url_elicitation_required_data_parses_url_entries() {
345        let data = serde_json::json!({
346            "elicitations": [
347                {
348                    "mode": "url",
349                    "message": "Auth",
350                    "url": "https://example.com/auth?elicitationId=el-1",
351                    "elicitationId": "el-1"
352                }
353            ]
354        });
355
356        let parsed: UrlElicitationRequiredData = serde_json::from_value(data).unwrap();
357        assert_eq!(parsed.elicitations.len(), 1);
358        assert!(matches!(
359            &parsed.elicitations[0],
360            ElicitRequestParams::UrlElicitationParams { elicitation_id, .. } if elicitation_id == "el-1"
361        ));
362    }
363
364    #[test]
365    fn parse_required_url_elicitations_filters_to_url_only() {
366        let error_data = rmcp::model::ErrorData {
367            code: URL_ELICITATION_REQUIRED,
368            message: "URL elicitation required".into(),
369            data: Some(serde_json::json!({
370                "elicitations": [
371                    {
372                        "mode": "url",
373                        "message": "Auth",
374                        "url": "https://example.com/auth",
375                        "elicitationId": "el-1"
376                    },
377                    {
378                        "mode": "form",
379                        "message": "Pick a color",
380                        "requestedSchema": { "type": "object", "properties": {} }
381                    }
382                ]
383            })),
384        };
385
386        let result = parse_required_url_elicitations(&error_data).unwrap();
387        assert_eq!(result.len(), 1);
388        assert!(matches!(
389            &result[0],
390            ElicitRequestParams::UrlElicitationParams { elicitation_id, .. } if elicitation_id == "el-1"
391        ));
392    }
393}
aether_core/mcp/run_mcp_task.rs

aether_core/mcp/
run_mcp_task.rs
