mcp_multiplexer/
server.rs

1use crate::config::Config;
2use crate::model::ToolInfo;
3use crate::search::search;
4use crate::stats::Stats;
5use crate::upstream::Upstreams;
6use rmcp::handler::server::router::tool::ToolRouter;
7use rmcp::handler::server::tool::ToolCallContext;
8use rmcp::handler::server::wrapper::Parameters;
9use rmcp::model::*;
10use rmcp::service::RequestContext;
11use rmcp::{RoleServer, ServerHandler, tool, tool_handler, tool_router};
12use schemars::JsonSchema;
13use serde::Serialize;
14use std::sync::Arc;
15
16#[derive(Serialize)]
17pub struct ServerSummary {
18    pub name: String,
19    pub status: String,
20    pub tool_count: usize,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub instructions: Option<String>,
23}
24
25#[derive(Serialize)]
26pub struct ToolSummary {
27    pub name: String,
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub description: Option<String>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub annotations: Option<serde_json::Value>,
32}
33
34#[derive(Serialize)]
35pub struct RefreshResult {
36    pub name: String,
37    pub status: String,
38    #[serde(skip_serializing_if = "Option::is_none")]
39    pub tool_count: Option<usize>,
40    #[serde(skip_serializing_if = "Option::is_none")]
41    pub error: Option<String>,
42}
43
44#[derive(serde::Deserialize, JsonSchema)]
45pub struct ServerArg {
46    pub server: String,
47}
48
49#[derive(serde::Deserialize, JsonSchema)]
50pub struct SearchArgs {
51    pub query: String,
52    pub server: Option<String>,
53    pub limit: Option<usize>,
54}
55
56#[derive(serde::Deserialize, JsonSchema)]
57pub struct RefreshArgs {
58    pub server: Option<String>,
59}
60
61#[derive(serde::Deserialize, JsonSchema)]
62pub struct DescribeArgs {
63    pub server: String,
64    pub tool: String,
65}
66
67#[derive(serde::Deserialize, JsonSchema)]
68pub struct CallArgs {
69    pub server: String,
70    pub tool: String,
71    pub arguments: Option<serde_json::Map<String, serde_json::Value>>,
72}
73
74#[derive(serde::Deserialize, JsonSchema)]
75pub struct AuthorizeArgs {
76    pub server: String,
77    /// Headless completion: the final redirect URL (http://127.0.0.1:.../callback?code=...)
78    pub pasted_url: Option<String>,
79}
80
81#[derive(Clone)]
82pub struct Aggregator {
83    cfg: Config,
84    ups: Arc<Upstreams>,
85    tool_router: ToolRouter<Self>,
86    stats: Arc<Stats>,
87}
88
89fn internal(e: impl std::fmt::Display) -> ErrorData {
90    ErrorData::internal_error(e.to_string(), None)
91}
92
93fn json_len<T: serde::Serialize>(v: &T) -> u64 {
94    serde_json::to_string(v)
95        .map(|s| s.len() as u64)
96        .unwrap_or(0)
97}
98
99impl Aggregator {
100    pub fn new(cfg: Config, ups: Arc<Upstreams>) -> Aggregator {
101        let tool_router = Self::tool_router();
102        let meta_bytes = tool_router.list_all().iter().map(json_len).sum();
103        Aggregator {
104            cfg,
105            ups,
106            tool_router,
107            stats: Stats::new(meta_bytes),
108        }
109    }
110
111    pub async fn list_servers(&self) -> anyhow::Result<Vec<ServerSummary>> {
112        let mut out = Vec::new();
113        for name in self.ups.server_names() {
114            out.push(ServerSummary {
115                // in-memory index only: list_servers must never connect (lazy startup)
116                tool_count: self.ups.cached_tools(&name).len(),
117                status: self.ups.status(&name).to_string(),
118                instructions: self.ups.instructions(&name),
119                name,
120            });
121        }
122        Ok(out)
123    }
124
125    pub async fn list_tools(&self, server: String) -> anyhow::Result<Vec<ToolSummary>> {
126        Ok(self
127            .ups
128            .tools(&server)
129            .await?
130            .into_iter()
131            .map(|t| ToolSummary {
132                name: t.name,
133                description: t.description,
134                annotations: t.annotations,
135            })
136            .collect())
137    }
138
139    pub async fn search_tools(
140        &self,
141        query: String,
142        server: Option<String>,
143        limit: usize,
144    ) -> anyhow::Result<Vec<(String, ToolInfo)>> {
145        let names = match server.as_deref() {
146            Some(s) => {
147                if !self.ups.server_names().iter().any(|n| n == s) {
148                    anyhow::bail!(
149                        "unknown server {s:?}; available: {}",
150                        self.ups.server_names().join(", ")
151                    );
152                }
153                vec![s.to_string()]
154            }
155            None => self.ups.server_names(),
156        };
157        // fetch concurrently: a cold/hanging server must not serialize N x connect-timeout into a search
158        let mut set = tokio::task::JoinSet::new();
159        for name in names {
160            let ups = self.ups.clone();
161            set.spawn(async move { (name.clone(), ups.tools(&name).await.unwrap_or_default()) });
162        }
163        let mut index = Vec::new();
164        while let Some(pair) = set.join_next().await {
165            index.push(pair?);
166        }
167        Ok(search(&index, &query, server.as_deref(), limit))
168    }
169
170    /// Reconnect one or all servers and rebuild the tool index. Per-server
171    /// errors are reported, not propagated — one dead server must not block
172    /// refreshing the others.
173    pub async fn refresh_tools(
174        &self,
175        server: Option<String>,
176    ) -> anyhow::Result<Vec<RefreshResult>> {
177        let names = match server.as_deref() {
178            Some(s) => {
179                if !self.ups.server_names().iter().any(|n| n == s) {
180                    anyhow::bail!(
181                        "unknown server {s:?}; available: {}",
182                        self.ups.server_names().join(", ")
183                    );
184                }
185                vec![s.to_string()]
186            }
187            None => self.ups.server_names(),
188        };
189        let mut set = tokio::task::JoinSet::new();
190        for name in names {
191            let ups = self.ups.clone();
192            set.spawn(async move {
193                match ups.refresh(&name).await {
194                    Ok(()) => RefreshResult {
195                        tool_count: Some(ups.cached_tools(&name).len()),
196                        status: "ok".into(),
197                        error: None,
198                        name,
199                    },
200                    Err(e) => RefreshResult {
201                        tool_count: None,
202                        status: "error".into(),
203                        error: Some(e.to_string()),
204                        name,
205                    },
206                }
207            });
208        }
209        let mut out = Vec::new();
210        while let Some(r) = set.join_next().await {
211            out.push(r?);
212        }
213        out.sort_by(|a, b| a.name.cmp(&b.name));
214        Ok(out)
215    }
216
217    pub async fn describe_tool(&self, server: String, tool: String) -> anyhow::Result<ToolInfo> {
218        if let Some(sc) = self.cfg.mcp_servers.get(&server)
219            && !sc.is_allowed(&tool)
220        {
221            anyhow::bail!("tool {tool:?} on server {server:?} is blocked by config");
222        }
223        self.ups
224            .tools(&server)
225            .await?
226            .into_iter()
227            .find(|t| t.name == tool)
228            .ok_or_else(|| anyhow::anyhow!("unknown tool {tool:?} on server {server:?}"))
229    }
230
231    pub async fn call_tool(
232        &self,
233        server: String,
234        tool: String,
235        arguments: Option<serde_json::Map<String, serde_json::Value>>,
236    ) -> Result<CallToolResult, ErrorData> {
237        self.ups
238            .call(&server, &tool, arguments)
239            .await
240            .map_err(internal)
241    }
242
243    pub async fn authorize_server(
244        &self,
245        server: String,
246        pasted_url: Option<String>,
247    ) -> anyhow::Result<String> {
248        match pasted_url {
249            Some(u) => {
250                self.ups.oauth_complete(&server, &u).await?;
251                self.ups.refresh(&server).await?;
252                Ok(format!("authorized; tool index refreshed for {server}"))
253            }
254            None => match self.ups.oauth_begin(&server).await? {
255                None => Ok(format!("{server} is already authorized")),
256                Some(url) => Ok(format!(
257                    "Authorization required for {server}.\n\
258                     1. Open in a browser: {url}\n\
259                     2. Approve access — the redirect completes automatically.\n\
260                     3. Retry your call.\n\
261                     Headless? Open the URL anywhere, then call authorize_server again \
262                     with pasted_url set to the final redirect URL."
263                )),
264            },
265        }
266    }
267
268    /// Split a `server__tool` name into (server, tool) by longest matching
269    /// expose-server prefix — server names may themselves contain `__`.
270    pub fn route_exposed<'n>(&self, name: &'n str) -> Option<(&str, &'n str)> {
271        self.cfg
272            .mcp_servers
273            .iter()
274            .filter(|(_, sc)| sc.expose)
275            .filter_map(|(n, _)| {
276                name.strip_prefix(&format!("{n}__"))
277                    .map(|t| (n.as_str(), t))
278            })
279            .max_by_key(|(n, _)| n.len())
280    }
281}
282
283#[tool_router]
284impl Aggregator {
285    #[tool(
286        name = "list_servers",
287        description = "List connected MCP servers: name, status, tool count, instructions"
288    )]
289    async fn list_servers_tool(&self) -> Result<String, ErrorData> {
290        self.stats.record_meta("list_servers").await;
291        let v = self.list_servers().await.map_err(internal)?;
292        serde_json::to_string_pretty(&v).map_err(internal)
293    }
294
295    #[tool(
296        name = "list_tools",
297        description = "List a server's tools: name, one-line description, annotations. No schemas."
298    )]
299    async fn list_tools_tool(
300        &self,
301        Parameters(p): Parameters<ServerArg>,
302    ) -> Result<String, ErrorData> {
303        self.stats.record_meta("list_tools").await;
304        let v = self.list_tools(p.server).await.map_err(internal)?;
305        serde_json::to_string_pretty(&v).map_err(internal)
306    }
307
308    #[tool(
309        name = "search_tools",
310        description = "Search tools across servers by name/description. Returns matches WITH full input schemas."
311    )]
312    async fn search_tools_tool(
313        &self,
314        Parameters(p): Parameters<SearchArgs>,
315    ) -> Result<String, ErrorData> {
316        self.stats.record_meta("search_tools").await;
317        let v = self
318            .search_tools(p.query, p.server, p.limit.unwrap_or(5).min(25))
319            .await
320            .map_err(internal)?;
321        self.stats
322            .record_served(v.iter().map(|(_, t)| json_len(&t.schema)).sum())
323            .await;
324        serde_json::to_string_pretty(&v).map_err(internal)
325    }
326
327    #[tool(
328        name = "refresh_tools",
329        description = "Reconnect one server (or all) and rebuild its tool index. Use when a server's tools changed."
330    )]
331    async fn refresh_tools_tool(
332        &self,
333        Parameters(p): Parameters<RefreshArgs>,
334    ) -> Result<String, ErrorData> {
335        self.stats.record_meta("refresh_tools").await;
336        let v = self.refresh_tools(p.server).await.map_err(internal)?;
337        serde_json::to_string_pretty(&v).map_err(internal)
338    }
339
340    #[tool(
341        name = "describe_tool",
342        description = "Get the full input schema of one exact tool"
343    )]
344    async fn describe_tool_tool(
345        &self,
346        Parameters(p): Parameters<DescribeArgs>,
347    ) -> Result<String, ErrorData> {
348        self.stats.record_meta("describe_tool").await;
349        let v = self
350            .describe_tool(p.server, p.tool)
351            .await
352            .map_err(internal)?;
353        self.stats.record_served(json_len(&v.schema)).await;
354        serde_json::to_string_pretty(&v).map_err(internal)
355    }
356
357    #[tool(
358        name = "call_tool",
359        description = "Call a tool on an upstream MCP server. Arguments must match its schema."
360    )]
361    async fn call_tool_tool(
362        &self,
363        Parameters(p): Parameters<CallArgs>,
364    ) -> Result<CallToolResult, ErrorData> {
365        self.stats.record_meta("call_tool").await;
366        self.stats.record_proxied().await;
367        self.call_tool(p.server, p.tool, p.arguments).await
368    }
369
370    #[tool(
371        name = "authorize_server",
372        description = "OAuth-authorize a server. Without pasted_url: returns the authorization URL or 'already authorized'. With pasted_url: completes a headless flow from the final redirect URL."
373    )]
374    async fn authorize_server_tool(
375        &self,
376        Parameters(p): Parameters<AuthorizeArgs>,
377    ) -> Result<String, ErrorData> {
378        self.stats.record_meta("authorize_server").await;
379        self.authorize_server(p.server, p.pasted_url)
380            .await
381            .map_err(internal)
382    }
383}
384
385#[tool_handler]
386impl ServerHandler for Aggregator {
387    fn get_info(&self) -> ServerInfo {
388        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
389            .with_instructions("Multiplexed MCP servers. Use list_servers → list_tools/search_tools → describe_tool → call_tool. refresh_tools re-indexes after upstream tool changes. authorize_server handles OAuth login. Tools named server__tool are directly exposed.")
390    }
391
392    async fn list_tools(
393        &self,
394        _req: Option<PaginatedRequestParams>,
395        _ctx: RequestContext<RoleServer>,
396    ) -> Result<ListToolsResult, ErrorData> {
397        let mut tools = self.tool_router.list_all();
398        for (name, sc) in &self.cfg.mcp_servers {
399            if !sc.expose {
400                continue;
401            }
402            if let Ok(list) = self.ups.tools(name).await {
403                for t in list {
404                    let schema = match t.schema {
405                        serde_json::Value::Object(m) => Arc::new(m),
406                        _ => Arc::new(serde_json::Map::new()),
407                    };
408                    let mut tool = Tool::new_with_raw(
409                        format!("{name}__{}", t.name),
410                        t.description.map(std::borrow::Cow::Owned),
411                        schema,
412                    );
413                    tool.annotations = t.annotations.and_then(|v| serde_json::from_value(v).ok());
414                    tools.push(tool);
415                }
416            }
417        }
418        Ok(ListToolsResult::with_all_items(tools))
419    }
420
421    async fn call_tool(
422        &self,
423        req: CallToolRequestParams,
424        ctx: RequestContext<RoleServer>,
425    ) -> Result<CallToolResponse, ErrorData> {
426        if let Some((server, tool)) = self.route_exposed(&req.name) {
427            let tool = tool.to_string();
428            self.stats.record_proxied().await;
429            return self
430                .ups
431                .call(server, &tool, req.arguments)
432                .await
433                .map(CallToolResponse::from)
434                .map_err(internal);
435        }
436        let tcc = ToolCallContext::new(self, req, ctx);
437        self.tool_router.call(tcc).await
438    }
439}
mcp_multiplexer/server.rs

mcp_multiplexer/
server.rs
