aether_core/mcp/
mcp_handle.rs

1use super::run_mcp_task::ManagerCommand;
2use futures::{Stream, future::join_all};
3use mcp_utils::client::{CallToolError, CallToolOptions, McpError, McpSnapshot, ToolCallEvent, ToolRoute, call_tool};
4use rmcp::model::{GetPromptRequestParams, GetPromptResult, Prompt};
5use serde_json::{Map, Value};
6use std::{pin::Pin, sync::Arc};
7use thiserror::Error;
8use tokio::sync::{mpsc, oneshot, watch};
9
10pub type ToolCallStream = Pin<Box<dyn Stream<Item = ToolCallEvent> + Send>>;
11
12#[derive(Clone)]
13pub struct McpHandle {
14    control_tx: mpsc::Sender<ManagerCommand>,
15    snapshot_rx: watch::Receiver<Arc<McpSnapshot>>,
16}
17
18#[derive(Debug, Error)]
19pub enum McpHandleError {
20    #[error(transparent)]
21    Route(#[from] McpError),
22    #[error("MCP manager is unavailable")]
23    ManagerUnavailable,
24    #[error("failed to list prompts for {server}: {message}")]
25    PromptList { server: String, message: String },
26    #[error("failed to get prompt '{prompt}' from {server}: {message}")]
27    PromptGet { server: String, prompt: String, message: String },
28}
29
30impl McpHandle {
31    pub(super) fn new(
32        control_tx: mpsc::Sender<ManagerCommand>,
33        snapshot_rx: watch::Receiver<Arc<McpSnapshot>>,
34    ) -> Self {
35        Self { control_tx, snapshot_rx }
36    }
37
38    pub fn snapshot(&self) -> Arc<McpSnapshot> {
39        self.snapshot_rx.borrow().clone()
40    }
41
42    pub fn subscribe(&self) -> watch::Receiver<Arc<McpSnapshot>> {
43        self.snapshot_rx.clone()
44    }
45
46    pub fn call(&self, route: ToolRoute, arguments: Map<String, Value>, options: CallToolOptions) -> ToolCallStream {
47        match self.snapshot().resolve(route, arguments) {
48            Ok((client, params)) => Box::pin(call_tool(client, params, options)),
49            Err(error) => Box::pin(futures::stream::once(async move {
50                ToolCallEvent::Complete(Err(CallToolError::Unavailable {
51                    message: format!("Failed to resolve tool: {error}"),
52                }))
53            })),
54        }
55    }
56
57    pub fn call_model_visible(
58        &self,
59        namespaced_name: String,
60        arguments_json: &str,
61        options: CallToolOptions,
62    ) -> ToolCallStream {
63        let parsed = serde_json::from_str::<Value>(arguments_json)
64            .map_err(McpError::from)
65            .and_then(|value| {
66                value
67                    .as_object()
68                    .cloned()
69                    .ok_or_else(|| McpError::JsonError("tool arguments must be a JSON object".to_string()))
70            })
71            .map(|arguments| (ToolRoute::ModelVisible { namespaced_name }, arguments));
72        match parsed {
73            Ok((route, arguments)) => self.call(route, arguments, options),
74            Err(error) => failed_call(error),
75        }
76    }
77
78    pub async fn list_prompts(&self) -> Result<Vec<Prompt>, McpHandleError> {
79        let futures = self.snapshot().clients_with_prompts().into_iter().map(|(server, client)| async move {
80            let response = client
81                .list_prompts(None)
82                .await
83                .map_err(|error| McpHandleError::PromptList { server: server.clone(), message: error.to_string() })?;
84            Ok::<_, McpHandleError>(
85                response
86                    .prompts
87                    .into_iter()
88                    .map(|prompt| {
89                        Prompt::new(format!("{server}__{}", prompt.name), prompt.description, prompt.arguments)
90                    })
91                    .collect::<Vec<_>>(),
92            )
93        });
94        let mut prompts = Vec::new();
95        for result in join_all(futures).await {
96            prompts.extend(result?);
97        }
98        Ok(prompts)
99    }
100
101    pub async fn get_prompt(
102        &self,
103        name: &str,
104        arguments: Option<Map<String, Value>>,
105    ) -> Result<GetPromptResult, McpHandleError> {
106        let (prompt, client) = self.snapshot().client_for_prompt(name)?;
107        let server = client.service().server_name().to_string();
108        let mut request = GetPromptRequestParams::new(prompt.clone());
109        if let Some(arguments) = arguments {
110            request = request.with_arguments(arguments);
111        }
112        client.get_prompt(request).await.map_err(|error| McpHandleError::PromptGet {
113            server,
114            prompt,
115            message: error.to_string(),
116        })
117    }
118
119    pub async fn authenticate_server(&self, name: &str) -> Result<(), McpHandleError> {
120        let (tx, rx) = oneshot::channel();
121        self.control_tx
122            .send(ManagerCommand::AuthenticateServer { name: name.to_string(), tx })
123            .await
124            .map_err(|_| McpHandleError::ManagerUnavailable)?;
125        rx.await.map_err(|_| McpHandleError::ManagerUnavailable)?.map_err(McpHandleError::Route)
126    }
127}
128
129fn failed_call(error: McpError) -> ToolCallStream {
130    Box::pin(futures::stream::once(async move {
131        ToolCallEvent::Complete(Err(CallToolError::Unavailable { message: format!("Failed to resolve tool: {error}") }))
132    }))
133}
aether_core/mcp/mcp_handle.rs

aether_core/mcp/
mcp_handle.rs
