Skip to main content

rtb_mcp/
server.rs

1//! [`McpServer`] — the core type that turns `BUILTIN_COMMANDS` into
2//! an MCP-tools-only `rmcp` `ServerHandler`.
3
4use 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
21/// MCP server that exposes every `mcp_exposed` [`Command`] as an
22/// MCP tool over the supplied [`Transport`].
23///
24/// Construction walks [`BUILTIN_COMMANDS`] eagerly: the tool registry
25/// is built once at `new()` time, so subsequent `tools/list` and
26/// `tools/call` requests are pure dispatch. Each entry retains a
27/// pointer to its `BUILTIN_COMMANDS` factory (rather than a
28/// long-lived `Box<dyn Command>`) so `Command::run` always sees a
29/// fresh instance — this matches the per-invocation lifecycle of
30/// CLI execution and avoids trait-object lifetime entanglement
31/// across the `tools/call` await boundary.
32pub struct McpServer {
33    app: App,
34    tools: Arc<Vec<RegisteredTool>>,
35    transport: Transport,
36}
37
38/// A single MCP-exposed command, captured at registry-build time.
39#[derive(Clone)]
40struct RegisteredTool {
41    name: &'static str,
42    about: &'static str,
43    aliases: &'static [&'static str],
44    /// Schema as JSON object. Defaults to `{"type":"object"}` —
45    /// the minimum a JSON Schema validator needs to accept any
46    /// arguments object (or none).
47    schema: serde_json::Map<String, serde_json::Value>,
48    /// Pointer back into `BUILTIN_COMMANDS` so we can build a fresh
49    /// `Box<dyn Command>` per invocation.
50    factory: fn() -> Box<dyn Command>,
51}
52
53impl McpServer {
54    /// Build a new server. Walks [`BUILTIN_COMMANDS`] eagerly,
55    /// filters by [`Command::mcp_exposed`], and freezes the
56    /// registry. The `transport` choice is honoured by [`Self::serve`].
57    #[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                    // Non-object schemas are invalid MCP — fall back
70                    // to the empty object so the tool still lists.
71                    let mut map = serde_json::Map::new();
72                    map.insert("type".into(), serde_json::Value::String("object".into()));
73                    let _ = other; // intentionally discarded
74                    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    /// Number of registered MCP tools.
94    #[must_use]
95    pub fn tool_count(&self) -> usize {
96        self.tools.len()
97    }
98
99    /// Iterator over `(name, about, schema)` triples for every
100    /// registered tool — used by `mcp list` to print the manifest.
101    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    /// Run the same dispatch logic that backs `tools/call`. Returns
106    /// `Ok(())` when the named tool's `Command::run` succeeded,
107    /// `Err(McpError::Command)` when it ran but failed, and
108    /// `Err(McpError::Protocol)` when no tool with the given name
109    /// is registered.
110    ///
111    /// This bypass exists for unit testing — the rmcp service loop
112    /// reaches the same logic via `ServerHandler::call_tool`.
113    ///
114    /// # Errors
115    ///
116    /// See variants above.
117    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    /// Run the server until the supplied transport closes or the
133    /// `app.shutdown` token fires.
134    ///
135    /// # Errors
136    ///
137    /// Returns [`McpError::Transport`] on bind / accept failure or
138    /// when an unsupported transport variant is selected.
139    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    /// Run the rmcp service against a caller-supplied `AsyncRead` +
156    /// `AsyncWrite` pair. The stdio transport delegates here; tests
157    /// pass a `tokio::io::duplex` half so they don't have to take
158    /// over the process's real stdin/stdout.
159    ///
160    /// # Errors
161    ///
162    /// [`McpError::Transport`] if `rmcp` fails to bring the service
163    /// up or the underlying transport surfaces an error.
164    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/// `rmcp::ServerHandler` implementation backed by an [`McpServer`]'s
187/// tool registry.
188///
189/// Held by value inside the running rmcp service. Cloning is cheap —
190/// every field is `Arc` or `Clone`-cheap — so the handler can survive
191/// the move into `serve_server` while leaving the `McpServer` builder
192/// API ergonomic.
193#[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        // rmcp 1.x marks ServerInfo / ServerCapabilities /
226        // Implementation as #[non_exhaustive]; constructors +
227        // field-assign replace the old struct-literal form.
228        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    // The following are no-op hooks we override only to silence the
265    // crate-default `tracing::info!("client initialized")` (which
266    // would corrupt the stdio transport by writing to the same
267    // stdout stream the protocol uses).
268    async fn on_initialized(&self, _context: NotificationContext<RoleServer>) {}
269}
270
271// `serve_server` is re-exported by `rmcp` and used here only via
272// `ServiceExt::serve`. Reference it once so a future rmcp release
273// that drops the symbol fails the build loudly instead of leaving
274// stale documentation behind.
275#[doc(hidden)]
276#[allow(dead_code)]
277const fn _link_check() {
278    let _ = serve_server::<McpHandler, (tokio::io::Stdin, tokio::io::Stdout), _, _>;
279}