diff --git a/Cargo.lock b/Cargo.lock index d169124e83..6e702a92f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2775,7 +2775,6 @@ dependencies = [ "jsonschema", "mcp-client", "mcp-core", - "mcp-server", "nix 0.30.1", "once_cell", "rand 0.8.5", @@ -2829,7 +2828,6 @@ dependencies = [ "lopdf", "lru", "mcp-core", - "mcp-server", "mpatch", "oauth2", "once_cell", @@ -2887,7 +2885,6 @@ dependencies = [ "goose-mcp", "http 1.2.0", "mcp-core", - "mcp-server", "reqwest 0.12.12", "rmcp", "schemars", @@ -4110,26 +4107,6 @@ dependencies = [ "utoipa", ] -[[package]] -name = "mcp-server" -version = "0.1.0" -dependencies = [ - "anyhow", - "futures", - "mcp-core", - "pin-project", - "rmcp", - "serde", - "serde_json", - "thiserror 1.0.69", - "tokio", - "tower 0.4.13", - "tower-service", - "tracing", - "tracing-appender", - "tracing-subscriber", -] - [[package]] name = "md-5" version = "0.10.6" diff --git a/crates/goose-cli/Cargo.toml b/crates/goose-cli/Cargo.toml index e7e3e0e20a..fb16ced65e 100644 --- a/crates/goose-cli/Cargo.toml +++ b/crates/goose-cli/Cargo.toml @@ -19,7 +19,6 @@ goose = { path = "../goose" } goose-bench = { path = "../goose-bench" } goose-mcp = { path = "../goose-mcp" } mcp-client = { path = "../mcp-client" } -mcp-server = { path = "../mcp-server" } mcp-core = { path = "../mcp-core" } rmcp = { workspace = true } agent-client-protocol = "0.4.0" diff --git a/crates/goose-cli/src/logging.rs b/crates/goose-cli/src/logging.rs index 99a39a0c2f..974c0a6c32 100644 --- a/crates/goose-cli/src/logging.rs +++ b/crates/goose-cli/src/logging.rs @@ -119,8 +119,6 @@ fn setup_logging_internal( let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| { // Set default levels for different modules EnvFilter::new("") - // Set mcp-server module to DEBUG - .add_directive("mcp_server=debug".parse().unwrap()) // Set mcp-client to DEBUG .add_directive("mcp_client=debug".parse().unwrap()) // Set goose module to DEBUG diff --git a/crates/goose-mcp/Cargo.toml b/crates/goose-mcp/Cargo.toml index a728ad2be8..e187832901 100644 --- a/crates/goose-mcp/Cargo.toml +++ b/crates/goose-mcp/Cargo.toml @@ -13,7 +13,6 @@ workspace = true [dependencies] goose = { path = "../goose" } mcp-core = { path = "../mcp-core" } -mcp-server = { path = "../mcp-server" } rmcp = { version = "0.6.0", features = ["server", "client", "transport-io", "macros"] } anyhow = "1.0.94" tokio = { version = "1", features = ["full"] } diff --git a/crates/goose-server/Cargo.toml b/crates/goose-server/Cargo.toml index de4dbba76c..4b84ecd8c3 100644 --- a/crates/goose-server/Cargo.toml +++ b/crates/goose-server/Cargo.toml @@ -14,7 +14,6 @@ workspace = true goose = { path = "../goose" } mcp-core = { path = "../mcp-core" } goose-mcp = { path = "../goose-mcp" } -mcp-server = { path = "../mcp-server" } rmcp = { workspace = true } schemars = "1.0" axum = { version = "0.8.1", features = ["ws", "macros"] } diff --git a/crates/goose-server/src/logging.rs b/crates/goose-server/src/logging.rs index 026176538d..c4ffd92870 100644 --- a/crates/goose-server/src/logging.rs +++ b/crates/goose-server/src/logging.rs @@ -57,8 +57,6 @@ pub fn setup_logging(name: Option<&str>) -> Result<()> { let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| { // Set default levels for different modules EnvFilter::new("") - // Set mcp-server module to DEBUG - .add_directive("mcp_server=debug".parse().unwrap()) // Set mcp-client to DEBUG .add_directive("mcp_client=debug".parse().unwrap()) // Set goose module to DEBUG diff --git a/crates/mcp-server/Cargo.toml b/crates/mcp-server/Cargo.toml deleted file mode 100644 index 3829837274..0000000000 --- a/crates/mcp-server/Cargo.toml +++ /dev/null @@ -1,23 +0,0 @@ -[package] -name = "mcp-server" -version = "0.1.0" -edition = "2021" - -[lints] -workspace = true - -[dependencies] -anyhow = "1.0.94" -thiserror = "1.0" -mcp-core = { path = "../mcp-core" } -rmcp = { workspace = true } -serde = { version = "1.0.216", features = ["derive"] } -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" -tracing-subscriber = { version = "0.3", features = ["env-filter"] } -tracing-appender = "0.2" diff --git a/crates/mcp-server/README.md b/crates/mcp-server/README.md deleted file mode 100644 index 1e4f061765..0000000000 --- a/crates/mcp-server/README.md +++ /dev/null @@ -1,7 +0,0 @@ -### 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 deleted file mode 100644 index d4cd59e363..0000000000 --- a/crates/mcp-server/src/errors.rs +++ /dev/null @@ -1,106 +0,0 @@ -use std::borrow::Cow; - -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), - - #[error("Tool not found: {0}")] - ToolNotFound(String), - - #[error("Resource not found: {0}")] - ResourceNotFound(String), - - #[error("Not found: {0}")] - PromptNotFound(String), -} - -impl From for rmcp::model::ErrorData { - fn from(err: RouterError) -> Self { - use rmcp::model::*; - match err { - RouterError::MethodNotFound(msg) => ErrorData { - code: ErrorCode::METHOD_NOT_FOUND, - message: Cow::from(msg), - data: None, - }, - RouterError::InvalidParams(msg) => ErrorData { - code: ErrorCode::INVALID_PARAMS, - message: Cow::from(msg), - data: None, - }, - RouterError::Internal(msg) => ErrorData { - code: ErrorCode::INTERNAL_ERROR, - message: Cow::from(msg), - data: None, - }, - RouterError::ToolNotFound(msg) => ErrorData { - code: ErrorCode::INVALID_REQUEST, - message: Cow::from(msg), - data: None, - }, - RouterError::ResourceNotFound(msg) => ErrorData { - code: ErrorCode::INVALID_REQUEST, - message: Cow::from(msg), - data: None, - }, - RouterError::PromptNotFound(msg) => ErrorData { - code: ErrorCode::INVALID_REQUEST, - message: Cow::from(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/lib.rs b/crates/mcp-server/src/lib.rs deleted file mode 100644 index 4b4fd89108..0000000000 --- a/crates/mcp-server/src/lib.rs +++ /dev/null @@ -1,289 +0,0 @@ -use std::{ - pin::Pin, - task::{Context, Poll}, -}; - -use futures::{Future, Stream}; -use pin_project::pin_project; -use rmcp::model::{ - ErrorData, JsonRpcError, JsonRpcMessage, JsonRpcResponse, JsonRpcVersion2_0, RequestId, -}; -use router::McpRequest; -use tokio::{ - io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader}, - sync::mpsc, -}; -use tower_service::Service; - -mod errors; -pub use errors::{BoxError, RouterError, ServerError, TransportError}; - -pub mod router; -pub use router::Router; - -/// A transport layer that handles JSON-RPC messages over byte -#[pin_project] -pub struct ByteTransport { - // Reader is a BufReader on the underlying stream (stdin or similar) buffering - // the underlying data across poll calls, we clear one line (\n) during each - // iteration of poll_next from this buffer - #[pin] - reader: BufReader, - #[pin] - writer: W, -} - -impl ByteTransport -where - R: AsyncRead, - W: AsyncWrite, -{ - pub fn new(reader: R, writer: W) -> Self { - Self { - // Default BufReader capacity is 8 * 1024, increase this to 2MB to the file size limit - // allows the buffer to have the capacity to read very large calls - reader: BufReader::with_capacity(2 * 1024 * 1024, 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 = this.reader.as_mut(); - 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)))), - }; - // Log incoming message here before serde conversion to - // track incomplete chunks which are not valid JSON - tracing::info!(json = %line, "incoming message"); - - // 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 + Send + 'static, - W: AsyncWrite + Unpin + Send + 'static, - { - 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) => { - let request_json = serde_json::to_string(&request) - .unwrap_or_else(|_| "Failed to serialize request".to_string()); - - tracing::info!( - method = ?request.request.method, - json = %request_json, - "Received request" - ); - - // Process the request using our service - let (notify_tx, mut notify_rx) = mpsc::channel(256); - let mcp_request = McpRequest { - request, - notifier: notify_tx, - }; - - let transport_fut = tokio::spawn(async move { - while let Some(notification) = notify_rx.recv().await { - if transport.write_message(notification).await.is_err() { - break; - } - } - transport - }); - - let response = match service.call(mcp_request).await { - Ok(resp) => resp, - Err(e) => { - let error_msg = e.into().to_string(); - tracing::error!(error = %error_msg, "Request processing failed"); - - // Return an error response instead of a regular response - return Err(ServerError::Transport(TransportError::Protocol( - error_msg, - ))); - } - }; - - transport = match transport_fut.await { - Ok(transport) => transport, - Err(e) => { - tracing::error!(error = %e, "Failed to spawn transport task"); - return Err(ServerError::Transport(TransportError::Io( - e.into(), - ))); - } - }; - - // 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))); - } - } - JsonRpcMessage::Response(_) - | JsonRpcMessage::Notification(_) - | JsonRpcMessage::Error(_) => { - // Ignore responses, notifications, batch messages and error messages for now - continue; - } - } - } - Err(e) => { - // Convert transport error to JSON-RPC error response - let error_data = match e { - TransportError::Json(_) | TransportError::InvalidMessage(_) => ErrorData { - code: rmcp::model::ErrorCode::PARSE_ERROR, - message: e.to_string().into(), - data: None, - }, - TransportError::Protocol(_) => ErrorData { - code: rmcp::model::ErrorCode::INVALID_REQUEST, - message: e.to_string().into(), - data: None, - }, - _ => ErrorData { - code: rmcp::model::ErrorCode::INTERNAL_ERROR, - message: e.to_string().into(), - data: None, - }, - }; - - let error_response = JsonRpcMessage::Error(JsonRpcError { - jsonrpc: JsonRpcVersion2_0, - id: RequestId::Number(0), // Use a default ID for transport errors - error: error_data, - }); - - if let Err(e) = transport.write_message(error_response).await { - return Err(ServerError::Transport(TransportError::Io(e))); - } - } - } - } - - Ok(()) - } -} - -// Define a specific service implementation that we need for any -// Any router implements this -pub trait BoundedService: - Service< - McpRequest, - Response = JsonRpcResponse, - Error = BoxError, - Future = Pin> + Send>>, - > + Send - + 'static -{ -} - -// Implement it for any type that meets the bounds -impl BoundedService for T where - T: Service< - McpRequest, - Response = JsonRpcResponse, - Error = BoxError, - Future = Pin> + Send>>, - > + Send - + 'static -{ -} diff --git a/crates/mcp-server/src/main.rs b/crates/mcp-server/src/main.rs deleted file mode 100644 index dd2cdfcc81..0000000000 --- a/crates/mcp-server/src/main.rs +++ /dev/null @@ -1,248 +0,0 @@ -use anyhow::Result; -use mcp_core::handler::{PromptError, ResourceError}; -use mcp_core::protocol::ServerCapabilities; -use mcp_server::router::{CapabilitiesBuilder, RouterService}; -use mcp_server::{ByteTransport, Router, Server}; -use rmcp::model::{ - Content, ErrorCode, ErrorData, JsonRpcMessage, Prompt, PromptArgument, RawResource, Resource, - Tool, ToolAnnotations, -}; -use rmcp::object; -use serde_json::Value; -use std::borrow::Cow; -use std::{future::Future, pin::Pin, sync::Arc}; -use tokio::sync::mpsc; -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)] -struct CounterRouter { - counter: Arc>, -} - -impl CounterRouter { - fn new() -> Self { - Self { - counter: Arc::new(Mutex::new(0)), - } - } - - 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) - } - - fn _create_resource_text(&self, uri: &str, name: &str) -> Resource { - Resource::new(RawResource::new(uri, name), None) - } -} - -impl Router for CounterRouter { - fn name(&self) -> String { - "counter".to_string() - } - - fn instructions(&self) -> String { - "This server provides a counter tool that can increment and decrement values. The counter starts at 0 and can be modified using the 'increment' and 'decrement' tools. Use 'get_value' to check the current count.".to_string() - } - - fn capabilities(&self) -> ServerCapabilities { - CapabilitiesBuilder::new() - .with_tools(false) - .with_resources(false, false) - .with_prompts(false) - .build() - } - - fn list_tools(&self) -> Vec { - vec![ - Tool::new( - "increment".to_string(), - "Increment the counter by 1".to_string(), - object!({ - "type": "object", - "properties": {}, - "required": [] - }), - ) - .annotate(ToolAnnotations { - title: Some("Increment Tool".to_string()), - read_only_hint: Some(false), - destructive_hint: Some(false), - idempotent_hint: Some(false), - open_world_hint: Some(false), - }), - Tool::new( - "decrement".to_string(), - "Decrement the counter by 1".to_string(), - object!({ - "type": "object", - "properties": {}, - "required": [] - }), - ) - .annotate(ToolAnnotations { - title: Some("Decrement Tool".to_string()), - read_only_hint: Some(false), - destructive_hint: Some(false), - idempotent_hint: Some(false), - open_world_hint: Some(false), - }), - Tool::new( - "get_value".to_string(), - "Get the current counter value".to_string(), - object!({ - "type": "object", - "properties": {}, - "required": [] - }), - ) - .annotate(ToolAnnotations { - title: Some("Get Value Tool".to_string()), - read_only_hint: Some(true), - destructive_hint: Some(false), - idempotent_hint: Some(false), - open_world_hint: Some(false), - }), - ] - } - - fn call_tool( - &self, - tool_name: &str, - _arguments: Value, - _notifier: mpsc::Sender, - ) -> Pin, ErrorData>> + 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(vec![Content::text(value.to_string())]) - } - "decrement" => { - let value = this.decrement().await?; - Ok(vec![Content::text(value.to_string())]) - } - "get_value" => { - let value = this.get_value().await?; - Ok(vec![Content::text(value.to_string())]) - } - _ => Err(ErrorData { - code: ErrorCode::INVALID_REQUEST, - message: Cow::from(format!("Tool {} not found", tool_name)), - data: None, - }), - } - }) - } - - fn list_resources(&self) -> Vec { - vec![ - self._create_resource_text("str:////Users/to/some/path/", "cwd"), - self._create_resource_text("memo://insights", "memo-name"), - ] - } - - fn read_resource( - &self, - uri: &str, - ) -> Pin> + Send + 'static>> { - let uri = uri.to_string(); - Box::pin(async move { - match uri.as_str() { - "str:////Users/to/some/path/" => { - let cwd = "/Users/to/some/path/"; - Ok(cwd.to_string()) - } - "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 - ))), - } - }) - } - - fn list_prompts(&self) -> Vec { - vec![Prompt::new( - "example_prompt", - Some("This is an example prompt that takes one required argument, message"), - Some(vec![PromptArgument { - name: "message".to_string(), - description: Some("A message to put in the prompt".to_string()), - required: Some(true), - }]), - )] - } - - fn get_prompt( - &self, - prompt_name: &str, - ) -> Pin> + Send + 'static>> { - let prompt_name = prompt_name.to_string(); - Box::pin(async move { - match prompt_name.as_str() { - "example_prompt" => { - let prompt = "This is an example prompt with your message here: '{message}'"; - Ok(prompt.to_string()) - } - _ => Err(PromptError::NotFound(format!( - "Prompt {} not found", - prompt_name - ))), - } - }) - } -} - -#[tokio::main] -async fn main() -> Result<()> { - // Set up file appender for logging - 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"); - - // 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?) -} diff --git a/crates/mcp-server/src/router.rs b/crates/mcp-server/src/router.rs deleted file mode 100644 index 02dd318f7a..0000000000 --- a/crates/mcp-server/src/router.rs +++ /dev/null @@ -1,425 +0,0 @@ -use std::{ - future::Future, - pin::Pin, - task::{Context, Poll}, -}; - -type PromptFuture = Pin> + Send + 'static>>; -use mcp_core::{ - handler::{PromptError, ResourceError}, - protocol::{ - CallToolResult, Implementation, InitializeResult, ListPromptsResult, ListResourcesResult, - ListToolsResult, PromptsCapability, ReadResourceResult, ResourcesCapability, - ServerCapabilities, ToolsCapability, - }, -}; -use rmcp::model::{ - Content, ErrorData, GetPromptResult, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, - JsonRpcVersion2_0, Prompt, PromptMessage, PromptMessageRole, RequestId, Resource, - ResourceContents, -}; -use serde_json::Value; -use tokio::sync::mpsc; -use tower_service::Service; - -use crate::{BoxError, RouterError}; - -/// Builder for configuring and constructing capabilities -pub struct CapabilitiesBuilder { - tools: Option, - prompts: Option, - resources: Option, -} - -impl Default for CapabilitiesBuilder { - fn default() -> Self { - Self::new() - } -} - -impl CapabilitiesBuilder { - pub fn new() -> Self { - Self { - tools: None, - prompts: None, - resources: None, - } - } - - /// Add multiple tools to the router - pub fn with_tools(mut self, list_changed: bool) -> Self { - self.tools = Some(ToolsCapability { - list_changed: Some(list_changed), - }); - 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) -> ServerCapabilities { - // Create capabilities based on what's configured - ServerCapabilities { - tools: self.tools, - prompts: self.prompts, - resources: self.resources, - } - } -} - -pub trait Router: Send + Sync + 'static { - fn name(&self) -> String; - // in the protocol, instructions are optional but we make it required - fn instructions(&self) -> String; - fn capabilities(&self) -> ServerCapabilities; - fn list_tools(&self) -> Vec; - fn call_tool( - &self, - tool_name: &str, - arguments: Value, - notifier: mpsc::Sender, - ) -> Pin, ErrorData>> + Send + 'static>>; - fn list_resources(&self) -> Vec; - fn read_resource( - &self, - uri: &str, - ) -> Pin> + Send + 'static>>; - fn list_prompts(&self) -> Vec; - fn get_prompt(&self, prompt_name: &str) -> PromptFuture; - - // Helper method to create base response - fn create_response(&self, id: RequestId) -> JsonRpcResponse { - JsonRpcResponse { - jsonrpc: JsonRpcVersion2_0, - id, - result: serde_json::Map::new(), - } - } - - // Helper method to set result on response - fn set_result( - &self, - response: &mut JsonRpcResponse, - result: T, - ) -> Result<(), RouterError> { - let value = serde_json::to_value(result) - .map_err(|e| RouterError::Internal(format!("JSON serialization error: {}", e)))?; - - if let Some(obj) = value.as_object() { - response.result = obj.clone(); - } else { - return Err(RouterError::Internal("Result must be a JSON object".into())); - } - - Ok(()) - } - - fn handle_initialize( - &self, - req: JsonRpcRequest, - ) -> impl Future> + Send { - async move { - let result = InitializeResult { - protocol_version: "2025-03-26".to_string(), - capabilities: self.capabilities().clone(), - server_info: Implementation { - name: self.name(), - version: env!("CARGO_PKG_VERSION").to_string(), - }, - instructions: Some(self.instructions()), - }; - - let mut response = self.create_response(req.id); - self.set_result(&mut response, result)?; - Ok(response) - } - } - - fn handle_tools_list( - &self, - req: JsonRpcRequest, - ) -> impl Future> + Send { - async move { - let tools = self.list_tools(); - - let result = ListToolsResult { - tools, - next_cursor: None, - }; - let mut response = self.create_response(req.id); - self.set_result(&mut response, result)?; - Ok(response) - } - } - - fn handle_tools_call( - &self, - req: JsonRpcRequest, - notifier: mpsc::Sender, - ) -> impl Future> + Send { - async move { - let params = &req.request.params; - - 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 result = match self.call_tool(name, arguments, notifier).await { - Ok(result) => CallToolResult { - content: result, - is_error: None, - }, - Err(err) => CallToolResult { - content: vec![Content::text(err.to_string())], - is_error: Some(true), - }, - }; - - let mut response = self.create_response(req.id); - self.set_result(&mut response, result)?; - Ok(response) - } - } - - fn handle_resources_list( - &self, - req: JsonRpcRequest, - ) -> impl Future> + Send { - async move { - let resources = self.list_resources(); - - let result = ListResourcesResult { - resources, - next_cursor: None, - }; - let mut response = self.create_response(req.id); - self.set_result(&mut response, result)?; - Ok(response) - } - } - - fn handle_resources_read( - &self, - req: JsonRpcRequest, - ) -> impl Future> + Send { - async move { - let params = &req.request.params; - - 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, - meta: None, - }], - }; - - let mut response = self.create_response(req.id); - self.set_result(&mut response, result)?; - Ok(response) - } - } - - fn handle_prompts_list( - &self, - req: JsonRpcRequest, - ) -> impl Future> + Send { - async move { - let prompts = self.list_prompts(); - - let result = ListPromptsResult { prompts }; - - let mut response = self.create_response(req.id); - self.set_result(&mut response, result)?; - Ok(response) - } - } - - fn handle_prompts_get( - &self, - req: JsonRpcRequest, - ) -> impl Future> + Send { - async move { - // Validate and extract parameters - let params = &req.request.params; - - // Extract "name" field - let prompt_name = params - .get("name") - .and_then(Value::as_str) - .ok_or_else(|| RouterError::InvalidParams("Missing prompt name".into()))?; - - // Extract "arguments" field - let arguments = params - .get("arguments") - .and_then(Value::as_object) - .ok_or_else(|| RouterError::InvalidParams("Missing arguments object".into()))?; - - // Fetch the prompt definition first - let prompt = self - .list_prompts() - .into_iter() - .find(|p| p.name == prompt_name) - .ok_or_else(|| { - RouterError::PromptNotFound(format!("Prompt '{}' not found", prompt_name)) - })?; - - // Validate required arguments - if let Some(args) = &prompt.arguments { - for arg in args { - if arg.required.is_some() - && arg.required.unwrap() - && (!arguments.contains_key(&arg.name) - || arguments - .get(&arg.name) - .and_then(Value::as_str) - .is_none_or(str::is_empty)) - { - return Err(RouterError::InvalidParams(format!( - "Missing required argument: '{}'", - arg.name - ))); - } - } - } - - // Now get the prompt content - let description = self - .get_prompt(prompt_name) - .await - .map_err(|e| RouterError::Internal(e.to_string()))?; - - // Validate prompt arguments for potential security issues from user text input - // Checks: - // - Prompt must be less than 10000 total characters - // - Argument keys must be less than 1000 characters - // - Argument values must be less than 1000 characters - // - Dangerous patterns, eg "../", "//", "\\\\", "