viam_mcp_sdk/lib.rs
1// MCP SDK for building MCP server components
2pub mod json_rpc;
3pub mod command;
4pub mod server;
5pub mod http;
6pub mod tool;
7
8// Re-export common types
9pub use command::{CommandHandler, HandlerResult, ServerStateChanges};
10pub use server::McpServer;
11pub use tool::{ToolHandler, ToolRegistry};
12
13// Generate bindings for the MCP server world
14#[allow(clippy::missing_safety_doc)]
15mod bindings {
16 wit_bindgen::generate!({
17 path: "./wit",
18 world: "mcp-sdk",
19 generate_all
20 // with: {
21 // "wasi:cli/run@0.2.0": generate,
22 // "wasi:io/streams@0.2.0": generate_all,
23 // }
24 });
25}
26
27/// Macro for creating an MCP server component
28///
29/// This macro simplifies the creation of MCP server components by automatically
30/// generating the necessary handlers and message processing logic.
31///
32/// # Example
33///
34/// ```rust
35/// use mcp_sdk::mcp_server_component;
36/// use serde_json::Value;
37///
38/// struct MyTool;
39///
40/// impl ToolHandler for MyTool {
41/// fn handle(&self, params: &Value) -> Result<Value, JsonRpcError> {
42/// // Handle tool call
43/// Ok(json!({ "result": "success" }))
44/// }
45///
46/// fn name(&self) -> &'static str {
47/// "my_tool"
48/// }
49///
50/// fn description(&self) -> &'static str {
51/// "A custom tool"
52/// }
53///
54/// fn input_schema(&self) -> Value {
55/// json!({
56/// "type": "object",
57/// "properties": {
58/// "param1": {
59/// "type": "string"
60/// }
61/// }
62/// })
63/// }
64/// }
65///
66/// mcp_server_component! {
67/// tools: [
68/// MyTool
69/// ]
70/// }
71/// ```
72#[macro_export]
73macro_rules! mcp_server_component {
74 (
75 $(
76 $method:expr => $handler:expr
77 ),*
78 $(,)?
79 ) => {
80 use $crate::server::McpServer;
81 use serde_json::Value;
82
83 // Initialize server with default handlers
84 static mut SERVER: Option<McpServer> = None;
85
86 // Initialize the server
87 fn init_server() -> &'static McpServer {
88 unsafe {
89 if SERVER.is_none() {
90 let mut server = McpServer::new();
91
92 // Register custom handlers
93 $(
94 server.register_handler($method, Box::new($handler));
95 )*
96
97 SERVER = Some(server);
98 }
99
100 SERVER.as_ref().unwrap()
101 }
102 }
103
104 // Process a JSON-RPC message
105 pub fn process_message(message: &str) -> String {
106 let server = init_server();
107 match server.process_message(message) {
108 Ok(response) => response,
109 Err(err) => err,
110 }
111 }
112
113 // Run the server with stdio transport
114 pub fn run() {
115 use std::io::{self, BufRead, Write};
116
117 let stdin = io::stdin();
118 let mut stdout = io::stdout();
119
120 // Initialize the server
121 let _server = init_server();
122
123 // Buffer to collect multiline JSON
124 let mut buffer = String::new();
125 let mut brace_count = 0;
126 let mut in_json = false;
127
128 // Process messages from stdin
129 for line in stdin.lock().lines() {
130 if let Ok(line_content) = line {
131 // Count opening and closing braces to detect complete JSON objects
132 for c in line_content.chars() {
133 if c == '{' {
134 if !in_json {
135 in_json = true;
136 }
137 brace_count += 1;
138 } else if c == '}' {
139 brace_count -= 1;
140 }
141 }
142
143 // Add the line to our buffer
144 buffer.push_str(&line_content);
145 buffer.push('\n');
146
147 // If we have a complete JSON object, process it
148 if in_json && brace_count == 0 {
149 let response = process_message(&buffer);
150 writeln!(stdout, "{}", response).unwrap();
151 stdout.flush().unwrap();
152
153 // Reset the buffer
154 buffer.clear();
155 in_json = false;
156 }
157 }
158 }
159 }
160 };
161
162 (
163 tools: [
164 $($tool_type:ty),* $(,)?
165 ]
166 $(,)?
167 ) => {
168 use $crate::server::McpServer;
169 use $crate::tool::{ToolRegistry, ToolsCallHandler, ListToolsHandler, InitializeHandler, InitializedHandler};
170 use $crate::command::{CommandHandler, HandlerResult};
171 use std::sync::Arc;
172 use serde_json::Value;
173
174 // Initialize server with tool handlers
175 static mut SERVER: Option<McpServer> = None;
176 static mut TOOL_REGISTRY: Option<Arc<ToolRegistry>> = None;
177
178 // Initialize the tool registry
179 fn init_tool_registry() -> Arc<ToolRegistry> {
180 unsafe {
181 if TOOL_REGISTRY.is_none() {
182 let mut registry = ToolRegistry::new();
183 $(
184 registry.register(Arc::new(<$tool_type>::default()));
185 )*
186
187 TOOL_REGISTRY = Some(Arc::new(registry));
188 }
189
190 Arc::clone(TOOL_REGISTRY.as_ref().unwrap())
191 }
192 }
193
194 // Initialize the server
195 fn init_server() -> &'static McpServer {
196 unsafe {
197 if SERVER.is_none() {
198 let mut server = McpServer::new();
199
200 // Get tool registry
201 let tool_registry = init_tool_registry();
202
203 // Create handlers directly as Box<dyn CommandHandler>
204 let tools_call_handler = Box::new(ToolsCallHandler::new(Arc::clone(&tool_registry)));
205 let list_tools_handler = Box::new(ListToolsHandler::new(Arc::clone(&tool_registry)));
206 let initialize_handler = Box::new(InitializeHandler::default());
207 let initialized_handler = Box::new(InitializedHandler::default());
208
209 // Register the tool handlers using their name() method
210 server.register_handler(tools_call_handler.name(), tools_call_handler);
211 server.register_handler(list_tools_handler.name(), list_tools_handler);
212 server.register_handler(initialize_handler.name(), initialize_handler);
213 server.register_handler(initialized_handler.name(), initialized_handler);
214
215 SERVER = Some(server);
216 }
217
218 SERVER.as_ref().unwrap()
219 }
220 }
221
222 // Process a JSON-RPC message
223 pub fn process_message(message: &str) -> String {
224 let server = init_server();
225 match server.process_message(message) {
226 Ok(response) => response,
227 Err(err) => err,
228 }
229 }
230
231 // Run the server with stdio transport
232 pub fn run() {
233 use std::io::{self, BufRead, Write};
234
235 let stdin = io::stdin();
236 let mut stdout = io::stdout();
237
238 // Initialize the server
239 let _server = init_server();
240
241 // Buffer to collect multiline JSON
242 let mut buffer = String::new();
243 let mut brace_count = 0;
244 let mut in_json = false;
245
246 // Process messages from stdin
247 for line in stdin.lock().lines() {
248 if let Ok(line_content) = line {
249 // Count opening and closing braces to detect complete JSON objects
250 for c in line_content.chars() {
251 if c == '{' {
252 if !in_json {
253 in_json = true;
254 }
255 brace_count += 1;
256 } else if c == '}' {
257 brace_count -= 1;
258 }
259 }
260
261 // Add the line to our buffer
262 buffer.push_str(&line_content);
263 buffer.push('\n');
264
265 // If we have a complete JSON object, process it
266 if in_json && brace_count == 0 {
267 let response = process_message(&buffer);
268 writeln!(stdout, "{}", response).unwrap();
269 stdout.flush().unwrap();
270
271 // Reset the buffer
272 buffer.clear();
273 in_json = false;
274 }
275 }
276 }
277 }
278 };
279
280 (
281 tools: [
282 $($tool_type:ty),* $(,)?
283 ],
284 overrides: {
285 $(
286 $override_name:ident: $override_type:ty
287 ),* $(,)?
288 }
289 $(,)?
290 ) => {
291 use $crate::server::McpServer;
292 use $crate::tool::{ToolRegistry, ToolsCallHandler, ListToolsHandler, InitializeHandler, InitializedHandler};
293 use $crate::command::{CommandHandler, HandlerResult};
294 use std::sync::Arc;
295 use serde_json::Value;
296
297 // Initialize server with tool handlers
298 static mut SERVER: Option<McpServer> = None;
299 static mut TOOL_REGISTRY: Option<Arc<ToolRegistry>> = None;
300
301 // Initialize the tool registry
302 fn init_tool_registry() -> Arc<ToolRegistry> {
303 unsafe {
304 if TOOL_REGISTRY.is_none() {
305 let mut registry = ToolRegistry::new();
306 $(
307 registry.register(Arc::new(<$tool_type>::default()));
308 )*
309
310 TOOL_REGISTRY = Some(Arc::new(registry));
311 }
312
313 Arc::clone(TOOL_REGISTRY.as_ref().unwrap())
314 }
315 }
316
317 // Initialize the server
318 fn init_server() -> &'static McpServer {
319 unsafe {
320 if SERVER.is_none() {
321 let mut server = McpServer::new();
322
323 // Get tool registry
324 let tool_registry = init_tool_registry();
325
326 // Create default handlers directly as Box<dyn CommandHandler>
327 let tools_call_handler = Box::new(ToolsCallHandler::new(Arc::clone(&tool_registry)));
328 let list_tools_handler = Box::new(ListToolsHandler::new(Arc::clone(&tool_registry)));
329 let initialize_handler = Box::new(InitializeHandler::default());
330 let initialized_handler = Box::new(InitializedHandler::default());
331
332 // Register the default handlers using their name() method
333 server.register_handler(tools_call_handler.name(), tools_call_handler);
334 server.register_handler(list_tools_handler.name(), list_tools_handler);
335 server.register_handler(initialize_handler.name(), initialize_handler);
336 server.register_handler(initialized_handler.name(), initialized_handler);
337
338 // Create and register override handlers
339 $(
340 let override_handler = Box::new(<$override_type>::default());
341 server.register_handler(override_handler.name(), override_handler);
342 )*
343
344 SERVER = Some(server);
345 }
346
347 SERVER.as_ref().unwrap()
348 }
349 }
350
351 // Process a JSON-RPC message
352 pub fn process_message(message: &str) -> String {
353 let server = init_server();
354 match server.process_message(message) {
355 Ok(response) => response,
356 Err(err) => err,
357 }
358 }
359
360 // Run the server with stdio transport
361 pub fn run() {
362 use std::io::{self, BufRead, Write};
363
364 let stdin = io::stdin();
365 let mut stdout = io::stdout();
366
367 // Initialize the server
368 let _server = init_server();
369
370 // Buffer to collect multiline JSON
371 let mut buffer = String::new();
372 let mut brace_count = 0;
373 let mut in_json = false;
374
375 // Process messages from stdin
376 for line in stdin.lock().lines() {
377 if let Ok(line_content) = line {
378 // Count opening and closing braces to detect complete JSON objects
379 for c in line_content.chars() {
380 if c == '{' {
381 if !in_json {
382 in_json = true;
383 }
384 brace_count += 1;
385 } else if c == '}' {
386 brace_count -= 1;
387 }
388 }
389
390 // Add the line to our buffer
391 buffer.push_str(&line_content);
392 buffer.push('\n');
393
394 // If we have a complete JSON object, process it
395 if in_json && brace_count == 0 {
396 let response = process_message(&buffer);
397 writeln!(stdout, "{}", response).unwrap();
398 stdout.flush().unwrap();
399
400 // Reset the buffer
401 buffer.clear();
402 in_json = false;
403 }
404 }
405 }
406 }
407 };
408}