From 45dafaabe1e03201c15c88d08a33706b1a9e0b49 Mon Sep 17 00:00:00 2001 From: Collin Brittain Date: Thu, 6 Feb 2025 16:25:39 -0600 Subject: [PATCH] WIP - notification and request dispatch and handling --- Cargo.lock | 9 +++ Cargo.toml | 2 + crates/mcp-client-sse/Cargo.toml | 9 +++ crates/mcp-client-sse/src/client.rs | 14 ++++ crates/mcp-client-sse/src/lib.rs | 1 + crates/mcp-core/src/protocol.rs | 28 +++---- crates/mcp-types/build.rs | 2 +- crates/mcp-types/src/lib.rs | 4 +- crates/mcp-types/src/v2024_11_05/messages.rs | 80 +++++++++++++++++++ crates/mcp-types/src/v2024_11_05/mod.rs | 2 + .../{v2024_11_05.rs => v2024_11_05/types.rs} | 0 11 files changed, 135 insertions(+), 16 deletions(-) create mode 100644 crates/mcp-client-sse/Cargo.toml create mode 100644 crates/mcp-client-sse/src/client.rs create mode 100644 crates/mcp-client-sse/src/lib.rs create mode 100644 crates/mcp-types/src/v2024_11_05/messages.rs create mode 100644 crates/mcp-types/src/v2024_11_05/mod.rs rename crates/mcp-types/src/{v2024_11_05.rs => v2024_11_05/types.rs} (100%) diff --git a/Cargo.lock b/Cargo.lock index fd4689b..e163843 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -932,6 +932,15 @@ version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +[[package]] +name = "mcp-client-sse" +version = "0.1.0" +dependencies = [ + "mcp-core", + "mcp-transport-sse", + "mcp-types", +] + [[package]] name = "mcp-core" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 5b49780..4327ffd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,11 +1,13 @@ [workspace] members = [ + "crates/mcp-client-sse", "crates/mcp-core", "crates/mcp-transport-sse", "crates/mcp-transport-stdio", "crates/mcp-types", ] default-members = [ + "crates/mcp-client-sse", "crates/mcp-core", "crates/mcp-transport-sse", "crates/mcp-transport-stdio", diff --git a/crates/mcp-client-sse/Cargo.toml b/crates/mcp-client-sse/Cargo.toml new file mode 100644 index 0000000..28f438a --- /dev/null +++ b/crates/mcp-client-sse/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "mcp-client-sse" +version = "0.1.0" +edition = "2021" + +[dependencies] +mcp-core = { path = "../mcp-core" } +mcp-transport-sse = { path = "../mcp-transport-sse" } +mcp-types = { path = "../mcp-types" } diff --git a/crates/mcp-client-sse/src/client.rs b/crates/mcp-client-sse/src/client.rs new file mode 100644 index 0000000..5554874 --- /dev/null +++ b/crates/mcp-client-sse/src/client.rs @@ -0,0 +1,14 @@ +use mcp_core::protocol::Protocol; +use mcp_transport_sse::{client::SSEClientTransport, error::SSETransportError}; + +pub struct SSEClient { + protocol: Protocol, +} + +impl SSEClient { + pub fn new(transport: SSEClientTransport) -> Self { + Self { + protocol: Protocol::new(Box::new(transport)), + } + } +} diff --git a/crates/mcp-client-sse/src/lib.rs b/crates/mcp-client-sse/src/lib.rs new file mode 100644 index 0000000..b9babe5 --- /dev/null +++ b/crates/mcp-client-sse/src/lib.rs @@ -0,0 +1 @@ +pub mod client; diff --git a/crates/mcp-core/src/protocol.rs b/crates/mcp-core/src/protocol.rs index 95664e0..5346390 100644 --- a/crates/mcp-core/src/protocol.rs +++ b/crates/mcp-core/src/protocol.rs @@ -1,17 +1,17 @@ use std::collections::HashMap; use mcp_types::{ - JSONRPCError, JSONRPCNotification, JSONRPCNotificationParams, JSONRPCRequest, - JSONRPCRequestParams, JSONRPCRequestParamsMeta, JSONRPCResponse, Notification, Request, - RequestId, + ClientCapabilities, Implementation, InitializeRequest, InitializeRequestParams, JSONRPCError, + JSONRPCNotification, JSONRPCNotificationParams, JSONRPCRequest, JSONRPCRequestParams, + JSONRPCRequestParamsMeta, JSONRPCResponse, MessageSchema, Notification, Request, RequestId, + Result as MCPResult, LATEST_PROTOCOL_VERSION, }; use crate::transport::Transport; pub struct Protocol> + Send + Sync + 'static> { - notification_handlers: HashMap>, - request_handlers: - HashMap Result>>, + notification_handlers: HashMap Result<(), E>>>, + request_handlers: HashMap Result>>, #[cfg(not(feature = "uuid"))] request_id: i64, transport: Box>, @@ -31,20 +31,20 @@ where } } - pub fn register_notification_handler( + pub fn set_notification_handler( &mut self, - method: String, - handler: Box, + schema: S, + handler: Box Result<(), E>>, ) { - self.notification_handlers.insert(method, handler); + self.notification_handlers.insert(schema.method(), handler); } - pub fn register_request_handler( + pub fn set_request_handler( &mut self, - method: String, - handler: Box Result>, + schema: S, + handler: Box Result>, ) { - self.request_handlers.insert(method, handler); + self.request_handlers.insert(schema.method(), handler); } pub async fn send_request(&mut self, request: Request) -> Result<(), E> { diff --git a/crates/mcp-types/build.rs b/crates/mcp-types/build.rs index 8d83e0d..b6cc9c7 100644 --- a/crates/mcp-types/build.rs +++ b/crates/mcp-types/build.rs @@ -20,7 +20,7 @@ fn generate(version: &str) { let contents = syn::parse2::(type_space.to_stream()).unwrap(); let contents = prettyplease::unparse(&contents); - let file_name = format!("src/v{}.rs", version.replace("-", "_")); + let file_name = format!("src/v{}/types.rs", version.replace("-", "_")); let mut out_file = Path::new(&env::var("CARGO_MANIFEST_DIR").unwrap()).to_path_buf(); out_file.push(file_name); diff --git a/crates/mcp-types/src/lib.rs b/crates/mcp-types/src/lib.rs index 39a563a..88a5768 100644 --- a/crates/mcp-types/src/lib.rs +++ b/crates/mcp-types/src/lib.rs @@ -25,10 +25,12 @@ mod v2024_11_05; /// Re-export of the types module. #[cfg(feature = "2024_11_05")] -pub use v2024_11_05::{ +pub use v2024_11_05::types::{ JsonrpcError as JSONRPCError, JsonrpcErrorError as JSONRPCInnerError, JsonrpcMessage as JSONRPCMessage, JsonrpcNotification as JSONRPCNotification, JsonrpcNotificationParams as JSONRPCNotificationParams, JsonrpcRequest as JSONRPCRequest, JsonrpcRequestParams as JSONRPCRequestParams, JsonrpcRequestParamsMeta as JSONRPCRequestParamsMeta, JsonrpcResponse as JSONRPCResponse, *, }; + +pub use v2024_11_05::messages::*; diff --git a/crates/mcp-types/src/v2024_11_05/messages.rs b/crates/mcp-types/src/v2024_11_05/messages.rs new file mode 100644 index 0000000..6c7dc73 --- /dev/null +++ b/crates/mcp-types/src/v2024_11_05/messages.rs @@ -0,0 +1,80 @@ +use serde::de::DeserializeOwned; + +pub trait MessageSchema: DeserializeOwned { + fn method(&self) -> String; + fn params(&self) -> Option>; +} + +/// Macro for implementing the `MessageSchema` trait for a given type that has +/// a `method` field and a `params` field. +macro_rules! impl_message_schema { + ($type:ty) => { + impl MessageSchema for $type { + fn method(&self) -> String { + self.method.clone() + } + + fn params(&self) -> Option> { + serde_json::to_value(&self.params) + .ok() + .and_then(|v| v.as_object().cloned()) + } + } + }; +} + +// Implement the `MessageSchema` trait for all generated message types +use super::types::*; + +impl_message_schema!(CancelledNotification); +impl_message_schema!(InitializedNotification); +impl_message_schema!(LoggingMessageNotification); +impl_message_schema!(ProgressNotification); +impl_message_schema!(PromptListChangedNotification); +impl_message_schema!(ResourceListChangedNotification); +impl_message_schema!(ResourceUpdatedNotification); +impl_message_schema!(ToolListChangedNotification); + +// Implement the `MessageSchema` trait for all generated request types +impl_message_schema!(CallToolRequest); +impl_message_schema!(CompleteRequest); +impl_message_schema!(CreateMessageRequest); +impl_message_schema!(GetPromptRequest); +impl_message_schema!(InitializeRequest); +impl_message_schema!(ListPromptsRequest); +impl_message_schema!(ListResourcesRequest); +impl_message_schema!(ListResourceTemplatesRequest); +impl_message_schema!(ListRootsRequest); +impl_message_schema!(ListToolsRequest); +impl_message_schema!(PingRequest); +impl_message_schema!(ReadResourceRequest); +impl_message_schema!(SetLevelRequest); +impl_message_schema!(SubscribeRequest); +impl_message_schema!(UnsubscribeRequest); + +#[cfg(test)] +mod tests { + use crate::{InitializeRequest, InitializeRequestParams, LATEST_PROTOCOL_VERSION}; + + use super::*; + + #[test] + fn test_message_schema() { + let request = InitializeRequest { + method: "initialize".to_string(), + params: InitializeRequestParams { + capabilities: ClientCapabilities::default(), + client_info: Implementation { + name: "test".to_string(), + version: "1.0.0".to_string(), + }, + protocol_version: LATEST_PROTOCOL_VERSION.to_string(), + }, + }; + assert_eq!(request.method(), "initialize"); + assert_eq!( + request.params().unwrap().get("protocolVersion").unwrap(), + LATEST_PROTOCOL_VERSION + ); + } +} diff --git a/crates/mcp-types/src/v2024_11_05/mod.rs b/crates/mcp-types/src/v2024_11_05/mod.rs new file mode 100644 index 0000000..3c75990 --- /dev/null +++ b/crates/mcp-types/src/v2024_11_05/mod.rs @@ -0,0 +1,2 @@ +pub mod messages; +pub mod types; diff --git a/crates/mcp-types/src/v2024_11_05.rs b/crates/mcp-types/src/v2024_11_05/types.rs similarity index 100% rename from crates/mcp-types/src/v2024_11_05.rs rename to crates/mcp-types/src/v2024_11_05/types.rs