1use std::sync::Arc;
5
6use rmcp::handler::server::ServerHandler;
7use rmcp::model::{
8 CallToolRequestParams, CallToolResult, ContentBlock, Implementation, ListToolsResult,
9 PaginatedRequestParams, ServerCapabilities, ServerInfo, Tool,
10};
11use rmcp::service::{NotificationContext, RequestContext, RoleServer};
12use rmcp::transport::io::stdio;
13use rmcp::ErrorData as RmcpError;
14use rmcp::{serve_server, ServiceExt};
15use rtb_app::app::App;
16use rtb_app::command::{Command, BUILTIN_COMMANDS};
17
18use crate::error::{McpError, Result};
19use crate::transport::Transport;
20
21pub struct McpServer {
33 app: App,
34 tools: Arc<Vec<RegisteredTool>>,
35 transport: Transport,
36}
37
38#[derive(Clone)]
40struct RegisteredTool {
41 name: &'static str,
42 about: &'static str,
43 aliases: &'static [&'static str],
44 schema: serde_json::Map<String, serde_json::Value>,
48 factory: fn() -> Box<dyn Command>,
51}
52
53impl McpServer {
54 #[must_use]
58 pub fn new(app: App, transport: Transport) -> Self {
59 let mut tools = Vec::new();
60 for factory in BUILTIN_COMMANDS {
61 let cmd = factory();
62 if !cmd.mcp_exposed() {
63 continue;
64 }
65 let spec = cmd.spec();
66 let schema = match cmd.mcp_input_schema() {
67 Some(serde_json::Value::Object(map)) => map,
68 Some(other) => {
69 let mut map = serde_json::Map::new();
72 map.insert("type".into(), serde_json::Value::String("object".into()));
73 let _ = other; map
75 }
76 None => {
77 let mut map = serde_json::Map::new();
78 map.insert("type".into(), serde_json::Value::String("object".into()));
79 map
80 }
81 };
82 tools.push(RegisteredTool {
83 name: spec.name,
84 about: spec.about,
85 aliases: spec.aliases,
86 schema,
87 factory: *factory,
88 });
89 }
90 Self { app, tools: Arc::new(tools), transport }
91 }
92
93 #[must_use]
95 pub fn tool_count(&self) -> usize {
96 self.tools.len()
97 }
98
99 pub fn tool_manifest(&self) -> impl Iterator<Item = (&str, &str, serde_json::Value)> + '_ {
102 self.tools.iter().map(|t| (t.name, t.about, serde_json::Value::Object(t.schema.clone())))
103 }
104
105 pub async fn dispatch(&self, name: &str) -> Result<()> {
118 let tool = self
119 .tools
120 .iter()
121 .find(|t| t.name == name || t.aliases.contains(&name))
122 .ok_or_else(|| McpError::Protocol(format!("unknown MCP tool: {name}")))?;
123 let cmd = (tool.factory)();
124 match cmd.run(self.app.clone()).await {
125 Ok(()) => Ok(()),
126 Err(e) => {
127 Err(McpError::Command { command: tool.name.to_string(), message: e.to_string() })
128 }
129 }
130 }
131
132 pub async fn serve(self) -> Result<()> {
140 match self.transport.clone() {
141 Transport::Stdio => self.serve_stdio().await,
142 Transport::Sse { .. } | Transport::Http { .. } => Err(McpError::Transport(
143 "SSE / streamable HTTP transports are not yet implemented in rtb-mcp v0.1; \
144 use --transport stdio"
145 .to_string(),
146 )),
147 }
148 }
149
150 async fn serve_stdio(self) -> Result<()> {
151 let (stdin, stdout) = stdio();
152 self.serve_with_pipe(stdin, stdout).await
153 }
154
155 pub async fn serve_with_pipe<R, W>(self, read: R, write: W) -> Result<()>
165 where
166 R: tokio::io::AsyncRead + Send + Sync + Unpin + 'static,
167 W: tokio::io::AsyncWrite + Send + Sync + Unpin + 'static,
168 {
169 let shutdown = self.app.shutdown.clone();
170 let handler = McpHandler::from_server(&self);
171 let running =
172 handler.serve((read, write)).await.map_err(|e| McpError::Transport(e.to_string()))?;
173 let cancel = running.cancellation_token();
174 tokio::select! {
175 res = running.waiting() => {
176 res.map(|_| ()).map_err(|e| McpError::Transport(e.to_string()))
177 }
178 () = shutdown.cancelled() => {
179 cancel.cancel();
180 Ok(())
181 }
182 }
183 }
184}
185
186#[derive(Clone)]
194struct McpHandler {
195 app: App,
196 tools: Arc<Vec<RegisteredTool>>,
197 server_name: String,
198 server_version: String,
199}
200
201impl McpHandler {
202 fn from_server(server: &McpServer) -> Self {
203 Self {
204 app: server.app.clone(),
205 tools: server.tools.clone(),
206 server_name: server.app.metadata.name.clone(),
207 server_version: server.app.version.version.to_string(),
208 }
209 }
210
211 fn render_tools(&self) -> Vec<Tool> {
212 self.tools
213 .iter()
214 .map(|t| Tool::new(t.name.to_string(), t.about.to_string(), t.schema.clone()))
215 .collect()
216 }
217
218 fn find_tool(&self, name: &str) -> Option<RegisteredTool> {
219 self.tools.iter().find(|t| t.name == name || t.aliases.contains(&name)).cloned()
220 }
221}
222
223impl ServerHandler for McpHandler {
224 fn get_info(&self) -> ServerInfo {
225 let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build());
229 info.server_info =
230 Implementation::new(self.server_name.clone(), self.server_version.clone());
231 info
232 }
233
234 async fn list_tools(
235 &self,
236 _request: Option<PaginatedRequestParams>,
237 _context: RequestContext<RoleServer>,
238 ) -> std::result::Result<ListToolsResult, RmcpError> {
239 Ok(ListToolsResult { tools: self.render_tools(), next_cursor: None, meta: None })
240 }
241
242 async fn call_tool(
243 &self,
244 request: CallToolRequestParams,
245 _context: RequestContext<RoleServer>,
246 ) -> std::result::Result<CallToolResult, RmcpError> {
247 let Some(tool) = self.find_tool(request.name.as_ref()) else {
248 return Err(RmcpError::invalid_params(
249 format!("unknown MCP tool: {}", request.name),
250 None,
251 ));
252 };
253 let cmd = (tool.factory)();
254 match cmd.run(self.app.clone()).await {
255 Ok(()) => {
256 Ok(CallToolResult::success(vec![ContentBlock::text(format!("{} ok", tool.name))]))
257 }
258 Err(e) => {
259 Ok(CallToolResult::error(vec![ContentBlock::text(format!("{}: {}", tool.name, e))]))
260 }
261 }
262 }
263
264 async fn on_initialized(&self, _context: NotificationContext<RoleServer>) {}
269}
270
271#[doc(hidden)]
276#[allow(dead_code)]
277const fn _link_check() {
278 let _ = serve_server::<McpHandler, (tokio::io::Stdin, tokio::io::Stdout), _, _>;
279}