aether_core/mcp/
gateway_service.rs

1use super::McpHandle;
2use futures::StreamExt;
3use mcp_utils::client::{CallToolOptions, CancellationToken, ToolCallEvent, ToolRoute};
4use mcp_utils::tool_gateway::LIST_SERVERS_TOOL;
5use rmcp::model::{
6    CallToolRequestParams, CallToolResponse, CallToolResult, ErrorData, Implementation, ListToolsResult,
7    PaginatedRequestParams, ServerCapabilities, ServerInfo, Tool, ToolAnnotations,
8};
9use rmcp::{RoleServer, ServerHandler, service::RequestContext};
10use serde_json::{Map, json};
11use std::{
12    sync::Arc,
13    time::{Duration, Instant},
14};
15
16#[derive(Clone)]
17pub struct GatewayService {
18    handle: McpHandle,
19}
20
21impl GatewayService {
22    pub fn new(handle: McpHandle) -> Self {
23        Self { handle }
24    }
25
26    fn tools(&self) -> Vec<Tool> {
27        let snapshot = self.handle.snapshot();
28        let mut tools = snapshot
29            .catalog()
30            .tools()
31            .deferred
32            .into_iter()
33            .map(|tool| {
34                let definition = tool.definition();
35                let schema = definition.parameters.as_object().cloned().unwrap_or_default();
36                let mut result = Tool::new(definition.name.clone(), definition.description.clone(), Arc::new(schema));
37                if let Some(annotations) = &definition.annotations {
38                    let mut converted = ToolAnnotations::new();
39                    converted.title.clone_from(&annotations.title);
40                    converted.read_only_hint = annotations.read_only_hint;
41                    converted.destructive_hint = annotations.destructive_hint;
42                    converted.idempotent_hint = annotations.idempotent_hint;
43                    converted.open_world_hint = annotations.open_world_hint;
44                    result = result.with_annotations(converted);
45                }
46                result
47            })
48            .collect::<Vec<_>>();
49        tools.push(Tool::new(
50            LIST_SERVERS_TOOL,
51            "List connected MCP servers with deferred tools",
52            Arc::new(Map::new()),
53        ));
54        tools
55    }
56
57    fn list_servers(&self) -> CallToolResult {
58        let servers = self.handle.snapshot().catalog().discoverable_deferred_servers();
59        let value = serde_json::to_value(
60            servers
61                .iter()
62                .map(|server| json!({ "name": server.name, "description": server.description }))
63                .collect::<Vec<_>>(),
64        )
65        .expect("server descriptions serialize");
66        CallToolResult::structured(value)
67    }
68}
69
70impl ServerHandler for GatewayService {
71    fn get_info(&self) -> ServerInfo {
72        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
73            .with_server_info(Implementation::new("aether-deferred-tool-gateway", env!("CARGO_PKG_VERSION")))
74    }
75
76    async fn list_tools(
77        &self,
78        _request: Option<PaginatedRequestParams>,
79        _context: RequestContext<RoleServer>,
80    ) -> Result<ListToolsResult, ErrorData> {
81        Ok(ListToolsResult::with_all_items(self.tools()))
82    }
83
84    async fn call_tool(
85        &self,
86        request: CallToolRequestParams,
87        context: RequestContext<RoleServer>,
88    ) -> Result<CallToolResponse, ErrorData> {
89        if request.name == LIST_SERVERS_TOOL {
90            return Ok(self.list_servers().into());
91        }
92
93        let (server, tool) = request
94            .name
95            .split_once("__")
96            .ok_or_else(|| ErrorData::invalid_params("deferred tool names must use server__tool", None))?;
97
98        let arguments = request.arguments.unwrap_or_default();
99        let cancellation = CancellationToken::new();
100        let mut guard = CallCancellationGuard::new(cancellation.clone());
101        let started = Instant::now();
102        let mut events = self.handle.call(
103            ToolRoute::Deferred { server: server.to_string(), tool: tool.to_string() },
104            arguments,
105            CallToolOptions { timeout: Duration::from_mins(10), meta: request.meta, cancel: cancellation.clone() },
106        );
107        let peer = context.peer;
108        let disconnected = async move {
109            while !peer.is_transport_closed() {
110                tokio::time::sleep(Duration::from_millis(25)).await;
111            }
112        };
113        tokio::pin!(disconnected);
114
115        loop {
116            let event = tokio::select! {
117                event = events.next() => event,
118                () = context.ct.cancelled() => None,
119                () = &mut disconnected => None,
120            };
121            let Some(event) = event else {
122                cancellation.cancel();
123                tracing::info!(
124                    route = "deferred",
125                    server,
126                    tool,
127                    outcome = "cancelled",
128                    duration_ms = started.elapsed().as_millis(),
129                    "deferred MCP gateway call completed"
130                );
131                guard.disarm();
132                return Err(ErrorData::internal_error("deferred tool call was cancelled", None));
133            };
134            match event {
135                ToolCallEvent::Complete(result) | ToolCallEvent::TaskComplete { result, .. } => {
136                    let outcome = if result.is_ok() { "success" } else { "error" };
137                    tracing::info!(
138                        route = "deferred",
139                        server,
140                        tool,
141                        outcome,
142                        duration_ms = started.elapsed().as_millis(),
143                        "deferred MCP gateway call completed"
144                    );
145                    guard.disarm();
146                    return result.map(Into::into).map_err(|error| ErrorData::internal_error(error.to_string(), None));
147                }
148                ToolCallEvent::Cancelled { .. } => {
149                    tracing::info!(
150                        route = "deferred",
151                        server,
152                        tool,
153                        outcome = "cancelled",
154                        duration_ms = started.elapsed().as_millis(),
155                        "deferred MCP gateway call completed"
156                    );
157                    guard.disarm();
158                    return Err(ErrorData::internal_error("deferred tool call was cancelled", None));
159                }
160                ToolCallEvent::Progress(_) | ToolCallEvent::TaskCreated(_) | ToolCallEvent::TaskStatus(_) => {}
161            }
162        }
163    }
164}
165
166struct CallCancellationGuard {
167    token: CancellationToken,
168    armed: bool,
169}
170
171impl CallCancellationGuard {
172    fn new(token: CancellationToken) -> Self {
173        Self { token, armed: true }
174    }
175
176    fn disarm(&mut self) {
177        self.armed = false;
178    }
179}
180
181impl Drop for CallCancellationGuard {
182    fn drop(&mut self) {
183        if self.armed {
184            self.token.cancel();
185        }
186    }
187}
aether_core/mcp/gateway_service.rs

aether_core/mcp/
gateway_service.rs
