1use super::mcp_catalog::{CatalogCore, Tier, namespaced};
2use mcp_utils::client::{ClientError, ClientOptions, McpClient, ToolCall, ToolCallError, ToolCallOptions, Transport};
3use mcp_utils::server::{McpServer, ServerHandle, UnixSocket};
4use rmcp::model::{
5 CallToolRequestParams, CallToolResponse, CallToolResult, ErrorData, Implementation, ListToolsResult,
6 PaginatedRequestParams, ServerCapabilities, ServerConfig, Tool,
7};
8use rmcp::{RoleServer, ServerHandler, service::RequestContext};
9use serde::{Deserialize, Serialize};
10use serde_json::{Map, Value};
11use std::future::{Future, ready};
12use std::path::PathBuf;
13use std::sync::Arc;
14use std::time::{Duration, Instant};
15use thiserror::Error;
16
17pub const AETHER_MCP_IPC_SOCKET: &str = "AETHER_MCP_IPC_SOCKET";
18pub(crate) const PROGRESSIVE_DISCOVERY_INSTRUCTIONS: &str = include_str!("progressive_discovery_instructions.md");
19
20#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
21pub struct DeferredServer {
22 pub name: String,
23 pub description: String,
24 pub tools: Vec<DeferredTool>,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28pub struct DeferredTool {
29 pub name: String,
30 #[serde(default, skip_serializing_if = "Option::is_none")]
31 pub description: Option<String>,
32}
33
34pub struct DeferredToolsClient {
35 client: McpClient,
36}
37
38#[derive(Debug, Error)]
39pub enum DeferredToolsError {
40 #[error("failed to connect to the active Aether session: {0}")]
41 Connect(#[source] ClientError),
42 #[error("deferred tool discovery timed out after {} seconds", .0.as_secs())]
43 TimedOut(Duration),
44 #[error("failed to list deferred servers: {0}")]
45 ListServers(#[source] ToolCallError),
46 #[error("invalid deferred server listing: {0}")]
47 MalformedListing(#[source] serde_json::Error),
48 #[error("failed to list deferred tools: {0}")]
49 ListTools(#[source] ClientError),
50}
51
52impl DeferredToolsClient {
53 pub async fn connect(socket: PathBuf) -> Result<Self, DeferredToolsError> {
54 let options = ClientOptions::default().implementation(Implementation::new("aether", env!("CARGO_PKG_VERSION")));
55 let client = McpClient::connect("aether", Transport::Unix(socket), &options).await;
56 client.map(|client| Self { client }).map_err(DeferredToolsError::Connect)
57 }
58
59 pub async fn servers(&self) -> Result<Vec<DeferredServer>, DeferredToolsError> {
60 let options = ToolCallOptions::with_timeout(DISCOVERY_TIMEOUT);
61 let listed = self.client.call_tool(LIST_SERVERS_TOOL, Map::new(), options).result().await;
62 let listed = listed.map_err(|error| match error {
63 ToolCallError::TimedOut(timeout) => DeferredToolsError::TimedOut(timeout),
64 error => DeferredToolsError::ListServers(error),
65 })?;
66 serde_json::from_value(listed.structured_content.unwrap_or_default())
67 .map_err(DeferredToolsError::MalformedListing)
68 }
69
70 pub async fn tool(&self, server: &str, tool: &str) -> Result<Option<Tool>, DeferredToolsError> {
71 let tools = tokio::time::timeout(DISCOVERY_TIMEOUT, self.client.list_tools())
72 .await
73 .map_err(|_| DeferredToolsError::TimedOut(DISCOVERY_TIMEOUT))?
74 .map_err(DeferredToolsError::ListTools)?;
75 let name = namespaced(server, tool);
76 Ok(tools.into_iter().find(|definition| definition.name == name))
77 }
78
79 pub fn call(&self, server: &str, tool: &str, args: Map<String, Value>, timeout: Duration) -> ToolCall {
80 self.client.call_tool(&namespaced(server, tool), args, ToolCallOptions::with_timeout(timeout))
81 }
82
83 pub async fn close(&self) {
84 self.client.close().await;
85 }
86}
87
88pub(super) fn serve(catalog: CatalogCore, socket: UnixSocket) -> ServerHandle {
89 McpServer::new(DeferredToolsMcp { catalog }).serve_unix(socket)
90}
91
92const LIST_SERVERS_TOOL: &str = "_aether_list_servers";
93const CALL_TIMEOUT: Duration = Duration::from_mins(10);
94const DISCOVERY_TIMEOUT: Duration = Duration::from_secs(5);
95
96struct DeferredToolsMcp {
97 catalog: CatalogCore,
98}
99
100impl DeferredToolsMcp {
101 fn tools(&self) -> Vec<Tool> {
102 let list_servers =
103 Tool::new(LIST_SERVERS_TOOL, "List connected MCP servers with deferred tools", Arc::new(Map::new()));
104 let mut tools: Vec<Tool> = self.catalog.tools(Tier::Deferred).into_iter().map(|(_, tool)| tool).collect();
105 tools.push(list_servers);
106 tools
107 }
108
109 fn list_servers(&self) -> CallToolResult {
110 let servers = self
111 .catalog
112 .servers(Tier::Deferred)
113 .into_iter()
114 .map(|(server, tools)| DeferredServer {
115 name: server.name().to_string(),
116 description: server.description().unwrap_or_else(|| server.name().to_string()),
117 tools: tools
118 .into_iter()
119 .map(|tool| DeferredTool {
120 name: tool.name.to_string(),
121 description: tool.description.map(Into::into),
122 })
123 .collect(),
124 })
125 .collect::<Vec<_>>();
126 CallToolResult::structured(serde_json::to_value(servers).expect("deferred servers serialize to JSON"))
127 }
128}
129
130impl ServerHandler for DeferredToolsMcp {
131 fn get_info(&self) -> ServerConfig {
132 ServerConfig::new(ServerCapabilities::builder().enable_tools().build())
133 .with_server_info(Implementation::new("aether-deferred-tools", env!("CARGO_PKG_VERSION")))
134 }
135
136 fn list_tools(
137 &self,
138 _request: Option<PaginatedRequestParams>,
139 _context: RequestContext<RoleServer>,
140 ) -> impl Future<Output = Result<ListToolsResult, ErrorData>> + Send + '_ {
141 ready(Ok(ListToolsResult::with_all_items(self.tools())))
142 }
143
144 async fn call_tool(
145 &self,
146 request: CallToolRequestParams,
147 context: RequestContext<RoleServer>,
148 ) -> Result<CallToolResponse, ErrorData> {
149 if request.name == LIST_SERVERS_TOOL {
150 return Ok(self.list_servers().into());
151 }
152 let name = &request.name;
153 if self.catalog.tier_of(name) == Some(Tier::Direct) {
154 let message = format!("Tool '{name}' is exposed directly; call '{name}' as a regular tool instead");
155 return Err(ErrorData::invalid_params(message, None));
156 }
157
158 let options = ToolCallOptions { timeout: Some(CALL_TIMEOUT), meta: request.meta, cancel: context.ct.clone() };
159 let args = request.arguments.unwrap_or_default();
160 let Some(call) = self.catalog.call(Tier::Deferred, name, args, options) else {
161 return Err(ErrorData::invalid_params(format!("Tool not found: {name}"), None));
162 };
163 forward(call, name, &context).await
164 }
165}
166
167async fn forward(
168 call: ToolCall,
169 name: &str,
170 context: &RequestContext<RoleServer>,
171) -> Result<CallToolResponse, ErrorData> {
172 let started = Instant::now();
173 let progress_token = context.meta.get_progress_token();
174 let result = call
175 .result_with_progress(|mut progress| {
176 let (peer, token) = (context.peer.clone(), progress_token.clone());
177 async move {
178 if let Some(token) = token {
179 progress.progress_token = token;
180 let _ = peer.notify_progress(progress).await;
181 }
182 }
183 })
184 .await;
185
186 let outcome = match &result {
187 Ok(_) => "success",
188 Err(ToolCallError::Cancelled) => "cancelled",
189 Err(_) => "error",
190 };
191 tracing::info!(tool = name, outcome, duration_ms = started.elapsed().as_millis(), "MCP tool call forwarded");
192
193 result.map(Into::into).map_err(|error| ErrorData::internal_error(error.to_string(), None))
194}