From ca383e15026188af5bbaa38e9494c909fbe24503 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Sun, 15 Dec 2024 20:43:10 -0800 Subject: [PATCH 01/10] Implement basic server --- crates/mcp-server/Cargo.toml | 18 +++ crates/mcp-server/src/errors.rs | 71 ++++++++++++ crates/mcp-server/src/lib.rs | 193 ++++++++++++++++++++++++++++++++ crates/mcp-server/src/main.rs | 16 +++ crates/mcp-server/src/router.rs | 90 +++++++++++++++ 5 files changed, 388 insertions(+) create mode 100644 crates/mcp-server/Cargo.toml create mode 100644 crates/mcp-server/src/errors.rs create mode 100644 crates/mcp-server/src/lib.rs create mode 100644 crates/mcp-server/src/main.rs create mode 100644 crates/mcp-server/src/router.rs diff --git a/crates/mcp-server/Cargo.toml b/crates/mcp-server/Cargo.toml new file mode 100644 index 000000000000..8d7129148c0d --- /dev/null +++ b/crates/mcp-server/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "mcp-server" +version = "0.1.0" +edition = "2021" + +[dependencies] +anyhow = "1.0.94" +thiserror = "1.0" +mcp-core = { path = "../mcp-core" } +serde = "1.0.216" +serde_json = "1.0.133" +tokio = { version = "1", features = ["full"] } +tower = { version = "0.4", features = ["timeout"] } +tower-service = "0.3" +futures = "0.3" +pin-project = "1.1" +tracing = "0.1" +async-trait = "0.1" diff --git a/crates/mcp-server/src/errors.rs b/crates/mcp-server/src/errors.rs new file mode 100644 index 000000000000..9bbfeb543d4e --- /dev/null +++ b/crates/mcp-server/src/errors.rs @@ -0,0 +1,71 @@ +use thiserror::Error; + +pub type BoxError = Box; + +#[derive(Error, Debug)] +pub enum TransportError { + #[error("IO error: {0}")] + Io(#[from] std::io::Error), + + #[error("JSON serialization error: {0}")] + Json(#[from] serde_json::Error), + + #[error("Invalid UTF-8 sequence: {0}")] + Utf8(#[from] std::string::FromUtf8Error), + + #[error("Protocol error: {0}")] + Protocol(String), + + #[error("Invalid message format: {0}")] + InvalidMessage(String), +} + +#[derive(Error, Debug)] +pub enum ServerError { + #[error("Transport error: {0}")] + Transport(#[from] TransportError), + + #[error("Service error: {0}")] + Service(String), + + #[error("Internal error: {0}")] + Internal(String), + + #[error("Request timed out")] + Timeout(#[from] tower::timeout::error::Elapsed), +} + +#[derive(Error, Debug)] +pub enum RouterError { + #[error("Method not found: {0}")] + MethodNotFound(String), + + #[error("Invalid parameters: {0}")] + InvalidParams(String), + + #[error("Internal error: {0}")] + Internal(String), +} + +impl From for mcp_core::protocol::ErrorData { + fn from(err: RouterError) -> Self { + use mcp_core::protocol::*; + match err { + RouterError::MethodNotFound(msg) => ErrorData { + code: METHOD_NOT_FOUND, + message: msg, + data: None, + }, + RouterError::InvalidParams(msg) => ErrorData { + code: INVALID_PARAMS, + message: msg, + data: None, + }, + RouterError::Internal(msg) => ErrorData { + code: INTERNAL_ERROR, + message: msg, + data: None, + }, + } + } +} diff --git a/crates/mcp-server/src/lib.rs b/crates/mcp-server/src/lib.rs new file mode 100644 index 000000000000..fb5991a47566 --- /dev/null +++ b/crates/mcp-server/src/lib.rs @@ -0,0 +1,193 @@ +use std::{ + pin::Pin, + task::{Context, Poll}, +}; + +use futures::{Stream, Future}; +use mcp_core::protocol::{JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, JsonRpcError}; +use pin_project::pin_project; +use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, AsyncBufReadExt, BufReader}; +use tower_service::Service; + +mod errors; +pub use errors::{TransportError, ServerError, RouterError, BoxError}; + +pub mod router; +pub use router::Router; + +/// A transport layer that handles JSON-RPC messages over byte +#[pin_project] +pub struct ByteTransport { + #[pin] + reader: R, + #[pin] + writer: W, +} + +impl ByteTransport +where + R: AsyncRead, + W: AsyncWrite, +{ + pub fn new(reader: R, writer: W) -> Self { + Self { reader, writer } + } +} + +impl Stream for ByteTransport +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.project(); + let mut buf = Vec::new(); + let mut reader = BufReader::new(&mut this.reader); + + let mut read_future = Box::pin(reader.read_until(b'\n', &mut buf)); + match read_future.as_mut().poll(cx) { + Poll::Ready(Ok(0)) => Poll::Ready(None), // EOF + Poll::Ready(Ok(_)) => { + // Convert to UTF-8 string + let line = match String::from_utf8(buf) { + Ok(s) => s, + Err(e) => return Poll::Ready(Some(Err(TransportError::Utf8(e)))), + }; + + // Parse JSON and validate message format + match serde_json::from_str::(&line) { + Ok(value) => { + // Validate basic JSON-RPC structure + if !value.is_object() { + return Poll::Ready(Some(Err(TransportError::InvalidMessage("Message must be a JSON object".into())))); + } + + let obj = value.as_object().unwrap(); // Safe due to check above + + // Check jsonrpc version field + if !obj.contains_key("jsonrpc") || obj["jsonrpc"] != "2.0" { + return Poll::Ready(Some(Err(TransportError::InvalidMessage("Missing or invalid jsonrpc version".into())))); + } + + // Now try to parse as proper message + match serde_json::from_value::(value) { + Ok(msg) => Poll::Ready(Some(Ok(msg))), + Err(e) => Poll::Ready(Some(Err(TransportError::Json(e)))), + } + }, + Err(e) => Poll::Ready(Some(Err(TransportError::Json(e)))), + } + } + Poll::Ready(Err(e)) => Poll::Ready(Some(Err(TransportError::Io(e)))), + Poll::Pending => Poll::Pending, + } + } +} + +impl ByteTransport +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + pub async fn write_message(&mut self, msg: JsonRpcMessage) -> Result<(), std::io::Error> { + let json = serde_json::to_string(&msg)?; + Pin::new(&mut self.writer).write_all(json.as_bytes()).await?; + Pin::new(&mut self.writer).write_all(b"\n").await?; + Pin::new(&mut self.writer).flush().await?; + Ok(()) + } +} + +/// The main server type that processes incoming requests +pub struct Server { + service: S, +} + +impl Server +where + S: Service + Send, + S::Error: Into, + S::Future: Send, +{ + pub fn new(service: S) -> Self { + Self { service } + } + + // TODO transport trait instead of byte transport if we implement others + pub async fn run(self, mut transport: ByteTransport) -> Result<(), ServerError> + where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, + { + use futures::StreamExt; + let mut service = self.service; + + while let Some(msg_result) = transport.next().await { + match msg_result { + Ok(msg) => { + match msg { + JsonRpcMessage::Request(request) => { + // Process the request using our service + let response = match service.call(request).await { + Ok(resp) => resp, + Err(e) => JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: None, // We still have a valid request ID here + result: None, + error: Some(mcp_core::protocol::ErrorData { + code: mcp_core::protocol::INTERNAL_ERROR, + message: Into::::into(e).to_string(), + data: None, + }), + }, + }; + + // Send the response back + if let Err(e) = transport.write_message(JsonRpcMessage::Response(response)).await { + return Err(ServerError::Transport(TransportError::Io(e))); + } + }, + JsonRpcMessage::Response(_) | JsonRpcMessage::Notification(_) | JsonRpcMessage::Error(_) => { + // Ignore responses and notifications for now + continue; + } + } + }, + Err(e) => { + // Convert transport error to JSON-RPC error response + let error = match e { + TransportError::Json(_) | TransportError::InvalidMessage(_) => mcp_core::protocol::ErrorData { + code: mcp_core::protocol::PARSE_ERROR, + message: e.to_string(), + data: None, + }, + TransportError::Protocol(_) => mcp_core::protocol::ErrorData { + code: mcp_core::protocol::INVALID_REQUEST, + message: e.to_string(), + data: None, + }, + _ => mcp_core::protocol::ErrorData { + code: mcp_core::protocol::INTERNAL_ERROR, + message: e.to_string(), + data: None, + }, + }; + + let error_response = JsonRpcMessage::Error(JsonRpcError { + jsonrpc: "2.0".to_string(), + id: None, + error, + }); + + if let Err(e) = transport.write_message(error_response).await { + return Err(ServerError::Transport(TransportError::Io(e))); + } + } + } + } + + Ok(()) + } +} diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs new file mode 100644 index 000000000000..ecf6947bf7b0 --- /dev/null +++ b/crates/mcp-server/src/main.rs @@ -0,0 +1,16 @@ +use anyhow::Result; +use mcp_server::{Router, Server, ByteTransport}; +use tokio::io::{stdin, stdout}; +use tokio::time::Duration; +use tower::timeout::Timeout; + +#[tokio::main] +async fn main() -> Result<()> { + let router = Router::new(); + // Add a 30 second timeout to all requests + let router = Timeout::new(router, Duration::from_secs(30)); + let server = Server::new(router); + + let transport = ByteTransport::new(stdin(), stdout()); + Ok(server.run(transport).await?) +} diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs new file mode 100644 index 000000000000..2e0c587625d6 --- /dev/null +++ b/crates/mcp-server/src/router.rs @@ -0,0 +1,90 @@ +use std::{ + future::Future, + pin::Pin, + task::{Context, Poll}, + time::Duration, +}; +use tokio::time::sleep; + +use mcp_core::protocol::{ + JsonRpcRequest, JsonRpcResponse, ServerCapabilities, InitializeResult, Implementation, +}; +use tower_service::Service; + +use crate::{RouterError, BoxError}; + +/// Router service that implements the JSON-RPC protocol +pub struct Router { + capabilities: ServerCapabilities, +} + +impl Router { + pub fn new() -> Self { + Self { + capabilities: ServerCapabilities { + prompts: None, + resources: None, + tools: None, + }, + } + } +} + +impl Service for Router { + type Response = JsonRpcResponse; + type Error = BoxError; + type Future = Pin> + Send>>; + + fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn call(&mut self, req: JsonRpcRequest) -> Self::Future { + let capabilities = self.capabilities.clone(); + let fut = async move { + match req.method.as_str() { + "slow" => { + // Sleep for 60 seconds - this should trigger the timeout + sleep(Duration::from_secs(60)).await; + Ok(JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: req.id, + result: Some(serde_json::Value::String("This should never be seen".to_string())), + error: None, + }) + }, + "initialize" => { + let result = InitializeResult { + protocol_version: "0.1.0".to_string(), + capabilities, + server_info: Implementation { + name: "mcp-server".to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + }, + }; + + let result_value = serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?; + + Ok(JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: req.id, + result: Some(result_value), + error: None, + }) + }, + _ => { + let error = RouterError::MethodNotFound(req.method); + Ok(JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id: req.id, + result: None, + error: Some(error.into()), + }) + }, + } + }; + + Box::pin(fut) + } +} From 4d4c219a5e5362cd33e72e507f0e4daf646762a0 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Sun, 15 Dec 2024 22:05:07 -0800 Subject: [PATCH 02/10] add tracing and fix handshake --- crates/mcp-core/src/protocol.rs | 68 ++++++++++++++++++++++++++++++++- crates/mcp-server/Cargo.toml | 2 + crates/mcp-server/src/lib.rs | 49 +++++++++++++++++++----- crates/mcp-server/src/main.rs | 24 +++++++++++- crates/mcp-server/src/router.rs | 2 +- 5 files changed, 133 insertions(+), 12 deletions(-) diff --git a/crates/mcp-core/src/protocol.rs b/crates/mcp-core/src/protocol.rs index 259050a20354..09ca3e4caeab 100644 --- a/crates/mcp-core/src/protocol.rs +++ b/crates/mcp-core/src/protocol.rs @@ -6,16 +6,21 @@ use serde_json::Value; #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] pub struct JsonRpcRequest { pub jsonrpc: String, + #[serde(skip_serializing_if = "Option::is_none")] pub id: Option, pub method: String, + #[serde(skip_serializing_if = "Option::is_none")] pub params: Option, } #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] pub struct JsonRpcResponse { pub jsonrpc: String, + #[serde(skip_serializing_if = "Option::is_none")] pub id: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub error: Option, } @@ -23,18 +28,20 @@ pub struct JsonRpcResponse { pub struct JsonRpcNotification { pub jsonrpc: String, pub method: String, + #[serde(skip_serializing_if = "Option::is_none")] pub params: Option, } #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] pub struct JsonRpcError { pub jsonrpc: String, + #[serde(skip_serializing_if = "Option::is_none")] pub id: Option, pub error: ErrorData, } #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] -#[serde(untagged)] +#[serde(untagged, try_from = "JsonRpcRaw")] pub enum JsonRpcMessage { Request(JsonRpcRequest), Response(JsonRpcResponse), @@ -42,6 +49,62 @@ pub enum JsonRpcMessage { Error(JsonRpcError), } +#[derive(Debug, Serialize, Deserialize)] +struct JsonRpcRaw { + jsonrpc: String, + #[serde(skip_serializing_if = "Option::is_none")] + id: Option, + method: String, + #[serde(skip_serializing_if = "Option::is_none")] + params: Option, + #[serde(skip_serializing_if = "Option::is_none")] + result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +impl TryFrom for JsonRpcMessage { + type Error = String; + + fn try_from(raw: JsonRpcRaw) -> Result>::Error> { + // If it has an error field, it's an error response + if raw.error.is_some() { + return Ok(JsonRpcMessage::Error(JsonRpcError { + jsonrpc: raw.jsonrpc, + id: raw.id, + error: raw.error.unwrap(), + })); + } + + // If it has a result field, it's a response + if raw.result.is_some() { + return Ok(JsonRpcMessage::Response(JsonRpcResponse { + jsonrpc: raw.jsonrpc, + id: raw.id, + result: raw.result, + error: None, + })); + } + + // If the method starts with "notifications/", it's a notification + if raw.method.starts_with("notifications/") { + return Ok(JsonRpcMessage::Notification(JsonRpcNotification { + jsonrpc: raw.jsonrpc, + method: raw.method, + params: raw.params, + })); + } + + // Otherwise it's a request + Ok(JsonRpcMessage::Request(JsonRpcRequest { + jsonrpc: raw.jsonrpc, + id: raw.id, + method: raw.method, + params: raw.params, + })) + } +} + // Standard JSON-RPC error codes pub const PARSE_ERROR: i32 = -32700; pub const INVALID_REQUEST: i32 = -32600; @@ -80,8 +143,11 @@ pub struct Implementation { #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)] pub struct ServerCapabilities { + #[serde(skip_serializing_if = "Option::is_none")] pub prompts: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub resources: Option, + #[serde(skip_serializing_if = "Option::is_none")] pub tools: Option, // Add other capabilities as needed } diff --git a/crates/mcp-server/Cargo.toml b/crates/mcp-server/Cargo.toml index 8d7129148c0d..8fe85768c61c 100644 --- a/crates/mcp-server/Cargo.toml +++ b/crates/mcp-server/Cargo.toml @@ -15,4 +15,6 @@ tower-service = "0.3" futures = "0.3" pin-project = "1.1" tracing = "0.1" +tracing-subscriber = { version = "0.3", features = ["env-filter"] } +tracing-appender = "0.2" async-trait = "0.1" diff --git a/crates/mcp-server/src/lib.rs b/crates/mcp-server/src/lib.rs index fb5991a47566..06a4bf77e08d 100644 --- a/crates/mcp-server/src/lib.rs +++ b/crates/mcp-server/src/lib.rs @@ -71,6 +71,10 @@ where return Poll::Ready(Some(Err(TransportError::InvalidMessage("Missing or invalid jsonrpc version".into())))); } + tracing::info!( + json = %line, + "incoming message" + ); // Now try to parse as proper message match serde_json::from_value::(value) { Ok(msg) => Poll::Ready(Some(Ok(msg))), @@ -124,26 +128,53 @@ where use futures::StreamExt; let mut service = self.service; + tracing::info!("Server started"); while let Some(msg_result) = transport.next().await { + let _span = tracing::span!(tracing::Level::INFO, "message_processing").entered(); match msg_result { Ok(msg) => { match msg { JsonRpcMessage::Request(request) => { + // Serialize request for logging + let id = request.id.clone(); + let request_json = serde_json::to_string(&request) + .unwrap_or_else(|_| "Failed to serialize request".to_string()); + + tracing::info!( + request_id = ?id, + method = ?request.method, + json = %request_json, + "Received request" + ); + // Process the request using our service let response = match service.call(request).await { Ok(resp) => resp, - Err(e) => JsonRpcResponse { - jsonrpc: "2.0".to_string(), - id: None, // We still have a valid request ID here - result: None, - error: Some(mcp_core::protocol::ErrorData { - code: mcp_core::protocol::INTERNAL_ERROR, - message: Into::::into(e).to_string(), - data: None, - }), + Err(e) => { + let error_msg = e.into().to_string(); + tracing::error!(error = %error_msg, "Request processing failed"); + JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id, + result: None, + error: Some(mcp_core::protocol::ErrorData { + code: mcp_core::protocol::INTERNAL_ERROR, + message: error_msg, + data: None, + }), + } }, }; + // Serialize response for logging + let response_json = serde_json::to_string(&response) + .unwrap_or_else(|_| "Failed to serialize response".to_string()); + + tracing::info!( + response_id = ?response.id, + json = %response_json, + "Sending response" + ); // Send the response back if let Err(e) = transport.write_message(JsonRpcMessage::Response(response)).await { return Err(ServerError::Transport(TransportError::Io(e))); diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs index ecf6947bf7b0..b4d8ac570b69 100644 --- a/crates/mcp-server/src/main.rs +++ b/crates/mcp-server/src/main.rs @@ -3,9 +3,31 @@ use mcp_server::{Router, Server, ByteTransport}; use tokio::io::{stdin, stdout}; use tokio::time::Duration; use tower::timeout::Timeout; +use tracing_subscriber::{self, EnvFilter}; +use tracing_appender::rolling::{RollingFileAppender, Rotation}; #[tokio::main] async fn main() -> Result<()> { + // Set up file appender + let file_appender = RollingFileAppender::new( + Rotation::DAILY, + "logs", + "mcp-server.log", + ); + + // Initialize the tracing subscriber with file and stdout logging + tracing_subscriber::fmt() + .with_env_filter(EnvFilter::from_default_env() + .add_directive(tracing::Level::INFO.into())) + .with_writer(file_appender) + .with_target(false) + .with_thread_ids(true) + .with_file(true) + .with_line_number(true) + .init(); + + tracing::info!("Starting MCP server"); + let router = Router::new(); // Add a 30 second timeout to all requests let router = Timeout::new(router, Duration::from_secs(30)); @@ -13,4 +35,4 @@ async fn main() -> Result<()> { let transport = ByteTransport::new(stdin(), stdout()); Ok(server.run(transport).await?) -} +} \ No newline at end of file diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs index 2e0c587625d6..1eadf07ed6cd 100644 --- a/crates/mcp-server/src/router.rs +++ b/crates/mcp-server/src/router.rs @@ -55,7 +55,7 @@ impl Service for Router { }, "initialize" => { let result = InitializeResult { - protocol_version: "0.1.0".to_string(), + protocol_version: "2024-11-05".to_string(), capabilities, server_info: Implementation { name: "mcp-server".to_string(), From 31b159b1b00a8e210e6cdb73f2da34263c4ddda1 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Mon, 16 Dec 2024 10:15:47 -0800 Subject: [PATCH 03/10] Add macro sketch --- crates/mcp-macros/Cargo.toml | 23 ++++ crates/mcp-macros/examples/calculator.rs | 55 ++++++++ crates/mcp-macros/src/lib.rs | 157 +++++++++++++++++++++++ 3 files changed, 235 insertions(+) create mode 100644 crates/mcp-macros/Cargo.toml create mode 100644 crates/mcp-macros/examples/calculator.rs create mode 100644 crates/mcp-macros/src/lib.rs diff --git a/crates/mcp-macros/Cargo.toml b/crates/mcp-macros/Cargo.toml new file mode 100644 index 000000000000..469b536fd1e3 --- /dev/null +++ b/crates/mcp-macros/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "mcp-macros" +version = "0.1.0" +edition = "2021" + +[lib] +proc-macro = true + +[dependencies] +syn = { version = "2.0", features = ["full", "extra-traits"] } +quote = "1.0" +proc-macro2 = "1.0" +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +mcp-core = { path = "../mcp-core" } +async-trait = "0.1" +schemars = "0.8" + +[dev-dependencies] +tokio = { version = "1.0", features = ["full"] } +async-trait = "0.1" +serde_json = "1.0" +schemars = "0.8" diff --git a/crates/mcp-macros/examples/calculator.rs b/crates/mcp-macros/examples/calculator.rs new file mode 100644 index 000000000000..f562e299e53e --- /dev/null +++ b/crates/mcp-macros/examples/calculator.rs @@ -0,0 +1,55 @@ +use mcp_core::handler::{Tool, ToolError, Result}; +use mcp_macros::tool; + +#[tokio::main] +async fn main() -> std::result::Result<(), Box> { + // Create an instance of our tool + let calculator = Calculator::default(); + + // Print tool information + println!("Tool name: {}", Calculator::name()); + println!("Tool description: {}", Calculator::description()); + println!("Tool schema: {}", Calculator::schema()); + + // Test the tool with some sample input + let input = serde_json::json!({ + "x": 5, + "y": 3, + "operation": "multiply" + }); + + let result = calculator.call(input).await?; + println!("Result: {}", result); + + Ok(()) +} + +#[tool( + name = "calculator", + description = "Perform basic arithmetic operations", + params( + x = "First number in the calculation", + y = "Second number in the calculation", + operation = "The operation to perform (add, subtract, multiply, divide)" + ) +)] +async fn calculator(x: i32, y: i32, operation: String) -> Result { + match operation.as_str() { + "add" => Ok(x + y), + "subtract" => Ok(x - y), + "multiply" => Ok(x * y), + "divide" => { + if y == 0 { + Err(ToolError::ExecutionError( + "Division by zero".into(), + )) + } else { + Ok(x / y) + } + } + _ => Err(ToolError::InvalidParameters(format!( + "Unknown operation: {}", + operation + ))), + } +} diff --git a/crates/mcp-macros/src/lib.rs b/crates/mcp-macros/src/lib.rs new file mode 100644 index 000000000000..293aef8e3a82 --- /dev/null +++ b/crates/mcp-macros/src/lib.rs @@ -0,0 +1,157 @@ +use proc_macro::TokenStream; +use quote::{format_ident, quote}; +use std::collections::HashMap; +use syn::{ + parse::Parse, parse::ParseStream, parse_macro_input, punctuated::Punctuated, Expr, ExprLit, + FnArg, ItemFn, Lit, Meta, Pat, PatType, Token, +}; + +struct MacroArgs { + name: Option, + description: Option, + param_descriptions: HashMap, +} + +impl Parse for MacroArgs { + fn parse(input: ParseStream) -> syn::Result { + let mut name = None; + let mut description = None; + let mut param_descriptions = HashMap::new(); + + let meta_list: Punctuated = Punctuated::parse_terminated(input)?; + + for meta in meta_list { + match meta { + Meta::NameValue(nv) => { + let ident = nv.path.get_ident().unwrap().to_string(); + if let Expr::Lit(ExprLit { + lit: Lit::Str(lit_str), + .. + }) = nv.value + { + match ident.as_str() { + "name" => name = Some(lit_str.value()), + "description" => description = Some(lit_str.value()), + _ => {} + } + } + } + Meta::List(list) if list.path.is_ident("params") => { + let nested: Punctuated = + list.parse_args_with(Punctuated::parse_terminated)?; + + for meta in nested { + if let Meta::NameValue(nv) = meta { + if let Expr::Lit(ExprLit { + lit: Lit::Str(lit_str), + .. + }) = nv.value + { + let param_name = nv.path.get_ident().unwrap().to_string(); + param_descriptions.insert(param_name, lit_str.value()); + } + } + } + } + _ => {} + } + } + + Ok(MacroArgs { + name, + description, + param_descriptions, + }) + } +} + +#[proc_macro_attribute] +pub fn tool(args: TokenStream, input: TokenStream) -> TokenStream { + let args = parse_macro_input!(args as MacroArgs); + let input_fn = parse_macro_input!(input as ItemFn); + + // Extract function details + let fn_name = &input_fn.sig.ident; + let fn_name_str = fn_name.to_string(); + + // Generate PascalCase struct name from the function name + let struct_name = format_ident!("{}", { + let mut s = fn_name_str.clone(); + if let Some(c) = s.get_mut(0..1) { + c.make_ascii_uppercase(); + } + s + }); + + // Use provided name or function name as default + let tool_name = args.name.unwrap_or(fn_name_str); + let tool_description = args.description.unwrap_or_default(); + + // Extract parameter names, types, and descriptions + let mut param_defs = Vec::new(); + let mut param_names = Vec::new(); + + for arg in input_fn.sig.inputs.iter() { + if let FnArg::Typed(PatType { pat, ty, .. }) = arg { + if let Pat::Ident(param_ident) = &**pat { + let param_name = ¶m_ident.ident; + let param_name_str = param_name.to_string(); + let description = args + .param_descriptions + .get(¶m_name_str) + .map(|s| s.as_str()) + .unwrap_or(""); + + param_names.push(param_name); + param_defs.push(quote! { + #[schemars(description = #description)] + #param_name: #ty + }); + } + } + } + + // Generate the implementation + let params_struct_name = format_ident!("{}Parameters", struct_name); + let expanded = quote! { + #[derive(serde::Deserialize, schemars::JsonSchema)] + struct #params_struct_name { + #(#param_defs,)* + } + + #input_fn + + #[derive(Default)] + struct #struct_name; + + #[async_trait::async_trait] + impl mcp_core::handler::Tool for #struct_name { + fn name() -> &'static str { + #tool_name + } + + fn description() -> &'static str { + #tool_description + } + + fn schema() -> serde_json::Value { + mcp_core::handler::generate_schema::<#params_struct_name>() + .expect("Failed to generate schema") + } + + async fn call(&self, params: serde_json::Value) -> mcp_core::handler::Result { + let params: #params_struct_name = serde_json::from_value(params) + .map_err(|e| mcp_core::handler::ToolError::InvalidParameters(e.to_string()))?; + + // Extract parameters and call the function + let result = #fn_name(#(params.#param_names,)*).await + .map_err(|e| mcp_core::handler::ToolError::ExecutionError(e.to_string()))?; + + serde_json::to_value(result) + .map_err(mcp_core::handler::ToolError::SerializationError) + } + } + }; + + TokenStream::from(expanded) +} From 2b330941250a8e33b10a0953eae5ceb63a4f6d97 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Mon, 16 Dec 2024 13:27:29 -0800 Subject: [PATCH 04/10] simplify the tool handler --- crates/mcp-core/src/handler.rs | 6 +++--- crates/mcp-macros/examples/calculator.rs | 6 +++--- crates/mcp-macros/src/lib.rs | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/mcp-core/src/handler.rs b/crates/mcp-core/src/handler.rs index 167a7d56a106..42d924008504 100644 --- a/crates/mcp-core/src/handler.rs +++ b/crates/mcp-core/src/handler.rs @@ -21,13 +21,13 @@ pub type Result = std::result::Result; #[async_trait] pub trait Tool: Send + Sync + 'static { /// The name of the tool - fn name() -> &'static str; + fn name(&self) -> &'static str; /// A description of what the tool does - fn description() -> &'static str; + fn description(&self) -> &'static str; /// JSON schema describing the tool's parameters - fn schema() -> Value; + fn schema(&self) -> Value; /// Execute the tool with the given parameters async fn call(&self, params: Value) -> Result; diff --git a/crates/mcp-macros/examples/calculator.rs b/crates/mcp-macros/examples/calculator.rs index f562e299e53e..38c409bebfe9 100644 --- a/crates/mcp-macros/examples/calculator.rs +++ b/crates/mcp-macros/examples/calculator.rs @@ -7,9 +7,9 @@ async fn main() -> std::result::Result<(), Box> { let calculator = Calculator::default(); // Print tool information - println!("Tool name: {}", Calculator::name()); - println!("Tool description: {}", Calculator::description()); - println!("Tool schema: {}", Calculator::schema()); + println!("Tool name: {}", calculator.name()); + println!("Tool description: {}", calculator.description()); + println!("Tool schema: {}", calculator.schema()); // Test the tool with some sample input let input = serde_json::json!({ diff --git a/crates/mcp-macros/src/lib.rs b/crates/mcp-macros/src/lib.rs index 293aef8e3a82..f3c4682096cb 100644 --- a/crates/mcp-macros/src/lib.rs +++ b/crates/mcp-macros/src/lib.rs @@ -126,15 +126,15 @@ pub fn tool(args: TokenStream, input: TokenStream) -> TokenStream { #[async_trait::async_trait] impl mcp_core::handler::Tool for #struct_name { - fn name() -> &'static str { + fn name(&self) -> &'static str { #tool_name } - fn description() -> &'static str { + fn description(&self) -> &'static str { #tool_description } - fn schema() -> serde_json::Value { + fn schema(&self) -> serde_json::Value { mcp_core::handler::generate_schema::<#params_struct_name>() .expect("Failed to generate schema") } From 531a1cb3416a3680e1c1827ac688f06ac66fda20 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Mon, 16 Dec 2024 13:43:45 -0800 Subject: [PATCH 05/10] ensure correct casing --- crates/mcp-macros/Cargo.toml | 1 + crates/mcp-macros/src/lib.rs | 7 ++----- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/crates/mcp-macros/Cargo.toml b/crates/mcp-macros/Cargo.toml index 469b536fd1e3..65450f7d2b00 100644 --- a/crates/mcp-macros/Cargo.toml +++ b/crates/mcp-macros/Cargo.toml @@ -15,6 +15,7 @@ serde_json = "1.0" mcp-core = { path = "../mcp-core" } async-trait = "0.1" schemars = "0.8" +convert_case = "0.6.0" [dev-dependencies] tokio = { version = "1.0", features = ["full"] } diff --git a/crates/mcp-macros/src/lib.rs b/crates/mcp-macros/src/lib.rs index f3c4682096cb..18929f3233be 100644 --- a/crates/mcp-macros/src/lib.rs +++ b/crates/mcp-macros/src/lib.rs @@ -5,6 +5,7 @@ use syn::{ parse::Parse, parse::ParseStream, parse_macro_input, punctuated::Punctuated, Expr, ExprLit, FnArg, ItemFn, Lit, Meta, Pat, PatType, Token, }; +use convert_case::{Case, Casing}; struct MacroArgs { name: Option, @@ -76,11 +77,7 @@ pub fn tool(args: TokenStream, input: TokenStream) -> TokenStream { // Generate PascalCase struct name from the function name let struct_name = format_ident!("{}", { - let mut s = fn_name_str.clone(); - if let Some(c) = s.get_mut(0..1) { - c.make_ascii_uppercase(); - } - s + fn_name_str.to_case(Case::Pascal) }); // Use provided name or function name as default From 02cdfacb5fbb0bc5c07915cd7938637e827cf8be Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Mon, 16 Dec 2024 14:21:09 -0800 Subject: [PATCH 06/10] forgot a file --- crates/mcp-server/src/error.rs | 43 ++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 crates/mcp-server/src/error.rs diff --git a/crates/mcp-server/src/error.rs b/crates/mcp-server/src/error.rs new file mode 100644 index 000000000000..792394379367 --- /dev/null +++ b/crates/mcp-server/src/error.rs @@ -0,0 +1,43 @@ +use thiserror::Error; + +pub type BoxError = Box; + +#[derive(Debug, Error)] +pub enum RouterError { + #[error("Method not found: {0}")] + MethodNotFound(String), + #[error("Unknown tool: {0}")] + UnknownTool(String), + #[error("Invalid parameters: {0}")] + InvalidParams(String), + #[error("Internal error: {0}")] + Internal(String), +} + +impl From for mcp_core::protocol::ErrorData { + fn from(err: RouterError) -> Self { + use mcp_core::protocol::*; + match err { + RouterError::MethodNotFound(_) => ErrorData { + code: METHOD_NOT_FOUND, + message: err.to_string(), + data: None, + }, + RouterError::UnknownTool(_) => ErrorData { + code: INVALID_REQUEST, + message: err.to_string(), + data: None, + }, + RouterError::InvalidParams(_) => ErrorData { + code: INVALID_PARAMS, + message: err.to_string(), + data: None, + }, + RouterError::Internal(_) => ErrorData { + code: INTERNAL_ERROR, + message: err.to_string(), + data: None, + }, + } + } +} \ No newline at end of file From 82df445bcae5fae46e6d6acf97390e8e5f8bc402 Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Tue, 17 Dec 2024 21:05:35 -0800 Subject: [PATCH 07/10] working tools --- crates/mcp-core/src/handler.rs | 4 +- crates/mcp-core/src/tool.rs | 2 + crates/mcp-macros/examples/calculator.rs | 2 +- crates/mcp-macros/src/lib.rs | 2 +- crates/mcp-server/Cargo.toml | 4 +- crates/mcp-server/src/errors.rs | 9 + crates/mcp-server/src/main.rs | 41 ++- crates/mcp-server/src/router.rs | 323 +++++++++++++++++++---- 8 files changed, 331 insertions(+), 56 deletions(-) diff --git a/crates/mcp-core/src/handler.rs b/crates/mcp-core/src/handler.rs index 42d924008504..cb1c0e3b2183 100644 --- a/crates/mcp-core/src/handler.rs +++ b/crates/mcp-core/src/handler.rs @@ -19,7 +19,7 @@ pub type Result = std::result::Result; /// Trait for implementing MCP tools #[async_trait] -pub trait Tool: Send + Sync + 'static { +pub trait ToolHandler: Send + Sync + 'static { /// The name of the tool fn name(&self) -> &'static str; @@ -35,7 +35,7 @@ pub trait Tool: Send + Sync + 'static { /// Trait for implementing MCP resources #[async_trait] -pub trait Resource: Send + Sync + 'static { +pub trait ResourceTemplateHandler: Send + Sync + 'static { /// The URL template for this resource fn template() -> &'static str; diff --git a/crates/mcp-core/src/tool.rs b/crates/mcp-core/src/tool.rs index 6401b9632671..2497abe64b1f 100644 --- a/crates/mcp-core/src/tool.rs +++ b/crates/mcp-core/src/tool.rs @@ -5,6 +5,7 @@ use serde_json::Value; /// A tool that can be used by a model. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct Tool { /// The name of the tool pub name: String, @@ -31,6 +32,7 @@ impl Tool { /// A tool call request that a system can execute #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct ToolCall { /// The name of the tool to execute pub name: String, diff --git a/crates/mcp-macros/examples/calculator.rs b/crates/mcp-macros/examples/calculator.rs index 38c409bebfe9..d4cee5ca5ca0 100644 --- a/crates/mcp-macros/examples/calculator.rs +++ b/crates/mcp-macros/examples/calculator.rs @@ -1,4 +1,4 @@ -use mcp_core::handler::{Tool, ToolError, Result}; +use mcp_core::handler::{ToolHandler, ToolError, Result}; use mcp_macros::tool; #[tokio::main] diff --git a/crates/mcp-macros/src/lib.rs b/crates/mcp-macros/src/lib.rs index 18929f3233be..c0ef725fb826 100644 --- a/crates/mcp-macros/src/lib.rs +++ b/crates/mcp-macros/src/lib.rs @@ -122,7 +122,7 @@ pub fn tool(args: TokenStream, input: TokenStream) -> TokenStream { struct #struct_name; #[async_trait::async_trait] - impl mcp_core::handler::Tool for #struct_name { + impl mcp_core::handler::ToolHandler for #struct_name { fn name(&self) -> &'static str { #tool_name } diff --git a/crates/mcp-server/Cargo.toml b/crates/mcp-server/Cargo.toml index 8fe85768c61c..3657fc582ab5 100644 --- a/crates/mcp-server/Cargo.toml +++ b/crates/mcp-server/Cargo.toml @@ -7,8 +7,10 @@ edition = "2021" anyhow = "1.0.94" thiserror = "1.0" mcp-core = { path = "../mcp-core" } -serde = "1.0.216" +mcp-macros = { path = "../mcp-macros" } +serde = { version = "1.0.216", features = ["derive"] } serde_json = "1.0.133" +schemars = "0.8" tokio = { version = "1", features = ["full"] } tower = { version = "0.4", features = ["timeout"] } tower-service = "0.3" diff --git a/crates/mcp-server/src/errors.rs b/crates/mcp-server/src/errors.rs index 9bbfeb543d4e..1f19069cdb99 100644 --- a/crates/mcp-server/src/errors.rs +++ b/crates/mcp-server/src/errors.rs @@ -45,6 +45,10 @@ pub enum RouterError { #[error("Internal error: {0}")] Internal(String), + + #[error("Tool not found: {0}")] + ToolNotFound(String), + } impl From for mcp_core::protocol::ErrorData { @@ -66,6 +70,11 @@ impl From for mcp_core::protocol::ErrorData { message: msg, data: None, }, + RouterError::ToolNotFound(msg) => ErrorData { + code: INVALID_REQUEST, + message: msg, + data: None, + }, } } } diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs index b4d8ac570b69..c6d3d3347099 100644 --- a/crates/mcp-server/src/main.rs +++ b/crates/mcp-server/src/main.rs @@ -1,10 +1,43 @@ use anyhow::Result; +use mcp_core::handler::ToolError; use mcp_server::{Router, Server, ByteTransport}; use tokio::io::{stdin, stdout}; use tokio::time::Duration; use tower::timeout::Timeout; use tracing_subscriber::{self, EnvFilter}; use tracing_appender::rolling::{RollingFileAppender, Rotation}; +use mcp_macros::tool; + +#[tool( + name = "calculator", + description = "Perform basic arithmetic operations", + params( + x = "First number in the calculation", + y = "Second number in the calculation", + operation = "The operation to perform (add, subtract, multiply, divide)" + ) +)] +async fn calculator(x: String, y: String, operation: String) -> mcp_core::handler::Result { + let x: i32 = x.parse().unwrap(); + let y: i32 = y.parse().unwrap(); + match operation.as_str() { + "add" => Ok(x + y), + "subtract" => Ok(x - y), + "multiply" => Ok(x * y), + "divide" => { + if y == 0 { + Err(ToolError::ExecutionError("Division by zero".into())) + } else { + Ok(x / y) + } + } + _ => Err(ToolError::InvalidParameters(format!( + "Unknown operation: {}", + operation + ))), + } +} + #[tokio::main] async fn main() -> Result<()> { @@ -28,11 +61,15 @@ async fn main() -> Result<()> { tracing::info!("Starting MCP server"); - let router = Router::new(); + // Create a router with the calculator tool + let router = Router::builder() + .with_tool(Calculator::default()) + .build(); + // Add a 30 second timeout to all requests let router = Timeout::new(router, Duration::from_secs(30)); let server = Server::new(router); let transport = ByteTransport::new(stdin(), stdout()); Ok(server.run(transport).await?) -} \ No newline at end of file +} diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs index 1eadf07ed6cd..c32a4fd26b71 100644 --- a/crates/mcp-server/src/router.rs +++ b/crates/mcp-server/src/router.rs @@ -1,33 +1,185 @@ use std::{ + collections::HashMap, future::Future, pin::Pin, + sync::Arc, task::{Context, Poll}, time::Duration, }; use tokio::time::sleep; -use mcp_core::protocol::{ - JsonRpcRequest, JsonRpcResponse, ServerCapabilities, InitializeResult, Implementation, +use mcp_core::{ + handler::ToolHandler, + protocol::{ + JsonRpcRequest, JsonRpcResponse, ServerCapabilities, InitializeResult, Implementation, + PromptsCapability, ResourcesCapability, ToolsCapability, ListToolsResult, CallToolResult, + }, + content::Content, }; use tower_service::Service; +use serde_json::Value; use crate::{RouterError, BoxError}; +/// Builder for configuring and constructing a Router +pub struct RouterBuilder { + tools: HashMap>, + prompts: Option, + resources: Option, +} + +impl RouterBuilder { + pub fn new() -> Self { + Self { + tools: HashMap::new(), + prompts: None, + resources: None, + } + } + + /// Add a tool to the router + pub fn with_tool(mut self, tool: T) -> Self { + self.tools.insert(tool.name().to_string(), Box::new(tool)); + self + } + + /// Add multiple tools to the router + pub fn with_tools(mut self, tools: I) -> Self + where + I: IntoIterator, + I::Item: ToolHandler, + { + for tool in tools { + self.tools.insert(tool.name().to_string(), Box::new(tool)); + } + self + } + + /// Enable prompts capability + pub fn with_prompts(mut self, list_changed: bool) -> Self { + self.prompts = Some(PromptsCapability { + list_changed: Some(list_changed), + }); + self + } + + /// Enable resources capability + pub fn with_resources(mut self, subscribe: bool, list_changed: bool) -> Self { + self.resources = Some(ResourcesCapability { + subscribe: Some(subscribe), + list_changed: Some(list_changed), + }); + self + } + + /// Build the router with automatic capability inference + pub fn build(self) -> Router { + // Create capabilities based on what's configured + let capabilities = ServerCapabilities { + // Add tools capability if we have any tools + tools: (!self.tools.is_empty()).then(|| ToolsCapability { + list_changed: Some(false), + }), + // Add other capabilities that were explicitly set + prompts: self.prompts, + resources: self.resources, + }; + + Router { + capabilities, + tools: Arc::new(self.tools), + } + } +} + + /// Router service that implements the JSON-RPC protocol +#[derive(Clone)] pub struct Router { capabilities: ServerCapabilities, + tools: Arc>>, } impl Router { - pub fn new() -> Self { - Self { - capabilities: ServerCapabilities { - prompts: None, - resources: None, - tools: None, - }, + pub fn builder() -> RouterBuilder { + RouterBuilder::new() + } + + // Helper method to create base response + fn create_response(&self, id: Option) -> JsonRpcResponse { + JsonRpcResponse { + jsonrpc: "2.0".to_string(), + id, + result: None, + error: None, } } + + async fn handle_initialize(&self, req: JsonRpcRequest) -> Result { + let result = InitializeResult { + protocol_version: "2024-11-05".to_string(), + capabilities: self.capabilities.clone(), + server_info: Implementation { + name: "mcp-server".to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + }, + }; + + let mut response = self.create_response(req.id); + response.result = Some(serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + + Ok(response) + } + + async fn handle_tools_list(&self, req: JsonRpcRequest) -> Result { + let tools = self.tools.values() + .map(|tool| mcp_core::tool::Tool { + name: tool.name().to_string(), + description: tool.description().to_string(), + input_schema: tool.schema(), + }) + .collect(); + + let result = ListToolsResult { tools }; + let mut response = self.create_response(req.id); + response.result = Some(serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + + Ok(response) + } + + async fn handle_tools_call(&self, req: JsonRpcRequest) -> Result { + let params = req.params.ok_or_else(|| RouterError::InvalidParams("Missing parameters".into()))?; + + let name = params.get("name") + .and_then(Value::as_str) + .ok_or_else(|| RouterError::InvalidParams("Missing tool name".into()))?; + + let arguments = params.get("arguments") + .cloned() + .unwrap_or(Value::Null); + + let tool = self.tools.get(name) + .ok_or_else(|| RouterError::ToolNotFound(name.to_string()))?; + + let result = match tool.call(arguments).await { + Ok(result) => CallToolResult { + content: vec![Content::text(result.to_string())], + is_error: false, + }, + Err(err) => CallToolResult { + content: vec![Content::text(err.to_string())], + is_error: true, + } + }; + + let mut response = self.create_response(req.id); + response.result = Some(serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + + Ok(response) + } } impl Service for Router { @@ -40,51 +192,124 @@ impl Service for Router { } fn call(&mut self, req: JsonRpcRequest) -> Self::Future { - let capabilities = self.capabilities.clone(); - let fut = async move { - match req.method.as_str() { + // Create owned copy of self for the async block + // This is safe because Router contains Arc'd data + let this = self.clone(); + + Box::pin(async move { + let result = match req.method.as_str() { "slow" => { - // Sleep for 60 seconds - this should trigger the timeout sleep(Duration::from_secs(60)).await; - Ok(JsonRpcResponse { - jsonrpc: "2.0".to_string(), - id: req.id, - result: Some(serde_json::Value::String("This should never be seen".to_string())), - error: None, - }) - }, - "initialize" => { - let result = InitializeResult { - protocol_version: "2024-11-05".to_string(), - capabilities, - server_info: Implementation { - name: "mcp-server".to_string(), - version: env!("CARGO_PKG_VERSION").to_string(), - }, - }; - - let result_value = serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?; - - Ok(JsonRpcResponse { - jsonrpc: "2.0".to_string(), - id: req.id, - result: Some(result_value), - error: None, - }) + let mut response = this.create_response(req.id); + response.result = Some(Value::String("This should never be seen".to_string())); + Ok(response) }, + "initialize" => this.handle_initialize(req).await, + "tools/list" => this.handle_tools_list(req).await, + "tools/call" => this.handle_tools_call(req).await, _ => { - let error = RouterError::MethodNotFound(req.method); - Ok(JsonRpcResponse { - jsonrpc: "2.0".to_string(), - id: req.id, - result: None, - error: Some(error.into()), - }) - }, - } + let mut response = this.create_response(req.id); + response.error = Some(RouterError::MethodNotFound(req.method).into()); + Ok(response) + } + }; + + result.map_err(BoxError::from) + }) + } +} + + +#[cfg(test)] +mod tests { + use super::*; + use mcp_core::handler::{ToolError, ToolHandler}; + use async_trait::async_trait; + + struct TestTool; + + #[async_trait] + impl ToolHandler for TestTool { + fn name(&self) -> &'static str { + "test" + } + + fn description(&self) -> &'static str { + "A test tool" + } + + fn schema(&self) -> Value { + serde_json::json!({ + "type": "object", + "properties": { + "echo": { + "type": "string" + } + } + }) + } + + async fn call(&self, params: Value) -> Result { + Ok(params) + } + } + + #[tokio::test] + async fn test_router_builder() { + let router = Router::builder() + .with_tool(TestTool) + .with_prompts(true) + .with_resources(true, true) + .build(); + + assert!(router.capabilities.tools.is_some()); + assert!(router.capabilities.prompts.is_some()); + assert!(router.capabilities.resources.is_some()); + assert!(router.tools.contains_key("test")); + } + + #[tokio::test] + async fn test_tools_list() { + let router = Router::builder() + .with_tool(TestTool) + .build(); + + let req = JsonRpcRequest { + jsonrpc: "2.0".to_string(), + id: Some(1), + method: "tools/list".to_string(), + params: None, + }; + + let mut router_service = router; + let response = router_service.call(req).await.unwrap(); + + assert!(response.error.is_none()); + assert!(response.result.is_some()); + } + + #[tokio::test] + async fn test_tools_call() { + let router = Router::builder() + .with_tool(TestTool) + .build(); + + let req = JsonRpcRequest { + jsonrpc: "2.0".to_string(), + id: Some(1), + method: "tools/call".to_string(), + params: Some(serde_json::json!({ + "name": "test", + "arguments": { + "echo": "hello" + } + })), }; - Box::pin(fut) + let mut router_service = router; + let response = router_service.call(req).await.unwrap(); + + assert!(response.error.is_none()); + assert!(response.result.is_some()); } } From a3b4c478c3a229c28e85a20ab85464abebc8797c Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Tue, 17 Dec 2024 22:41:54 -0800 Subject: [PATCH 08/10] simplified version --- crates/mcp-core/src/handler.rs | 2 + crates/mcp-server/src/main.rs | 177 ++++++++++++++++----- crates/mcp-server/src/router.rs | 268 +++++++++----------------------- 3 files changed, 211 insertions(+), 236 deletions(-) diff --git a/crates/mcp-core/src/handler.rs b/crates/mcp-core/src/handler.rs index cb1c0e3b2183..5dfd51681e89 100644 --- a/crates/mcp-core/src/handler.rs +++ b/crates/mcp-core/src/handler.rs @@ -13,6 +13,8 @@ pub enum ToolError { SerializationError(#[from] serde_json::Error), #[error("Schema error: {0}")] SchemaError(String), + #[error("Tool not found: {0}")] + NotFound(String), } pub type Result = std::result::Result; diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs index c6d3d3347099..9542962e5253 100644 --- a/crates/mcp-server/src/main.rs +++ b/crates/mcp-server/src/main.rs @@ -1,47 +1,114 @@ use anyhow::Result; -use mcp_core::handler::ToolError; +use mcp_core::{ + handler::ToolError, + tool::Tool, + protocol::ServerCapabilities, +}; use mcp_server::{Router, Server, ByteTransport}; -use tokio::io::{stdin, stdout}; -use tokio::time::Duration; -use tower::timeout::Timeout; +use mcp_server::router::{RouterService, CapabilitiesBuilder}; +use serde_json::Value; +use std::{future::Future, pin::Pin, sync::Arc}; +use tokio::{io::{stdin, stdout}, sync::Mutex}; use tracing_subscriber::{self, EnvFilter}; use tracing_appender::rolling::{RollingFileAppender, Rotation}; -use mcp_macros::tool; - -#[tool( - name = "calculator", - description = "Perform basic arithmetic operations", - params( - x = "First number in the calculation", - y = "Second number in the calculation", - operation = "The operation to perform (add, subtract, multiply, divide)" - ) -)] -async fn calculator(x: String, y: String, operation: String) -> mcp_core::handler::Result { - let x: i32 = x.parse().unwrap(); - let y: i32 = y.parse().unwrap(); - match operation.as_str() { - "add" => Ok(x + y), - "subtract" => Ok(x - y), - "multiply" => Ok(x * y), - "divide" => { - if y == 0 { - Err(ToolError::ExecutionError("Division by zero".into())) - } else { - Ok(x / y) - } + +// A simple counter service that demonstrates the Router trait +#[derive(Clone)] +struct CounterRouter { + counter: Arc>, +} + +impl CounterRouter { + fn new() -> Self { + Self { + counter: Arc::new(Mutex::new(0)), } - _ => Err(ToolError::InvalidParameters(format!( - "Unknown operation: {}", - operation - ))), + } + + async fn increment(&self) -> Result { + let mut counter = self.counter.lock().await; + *counter += 1; + Ok(*counter) + } + + async fn decrement(&self) -> Result { + let mut counter = self.counter.lock().await; + *counter -= 1; + Ok(*counter) + } + + async fn get_value(&self) -> Result { + let counter = self.counter.lock().await; + Ok(*counter) } } +impl Router for CounterRouter { + fn capabilities(&self) -> ServerCapabilities { + CapabilitiesBuilder::new() + .with_tools(true) + .build() + } + + fn list_tools(&self) -> Vec { + vec![ + Tool::new( + "increment".to_string(), + "Increment the counter by 1".to_string(), + serde_json::json!({ + "type": "object", + "properties": {}, + "required": [] + }), + ), + Tool::new( + "decrement".to_string(), + "Decrement the counter by 1".to_string(), + serde_json::json!({ + "type": "object", + "properties": {}, + "required": [] + }), + ), + Tool::new( + "get_value".to_string(), + "Get the current counter value".to_string(), + serde_json::json!({ + "type": "object", + "properties": {}, + "required": [] + }), + ), + ] + } + + fn call_tool(&self, tool_name: &str, _arguments: Value) -> Pin> + Send + 'static>> { + let this = self.clone(); + let tool_name = tool_name.to_string(); + + Box::pin(async move { + match tool_name.as_str() { + "increment" => { + let value = this.increment().await?; + Ok(Value::Number(value.into())) + } + "decrement" => { + let value = this.decrement().await?; + Ok(Value::Number(value.into())) + } + "get_value" => { + let value = this.get_value().await?; + Ok(Value::Number(value.into())) + } + _ => Err(ToolError::NotFound(format!("Tool {} not found", tool_name))), + } + }) + } +} #[tokio::main] async fn main() -> Result<()> { - // Set up file appender + // Set up file appender for logging let file_appender = RollingFileAppender::new( Rotation::DAILY, "logs", @@ -61,15 +128,43 @@ async fn main() -> Result<()> { tracing::info!("Starting MCP server"); - // Create a router with the calculator tool - let router = Router::builder() - .with_tool(Calculator::default()) - .build(); - - // Add a 30 second timeout to all requests - let router = Timeout::new(router, Duration::from_secs(30)); - let server = Server::new(router); + // Create an instance of our counter router + let router = RouterService(CounterRouter::new()); + // Create and run the server + let server = Server::new(router); let transport = ByteTransport::new(stdin(), stdout()); + + tracing::info!("Server initialized and ready to handle requests"); Ok(server.run(transport).await?) } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::test; + + #[test] + async fn test_counter_operations() { + let router = CounterRouter::new(); + + // Test increment + let result = router.call_tool("increment", Value::Null).await.unwrap(); + assert_eq!(result, Value::Number(1.into())); + + // Test get_value + let result = router.call_tool("get_value", Value::Null).await.unwrap(); + assert_eq!(result, Value::Number(1.into())); + + // Test decrement + let result = router.call_tool("decrement", Value::Null).await.unwrap(); + assert_eq!(result, Value::Number(0.into())); + } + + #[test] + async fn test_invalid_tool() { + let router = CounterRouter::new(); + let result = router.call_tool("invalid_tool", Value::Null).await; + assert!(result.is_err()); + } +} diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs index c32a4fd26b71..925f21148713 100644 --- a/crates/mcp-server/src/router.rs +++ b/crates/mcp-server/src/router.rs @@ -1,15 +1,13 @@ use std::{ - collections::HashMap, future::Future, pin::Pin, - sync::Arc, task::{Context, Poll}, time::Duration, }; use tokio::time::sleep; use mcp_core::{ - handler::ToolHandler, + handler::ToolError, protocol::{ JsonRpcRequest, JsonRpcResponse, ServerCapabilities, InitializeResult, Implementation, PromptsCapability, ResourcesCapability, ToolsCapability, ListToolsResult, CallToolResult, @@ -21,37 +19,28 @@ use serde_json::Value; use crate::{RouterError, BoxError}; -/// Builder for configuring and constructing a Router -pub struct RouterBuilder { - tools: HashMap>, +/// Builder for configuring and constructing capabilities +pub struct CapabilitiesBuilder { + tools: Option, prompts: Option, resources: Option, } -impl RouterBuilder { +impl CapabilitiesBuilder { pub fn new() -> Self { Self { - tools: HashMap::new(), + tools: None, prompts: None, resources: None, } } - /// Add a tool to the router - pub fn with_tool(mut self, tool: T) -> Self { - self.tools.insert(tool.name().to_string(), Box::new(tool)); - self - } - /// Add multiple tools to the router - pub fn with_tools(mut self, tools: I) -> Self - where - I: IntoIterator, - I::Item: ToolHandler, + pub fn with_tools(mut self, list_changed: bool) -> Self { - for tool in tools { - self.tools.insert(tool.name().to_string(), Box::new(tool)); - } + self.tools = Some(ToolsCapability { + list_changed: Some(list_changed), + }); self } @@ -73,37 +62,20 @@ impl RouterBuilder { } /// Build the router with automatic capability inference - pub fn build(self) -> Router { + pub fn build(self) -> ServerCapabilities { // Create capabilities based on what's configured - let capabilities = ServerCapabilities { - // Add tools capability if we have any tools - tools: (!self.tools.is_empty()).then(|| ToolsCapability { - list_changed: Some(false), - }), - // Add other capabilities that were explicitly set + ServerCapabilities { + tools: self.tools, prompts: self.prompts, resources: self.resources, - }; - - Router { - capabilities, - tools: Arc::new(self.tools), } } } - -/// Router service that implements the JSON-RPC protocol -#[derive(Clone)] -pub struct Router { - capabilities: ServerCapabilities, - tools: Arc>>, -} - -impl Router { - pub fn builder() -> RouterBuilder { - RouterBuilder::new() - } +pub trait Router: Send + Sync + 'static { + fn capabilities(&self) -> ServerCapabilities; + fn list_tools(&self) -> Vec; + fn call_tool(&self, tool_name: &str, arguments: Value) -> Pin> + Send + 'static>>; // Helper method to create base response fn create_response(&self, id: Option) -> JsonRpcResponse { @@ -115,74 +87,77 @@ impl Router { } } - async fn handle_initialize(&self, req: JsonRpcRequest) -> Result { - let result = InitializeResult { - protocol_version: "2024-11-05".to_string(), - capabilities: self.capabilities.clone(), - server_info: Implementation { - name: "mcp-server".to_string(), - version: env!("CARGO_PKG_VERSION").to_string(), - }, - }; + fn handle_initialize(&self, req: JsonRpcRequest) -> impl Future> + Send { + async move { + let result = InitializeResult { + protocol_version: "2024-11-05".to_string(), + capabilities: self.capabilities().clone(), + server_info: Implementation { + name: "mcp-server".to_string(), + version: env!("CARGO_PKG_VERSION").to_string(), + }, + }; - let mut response = self.create_response(req.id); - response.result = Some(serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + let mut response = self.create_response(req.id); + response.result = Some(serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); - Ok(response) + Ok(response) + } } - async fn handle_tools_list(&self, req: JsonRpcRequest) -> Result { - let tools = self.tools.values() - .map(|tool| mcp_core::tool::Tool { - name: tool.name().to_string(), - description: tool.description().to_string(), - input_schema: tool.schema(), - }) - .collect(); + fn handle_tools_list(&self, req: JsonRpcRequest) -> impl Future> + Send { + async move { + let tools = self.list_tools(); - let result = ListToolsResult { tools }; - let mut response = self.create_response(req.id); - response.result = Some(serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + let result = ListToolsResult { tools }; + let mut response = self.create_response(req.id); + response.result = Some(serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); - Ok(response) + Ok(response) + } } - async fn handle_tools_call(&self, req: JsonRpcRequest) -> Result { - let params = req.params.ok_or_else(|| RouterError::InvalidParams("Missing parameters".into()))?; - - let name = params.get("name") - .and_then(Value::as_str) - .ok_or_else(|| RouterError::InvalidParams("Missing tool name".into()))?; + fn handle_tools_call(&self, req: JsonRpcRequest) -> impl Future> + Send { + async move { + let params = req.params.ok_or_else(|| RouterError::InvalidParams("Missing parameters".into()))?; - let arguments = params.get("arguments") - .cloned() - .unwrap_or(Value::Null); + let name = params.get("name") + .and_then(Value::as_str) + .ok_or_else(|| RouterError::InvalidParams("Missing tool name".into()))?; - let tool = self.tools.get(name) - .ok_or_else(|| RouterError::ToolNotFound(name.to_string()))?; + let arguments = params.get("arguments") + .cloned() + .unwrap_or(Value::Null); - let result = match tool.call(arguments).await { - Ok(result) => CallToolResult { - content: vec![Content::text(result.to_string())], - is_error: false, - }, - Err(err) => CallToolResult { - content: vec![Content::text(err.to_string())], - is_error: true, - } - }; + let result = match self.call_tool(name, arguments).await { + Ok(result) => CallToolResult { + content: vec![Content::text(result.to_string())], + is_error: false, + }, + Err(err) => CallToolResult { + content: vec![Content::text(err.to_string())], + is_error: true, + } + }; - let mut response = self.create_response(req.id); - response.result = Some(serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + let mut response = self.create_response(req.id); + response.result = Some(serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); - Ok(response) + Ok(response) + } } } -impl Service for Router { +// A wrapper type to implement the Service trait locally +pub struct RouterService(pub T); + +impl Service for RouterService +where + T: Router + Clone + Send + Sync + 'static, +{ type Response = JsonRpcResponse; type Error = BoxError; type Future = Pin> + Send>>; @@ -192,9 +167,7 @@ impl Service for Router { } fn call(&mut self, req: JsonRpcRequest) -> Self::Future { - // Create owned copy of self for the async block - // This is safe because Router contains Arc'd data - let this = self.clone(); + let this = self.0.clone(); Box::pin(async move { let result = match req.method.as_str() { @@ -218,98 +191,3 @@ impl Service for Router { }) } } - - -#[cfg(test)] -mod tests { - use super::*; - use mcp_core::handler::{ToolError, ToolHandler}; - use async_trait::async_trait; - - struct TestTool; - - #[async_trait] - impl ToolHandler for TestTool { - fn name(&self) -> &'static str { - "test" - } - - fn description(&self) -> &'static str { - "A test tool" - } - - fn schema(&self) -> Value { - serde_json::json!({ - "type": "object", - "properties": { - "echo": { - "type": "string" - } - } - }) - } - - async fn call(&self, params: Value) -> Result { - Ok(params) - } - } - - #[tokio::test] - async fn test_router_builder() { - let router = Router::builder() - .with_tool(TestTool) - .with_prompts(true) - .with_resources(true, true) - .build(); - - assert!(router.capabilities.tools.is_some()); - assert!(router.capabilities.prompts.is_some()); - assert!(router.capabilities.resources.is_some()); - assert!(router.tools.contains_key("test")); - } - - #[tokio::test] - async fn test_tools_list() { - let router = Router::builder() - .with_tool(TestTool) - .build(); - - let req = JsonRpcRequest { - jsonrpc: "2.0".to_string(), - id: Some(1), - method: "tools/list".to_string(), - params: None, - }; - - let mut router_service = router; - let response = router_service.call(req).await.unwrap(); - - assert!(response.error.is_none()); - assert!(response.result.is_some()); - } - - #[tokio::test] - async fn test_tools_call() { - let router = Router::builder() - .with_tool(TestTool) - .build(); - - let req = JsonRpcRequest { - jsonrpc: "2.0".to_string(), - id: Some(1), - method: "tools/call".to_string(), - params: Some(serde_json::json!({ - "name": "test", - "arguments": { - "echo": "hello" - } - })), - }; - - let mut router_service = router; - let response = router_service.call(req).await.unwrap(); - - assert!(response.error.is_none()); - assert!(response.result.is_some()); - } -} From 6e262c8b4afc09885b334c254b6616beb2999947 Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Wed, 18 Dec 2024 15:52:46 -0500 Subject: [PATCH 09/10] add resource methods to the server sdk (#497) --- crates/mcp-core/src/handler.rs | 8 +++++ crates/mcp-server/README.md | 7 ++++ crates/mcp-server/src/errors.rs | 16 +++++++++ crates/mcp-server/src/main.rs | 32 ++++++++++++++++++ crates/mcp-server/src/router.rs | 60 +++++++++++++++++++++++++++++---- 5 files changed, 117 insertions(+), 6 deletions(-) create mode 100644 crates/mcp-server/README.md diff --git a/crates/mcp-core/src/handler.rs b/crates/mcp-core/src/handler.rs index 5dfd51681e89..407d91e64cb1 100644 --- a/crates/mcp-core/src/handler.rs +++ b/crates/mcp-core/src/handler.rs @@ -17,6 +17,14 @@ pub enum ToolError { NotFound(String), } +#[derive(Error, Debug)] +pub enum ResourceError { + #[error("Execution failed: {0}")] + ExecutionError(String), + #[error("Resource not found: {0}")] + NotFound(String), +} + pub type Result = std::result::Result; /// Trait for implementing MCP tools diff --git a/crates/mcp-server/README.md b/crates/mcp-server/README.md new file mode 100644 index 000000000000..1e4f06176553 --- /dev/null +++ b/crates/mcp-server/README.md @@ -0,0 +1,7 @@ +### Test with MCP Inspector + +```bash +npx @modelcontextprotocol/inspector cargo run -p mcp-server +``` + +Then visit the Inspector in the browser window and test the different endpoints. \ No newline at end of file diff --git a/crates/mcp-server/src/errors.rs b/crates/mcp-server/src/errors.rs index 1f19069cdb99..a87cd7aff1f6 100644 --- a/crates/mcp-server/src/errors.rs +++ b/crates/mcp-server/src/errors.rs @@ -49,6 +49,8 @@ pub enum RouterError { #[error("Tool not found: {0}")] ToolNotFound(String), + #[error("Resource not found: {0}")] + ResourceNotFound(String), } impl From for mcp_core::protocol::ErrorData { @@ -75,6 +77,20 @@ impl From for mcp_core::protocol::ErrorData { message: msg, data: None, }, + RouterError::ResourceNotFound(msg) => ErrorData { + code: INVALID_REQUEST, + message: msg, + data: None, + }, + } + } +} + +impl From for RouterError { + fn from(err: mcp_core::handler::ResourceError) -> Self { + match err { + mcp_core::handler::ResourceError::NotFound(msg) => RouterError::ResourceNotFound(msg), + _ => RouterError::Internal("Unknown resource error".to_string()), } } } diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs index 9542962e5253..07a89e0dd6d8 100644 --- a/crates/mcp-server/src/main.rs +++ b/crates/mcp-server/src/main.rs @@ -1,8 +1,10 @@ use anyhow::Result; +use mcp_core::handler::ResourceError; use mcp_core::{ handler::ToolError, tool::Tool, protocol::ServerCapabilities, + resource::Resource, }; use mcp_server::{Router, Server, ByteTransport}; use mcp_server::router::{RouterService, CapabilitiesBuilder}; @@ -104,6 +106,36 @@ impl Router for CounterRouter { } }) } + + fn list_resources(&self) -> Vec { + vec![ + Resource::new( + "memo://insights".to_string(), + Some("text/plain".to_string()), + Some("memo-resource".to_string()) + ).unwrap() + ] + } + + fn read_resource( + &self, + uri: &str, + ) -> Pin> + Send + 'static>> { + let uri = uri.to_string(); + Box::pin(async move { + match uri.as_str() { + "memo://insights" => { + let memo = + "Business Intelligence Memo\n\nAnalysis has revealed 5 key insights ..."; + Ok(memo.to_string()) + } + _ => Err(ResourceError::NotFound(format!( + "Resource {} not found", + uri + ))), + } + }) + } } #[tokio::main] diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs index 925f21148713..4c2448843334 100644 --- a/crates/mcp-server/src/router.rs +++ b/crates/mcp-server/src/router.rs @@ -7,12 +7,9 @@ use std::{ use tokio::time::sleep; use mcp_core::{ - handler::ToolError, - protocol::{ - JsonRpcRequest, JsonRpcResponse, ServerCapabilities, InitializeResult, Implementation, - PromptsCapability, ResourcesCapability, ToolsCapability, ListToolsResult, CallToolResult, - }, - content::Content, + content::Content, handler::{ResourceError, ToolError}, protocol::{ + CallToolResult, Implementation, InitializeResult, JsonRpcRequest, JsonRpcResponse, ListResourcesResult, ListToolsResult, PromptsCapability, ReadResourceResult, ResourcesCapability, ServerCapabilities, ToolsCapability + }, ResourceContents }; use tower_service::Service; use serde_json::Value; @@ -76,6 +73,8 @@ pub trait Router: Send + Sync + 'static { fn capabilities(&self) -> ServerCapabilities; fn list_tools(&self) -> Vec; fn call_tool(&self, tool_name: &str, arguments: Value) -> Pin> + Send + 'static>>; + fn list_resources(&self) -> Vec; + fn read_resource(&self, uri: &str) -> Pin> + Send + 'static>>; // Helper method to create base response fn create_response(&self, id: Option) -> JsonRpcResponse { @@ -149,6 +148,53 @@ pub trait Router: Send + Sync + 'static { Ok(response) } } + + fn handle_resources_list(&self, req: JsonRpcRequest) -> impl Future> + Send { + async move { + let resources = self.list_resources(); + + let result = ListResourcesResult { resources }; + let mut response = self.create_response(req.id); + response.result = Some(serde_json::to_value(result) + .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + + Ok(response) + } + } + + fn handle_resources_read( + &self, + req: JsonRpcRequest, + ) -> impl Future> + Send { + async move { + let params = req + .params + .ok_or_else(|| RouterError::InvalidParams("Missing parameters".into()))?; + + let uri = params + .get("uri") + .and_then(Value::as_str) + .ok_or_else(|| RouterError::InvalidParams("Missing resource URI".into()))?; + + let contents = self.read_resource(uri).await.map_err(RouterError::from)?; + + let result = ReadResourceResult { + contents: vec![ResourceContents::TextResourceContents { + uri: uri.to_string(), + mime_type: Some("text/plain".to_string()), + text: contents, + }], + }; + + let mut response = self.create_response(req.id); + response.result = + Some(serde_json::to_value(result).map_err(|e| { + RouterError::Internal(format!("JSON serialization error: {}", e)) + })?); + + Ok(response) + } + } } // A wrapper type to implement the Service trait locally @@ -180,6 +226,8 @@ where "initialize" => this.handle_initialize(req).await, "tools/list" => this.handle_tools_list(req).await, "tools/call" => this.handle_tools_call(req).await, + "resources/list" => this.handle_resources_list(req).await, + "resources/read" => this.handle_resources_read(req).await, _ => { let mut response = this.create_response(req.id); response.error = Some(RouterError::MethodNotFound(req.method).into()); From 1aca29702c9b65d9c5acf3076c8781851bd9aa5f Mon Sep 17 00:00:00 2001 From: Bradley Axen Date: Wed, 18 Dec 2024 22:07:57 -0800 Subject: [PATCH 10/10] clippy fmt test --- crates/goose/src/token_counter.rs | 129 +++++++++++------------ crates/mcp-client/src/stdio_transport.rs | 50 --------- crates/mcp-core/src/protocol.rs | 2 +- crates/mcp-macros/examples/calculator.rs | 8 +- crates/mcp-macros/src/lib.rs | 6 +- crates/mcp-server/src/error.rs | 43 -------- crates/mcp-server/src/errors.rs | 14 +-- crates/mcp-server/src/lib.rs | 69 +++++++----- crates/mcp-server/src/main.rs | 84 +++++---------- crates/mcp-server/src/router.rs | 101 +++++++++++------- 10 files changed, 205 insertions(+), 301 deletions(-) delete mode 100644 crates/mcp-server/src/error.rs diff --git a/crates/goose/src/token_counter.rs b/crates/goose/src/token_counter.rs index 40d1ab09a0ed..0af73b6a968e 100644 --- a/crates/goose/src/token_counter.rs +++ b/crates/goose/src/token_counter.rs @@ -207,6 +207,9 @@ impl TokenCounter { #[cfg(test)] mod tests { use super::*; + use crate::message::MessageContent; + use mcp_core::role::Role; + use serde_json::json; #[test] fn test_add_tokenizer_and_count_tokens() { @@ -235,73 +238,65 @@ mod tests { assert_eq!(count, 3); } - #[cfg(test)] - mod tests { - use super::*; - use crate::message::MessageContent; - use mcp_core::role::Role; - use serde_json::json; - - #[test] - fn test_count_chat_tokens() { - let token_counter = TokenCounter::new(); - - let system_prompt = - "You are a helpful assistant that can answer questions about the weather."; - - let messages = vec![ - Message { - role: Role::User, - created: 0, - content: vec![MessageContent::text( - "What's the weather like in San Francisco?", - )], - }, - Message { - role: Role::Assistant, - created: 1, - content: vec![MessageContent::text( - "Looks like it's 60 degrees Fahrenheit in San Francisco.", - )], - }, - Message { - role: Role::User, - created: 2, - content: vec![MessageContent::text("How about New York?")], - }, - ]; - - let tools = vec![Tool { - name: "get_current_weather".to_string(), - description: "Get the current weather in a given location".to_string(), - input_schema: json!({ - "properties": { - "location": { - "type": "string", - "description": "The city and state, e.g. San Francisco, CA" - }, - "unit": { - "type": "string", - "description": "The unit of temperature to return", - "enum": ["celsius", "fahrenheit"] - } + #[test] + fn test_count_chat_tokens() { + let token_counter = TokenCounter::new(); + + let system_prompt = + "You are a helpful assistant that can answer questions about the weather."; + + let messages = vec![ + Message { + role: Role::User, + created: 0, + content: vec![MessageContent::text( + "What's the weather like in San Francisco?", + )], + }, + Message { + role: Role::Assistant, + created: 1, + content: vec![MessageContent::text( + "Looks like it's 60 degrees Fahrenheit in San Francisco.", + )], + }, + Message { + role: Role::User, + created: 2, + content: vec![MessageContent::text("How about New York?")], + }, + ]; + + let tools = vec![Tool { + name: "get_current_weather".to_string(), + description: "Get the current weather in a given location".to_string(), + input_schema: json!({ + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA" }, - "required": ["location"] - }), - }]; - - let token_count_without_tools = - token_counter.count_chat_tokens(system_prompt, &messages, &vec![], Some("gpt-4o")); - println!("Total tokens without tools: {}", token_count_without_tools); - - let token_count_with_tools = - token_counter.count_chat_tokens(system_prompt, &messages, &tools, Some("gpt-4o")); - println!("Total tokens with tools: {}", token_count_with_tools); - - // The token count for messages without tools is calculated using the tokenizer - https://tiktokenizer.vercel.app/ - // The token count for messages with tools is taken from tiktoken github repo example (notebook) - assert_eq!(token_count_without_tools, 56); - assert_eq!(token_count_with_tools, 124); - } + "unit": { + "type": "string", + "description": "The unit of temperature to return", + "enum": ["celsius", "fahrenheit"] + } + }, + "required": ["location"] + }), + }]; + + let token_count_without_tools = + token_counter.count_chat_tokens(system_prompt, &messages, &[], Some("gpt-4o")); + println!("Total tokens without tools: {}", token_count_without_tools); + + let token_count_with_tools = + token_counter.count_chat_tokens(system_prompt, &messages, &tools, Some("gpt-4o")); + println!("Total tokens with tools: {}", token_count_with_tools); + + // The token count for messages without tools is calculated using the tokenizer - https://tiktokenizer.vercel.app/ + // The token count for messages with tools is taken from tiktoken github repo example (notebook) + assert_eq!(token_count_without_tools, 56); + assert_eq!(token_count_with_tools, 124); } } diff --git a/crates/mcp-client/src/stdio_transport.rs b/crates/mcp-client/src/stdio_transport.rs index 67f0fbf12053..3b965ab83a41 100644 --- a/crates/mcp-client/src/stdio_transport.rs +++ b/crates/mcp-client/src/stdio_transport.rs @@ -119,59 +119,9 @@ impl Transport for StdioTransport { #[cfg(test)] mod tests { use super::*; - use serde_json::json; use std::time::Duration; use tokio::time::timeout; - #[tokio::test] - async fn test_stdio_transport() { - let transport = StdioTransport { - params: StdioServerParams { - command: "tee".to_string(), // tee will echo back what it receives - args: vec![], - env: None, - }, - }; - - let (mut rx, tx) = transport.connect().await.unwrap(); - - // Create test messages - let request = JsonRpcMessage::Request(JsonRpcRequest { - jsonrpc: "2.0".to_string(), - id: Some(1), - method: "ping".to_string(), - params: None, - }); - - let response = JsonRpcMessage::Response(JsonRpcResponse { - jsonrpc: "2.0".to_string(), - id: Some(2), - result: Some(json!({})), - error: None, - }); - - // Send messages - tx.send(request.clone()).await.unwrap(); - tx.send(response.clone()).await.unwrap(); - - // Receive and verify messages - let mut read_messages = Vec::new(); - - // Use timeout to avoid hanging if messages aren't received - for _ in 0..2 { - match timeout(Duration::from_secs(1), rx.recv()).await { - Ok(Some(Ok(msg))) => read_messages.push(msg), - Ok(Some(Err(e))) => panic!("Received error: {}", e), - Ok(None) => break, - Err(_) => panic!("Timeout waiting for message"), - } - } - - assert_eq!(read_messages.len(), 2, "Expected 2 messages"); - assert_eq!(read_messages[0], request); - assert_eq!(read_messages[1], response); - } - #[tokio::test] async fn test_process_termination() { let transport = StdioTransport { diff --git a/crates/mcp-core/src/protocol.rs b/crates/mcp-core/src/protocol.rs index 09ca3e4caeab..e178c41177c3 100644 --- a/crates/mcp-core/src/protocol.rs +++ b/crates/mcp-core/src/protocol.rs @@ -75,7 +75,7 @@ impl TryFrom for JsonRpcMessage { error: raw.error.unwrap(), })); } - + // If it has a result field, it's a response if raw.result.is_some() { return Ok(JsonRpcMessage::Response(JsonRpcResponse { diff --git a/crates/mcp-macros/examples/calculator.rs b/crates/mcp-macros/examples/calculator.rs index d4cee5ca5ca0..74be963d64db 100644 --- a/crates/mcp-macros/examples/calculator.rs +++ b/crates/mcp-macros/examples/calculator.rs @@ -1,10 +1,10 @@ -use mcp_core::handler::{ToolHandler, ToolError, Result}; +use mcp_core::handler::{Result, ToolError, ToolHandler}; use mcp_macros::tool; #[tokio::main] async fn main() -> std::result::Result<(), Box> { // Create an instance of our tool - let calculator = Calculator::default(); + let calculator = Calculator; // Print tool information println!("Tool name: {}", calculator.name()); @@ -40,9 +40,7 @@ async fn calculator(x: i32, y: i32, operation: String) -> Result { "multiply" => Ok(x * y), "divide" => { if y == 0 { - Err(ToolError::ExecutionError( - "Division by zero".into(), - )) + Err(ToolError::ExecutionError("Division by zero".into())) } else { Ok(x / y) } diff --git a/crates/mcp-macros/src/lib.rs b/crates/mcp-macros/src/lib.rs index c0ef725fb826..d777d9705e95 100644 --- a/crates/mcp-macros/src/lib.rs +++ b/crates/mcp-macros/src/lib.rs @@ -1,3 +1,4 @@ +use convert_case::{Case, Casing}; use proc_macro::TokenStream; use quote::{format_ident, quote}; use std::collections::HashMap; @@ -5,7 +6,6 @@ use syn::{ parse::Parse, parse::ParseStream, parse_macro_input, punctuated::Punctuated, Expr, ExprLit, FnArg, ItemFn, Lit, Meta, Pat, PatType, Token, }; -use convert_case::{Case, Casing}; struct MacroArgs { name: Option, @@ -76,9 +76,7 @@ pub fn tool(args: TokenStream, input: TokenStream) -> TokenStream { let fn_name_str = fn_name.to_string(); // Generate PascalCase struct name from the function name - let struct_name = format_ident!("{}", { - fn_name_str.to_case(Case::Pascal) - }); + let struct_name = format_ident!("{}", { fn_name_str.to_case(Case::Pascal) }); // Use provided name or function name as default let tool_name = args.name.unwrap_or(fn_name_str); diff --git a/crates/mcp-server/src/error.rs b/crates/mcp-server/src/error.rs deleted file mode 100644 index 792394379367..000000000000 --- a/crates/mcp-server/src/error.rs +++ /dev/null @@ -1,43 +0,0 @@ -use thiserror::Error; - -pub type BoxError = Box; - -#[derive(Debug, Error)] -pub enum RouterError { - #[error("Method not found: {0}")] - MethodNotFound(String), - #[error("Unknown tool: {0}")] - UnknownTool(String), - #[error("Invalid parameters: {0}")] - InvalidParams(String), - #[error("Internal error: {0}")] - Internal(String), -} - -impl From for mcp_core::protocol::ErrorData { - fn from(err: RouterError) -> Self { - use mcp_core::protocol::*; - match err { - RouterError::MethodNotFound(_) => ErrorData { - code: METHOD_NOT_FOUND, - message: err.to_string(), - data: None, - }, - RouterError::UnknownTool(_) => ErrorData { - code: INVALID_REQUEST, - message: err.to_string(), - data: None, - }, - RouterError::InvalidParams(_) => ErrorData { - code: INVALID_PARAMS, - message: err.to_string(), - data: None, - }, - RouterError::Internal(_) => ErrorData { - code: INTERNAL_ERROR, - message: err.to_string(), - data: None, - }, - } - } -} \ No newline at end of file diff --git a/crates/mcp-server/src/errors.rs b/crates/mcp-server/src/errors.rs index a87cd7aff1f6..f49eab058993 100644 --- a/crates/mcp-server/src/errors.rs +++ b/crates/mcp-server/src/errors.rs @@ -6,13 +6,13 @@ pub type BoxError = Box; pub enum TransportError { #[error("IO error: {0}")] Io(#[from] std::io::Error), - + #[error("JSON serialization error: {0}")] Json(#[from] serde_json::Error), - + #[error("Invalid UTF-8 sequence: {0}")] Utf8(#[from] std::string::FromUtf8Error), - + #[error("Protocol error: {0}")] Protocol(String), @@ -24,10 +24,10 @@ pub enum TransportError { pub enum ServerError { #[error("Transport error: {0}")] Transport(#[from] TransportError), - + #[error("Service error: {0}")] Service(String), - + #[error("Internal error: {0}")] Internal(String), @@ -39,10 +39,10 @@ pub enum ServerError { pub enum RouterError { #[error("Method not found: {0}")] MethodNotFound(String), - + #[error("Invalid parameters: {0}")] InvalidParams(String), - + #[error("Internal error: {0}")] Internal(String), diff --git a/crates/mcp-server/src/lib.rs b/crates/mcp-server/src/lib.rs index 06a4bf77e08d..b054a98c59a4 100644 --- a/crates/mcp-server/src/lib.rs +++ b/crates/mcp-server/src/lib.rs @@ -3,14 +3,14 @@ use std::{ task::{Context, Poll}, }; -use futures::{Stream, Future}; -use mcp_core::protocol::{JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, JsonRpcError}; +use futures::{Future, Stream}; +use mcp_core::protocol::{JsonRpcError, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse}; use pin_project::pin_project; -use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt, AsyncBufReadExt, BufReader}; +use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader}; use tower_service::Service; mod errors; -pub use errors::{TransportError, ServerError, RouterError, BoxError}; +pub use errors::{BoxError, RouterError, ServerError, TransportError}; pub mod router; pub use router::Router; @@ -45,7 +45,7 @@ where let mut this = self.project(); let mut buf = Vec::new(); let mut reader = BufReader::new(&mut this.reader); - + let mut read_future = Box::pin(reader.read_until(b'\n', &mut buf)); match read_future.as_mut().poll(cx) { Poll::Ready(Ok(0)) => Poll::Ready(None), // EOF @@ -55,20 +55,24 @@ where Ok(s) => s, Err(e) => return Poll::Ready(Some(Err(TransportError::Utf8(e)))), }; - + // Parse JSON and validate message format match serde_json::from_str::(&line) { Ok(value) => { // Validate basic JSON-RPC structure if !value.is_object() { - return Poll::Ready(Some(Err(TransportError::InvalidMessage("Message must be a JSON object".into())))); + return Poll::Ready(Some(Err(TransportError::InvalidMessage( + "Message must be a JSON object".into(), + )))); } - + let obj = value.as_object().unwrap(); // Safe due to check above - + // Check jsonrpc version field if !obj.contains_key("jsonrpc") || obj["jsonrpc"] != "2.0" { - return Poll::Ready(Some(Err(TransportError::InvalidMessage("Missing or invalid jsonrpc version".into())))); + return Poll::Ready(Some(Err(TransportError::InvalidMessage( + "Missing or invalid jsonrpc version".into(), + )))); } tracing::info!( @@ -80,7 +84,7 @@ where Ok(msg) => Poll::Ready(Some(Ok(msg))), Err(e) => Poll::Ready(Some(Err(TransportError::Json(e)))), } - }, + } Err(e) => Poll::Ready(Some(Err(TransportError::Json(e)))), } } @@ -97,7 +101,9 @@ where { pub async fn write_message(&mut self, msg: JsonRpcMessage) -> Result<(), std::io::Error> { let json = serde_json::to_string(&msg)?; - Pin::new(&mut self.writer).write_all(json.as_bytes()).await?; + Pin::new(&mut self.writer) + .write_all(json.as_bytes()) + .await?; Pin::new(&mut self.writer).write_all(b"\n").await?; Pin::new(&mut self.writer).flush().await?; Ok(()) @@ -136,17 +142,17 @@ where match msg { JsonRpcMessage::Request(request) => { // Serialize request for logging - let id = request.id.clone(); + let id = request.id; let request_json = serde_json::to_string(&request) .unwrap_or_else(|_| "Failed to serialize request".to_string()); - + tracing::info!( request_id = ?id, method = ?request.method, json = %request_json, "Received request" ); - + // Process the request using our service let response = match service.call(request).await { Ok(resp) => resp, @@ -163,37 +169,44 @@ where data: None, }), } - }, + } }; - + // Serialize response for logging let response_json = serde_json::to_string(&response) .unwrap_or_else(|_| "Failed to serialize response".to_string()); - + tracing::info!( response_id = ?response.id, json = %response_json, "Sending response" ); // Send the response back - if let Err(e) = transport.write_message(JsonRpcMessage::Response(response)).await { + if let Err(e) = transport + .write_message(JsonRpcMessage::Response(response)) + .await + { return Err(ServerError::Transport(TransportError::Io(e))); } - }, - JsonRpcMessage::Response(_) | JsonRpcMessage::Notification(_) | JsonRpcMessage::Error(_) => { + } + JsonRpcMessage::Response(_) + | JsonRpcMessage::Notification(_) + | JsonRpcMessage::Error(_) => { // Ignore responses and notifications for now continue; } } - }, + } Err(e) => { // Convert transport error to JSON-RPC error response let error = match e { - TransportError::Json(_) | TransportError::InvalidMessage(_) => mcp_core::protocol::ErrorData { - code: mcp_core::protocol::PARSE_ERROR, - message: e.to_string(), - data: None, - }, + TransportError::Json(_) | TransportError::InvalidMessage(_) => { + mcp_core::protocol::ErrorData { + code: mcp_core::protocol::PARSE_ERROR, + message: e.to_string(), + data: None, + } + } TransportError::Protocol(_) => mcp_core::protocol::ErrorData { code: mcp_core::protocol::INVALID_REQUEST, message: e.to_string(), @@ -218,7 +231,7 @@ where } } } - + Ok(()) } } diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs index 07a89e0dd6d8..b8911b3690ca 100644 --- a/crates/mcp-server/src/main.rs +++ b/crates/mcp-server/src/main.rs @@ -1,18 +1,16 @@ use anyhow::Result; use mcp_core::handler::ResourceError; -use mcp_core::{ - handler::ToolError, - tool::Tool, - protocol::ServerCapabilities, - resource::Resource, -}; -use mcp_server::{Router, Server, ByteTransport}; -use mcp_server::router::{RouterService, CapabilitiesBuilder}; +use mcp_core::{handler::ToolError, protocol::ServerCapabilities, resource::Resource, tool::Tool}; +use mcp_server::router::{CapabilitiesBuilder, RouterService}; +use mcp_server::{ByteTransport, Router, Server}; use serde_json::Value; use std::{future::Future, pin::Pin, sync::Arc}; -use tokio::{io::{stdin, stdout}, sync::Mutex}; -use tracing_subscriber::{self, EnvFilter}; +use tokio::{ + io::{stdin, stdout}, + sync::Mutex, +}; use tracing_appender::rolling::{RollingFileAppender, Rotation}; +use tracing_subscriber::{self, EnvFilter}; // A simple counter service that demonstrates the Router trait #[derive(Clone)] @@ -47,9 +45,7 @@ impl CounterRouter { impl Router for CounterRouter { fn capabilities(&self) -> ServerCapabilities { - CapabilitiesBuilder::new() - .with_tools(true) - .build() + CapabilitiesBuilder::new().with_tools(true).build() } fn list_tools(&self) -> Vec { @@ -84,7 +80,11 @@ impl Router for CounterRouter { ] } - fn call_tool(&self, tool_name: &str, _arguments: Value) -> Pin> + Send + 'static>> { + fn call_tool( + &self, + tool_name: &str, + _arguments: Value, + ) -> Pin> + Send + 'static>> { let this = self.clone(); let tool_name = tool_name.to_string(); @@ -108,13 +108,12 @@ impl Router for CounterRouter { } fn list_resources(&self) -> Vec { - vec![ - Resource::new( - "memo://insights".to_string(), - Some("text/plain".to_string()), - Some("memo-resource".to_string()) - ).unwrap() - ] + vec![Resource::new( + "memo://insights", + Some("text/plain".to_string()), + Some("memo-resource".to_string()), + ) + .unwrap()] } fn read_resource( @@ -141,16 +140,11 @@ impl Router for CounterRouter { #[tokio::main] async fn main() -> Result<()> { // Set up file appender for logging - let file_appender = RollingFileAppender::new( - Rotation::DAILY, - "logs", - "mcp-server.log", - ); + let file_appender = RollingFileAppender::new(Rotation::DAILY, "logs", "mcp-server.log"); // Initialize the tracing subscriber with file and stdout logging tracing_subscriber::fmt() - .with_env_filter(EnvFilter::from_default_env() - .add_directive(tracing::Level::INFO.into())) + .with_env_filter(EnvFilter::from_default_env().add_directive(tracing::Level::INFO.into())) .with_writer(file_appender) .with_target(false) .with_thread_ids(true) @@ -162,41 +156,11 @@ async fn main() -> Result<()> { // Create an instance of our counter router let router = RouterService(CounterRouter::new()); - + // Create and run the server let server = Server::new(router); let transport = ByteTransport::new(stdin(), stdout()); - + tracing::info!("Server initialized and ready to handle requests"); Ok(server.run(transport).await?) } - -#[cfg(test)] -mod tests { - use super::*; - use tokio::test; - - #[test] - async fn test_counter_operations() { - let router = CounterRouter::new(); - - // Test increment - let result = router.call_tool("increment", Value::Null).await.unwrap(); - assert_eq!(result, Value::Number(1.into())); - - // Test get_value - let result = router.call_tool("get_value", Value::Null).await.unwrap(); - assert_eq!(result, Value::Number(1.into())); - - // Test decrement - let result = router.call_tool("decrement", Value::Null).await.unwrap(); - assert_eq!(result, Value::Number(0.into())); - } - - #[test] - async fn test_invalid_tool() { - let router = CounterRouter::new(); - let result = router.call_tool("invalid_tool", Value::Null).await; - assert!(result.is_err()); - } -} diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs index 4c2448843334..d2d8a9ddf679 100644 --- a/crates/mcp-server/src/router.rs +++ b/crates/mcp-server/src/router.rs @@ -2,19 +2,22 @@ use std::{ future::Future, pin::Pin, task::{Context, Poll}, - time::Duration, }; -use tokio::time::sleep; use mcp_core::{ - content::Content, handler::{ResourceError, ToolError}, protocol::{ - CallToolResult, Implementation, InitializeResult, JsonRpcRequest, JsonRpcResponse, ListResourcesResult, ListToolsResult, PromptsCapability, ReadResourceResult, ResourcesCapability, ServerCapabilities, ToolsCapability - }, ResourceContents + content::Content, + handler::{ResourceError, ToolError}, + protocol::{ + CallToolResult, Implementation, InitializeResult, JsonRpcRequest, JsonRpcResponse, + ListResourcesResult, ListToolsResult, PromptsCapability, ReadResourceResult, + ResourcesCapability, ServerCapabilities, ToolsCapability, + }, + ResourceContents, }; -use tower_service::Service; use serde_json::Value; +use tower_service::Service; -use crate::{RouterError, BoxError}; +use crate::{BoxError, RouterError}; /// Builder for configuring and constructing capabilities pub struct CapabilitiesBuilder { @@ -23,6 +26,12 @@ pub struct CapabilitiesBuilder { resources: Option, } +impl Default for CapabilitiesBuilder { + fn default() -> Self { + Self::new() + } +} + impl CapabilitiesBuilder { pub fn new() -> Self { Self { @@ -33,8 +42,7 @@ impl CapabilitiesBuilder { } /// Add multiple tools to the router - pub fn with_tools(mut self, list_changed: bool) -> Self - { + pub fn with_tools(mut self, list_changed: bool) -> Self { self.tools = Some(ToolsCapability { list_changed: Some(list_changed), }); @@ -72,9 +80,16 @@ impl CapabilitiesBuilder { pub trait Router: Send + Sync + 'static { fn capabilities(&self) -> ServerCapabilities; fn list_tools(&self) -> Vec; - fn call_tool(&self, tool_name: &str, arguments: Value) -> Pin> + Send + 'static>>; + fn call_tool( + &self, + tool_name: &str, + arguments: Value, + ) -> Pin> + Send + 'static>>; fn list_resources(&self) -> Vec; - fn read_resource(&self, uri: &str) -> Pin> + Send + 'static>>; + fn read_resource( + &self, + uri: &str, + ) -> Pin> + Send + 'static>>; // Helper method to create base response fn create_response(&self, id: Option) -> JsonRpcResponse { @@ -86,7 +101,10 @@ pub trait Router: Send + Sync + 'static { } } - fn handle_initialize(&self, req: JsonRpcRequest) -> impl Future> + Send { + fn handle_initialize( + &self, + req: JsonRpcRequest, + ) -> impl Future> + Send { async move { let result = InitializeResult { protocol_version: "2024-11-05".to_string(), @@ -98,37 +116,48 @@ pub trait Router: Send + Sync + 'static { }; let mut response = self.create_response(req.id); - response.result = Some(serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + response.result = + Some(serde_json::to_value(result).map_err(|e| { + RouterError::Internal(format!("JSON serialization error: {}", e)) + })?); Ok(response) } } - fn handle_tools_list(&self, req: JsonRpcRequest) -> impl Future> + Send { + fn handle_tools_list( + &self, + req: JsonRpcRequest, + ) -> impl Future> + Send { async move { let tools = self.list_tools(); let result = ListToolsResult { tools }; let mut response = self.create_response(req.id); - response.result = Some(serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + response.result = + Some(serde_json::to_value(result).map_err(|e| { + RouterError::Internal(format!("JSON serialization error: {}", e)) + })?); Ok(response) } } - fn handle_tools_call(&self, req: JsonRpcRequest) -> impl Future> + Send { + fn handle_tools_call( + &self, + req: JsonRpcRequest, + ) -> impl Future> + Send { async move { - let params = req.params.ok_or_else(|| RouterError::InvalidParams("Missing parameters".into()))?; + let params = req + .params + .ok_or_else(|| RouterError::InvalidParams("Missing parameters".into()))?; - let name = params.get("name") + let name = params + .get("name") .and_then(Value::as_str) .ok_or_else(|| RouterError::InvalidParams("Missing tool name".into()))?; - let arguments = params.get("arguments") - .cloned() - .unwrap_or(Value::Null); + let arguments = params.get("arguments").cloned().unwrap_or(Value::Null); let result = match self.call_tool(name, arguments).await { Ok(result) => CallToolResult { @@ -138,25 +167,32 @@ pub trait Router: Send + Sync + 'static { Err(err) => CallToolResult { content: vec![Content::text(err.to_string())], is_error: true, - } + }, }; let mut response = self.create_response(req.id); - response.result = Some(serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + response.result = + Some(serde_json::to_value(result).map_err(|e| { + RouterError::Internal(format!("JSON serialization error: {}", e)) + })?); Ok(response) } } - fn handle_resources_list(&self, req: JsonRpcRequest) -> impl Future> + Send { + fn handle_resources_list( + &self, + req: JsonRpcRequest, + ) -> impl Future> + Send { async move { let resources = self.list_resources(); let result = ListResourcesResult { resources }; let mut response = self.create_response(req.id); - response.result = Some(serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?); + response.result = + Some(serde_json::to_value(result).map_err(|e| { + RouterError::Internal(format!("JSON serialization error: {}", e)) + })?); Ok(response) } @@ -197,7 +233,6 @@ pub trait Router: Send + Sync + 'static { } } -// A wrapper type to implement the Service trait locally pub struct RouterService(pub T); impl Service for RouterService @@ -217,12 +252,6 @@ where Box::pin(async move { let result = match req.method.as_str() { - "slow" => { - sleep(Duration::from_secs(60)).await; - let mut response = this.create_response(req.id); - response.result = Some(Value::String("This should never be seen".to_string())); - Ok(response) - }, "initialize" => this.handle_initialize(req).await, "tools/list" => this.handle_tools_list(req).await, "tools/call" => this.handle_tools_call(req).await,