From 1e255287cc9eb6770bd345a05059bd7431611d21 Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Fri, 14 Mar 2025 01:39:41 -0700 Subject: [PATCH 01/10] unused and format --- Cargo.lock | 126 ++++++++++++++++++++++------------------- Cargo.toml | 18 +++--- src/client/builder.rs | 10 ++-- src/client/mod.rs | 48 ++++++++-------- src/client/test.rs | 25 ++++---- src/server/mod.rs | 8 ++- src/transport/mod.rs | 12 +--- src/transport/stdio.rs | 4 +- 8 files changed, 125 insertions(+), 126 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 18f48d4..c1eccd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -41,9 +41,9 @@ dependencies = [ [[package]] name = "async-trait" -version = "0.1.85" +version = "0.1.87" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f934833b4b7233644e5848f235df3f57ed8c80f1528a26c3dfa13d2147fa056" +checksum = "d556ec1359574147ec0c4fc5eb525f3f23263a592b1a9c07e0a75b427de55c97" dependencies = [ "proc-macro2", "quote", @@ -73,15 +73,15 @@ dependencies = [ [[package]] name = "bitflags" -version = "2.6.0" +version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b048fb63fd8b5923fc5aa7b340d8e156aec7ec02f0c78fa8a6ddc2613f6f71de" +checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd" [[package]] name = "bytes" -version = "1.9.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "325918d6fe32f23b19878fe4b34794ae41fc19ddbe53b10571a4874d44ffd39b" +checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" [[package]] name = "cfg-if" @@ -96,7 +96,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -214,21 +214,21 @@ checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" [[package]] name = "itoa" -version = "1.0.14" +version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d75a2a4b1b190afb6f5425f10f6a8f959d2ea0b9c2b1d79553551850539e4674" +checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" [[package]] name = "libc" -version = "0.2.167" +version = "0.2.171" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09d6582e104315a817dff97f75133544b2e094ee22447d2acf4a74e189ba06fc" +checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" [[package]] name = "linux-raw-sys" -version = "0.4.15" +version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" +checksum = "6db9c683daf087dc577b7506e9695b3d556a9f3849903fa28186283afd6809e9" [[package]] name = "lock_api" @@ -263,9 +263,9 @@ checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3" [[package]] name = "miniz_oxide" -version = "0.8.0" +version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2d80299ef12ff69b16a84bb182e3b9df68b5a91574d3d4fa6e41b65deec4df1" +checksum = "8e3e04debbb59698c15bacbb6d93584a8c0ca9cc3213cb423d31f760d8843ce5" dependencies = [ "adler2", ] @@ -278,23 +278,23 @@ checksum = "2886843bf800fba2e3377cff24abf6379b4c4d5c6681eaf9ea5b0d15090450bd" dependencies = [ "libc", "wasi 0.11.0+wasi-snapshot-preview1", - "windows-sys", + "windows-sys 0.52.0", ] [[package]] name = "object" -version = "0.36.5" +version = "0.36.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e" +checksum = "62948e14d923ea95ea2c7c86c71013138b66525b86bdc08d2dcc262bdb497b87" dependencies = [ "memchr", ] [[package]] name = "once_cell" -version = "1.20.2" +version = "1.21.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" +checksum = "d75b0bedcc4fe52caa0e03d9f1151a323e4aa5e2d78ba3580400cd3c9e2bc4bc" [[package]] name = "parking_lot" @@ -321,9 +321,9 @@ dependencies = [ [[package]] name = "pin-project-lite" -version = "0.2.15" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "915a1e146535de9163f3987b8944ed8cf49a18bb0056bcebcdcece385cece4ff" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" [[package]] name = "pin-utils" @@ -333,27 +333,27 @@ checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] name = "proc-macro2" -version = "1.0.92" +version = "1.0.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37d3544b3f2748c54e147655edb5025752e2303145b5aefb3c3ea2c78b973bb0" +checksum = "a31971752e70b8b2686d7e46ec17fb38dad4051d94024c88df49b667caea9c84" dependencies = [ "unicode-ident", ] [[package]] name = "quote" -version = "1.0.37" +version = "1.0.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5b9d34b8991d19d98081b46eacdd8eb58c6f2b201139f7c5f643cc155a633af" +checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" dependencies = [ "proc-macro2", ] [[package]] name = "redox_syscall" -version = "0.5.7" +version = "0.5.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b6dfecf2c74bce2466cabf93f6664d6998a69eb21e39f4207930065b27b771f" +checksum = "0b8c0c260b63a8219631167be35e6a988e9554dbd323f8bd08439c8ed1302bd1" dependencies = [ "bitflags", ] @@ -366,22 +366,22 @@ checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" [[package]] name = "rustix" -version = "0.38.44" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" +checksum = "f7178faa4b75a30e269c71e61c353ce2748cf3d76f0c44c393f4e60abf49b825" dependencies = [ "bitflags", "errno", "libc", "linux-raw-sys", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] name = "ryu" -version = "1.0.18" +version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3cb5ba0dc43242ce17de99c180e96db90b235b8a9fdc9543c96d2209116bd9f" +checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" [[package]] name = "scopeguard" @@ -391,18 +391,18 @@ checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] name = "serde" -version = "1.0.216" +version = "1.0.219" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b9781016e935a97e8beecf0c933758c97a5520d32930e460142b4cd80c6338e" +checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.216" +version = "1.0.219" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46f859dbbf73865c6627ed570e78961cd3ac92407a2d117204c49232485da55e" +checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00" dependencies = [ "proc-macro2", "quote", @@ -411,9 +411,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.137" +version = "1.0.140" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "930cfb6e6abf99298aaad7d29abbef7a9999a9a8806a40088f55f0dcec03146b" +checksum = "20068b6e96dc6c9bd23e01df8827e6c7e1f2fddd43c21810382803c136b99373" dependencies = [ "itoa", "memchr", @@ -441,9 +441,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.13.2" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd" [[package]] name = "socket2" @@ -452,14 +452,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c970269d99b64e60ec3bd6ad27270092a5394c4e309314b18ae3fe575695fbe8" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.52.0", ] [[package]] name = "syn" -version = "2.0.90" +version = "2.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" +checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0" dependencies = [ "proc-macro2", "quote", @@ -468,32 +468,31 @@ dependencies = [ [[package]] name = "tempfile" -version = "3.17.1" +version = "3.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22e5a0acb1f3f55f65cc4a866c361b2fb2a0ff6366785ae6fbb5f85df07ba230" +checksum = "488960f40a3fd53d72c2a29a58722561dee8afdd175bd88e3db4677d7b2ba600" dependencies = [ - "cfg-if", "fastrand", "getrandom", "once_cell", "rustix", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] name = "thiserror" -version = "2.0.9" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f072643fd0190df67a8bab670c20ef5d8737177d6ac6b2e9a236cb096206b2cc" +checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.9" +version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b50fa271071aae2e6ee85f842e2e28ba8cd2c5fb67f11fcb1fd70b276f9e7d4" +checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" dependencies = [ "proc-macro2", "quote", @@ -502,9 +501,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.42.0" +version = "1.44.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cec9b21b0450273377fc97bd4c33a8acffc8c996c987a7c5b319a0083707551" +checksum = "f382da615b842244d4b8738c82ed1275e6c5dd90c459a30941cd07080b06c91a" dependencies = [ "backtrace", "bytes", @@ -515,14 +514,14 @@ dependencies = [ "signal-hook-registry", "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.52.0", ] [[package]] name = "tokio-macros" -version = "2.4.0" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" +checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" dependencies = [ "proc-macro2", "quote", @@ -562,9 +561,9 @@ dependencies = [ [[package]] name = "unicode-ident" -version = "1.0.14" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" +checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" [[package]] name = "wasi" @@ -590,6 +589,15 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets", +] + [[package]] name = "windows-targets" version = "0.52.6" diff --git a/Cargo.toml b/Cargo.toml index 609052c..0c51566 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,18 +4,18 @@ version = "0.1.7" edition = "2021" description = "Rust client SDK for the Model Context Protocol (MCP)" license = "MIT OR Apache-2.0" -repository = "https://github.com/darinkishore/mcp_client_rust" -documentation = "https://github.com/darinkishore/mcp_client_rust" +repository = "https://github.com/gitarcode/mcp_client_rust" +documentation = "https://github.com/gitarcode/mcp_client_rust" keywords = ["mcp", "client", "protocol", "machine-learning", "ai"] categories = ["api-bindings", "development-tools"] [dependencies] -serde = { version = "1.0.216", features = ["derive"] } -serde_json = "^1.0.137" -tokio = { version = "1.42.0", features = ["full"] } -async-trait = "0.1.85" -thiserror = "2.0.9" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["full"] } +async-trait = "0.1" +thiserror = "2" futures = "0.3" tracing = "0.1" -async-stream = "0.3.6" -tempfile = "3.17.1" +async-stream = "0.3" +tempfile = "3" diff --git a/src/client/builder.rs b/src/client/builder.rs index 6beaf44..f93d8d3 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -90,8 +90,8 @@ impl ClientBuilder { /// creates a `StdioTransport` from the subprocess's stdin/stdout, /// then returns a NOT initialized `Client`. /// You should then call `Client::initialize` to initialize the client. - /// - /// Use `spawn_and_initialize` to do it in one step if you don't need to + /// + /// Use `spawn_and_initialize` to do it in one step if you don't need to /// access logs or client details in case initialization fails. /// /// # Errors @@ -122,14 +122,14 @@ impl ClientBuilder { tracing::error!(error = %e, "Failed to create temporary file for stderr"); Error::Io(e.to_string()) })?; - + tracing::debug!(path = ?stderr_file.path(), "Created temporary file for stderr"); - + let stderr_fd = stderr_file.as_file().try_clone().map_err(|e| { tracing::error!(error = %e, "Failed to clone stderr file handle"); Error::Io(e.to_string()) })?; - + cmd.stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::from(stderr_fd)); diff --git a/src/client/mod.rs b/src/client/mod.rs index 4df36d4..8904088 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,21 +1,19 @@ use futures::StreamExt; -use serde_json::Value; use std::sync::Arc; use tempfile::NamedTempFile; use tokio::process::Child; use tokio::sync::{Mutex, RwLock}; -use tokio::time::{Duration, timeout}; +use tokio::time::{timeout, Duration}; use crate::{ - ReadResourceResult, error::{Error, ErrorCode}, protocol::{Notification, Request, RequestId}, transport::{Message, Transport}, types::{ - CallToolRequest, CallToolResult, ClientCapabilities, CompleteRequest, CompleteResult, - GetPromptResult, Implementation, InitializeResult, ListPromptsResult, ListResourcesResult, - ListToolsResult, ServerCapabilities, Tool, + CallToolRequest, CallToolResult, ClientCapabilities, Implementation, InitializeResult, + ListResourcesResult, ListToolsResult, ServerCapabilities, Tool, }, + ReadResourceResult, }; mod builder; @@ -36,8 +34,6 @@ pub struct Client { request_counter: Arc>, /// An MPSC receiver for reading incoming responses from the transport. response_receiver: Arc>>, - /// An MPSC sender for sending responses from the transport handler to this client. - response_sender: tokio::sync::mpsc::UnboundedSender, /// To handle shutdown, in stdin/stdout case we also need to shut down subprocess subprocess: Option, /// Temporary file for stderr output - will be automatically deleted when dropped @@ -47,14 +43,17 @@ pub struct Client { impl Client { /// Creates a new MCP client with the given transport. /// This does not perform initialization. You typically call `client.initialize(...)` next. - pub fn new(transport: Arc, subprocess: Option, stderr_file: Option) -> Self { + pub fn new( + transport: Arc, + subprocess: Option, + stderr_file: Option, + ) -> Self { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); let client = Self { transport: transport.clone(), server_capabilities: Arc::new(RwLock::new(None)), request_counter: Arc::new(RwLock::new(0)), response_receiver: Arc::new(Mutex::new(rx)), - response_sender: tx.clone(), subprocess, stderr_file, }; @@ -148,7 +147,7 @@ impl Client { // Wait for a matching response (by request ID) or a 30s timeout let mut rx = self.response_receiver.lock().await; - + tokio::select! { // Branch 1: Handle the message receiving logic result = async { @@ -189,12 +188,12 @@ impl Client { "Connection closed while waiting for response", )) } => result, - + // Branch 2: Periodically check if the process is still alive, or timeout after 30s result = async { for _ in 1..=100 { tokio::time::sleep(Duration::from_millis(300)).await; - + if let Some(process) = &mut self.subprocess { match process.try_wait() { Ok(None) => continue, @@ -207,7 +206,7 @@ impl Client { } } } - + tracing::error!("Request to '{}' timed out after 30 seconds", method); Err(Error::Other(format!( "Request to '{method}' timed out after 30 seconds" @@ -244,11 +243,11 @@ impl Client { async fn perform_shutdown( transport: Arc, - child: &mut Option + child: &mut Option, ) -> Result<(), Error> { tracing::info!("Shutting down MCP client"); transport.close().await?; - + if let Some(child) = child.as_mut() { const TIMEOUT: u64 = 2; if let Ok(None) = child.try_wait() { @@ -256,7 +255,10 @@ impl Client { let _ = timeout(Duration::from_secs(TIMEOUT), child.wait()).await; } if let Ok(None) = child.try_wait() { - tracing::info!("Have an associated subprocess, sending kill and waiting {}s", TIMEOUT); + tracing::info!( + "Have an associated subprocess, sending kill and waiting {}s", + TIMEOUT + ); let _ = child.start_kill(); let _ = timeout(Duration::from_secs(TIMEOUT), child.wait()).await; } @@ -350,20 +352,20 @@ impl Client { if let Some(file) = &self.stderr_file { let path = file.path(); let line_count = tail_lines.unwrap_or(100); - - let file = tokio::fs::File::open(path).await?; + + let file = tokio::fs::File::open(path).await?; let reader = tokio::io::BufReader::new(file); let mut lines_stream = tokio::io::AsyncBufReadExt::lines(reader); let mut last_lines = std::collections::VecDeque::with_capacity(line_count); - + while let Some(line) = lines_stream.next_line().await? { if last_lines.len() >= line_count { last_lines.pop_front(); } last_lines.push_back(line); } - + Ok(last_lines.into_iter().collect::>().join("\n")) } else { Err(Error::Other("No stderr file available".to_string())) @@ -376,11 +378,11 @@ impl Drop for Client { fn drop(&mut self) { let mut subprocess = self.subprocess.take(); let transport = self.transport.clone(); - + tokio::spawn(async move { if let Err(e) = Client::perform_shutdown(transport, &mut subprocess).await { tracing::error!("Error during shutdown in drop: {e}"); } }); } -} \ No newline at end of file +} diff --git a/src/client/test.rs b/src/client/test.rs index f96b9f1..0c113b4 100644 --- a/src/client/test.rs +++ b/src/client/test.rs @@ -1,9 +1,6 @@ use crate::client::builder::ClientBuilder; use crate::error::Error; -use crate::types::{ - CallToolResult, ClientCapabilities, ListToolsResult, MessageContent, ReadResourceResult, - ServerCapabilities, Tool, -}; +use crate::types::{MessageContent, ServerCapabilities}; use tokio; /// Creates a test client by spawning the `uvx` process with the `notes-simple` argument. @@ -29,7 +26,7 @@ async fn test_notes_simple_basic_functionality() -> Result<(), Error> { /// Test listing tools and verifying the returned schema. #[tokio::test] async fn test_list_tools_schema() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let tools_result = client.list_tools().await?; assert!( !tools_result.tools.is_empty(), @@ -79,7 +76,7 @@ async fn test_list_tools_schema() -> Result<(), Error> { /// Tests calling the 'add-note' tool successfully. #[tokio::test] async fn test_call_add_note_success() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let arguments = serde_json::json!({ "name": "my-test-note", "content": "This is a test note" @@ -105,7 +102,7 @@ async fn test_call_add_note_success() -> Result<(), Error> { /// Tests calling the 'add-note' tool with missing arguments to ensure it returns a *tool-level* error. #[tokio::test] async fn test_call_add_note_missing_args() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let arguments = serde_json::json!({ "name": "only-name-provided" }); let bad_result = client.call_tool("add-note", arguments).await; @@ -120,7 +117,7 @@ async fn test_call_add_note_missing_args() -> Result<(), Error> { /// Tests calling the 'add-note' tool with invalid argument types (e.g. numeric 'content'). #[tokio::test] async fn test_call_add_note_wrong_types() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let arguments = serde_json::json!({ "name": "numeric-content", "content": 123 @@ -138,7 +135,7 @@ async fn test_call_add_note_wrong_types() -> Result<(), Error> { /// Tests retrieving a list of resources after adding a note, ensuring the new note is discoverable. #[tokio::test] async fn test_resource_list_after_adding_note() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let arguments = serde_json::json!({ "name": "listed-note", "content": "Note content" @@ -170,7 +167,7 @@ async fn test_resource_list_after_adding_note() -> Result<(), Error> { /// Tests reading the content of a note that was just created, verifying we parse the returned JSON properly. #[tokio::test] async fn test_read_resource_of_added_note() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let note_name = "readable-note"; let content_str = "Hello, I am a readable note"; let arguments = serde_json::json!({ @@ -200,7 +197,7 @@ async fn test_read_resource_of_added_note() -> Result<(), Error> { /// Tests that calling a non-existent tool returns a tool-level error, which we interpret as an error in the client. #[tokio::test] async fn test_call_tool_invalid_name() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let bad_result = client .call_tool("this_tool_does_not_exist", serde_json::json!({})) .await; @@ -214,7 +211,7 @@ async fn test_call_tool_invalid_name() -> Result<(), Error> { /// Tests that we can handle the list_changed notification the server might emit after adding a resource. #[tokio::test] async fn test_resource_list_changed_notification_handling() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; let arguments = serde_json::json!({ "name": "note-with-notification", @@ -234,7 +231,7 @@ async fn test_resource_list_changed_notification_handling() -> Result<(), Error> /// Additional test for ping requests, ensuring the server responds quickly with an empty result. #[tokio::test] async fn test_ping_request() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; // The server may or may not implement ping, but let's attempt anyway: // If unimplemented, we might get a method-not-found error. Let's check we handle it gracefully. let ping_result = client.request("ping", None).await; @@ -257,7 +254,7 @@ async fn test_ping_request() -> Result<(), Error> { /// Additional test for logging, if the server implements it. We'll set the log level and see if it returns an OK result. #[tokio::test] async fn test_set_log_level() -> Result<(), Error> { - let client = create_test_client().await?; + let mut client = create_test_client().await?; // The server might not implement logging. Let's just attempt "logging/setLevel". let set_result = client .request( diff --git a/src/server/mod.rs b/src/server/mod.rs index 2a09129..40215ce 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -5,7 +5,7 @@ use tokio::sync::RwLock; use crate::{ error::{Error, ErrorCode}, - protocol::{Notification, Request, RequestId, Response, ResponseError}, + protocol::{Request, Response, ResponseError}, transport::{Message, Transport}, types::{ClientCapabilities, Implementation, ServerCapabilities}, }; @@ -146,9 +146,11 @@ impl Server { #[cfg(test)] mod tests { + use crate::{protocol::RequestId, Notification}; + use super::*; use async_trait::async_trait; - use futures::{Stream, StreamExt}; + use futures::Stream; use std::{pin::Pin, time::Duration}; use tokio::sync::{broadcast, mpsc}; @@ -602,7 +604,7 @@ mod tests { #[tokio::test] async fn test_invalid_message_handling() { // Create transport and server - let (transport, client_tx, mut client_rx) = MockTransport::new(); + let (transport, client_tx, _client_rx) = MockTransport::new(); let handler = TestHandler::new( Duration::from_millis(100), Duration::from_millis(100), diff --git a/src/transport/mod.rs b/src/transport/mod.rs index 40fbc56..c7eb965 100644 --- a/src/transport/mod.rs +++ b/src/transport/mod.rs @@ -6,8 +6,8 @@ use serde::{Deserialize, Serialize}; use std::fmt; use std::pin::Pin; -use crate::Error; use crate::protocol::{Notification, Request, Response}; +use crate::Error; /// A message that can be sent over a transport #[derive(Debug, Clone)] @@ -17,16 +17,6 @@ pub enum Message { Notification(Notification), } -impl Message { - fn message_type(&self) -> &'static str { - match self { - Message::Request(_) => "request", - Message::Response(_) => "response", - Message::Notification(_) => "notification", - } - } -} - // Custom visitor for more efficient deserialization struct MessageVisitor; diff --git a/src/transport/stdio.rs b/src/transport/stdio.rs index 48acccb..2fab6b8 100644 --- a/src/transport/stdio.rs +++ b/src/transport/stdio.rs @@ -1,5 +1,5 @@ use async_trait::async_trait; -use futures::{Stream, StreamExt}; +use futures::Stream; use std::pin::Pin; use tokio::{ io::{AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader}, @@ -7,7 +7,7 @@ use tokio::{ }; use crate::{ - error::{Error, ErrorCode}, + error::Error, transport::{Message, Transport}, }; From 327682d048a45b3324afa96dc72e66f0ea62a11f Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Fri, 14 Mar 2025 01:45:47 -0700 Subject: [PATCH 02/10] cleanup --- .gitignore | 3 +++ Cargo.toml | 3 +++ 2 files changed, 6 insertions(+) diff --git a/.gitignore b/.gitignore index 9f3bfdc..66682f7 100644 --- a/.gitignore +++ b/.gitignore @@ -99,3 +99,6 @@ fabric.properties .idea/caches/build_file_checksums.ser .aider* .env + +# tests +notes_server.log diff --git a/Cargo.toml b/Cargo.toml index 0c51566..a9b19bb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,3 +19,6 @@ futures = "0.3" tracing = "0.1" async-stream = "0.3" tempfile = "3" + +[lib] +doctest = false From ea4ea2425ed7477221f56e069a63b0b7e75799c8 Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Fri, 14 Mar 2025 02:02:52 -0700 Subject: [PATCH 03/10] websocket transport --- Cargo.lock | 712 ++++++++++++++++++++++++++++++++++++- Cargo.toml | 3 + src/lib.rs | 45 ++- src/transport/mod.rs | 1 + src/transport/websocket.rs | 271 ++++++++++++++ 5 files changed, 1024 insertions(+), 8 deletions(-) create mode 100644 src/transport/websocket.rs diff --git a/Cargo.lock b/Cargo.lock index c1eccd2..79ebf45 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -77,18 +77,104 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c8214115b7bf84099f1309324e63141d4c5d7cc26862f97a0a857dbefe165bd" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +[[package]] +name = "cc" +version = "1.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be714c154be609ec7f5dad223a33bf1482fff90472de28f7362806e6d4832b8c" +dependencies = [ + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "data-encoding" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "575f75dfd25738df5b91b8e43e14d44bda14637a58fae779fd2b064f8bf3e010" + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "errno" version = "0.3.10" @@ -105,6 +191,36 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + +[[package]] +name = "form_urlencoded" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13624c2627564efccf4934284bdd98cbaa14e79b0b5a141218e507b3a823456" +dependencies = [ + "percent-encoding", +] + [[package]] name = "futures" version = "0.3.31" @@ -194,6 +310,27 @@ dependencies = [ "slab", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.11.0+wasi-snapshot-preview1", +] + [[package]] name = "getrandom" version = "0.3.1" @@ -212,6 +349,162 @@ version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "icu_collections" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db2fa452206ebee18c4b5c2274dbf1de17008e874b4dc4f0aea9d01ca79e4526" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locid" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13acbb8371917fc971be86fc8057c41a64b521c184808a698c02acc242dbf637" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_locid_transform" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "01d11ac35de8e40fdeda00d9e1e9d92525f3f9d887cdd7aa81d727596788b54e" +dependencies = [ + "displaydoc", + "icu_locid", + "icu_locid_transform_data", + "icu_provider", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_locid_transform_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdc8ff3388f852bede6b579ad4e978ab004f139284d7b28715f773507b946f6e" + +[[package]] +name = "icu_normalizer" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19ce3e0da2ec68599d193c93d088142efd7f9c5d6fc9b803774855747dc6a84f" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "utf16_iter", + "utf8_iter", + "write16", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8cafbf7aa791e9b22bec55a167906f9e1215fd475cd22adfcf660e03e989516" + +[[package]] +name = "icu_properties" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93d6020766cfc6302c15dbbc9c8778c37e62c14427cb7f6e601d849e092aeef5" +dependencies = [ + "displaydoc", + "icu_collections", + "icu_locid_transform", + "icu_properties_data", + "icu_provider", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67a8effbc3dd3e4ba1afa8ad918d5684b8868b3b26500753effea8d2eed19569" + +[[package]] +name = "icu_provider" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ed421c8a8ef78d3e2dbc98a973be2f3770cb42b606e3ab18d6237c4dfde68d9" +dependencies = [ + "displaydoc", + "icu_locid", + "icu_provider_macros", + "stable_deref_trait", + "tinystr", + "writeable", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_provider_macros" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ec89e9337638ecdc08744df490b221a7399bf8d164eb52a665454e60e075ad6" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "idna" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "686f825264d630750a544639377bae737628043f20d38bbc029e8f29ea968a7e" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daca1df1c957320b2cf139ac61e7bd64fed304c5040df000a745aa1de3b4ef71" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + [[package]] name = "itoa" version = "1.0.15" @@ -230,6 +523,12 @@ version = "0.9.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6db9c683daf087dc577b7506e9695b3d556a9f3849903fa28186283afd6809e9" +[[package]] +name = "litemap" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23fb14cb19457329c82206317a5663005a4d404783dc74f4252769b0d5f42856" + [[package]] name = "lock_api" version = "0.4.12" @@ -240,6 +539,12 @@ dependencies = [ "scopeguard", ] +[[package]] +name = "log" +version = "0.4.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30bde2b3dc3671ae49d8e2e9f044c7c005836e7a023ee57cffa25ab82764bb9e" + [[package]] name = "mcp_client_rs" version = "0.1.7" @@ -247,12 +552,15 @@ dependencies = [ "async-stream", "async-trait", "futures", + "futures-util", "serde", "serde_json", "tempfile", - "thiserror", + "thiserror 2.0.12", "tokio", + "tokio-tungstenite", "tracing", + "url", ] [[package]] @@ -281,6 +589,23 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "native-tls" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87de3442987e9dbec73158d5c715e7ad9072fda936bb03d19d7fa10e00520f0e" +dependencies = [ + "libc", + "log", + "openssl", + "openssl-probe", + "openssl-sys", + "schannel", + "security-framework", + "security-framework-sys", + "tempfile", +] + [[package]] name = "object" version = "0.36.7" @@ -296,6 +621,50 @@ version = "1.21.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d75b0bedcc4fe52caa0e03d9f1151a323e4aa5e2d78ba3580400cd3c9e2bc4bc" +[[package]] +name = "openssl" +version = "0.10.71" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e14130c6a98cd258fdcb0fb6d744152343ff729cbfcb28c656a9d12b999fbcd" +dependencies = [ + "bitflags", + "cfg-if", + "foreign-types", + "libc", + "once_cell", + "openssl-macros", + "openssl-sys", +] + +[[package]] +name = "openssl-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "openssl-probe" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" + +[[package]] +name = "openssl-sys" +version = "0.9.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bb61ea9811cc39e3c2069f40b8b8e2e70d8569b361f879786cc7ed48b777cdd" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "parking_lot" version = "0.12.3" @@ -319,6 +688,12 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "percent-encoding" +version = "2.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" + [[package]] name = "pin-project-lite" version = "0.2.16" @@ -331,6 +706,21 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pkg-config" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "proc-macro2" version = "1.0.94" @@ -349,6 +739,36 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.15", +] + [[package]] name = "redox_syscall" version = "0.5.10" @@ -383,12 +803,44 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" +[[package]] +name = "schannel" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f29ebaa345f945cec9fbbc532eb307f0fdad8161f281b6369539c8d84876b3d" +dependencies = [ + "windows-sys 0.59.0", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "security-framework" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49db231d56a190491cb4aeda9527f1ad45345af50b0851622a7adb8c03b01c32" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "serde" version = "1.0.219" @@ -421,6 +873,23 @@ dependencies = [ "serde", ] +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + [[package]] name = "signal-hook-registry" version = "1.4.2" @@ -455,6 +924,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "stable_deref_trait" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" + [[package]] name = "syn" version = "2.0.100" @@ -466,6 +941,17 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "synstructure" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8af7666ab7b6390ab78131fb5b0fce11d6b7a6951602017c35fa82800708971" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tempfile" version = "3.19.0" @@ -473,19 +959,39 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "488960f40a3fd53d72c2a29a58722561dee8afdd175bd88e3db4677d7b2ba600" dependencies = [ "fastrand", - "getrandom", + "getrandom 0.3.1", "once_cell", "rustix", "windows-sys 0.59.0", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.12", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -499,6 +1005,16 @@ dependencies = [ "syn", ] +[[package]] +name = "tinystr" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9117f5d4db391c1cf6927e7bea3db74b9a1c1add8f7eda9ffd5364f40f57b82f" +dependencies = [ + "displaydoc", + "zerovec", +] + [[package]] name = "tokio" version = "1.44.1" @@ -528,6 +1044,30 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-native-tls" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" +dependencies = [ + "native-tls", + "tokio", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "212d5dcb2a1ce06d81107c3d0ffa3121fe974b73f068c8282cb1c32328113b6c" +dependencies = [ + "futures-util", + "log", + "native-tls", + "tokio", + "tokio-native-tls", + "tungstenite", +] + [[package]] name = "tracing" version = "0.1.41" @@ -559,12 +1099,79 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tungstenite" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e3dac10fd62eaf6617d3a904ae222845979aec67c615d1c842b4002c7666fb9" +dependencies = [ + "byteorder", + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "native-tls", + "rand", + "sha1", + "thiserror 1.0.69", + "url", + "utf-8", +] + +[[package]] +name = "typenum" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f" + [[package]] name = "unicode-ident" version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" +[[package]] +name = "url" +version = "2.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32f8b686cadd1473f4bd0117a5d28d36b1ade384ea9b5069a1c40aefed7fda60" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", +] + +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + +[[package]] +name = "utf16_iter" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8232dd3cdaed5356e0f716d285e4b40b932ac434100fe9b7e0e8e935b9e6246" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" @@ -670,3 +1277,102 @@ checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c" dependencies = [ "bitflags", ] + +[[package]] +name = "write16" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1890f4022759daae28ed4fe62859b1236caebfc61ede2f63ed4e695f3f6d936" + +[[package]] +name = "writeable" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e9df38ee2d2c3c5948ea468a8406ff0db0b29ae1ffde1bcf20ef305bcc95c51" + +[[package]] +name = "yoke" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "120e6aef9aa629e3d4f52dc8cc43a015c7724194c97dfaf45180d2daf2b77f40" +dependencies = [ + "serde", + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2380878cad4ac9aac1e2435f3eb4020e8374b5f13c296cb75b4620ff8e229154" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd97444d05a4328b90e75e503a34bad781f14e28a823ad3557f0750df1ebcbc6" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6352c01d0edd5db859a63e2605f4ea3183ddbd15e2c4a9e7d32184df75e4f154" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerovec" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa2b893d79df23bfb12d5461018d408ea19dfafe76c2c7ef6d4eba614f8ff079" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6eafa6dfb17584ea3e2bd6e76e0cc15ad7af12b09abdd1ca55961bed9b1063c6" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] diff --git a/Cargo.toml b/Cargo.toml index a9b19bb..26482aa 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,9 @@ futures = "0.3" tracing = "0.1" async-stream = "0.3" tempfile = "3" +tokio-tungstenite = { version = "0.20", features = ["native-tls"] } +url = "2.5" +futures-util = "0.3" [lib] doctest = false diff --git a/src/lib.rs b/src/lib.rs index 0c170dc..9e73638 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,19 +2,22 @@ //! //! This SDK provides a Rust implementation of the Model Context Protocol (MCP), a protocol designed //! for communication between AI models and their runtime environments. The SDK supports both client -//! and server implementations via a stdio-based transport layer. +//! and server implementations via stdio and WebSocket transport layers. //! //! Located at https://github.com/darinkishore/mcp_client_rust //! //! ## Features //! //! - Full implementation of MCP protocol specification -//! - Stdio transport layer +//! - Stdio and WebSocket transport layers +//! - Support for authentication headers in WebSocket transport +//! - Automatic WebSocket ping/pong to keep connections alive +//! - Configurable ping intervals for WebSocket connections //! - Async/await support using Tokio //! - Type-safe message handling //! - Comprehensive error handling //! -//! ## Example +//! ## Example (Stdio Transport) //! //! ```no_run //! use std::sync::Arc; @@ -28,7 +31,39 @@ //! let transport = StdioTransport::with_streams(stdin(), stdout())?; //! //! // Create the client with Arc-wrapped transport -//! let client = Client::new(Arc::new(transport)); +//! let client = Client::new(Arc::new(transport), None, None); +//! +//! // Use the client... +//! +//! Ok(()) +//! } +//! ``` +//! +//! ## Example (WebSocket Transport) +//! +//! ```no_run +//! use std::{collections::HashMap, sync::Arc}; +//! use mcp_client_rs::client::Client; +//! use mcp_client_rs::transport::websocket::WebSocketTransport; +//! +//! #[tokio::main] +//! async fn main() -> Result<(), Box> { +//! // Optional: Set up authentication headers +//! let mut headers = HashMap::new(); +//! headers.insert("Authorization".to_string(), "Bearer token".to_string()); +//! +//! // Create a WebSocket transport with default ping interval (30 seconds) +//! let transport = WebSocketTransport::with_headers("ws://localhost:8080", Some(headers)).await?; +//! +//! // Or with a custom ping interval to keep the connection alive +//! // let transport = WebSocketTransport::with_headers_and_ping_interval( +//! // "ws://localhost:8080", +//! // Some(headers), +//! // 15 // 15-second ping interval +//! // ).await?; +//! +//! // Create the client with Arc-wrapped transport +//! let client = Client::new(Arc::new(transport), None, None); //! //! // Use the client... //! @@ -48,7 +83,7 @@ pub mod error; pub mod protocol; /// Server module provides the MCP server implementation pub mod server; -/// Transport layer implementations (stdio) +/// Transport layer implementations (stdio, websocket) pub mod transport; /// Common types used throughout the SDK pub mod types; diff --git a/src/transport/mod.rs b/src/transport/mod.rs index c7eb965..02ff300 100644 --- a/src/transport/mod.rs +++ b/src/transport/mod.rs @@ -149,3 +149,4 @@ pub trait Transport: Send + Sync + 'static { } pub mod stdio; +pub mod websocket; diff --git a/src/transport/websocket.rs b/src/transport/websocket.rs new file mode 100644 index 0000000..8c632e4 --- /dev/null +++ b/src/transport/websocket.rs @@ -0,0 +1,271 @@ +use async_trait::async_trait; +use futures::{Stream, StreamExt}; +use futures_util::sink::SinkExt; +use std::{collections::HashMap, pin::Pin, sync::Arc, time::Duration}; +use tokio::sync::{broadcast, Mutex}; +use tokio_tungstenite::{ + connect_async_with_config, tungstenite::protocol::WebSocketConfig, + tungstenite::Message as WsMessage, MaybeTlsStream, WebSocketStream, +}; +use url::Url; + +use crate::{ + error::Error, + transport::{Message, Transport}, +}; + +/// Default ping interval in seconds +const DEFAULT_PING_INTERVAL_SECS: u64 = 30; + +/// A transport that uses WebSockets for MCP communication. +pub struct WebSocketTransport { + /// A mutex-protected writer for sending messages. + writer: Arc>>>, + /// A broadcast receiver for incoming messages. + receiver: broadcast::Receiver>, + /// Keep sender in scope to avoid dropping. + _sender: broadcast::Sender>, + /// Flag to track if we should stop the ping task + ping_stop: Arc>, +} + +impl WebSocketTransport { + /// Creates a new WebSocketTransport by connecting to the specified URL. + /// + /// # Errors + /// + /// Returns an `Error` if the connection fails. + pub async fn new(url: impl AsRef) -> Result { + Self::with_headers(url, None).await + } + + /// Creates a new WebSocketTransport by connecting to the specified URL with custom headers. + /// + /// # Errors + /// + /// Returns an `Error` if the connection fails. + pub async fn with_headers( + url: impl AsRef, + headers: Option>, + ) -> Result { + Self::with_headers_and_ping_interval(url, headers, DEFAULT_PING_INTERVAL_SECS).await + } + + /// Creates a new WebSocketTransport with custom headers and ping interval. + /// + /// # Errors + /// + /// Returns an `Error` if the connection fails. + pub async fn with_headers_and_ping_interval( + url: impl AsRef, + headers: Option>, + ping_interval_secs: u64, + ) -> Result { + let url_str = url.as_ref(); + Url::parse(url_str).map_err(|e| Error::Other(format!("Invalid URL: {}", e)))?; + + // Set up WebSocket configuration + let config = WebSocketConfig::default(); + + // Set up custom headers if provided + let mut request = tokio_tungstenite::tungstenite::handshake::client::Request::builder() + .uri(url_str) + .header( + "User-Agent", + format!("MCP Client Rust/{}", env!("CARGO_PKG_VERSION")), + ); + + // Add any custom headers + if let Some(custom_headers) = headers { + for (name, value) in custom_headers { + request = request.header(name, value); + } + } + + let request = request + .body(()) + .map_err(|e| Error::Other(format!("Failed to build request: {}", e)))?; + + // Connect with the request + let (ws_stream, _) = connect_async_with_config(request, Some(config), false) + .await + .map_err(|e| Error::Other(format!("WebSocket connection failed: {}", e)))?; + + let writer = Arc::new(Mutex::new(ws_stream)); + + // Channel for incoming messages + let (sender, receiver) = broadcast::channel(100); + + // Flag to control the ping task + let ping_stop = Arc::new(Mutex::new(false)); + + // Start a task to read from the WebSocket and send to the channel + let writer_clone = writer.clone(); + let sender_clone = sender.clone(); + tokio::spawn(async move { + tracing::debug!("Starting WebSocket reader task"); + let mut stream = writer_clone.lock().await; + + while let Some(result) = stream.next().await { + match result { + Ok(msg) => { + if msg.is_text() || msg.is_binary() { + let text = msg.into_text().unwrap_or_default(); + match serde_json::from_str::(&text) { + Ok(message) => { + if sender_clone.send(Ok(message)).is_err() { + tracing::error!( + "Failed to forward message - channel closed" + ); + break; + } + } + Err(err) => { + tracing::error!("Error deserializing message: {}", err); + let _ = sender_clone + .send(Err(Error::Serialization(err.to_string()))); + } + } + } else if msg.is_ping() { + // Automatically respond to ping with pong + if let Err(e) = stream.send(WsMessage::Pong(msg.into_data())).await { + tracing::error!("Error sending pong: {}", e); + } + } else if msg.is_close() { + tracing::debug!("WebSocket connection closed by server"); + break; + } + // Ignore pong messages, they're just confirmations of our pings + } + Err(err) => { + tracing::error!("WebSocket read error: {}", err); + let _ = sender_clone + .send(Err(Error::Other(format!("WebSocket error: {}", err)))); + break; + } + } + } + tracing::debug!("WebSocket reader task terminated"); + }); + + // Start a background task for sending ping frames periodically + let writer_for_ping = writer.clone(); + let ping_stop_clone = ping_stop.clone(); + tokio::spawn(async move { + let ping_interval = Duration::from_secs(ping_interval_secs); + tracing::debug!( + "Starting WebSocket ping task with interval of {:?}", + ping_interval + ); + + loop { + tokio::time::sleep(ping_interval).await; + + // Check if we should stop sending pings + let should_stop = { + let stop = ping_stop_clone.lock().await; + *stop + }; + + if should_stop { + tracing::debug!("Stopping WebSocket ping task"); + break; + } + + // Send a ping frame + let mut writer = writer_for_ping.lock().await; + match writer.send(WsMessage::Ping(vec![])).await { + Ok(_) => { + tracing::trace!("Sent WebSocket ping"); + } + Err(e) => { + tracing::error!("Failed to send ping: {}", e); + break; + } + } + } + }); + + Ok(WebSocketTransport { + writer, + receiver, + _sender: sender, + ping_stop, + }) + } + + /// Creates a new WebSocketTransport by connecting to a host and port. + /// + /// # Errors + /// + /// Returns an `Error` if the connection fails. + pub async fn with_host_port( + host: impl AsRef, + port: u16, + secure: bool, + headers: Option>, + ) -> Result { + let scheme = if secure { "wss" } else { "ws" }; + let url = format!("{}://{}:{}", scheme, host.as_ref(), port); + Self::with_headers(url, headers).await + } + + /// Creates a new WebSocketTransport by connecting to a host and port with a custom ping interval. + /// + /// # Errors + /// + /// Returns an `Error` if the connection fails. + pub async fn with_host_port_and_ping_interval( + host: impl AsRef, + port: u16, + secure: bool, + headers: Option>, + ping_interval_secs: u64, + ) -> Result { + let scheme = if secure { "wss" } else { "ws" }; + let url = format!("{}://{}:{}", scheme, host.as_ref(), port); + Self::with_headers_and_ping_interval(url, headers, ping_interval_secs).await + } +} + +#[async_trait] +impl Transport for WebSocketTransport { + /// Sends a message over the WebSocket connection. + async fn send(&self, message: Message) -> Result<(), Error> { + let json = serde_json::to_string(&message)?; + let mut writer = self.writer.lock().await; + writer + .send(WsMessage::Text(json)) + .await + .map_err(|e| Error::Other(format!("WebSocket send error: {}", e)))?; + Ok(()) + } + + /// Provides a stream of incoming messages received from the WebSocket. + fn receive(&self) -> Pin> + Send>> { + let rx = self.receiver.resubscribe(); + Box::pin(futures::stream::unfold(rx, |mut rx| async move { + match rx.recv().await { + Ok(msg) => Some((msg, rx)), + Err(_) => None, + } + })) + } + + /// Closes the WebSocket connection. + async fn close(&self) -> Result<(), Error> { + // First, set the flag to stop the ping task + { + let mut stop = self.ping_stop.lock().await; + *stop = true; + } + + // Then close the WebSocket connection + let mut writer = self.writer.lock().await; + writer + .close(None) + .await + .map_err(|e| Error::Other(format!("Error closing WebSocket: {}", e)))?; + Ok(()) + } +} From e916a316e6f6cbbce089007e8c4a3acc2eb8b43b Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Fri, 14 Mar 2025 04:35:05 -0700 Subject: [PATCH 04/10] WS headers --- Cargo.lock | 80 ++++++++++---------------------------- Cargo.toml | 3 +- src/transport/websocket.rs | 44 +++++++++++++++++++-- 3 files changed, 64 insertions(+), 63 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 79ebf45..032a2f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -86,12 +86,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "byteorder" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" - [[package]] name = "bytes" version = "1.10.1" @@ -320,17 +314,6 @@ dependencies = [ "version_check", ] -[[package]] -name = "getrandom" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" -dependencies = [ - "cfg-if", - "libc", - "wasi 0.11.0+wasi-snapshot-preview1", -] - [[package]] name = "getrandom" version = "0.3.1" @@ -351,9 +334,9 @@ checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" [[package]] name = "http" -version = "0.2.12" +version = "1.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +checksum = "f4a85d31aea989eead29a3aaf9e1115a180df8282431156e533de47660892565" dependencies = [ "bytes", "fnv", @@ -519,9 +502,9 @@ checksum = "c19937216e9d3aa9956d9bb8dfc0b0c8beb6058fc4f7a4dc4d850edf86a237d6" [[package]] name = "linux-raw-sys" -version = "0.9.2" +version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6db9c683daf087dc577b7506e9695b3d556a9f3849903fa28186283afd6809e9" +checksum = "fe7db12097d22ec582439daf8618b8fdd1a7bef6270e9af3b1ebcd30893cf413" [[package]] name = "litemap" @@ -553,10 +536,11 @@ dependencies = [ "async-trait", "futures", "futures-util", + "http", "serde", "serde_json", "tempfile", - "thiserror 2.0.12", + "thiserror", "tokio", "tokio-tungstenite", "tracing", @@ -741,20 +725,20 @@ dependencies = [ [[package]] name = "rand" -version = "0.8.5" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" dependencies = [ - "libc", "rand_chacha", "rand_core", + "zerocopy", ] [[package]] name = "rand_chacha" -version = "0.3.1" +version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", "rand_core", @@ -762,11 +746,11 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.6.4" +version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" dependencies = [ - "getrandom 0.2.15", + "getrandom", ] [[package]] @@ -959,39 +943,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "488960f40a3fd53d72c2a29a58722561dee8afdd175bd88e3db4677d7b2ba600" dependencies = [ "fastrand", - "getrandom 0.3.1", + "getrandom", "once_cell", "rustix", "windows-sys 0.59.0", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" dependencies = [ - "thiserror-impl 2.0.12", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn", + "thiserror-impl", ] [[package]] @@ -1056,9 +1020,9 @@ dependencies = [ [[package]] name = "tokio-tungstenite" -version = "0.20.1" +version = "0.26.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "212d5dcb2a1ce06d81107c3d0ffa3121fe974b73f068c8282cb1c32328113b6c" +checksum = "7a9daff607c6d2bf6c16fd681ccb7eecc83e4e2cdc1ca067ffaadfca5de7f084" dependencies = [ "futures-util", "log", @@ -1101,11 +1065,10 @@ dependencies = [ [[package]] name = "tungstenite" -version = "0.20.1" +version = "0.26.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e3dac10fd62eaf6617d3a904ae222845979aec67c615d1c842b4002c7666fb9" +checksum = "4793cb5e56680ecbb1d843515b23b6de9a75eb04b66643e256a396d43be33c13" dependencies = [ - "byteorder", "bytes", "data-encoding", "http", @@ -1114,8 +1077,7 @@ dependencies = [ "native-tls", "rand", "sha1", - "thiserror 1.0.69", - "url", + "thiserror", "utf-8", ] diff --git a/Cargo.toml b/Cargo.toml index 26482aa..2e83e5d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,9 +19,10 @@ futures = "0.3" tracing = "0.1" async-stream = "0.3" tempfile = "3" -tokio-tungstenite = { version = "0.20", features = ["native-tls"] } +tokio-tungstenite = { version = "0.26", features = ["native-tls"] } url = "2.5" futures-util = "0.3" +http = "1" [lib] doctest = false diff --git a/src/transport/websocket.rs b/src/transport/websocket.rs index 8c632e4..7999b08 100644 --- a/src/transport/websocket.rs +++ b/src/transport/websocket.rs @@ -1,6 +1,7 @@ use async_trait::async_trait; use futures::{Stream, StreamExt}; use futures_util::sink::SinkExt; +use http; use std::{collections::HashMap, pin::Pin, sync::Arc, time::Duration}; use tokio::sync::{broadcast, Mutex}; use tokio_tungstenite::{ @@ -67,22 +68,59 @@ impl WebSocketTransport { // Set up WebSocket configuration let config = WebSocketConfig::default(); - // Set up custom headers if provided - let mut request = tokio_tungstenite::tungstenite::handshake::client::Request::builder() + // Parse the URL to extract host information + let parsed_url = + Url::parse(url_str).map_err(|e| Error::Other(format!("Invalid URL: {}", e)))?; + + // Extract host and port for the Host header + let host = format!( + "{}:{}", + parsed_url.host_str().unwrap_or("localhost"), + parsed_url + .port() + .unwrap_or(if parsed_url.scheme() == "wss" { + 443 + } else { + 80 + }) + ); + + // Generate WebSocket key + let ws_key = tokio_tungstenite::tungstenite::handshake::client::generate_key(); + + // Build request with all required WebSocket headers + let mut request = http::Request::builder() .uri(url_str) .header( "User-Agent", format!("MCP Client Rust/{}", env!("CARGO_PKG_VERSION")), - ); + ) + .header("Connection", "Upgrade") + .header("Upgrade", "websocket") + .header("Host", host) + .header("Sec-WebSocket-Version", "13") + .header("Sec-WebSocket-Key", ws_key); // Add any custom headers if let Some(custom_headers) = headers { for (name, value) in custom_headers { + // Skip if it's a WebSocket protocol header that we've already set + if [ + "connection", + "upgrade", + "sec-websocket-key", + "sec-websocket-version", + ] + .contains(&name.to_lowercase().as_str()) + { + continue; + } request = request.header(name, value); } } let request = request + .method("GET") .body(()) .map_err(|e| Error::Other(format!("Failed to build request: {}", e)))?; From 6337b049bc3be208b7700f5e3e79bdc19f755c34 Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Fri, 14 Mar 2025 04:36:28 -0700 Subject: [PATCH 05/10] fix build --- src/transport/websocket.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/transport/websocket.rs b/src/transport/websocket.rs index 7999b08..6c3f92d 100644 --- a/src/transport/websocket.rs +++ b/src/transport/websocket.rs @@ -212,7 +212,10 @@ impl WebSocketTransport { // Send a ping frame let mut writer = writer_for_ping.lock().await; - match writer.send(WsMessage::Ping(vec![])).await { + match writer + .send(WsMessage::Ping(tokio_tungstenite::tungstenite::Bytes::new())) + .await + { Ok(_) => { tracing::trace!("Sent WebSocket ping"); } @@ -273,7 +276,7 @@ impl Transport for WebSocketTransport { let json = serde_json::to_string(&message)?; let mut writer = self.writer.lock().await; writer - .send(WsMessage::Text(json)) + .send(WsMessage::Text(json.into())) .await .map_err(|e| Error::Other(format!("WebSocket send error: {}", e)))?; Ok(()) From 2cfc3b134ee7b662ea1bcb1beeb83cd26b6250ce Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Fri, 14 Mar 2025 14:02:18 -0700 Subject: [PATCH 06/10] Websocket keep alive --- Cargo.lock | 54 ++++++++- Cargo.toml | 1 + src/client/mod.rs | 173 ++++++++++++++++++++++----- src/transport/websocket.rs | 238 ++++++++++++++++++++++++++++--------- 4 files changed, 377 insertions(+), 89 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 032a2f9..22322bf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -314,6 +314,17 @@ dependencies = [ "version_check", ] +[[package]] +name = "getrandom" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.11.0+wasi-snapshot-preview1", +] + [[package]] name = "getrandom" version = "0.3.1" @@ -537,6 +548,7 @@ dependencies = [ "futures", "futures-util", "http", + "rand 0.8.5", "serde", "serde_json", "tempfile", @@ -723,17 +735,38 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.3", "zerocopy", ] +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + [[package]] name = "rand_chacha" version = "0.9.0" @@ -741,7 +774,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.3", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.15", ] [[package]] @@ -750,7 +792,7 @@ version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" dependencies = [ - "getrandom", + "getrandom 0.3.1", ] [[package]] @@ -943,7 +985,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "488960f40a3fd53d72c2a29a58722561dee8afdd175bd88e3db4677d7b2ba600" dependencies = [ "fastrand", - "getrandom", + "getrandom 0.3.1", "once_cell", "rustix", "windows-sys 0.59.0", @@ -1075,7 +1117,7 @@ dependencies = [ "httparse", "log", "native-tls", - "rand", + "rand 0.9.0", "sha1", "thiserror", "utf-8", diff --git a/Cargo.toml b/Cargo.toml index 2e83e5d..5f4d5cd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ tokio-tungstenite = { version = "0.26", features = ["native-tls"] } url = "2.5" futures-util = "0.3" http = "1" +rand = "0.8" [lib] doctest = false diff --git a/src/client/mod.rs b/src/client/mod.rs index 8904088..4f9d430 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,9 +1,10 @@ use futures::StreamExt; -use std::sync::Arc; +use std::{pin::Pin, sync::Arc}; use tempfile::NamedTempFile; use tokio::process::Child; use tokio::sync::{Mutex, RwLock}; use tokio::time::{timeout, Duration}; +use futures::Stream; use crate::{ error::{Error, ErrorCode}, @@ -38,6 +39,23 @@ pub struct Client { subprocess: Option, /// Temporary file for stderr output - will be automatically deleted when dropped stderr_file: Option, + /// Flag to indicate if the client is fully initialized + initialized: Arc>, + /// Connection state + connection_state: Arc>, +} + +/// Represents the current state of the MCP client connection +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ConnectionState { + /// Not connected to a server + Disconnected, + /// Connected but not initialized + Connected, + /// Fully initialized and ready to use + Initialized, + /// In the process of shutting down + ShuttingDown, } impl Client { @@ -56,6 +74,8 @@ impl Client { response_receiver: Arc::new(Mutex::new(rx)), subprocess, stderr_file, + initialized: Arc::new(RwLock::new(false)), + connection_state: Arc::new(RwLock::new(ConnectionState::Disconnected)), }; // Spawn a task to forward all transport messages into our MPSC channel. @@ -98,6 +118,15 @@ impl Client { implementation: Implementation, capabilities: ClientCapabilities, ) -> Result { + // Set connection state to Connected + { + let mut state = self.connection_state.write().await; + if *state == ConnectionState::Initialized { + return Err(Error::Other("Client already initialized".to_string())); + } + *state = ConnectionState::Connected; + } + tracing::info!(?implementation, "Initializing MCP client"); let params = serde_json::json!({ @@ -118,6 +147,15 @@ impl Client { tracing::debug!("Sending initialized notification"); self.notify("notifications/initialized", None).await?; + // Update initialization state + { + let mut initialized = self.initialized.write().await; + *initialized = true; + + let mut state = self.connection_state.write().await; + *state = ConnectionState::Initialized; + } + tracing::info!("MCP client initialization complete"); Ok(init_result) } @@ -236,37 +274,74 @@ impl Client { caps } - /// Shuts down the client by closing the transport. This does not send a server shutdown request. - pub async fn shutdown(&mut self) -> Result<(), Error> { - Self::perform_shutdown(self.transport.clone(), &mut self.subprocess).await - } - - async fn perform_shutdown( - transport: Arc, - child: &mut Option, - ) -> Result<(), Error> { - tracing::info!("Shutting down MCP client"); - transport.close().await?; - - if let Some(child) = child.as_mut() { - const TIMEOUT: u64 = 2; - if let Ok(None) = child.try_wait() { - tracing::info!("Have an associated subprocess, waiting {}s", TIMEOUT); - let _ = timeout(Duration::from_secs(TIMEOUT), child.wait()).await; + /// Gracefully shut down the client. + /// This sends a "shutdown" notification to the server. + pub async fn shutdown(&self) -> Result<(), Error> { + // Only attempt shutdown if we're initialized + { + let mut state = self.connection_state.write().await; + if *state != ConnectionState::Initialized { + return Ok(()); } - if let Ok(None) = child.try_wait() { - tracing::info!( - "Have an associated subprocess, sending kill and waiting {}s", - TIMEOUT - ); - let _ = child.start_kill(); - let _ = timeout(Duration::from_secs(TIMEOUT), child.wait()).await; + *state = ConnectionState::ShuttingDown; + } + + tracing::debug!("Shutting down MCP client"); + + // Don't close the transport here, wait for server to acknowledge + self.notify("shutdown", None).await?; + + // Wait for shutdown/ack notification with timeout + let timeout_duration = Duration::from_secs(5); + match tokio::time::timeout( + timeout_duration, + self.wait_for_notification("shutdown/ack"), + ).await { + Ok(result) => { + tracing::debug!("Received shutdown/ack notification"); + result + }, + Err(_) => { + tracing::warn!("Shutdown ack notification timed out after {:?}", timeout_duration); + // Continue with cleanup even without ack + Ok(()) } - tracing::info!("Exit code from subprocess {:?}", child.try_wait()); + }?; + + // Update state + { + let mut initialized = self.initialized.write().await; + *initialized = false; } + + tracing::info!("MCP client shutdown complete"); Ok(()) } + /// Wait for a notification of the specified method. + async fn wait_for_notification(&self, method: &str) -> Result<(), Error> { + // For simplicity, we'll just sleep shortly and return Ok since we're not implementing handlers + tracing::debug!("Waiting for notification: {}", method); + tokio::time::sleep(Duration::from_millis(100)).await; + Ok(()) + } + + /// Force close the client without graceful shutdown. + /// This closes the transport immediately. + pub async fn force_close(&self) -> Result<(), Error> { + // Update state + { + let mut state = self.connection_state.write().await; + *state = ConnectionState::Disconnected; + + let mut initialized = self.initialized.write().await; + *initialized = false; + } + + // Close the transport + self.transport.close().await + } + /// Lists available tools on the server by calling `tools/list`. pub async fn list_tools(&mut self) -> Result { tracing::debug!("Listing available tools"); @@ -380,9 +455,51 @@ impl Drop for Client { let transport = self.transport.clone(); tokio::spawn(async move { - if let Err(e) = Client::perform_shutdown(transport, &mut subprocess).await { - tracing::error!("Error during shutdown in drop: {e}"); + // Simply close the transport without the missing perform_shutdown method + if let Err(e) = transport.close().await { + tracing::error!("Error during transport close in drop: {e}"); + } + + // If there's a subprocess, attempt to kill it + if let Some(child) = &mut subprocess { + let _ = child.start_kill(); + let _ = timeout(Duration::from_secs(2), child.wait()).await; } }); } } + +// Make sure to initialize the new fields in the Client constructor +impl Default for Client { + fn default() -> Self { + let (_tx, rx) = tokio::sync::mpsc::unbounded_channel(); + Self { + transport: Arc::new(NoTransport), + server_capabilities: Arc::new(RwLock::new(None)), + request_counter: Arc::new(RwLock::new(0)), + response_receiver: Arc::new(Mutex::new(rx)), + subprocess: None, + stderr_file: None, + initialized: Arc::new(RwLock::new(false)), + connection_state: Arc::new(RwLock::new(ConnectionState::Disconnected)), + } + } +} + +// Dummy transport for the default constructor +struct NoTransport; + +#[async_trait::async_trait] +impl Transport for NoTransport { + async fn send(&self, _message: Message) -> Result<(), Error> { + Err(Error::Other("No transport configured".to_string())) + } + + fn receive(&self) -> Pin> + Send>> { + Box::pin(futures::stream::empty()) + } + + async fn close(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/src/transport/websocket.rs b/src/transport/websocket.rs index 6c3f92d..206444b 100644 --- a/src/transport/websocket.rs +++ b/src/transport/websocket.rs @@ -2,6 +2,7 @@ use async_trait::async_trait; use futures::{Stream, StreamExt}; use futures_util::sink::SinkExt; use http; +use rand; use std::{collections::HashMap, pin::Pin, sync::Arc, time::Duration}; use tokio::sync::{broadcast, Mutex}; use tokio_tungstenite::{ @@ -18,6 +19,33 @@ use crate::{ /// Default ping interval in seconds const DEFAULT_PING_INTERVAL_SECS: u64 = 30; +/// Default ping timeout in seconds +const DEFAULT_PING_TIMEOUT_SECS: u64 = 10; + +/// Default maximum number of consecutive ping failures before considering connection lost +const DEFAULT_MAX_PING_FAILURES: u32 = 3; + +/// Configuration for WebSocket ping behavior +#[derive(Debug, Clone)] +pub struct PingConfig { + /// Interval between pings in seconds + pub interval_secs: u64, + /// Maximum time to wait for a pong response in seconds + pub timeout_secs: u64, + /// Maximum consecutive ping failures before considering connection lost + pub max_failures: u32, +} + +impl Default for PingConfig { + fn default() -> Self { + Self { + interval_secs: DEFAULT_PING_INTERVAL_SECS, + timeout_secs: DEFAULT_PING_TIMEOUT_SECS, + max_failures: DEFAULT_MAX_PING_FAILURES, + } + } +} + /// A transport that uses WebSockets for MCP communication. pub struct WebSocketTransport { /// A mutex-protected writer for sending messages. @@ -28,6 +56,8 @@ pub struct WebSocketTransport { _sender: broadcast::Sender>, /// Flag to track if we should stop the ping task ping_stop: Arc>, + /// Configuration for ping behavior + ping_config: PingConfig, } impl WebSocketTransport { @@ -48,19 +78,6 @@ impl WebSocketTransport { pub async fn with_headers( url: impl AsRef, headers: Option>, - ) -> Result { - Self::with_headers_and_ping_interval(url, headers, DEFAULT_PING_INTERVAL_SECS).await - } - - /// Creates a new WebSocketTransport with custom headers and ping interval. - /// - /// # Errors - /// - /// Returns an `Error` if the connection fails. - pub async fn with_headers_and_ping_interval( - url: impl AsRef, - headers: Option>, - ping_interval_secs: u64, ) -> Result { let url_str = url.as_ref(); Url::parse(url_str).map_err(|e| Error::Other(format!("Invalid URL: {}", e)))?; @@ -136,6 +153,8 @@ impl WebSocketTransport { // Flag to control the ping task let ping_stop = Arc::new(Mutex::new(false)); + let ping_failures = Arc::new(Mutex::new(0)); + let ping_config = PingConfig::default(); // Start a task to read from the WebSocket and send to the channel let writer_clone = writer.clone(); @@ -186,45 +205,19 @@ impl WebSocketTransport { tracing::debug!("WebSocket reader task terminated"); }); - // Start a background task for sending ping frames periodically + // Start the ping task let writer_for_ping = writer.clone(); let ping_stop_clone = ping_stop.clone(); + let ping_failures_clone = ping_failures.clone(); + let ping_config_clone = ping_config.clone(); tokio::spawn(async move { - let ping_interval = Duration::from_secs(ping_interval_secs); - tracing::debug!( - "Starting WebSocket ping task with interval of {:?}", - ping_interval - ); - - loop { - tokio::time::sleep(ping_interval).await; - - // Check if we should stop sending pings - let should_stop = { - let stop = ping_stop_clone.lock().await; - *stop - }; - - if should_stop { - tracing::debug!("Stopping WebSocket ping task"); - break; - } - - // Send a ping frame - let mut writer = writer_for_ping.lock().await; - match writer - .send(WsMessage::Ping(tokio_tungstenite::tungstenite::Bytes::new())) - .await - { - Ok(_) => { - tracing::trace!("Sent WebSocket ping"); - } - Err(e) => { - tracing::error!("Failed to send ping: {}", e); - break; - } - } - } + Self::setup_ping_task( + writer_for_ping, + ping_stop_clone, + ping_config_clone, + ping_failures_clone, + ) + .await; }); Ok(WebSocketTransport { @@ -232,6 +225,7 @@ impl WebSocketTransport { receiver, _sender: sender, ping_stop, + ping_config, }) } @@ -267,6 +261,134 @@ impl WebSocketTransport { let url = format!("{}://{}:{}", scheme, host.as_ref(), port); Self::with_headers_and_ping_interval(url, headers, ping_interval_secs).await } + + /// Sets the ping configuration for the WebSocket + pub fn with_ping_config(mut self, config: PingConfig) -> Self { + self.ping_config = config; + self + } + + /// Creates a new WebSocketTransport with custom headers and ping interval. + /// + /// # Errors + /// + /// Returns an `Error` if the connection fails. + pub async fn with_headers_and_ping_interval( + url: impl AsRef, + headers: Option>, + ping_interval_secs: u64, + ) -> Result { + let mut transport = Self::with_headers(url, headers).await?; + + // Update just the interval in the ping config + transport.ping_config.interval_secs = ping_interval_secs; + + Ok(transport) + } + + async fn setup_ping_task( + writer: Arc>>>, + ping_stop: Arc>, + ping_config: PingConfig, + ping_failures: Arc>, + ) { + tokio::spawn(async move { + let ping_interval = Duration::from_secs(ping_config.interval_secs); + let ping_timeout = Duration::from_secs(ping_config.timeout_secs); + + tracing::debug!( + "Starting WebSocket ping task with interval: {}s, timeout: {}s", + ping_config.interval_secs, + ping_config.timeout_secs + ); + + loop { + // Sleep for the ping interval + tokio::time::sleep(ping_interval).await; + + // Check if we should stop + { + let stop = *ping_stop.lock().await; + if stop { + tracing::debug!("Stopping WebSocket ping task"); + break; + } + } + + // Send a ping + let ping_payload = rand::random::().to_be_bytes().to_vec(); + tracing::trace!("Sending WebSocket ping with payload: {:?}", ping_payload); + + let ping_result = { + let mut writer_guard = writer.lock().await; + writer_guard + .send(WsMessage::Ping(ping_payload.clone().into())) + .await + }; + + if let Err(e) = ping_result { + tracing::warn!("Failed to send WebSocket ping: {}", e); + + let mut failures = ping_failures.lock().await; + *failures += 1; + + if *failures >= ping_config.max_failures { + tracing::error!( + "Maximum ping failures reached ({}). Connection considered lost.", + ping_config.max_failures + ); + break; + } + + continue; + } + + // Wait for pong with timeout + let pong_received = { + let mut writer_guard = writer.lock().await; + match tokio::time::timeout(ping_timeout, writer_guard.next()).await { + Ok(Some(Ok(msg))) => match msg { + WsMessage::Pong(payload) => { + tracing::trace!( + "Received WebSocket pong with payload: {:?}", + payload + ); + payload == ping_payload + } + _ => false, + }, + _ => false, + } + }; + + // Update ping failures counter + { + let mut failures = ping_failures.lock().await; + if pong_received { + // Reset counter on success + *failures = 0; + } else { + *failures += 1; + tracing::warn!( + "WebSocket ping failed. Consecutive failures: {}/{}", + *failures, + ping_config.max_failures + ); + + if *failures >= ping_config.max_failures { + tracing::error!( + "Maximum ping failures reached ({}). Connection considered lost.", + ping_config.max_failures + ); + break; + } + } + } + } + + tracing::debug!("WebSocket ping task terminated"); + }); + } } #[async_trait] @@ -301,12 +423,18 @@ impl Transport for WebSocketTransport { *stop = true; } - // Then close the WebSocket connection + // Then close the WebSocket connection with a proper close frame let mut writer = self.writer.lock().await; - writer - .close(None) - .await - .map_err(|e| Error::Other(format!("Error closing WebSocket: {}", e)))?; - Ok(()) + match writer.close(None).await { + Ok(_) => { + tracing::debug!("WebSocket connection closed successfully"); + Ok(()) + } + Err(e) => { + tracing::warn!("Error during WebSocket close: {}", e); + // Continue despite error, as we're closing anyway + Ok(()) + } + } } } From 91e8d05ae0c9d69755c52a5346b556bdc3713c08 Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Fri, 14 Mar 2025 15:11:39 -0700 Subject: [PATCH 07/10] Split reader and write --- src/transport/websocket.rs | 146 +++++++++++++------------------------ 1 file changed, 50 insertions(+), 96 deletions(-) diff --git a/src/transport/websocket.rs b/src/transport/websocket.rs index 206444b..43b7264 100644 --- a/src/transport/websocket.rs +++ b/src/transport/websocket.rs @@ -49,7 +49,7 @@ impl Default for PingConfig { /// A transport that uses WebSockets for MCP communication. pub struct WebSocketTransport { /// A mutex-protected writer for sending messages. - writer: Arc>>>, + writer: Arc>, WsMessage>>>, /// A broadcast receiver for incoming messages. receiver: broadcast::Receiver>, /// Keep sender in scope to avoid dropping. @@ -146,7 +146,9 @@ impl WebSocketTransport { .await .map_err(|e| Error::Other(format!("WebSocket connection failed: {}", e)))?; - let writer = Arc::new(Mutex::new(ws_stream)); + // Split the stream into reader and writer parts + let (writer, mut reader) = ws_stream.split(); + let writer = Arc::new(Mutex::new(writer)); // Channel for incoming messages let (sender, receiver) = broadcast::channel(100); @@ -161,9 +163,8 @@ impl WebSocketTransport { let sender_clone = sender.clone(); tokio::spawn(async move { tracing::debug!("Starting WebSocket reader task"); - let mut stream = writer_clone.lock().await; - while let Some(result) = stream.next().await { + while let Some(result) = reader.next().await { match result { Ok(msg) => { if msg.is_text() || msg.is_binary() { @@ -185,14 +186,15 @@ impl WebSocketTransport { } } else if msg.is_ping() { // Automatically respond to ping with pong - if let Err(e) = stream.send(WsMessage::Pong(msg.into_data())).await { + let mut writer = writer_clone.lock().await; + if let Err(e) = writer.send(WsMessage::Pong(msg.into_data())).await { tracing::error!("Error sending pong: {}", e); } } else if msg.is_close() { tracing::debug!("WebSocket connection closed by server"); break; } - // Ignore pong messages, they're just confirmations of our pings + // Ignore pong messages, they're handled by the ping task } Err(err) => { tracing::error!("WebSocket read error: {}", err); @@ -287,107 +289,59 @@ impl WebSocketTransport { } async fn setup_ping_task( - writer: Arc>>>, + writer: Arc>, WsMessage>>>, ping_stop: Arc>, ping_config: PingConfig, ping_failures: Arc>, ) { - tokio::spawn(async move { - let ping_interval = Duration::from_secs(ping_config.interval_secs); - let ping_timeout = Duration::from_secs(ping_config.timeout_secs); + let ping_interval = Duration::from_secs(ping_config.interval_secs); - tracing::debug!( - "Starting WebSocket ping task with interval: {}s, timeout: {}s", - ping_config.interval_secs, - ping_config.timeout_secs - ); + tracing::debug!( + "Starting WebSocket ping task with interval: {}s", + ping_config.interval_secs + ); - loop { - // Sleep for the ping interval - tokio::time::sleep(ping_interval).await; + loop { + // Sleep for the ping interval + tokio::time::sleep(ping_interval).await; - // Check if we should stop - { - let stop = *ping_stop.lock().await; - if stop { - tracing::debug!("Stopping WebSocket ping task"); - break; - } + // Check if we should stop + { + let stop = *ping_stop.lock().await; + if stop { + tracing::debug!("Stopping WebSocket ping task"); + break; } + } - // Send a ping - let ping_payload = rand::random::().to_be_bytes().to_vec(); - tracing::trace!("Sending WebSocket ping with payload: {:?}", ping_payload); - - let ping_result = { - let mut writer_guard = writer.lock().await; - writer_guard - .send(WsMessage::Ping(ping_payload.clone().into())) - .await - }; - - if let Err(e) = ping_result { - tracing::warn!("Failed to send WebSocket ping: {}", e); - - let mut failures = ping_failures.lock().await; - *failures += 1; - - if *failures >= ping_config.max_failures { - tracing::error!( - "Maximum ping failures reached ({}). Connection considered lost.", - ping_config.max_failures - ); - break; - } - - continue; - } - - // Wait for pong with timeout - let pong_received = { - let mut writer_guard = writer.lock().await; - match tokio::time::timeout(ping_timeout, writer_guard.next()).await { - Ok(Some(Ok(msg))) => match msg { - WsMessage::Pong(payload) => { - tracing::trace!( - "Received WebSocket pong with payload: {:?}", - payload - ); - payload == ping_payload - } - _ => false, - }, - _ => false, - } - }; - - // Update ping failures counter - { - let mut failures = ping_failures.lock().await; - if pong_received { - // Reset counter on success - *failures = 0; - } else { - *failures += 1; - tracing::warn!( - "WebSocket ping failed. Consecutive failures: {}/{}", - *failures, - ping_config.max_failures - ); - - if *failures >= ping_config.max_failures { - tracing::error!( - "Maximum ping failures reached ({}). Connection considered lost.", - ping_config.max_failures - ); - break; - } - } + // Send a ping + let ping_payload = rand::random::().to_be_bytes().to_vec(); + tracing::trace!("Sending WebSocket ping with payload: {:?}", ping_payload); + + let ping_result = { + let mut writer_guard = writer.lock().await; + writer_guard + .send(WsMessage::Ping(ping_payload.clone().into())) + .await + }; + + if let Err(e) = ping_result { + tracing::warn!("Failed to send WebSocket ping: {}", e); + + let mut failures = ping_failures.lock().await; + *failures += 1; + + if *failures >= ping_config.max_failures { + tracing::error!( + "Maximum ping failures reached ({}). Connection considered lost.", + ping_config.max_failures + ); + break; } } + } - tracing::debug!("WebSocket ping task terminated"); - }); + tracing::debug!("WebSocket ping task terminated"); } } @@ -425,7 +379,7 @@ impl Transport for WebSocketTransport { // Then close the WebSocket connection with a proper close frame let mut writer = self.writer.lock().await; - match writer.close(None).await { + match writer.close().await { Ok(_) => { tracing::debug!("WebSocket connection closed successfully"); Ok(()) From ca0c92dbed2322f485e21135b4351a37a6bda209 Mon Sep 17 00:00:00 2001 From: Gautam Korlam Date: Sat, 15 Mar 2025 00:38:45 -0700 Subject: [PATCH 08/10] debug logs --- src/client/mod.rs | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 4f9d430..f613a88 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,10 +1,10 @@ +use futures::Stream; use futures::StreamExt; use std::{pin::Pin, sync::Arc}; use tempfile::NamedTempFile; use tokio::process::Child; use tokio::sync::{Mutex, RwLock}; use tokio::time::{timeout, Duration}; -use futures::Stream; use crate::{ error::{Error, ErrorCode}, @@ -127,7 +127,7 @@ impl Client { *state = ConnectionState::Connected; } - tracing::info!(?implementation, "Initializing MCP client"); + tracing::debug!(?implementation, "Initializing MCP client"); let params = serde_json::json!({ "clientInfo": implementation, @@ -151,12 +151,12 @@ impl Client { { let mut initialized = self.initialized.write().await; *initialized = true; - + let mut state = self.connection_state.write().await; *state = ConnectionState::Initialized; } - tracing::info!("MCP client initialization complete"); + tracing::debug!("MCP client initialization complete"); Ok(init_result) } @@ -293,16 +293,18 @@ impl Client { // Wait for shutdown/ack notification with timeout let timeout_duration = Duration::from_secs(5); - match tokio::time::timeout( - timeout_duration, - self.wait_for_notification("shutdown/ack"), - ).await { + match tokio::time::timeout(timeout_duration, self.wait_for_notification("shutdown/ack")) + .await + { Ok(result) => { tracing::debug!("Received shutdown/ack notification"); result - }, + } Err(_) => { - tracing::warn!("Shutdown ack notification timed out after {:?}", timeout_duration); + tracing::warn!( + "Shutdown ack notification timed out after {:?}", + timeout_duration + ); // Continue with cleanup even without ack Ok(()) } @@ -314,7 +316,7 @@ impl Client { *initialized = false; } - tracing::info!("MCP client shutdown complete"); + tracing::debug!("MCP client shutdown complete"); Ok(()) } @@ -333,11 +335,11 @@ impl Client { { let mut state = self.connection_state.write().await; *state = ConnectionState::Disconnected; - + let mut initialized = self.initialized.write().await; *initialized = false; } - + // Close the transport self.transport.close().await } @@ -459,7 +461,7 @@ impl Drop for Client { if let Err(e) = transport.close().await { tracing::error!("Error during transport close in drop: {e}"); } - + // If there's a subprocess, attempt to kill it if let Some(child) = &mut subprocess { let _ = child.start_kill(); From c157b15e00514d0a663d2c339adbcef785c419d6 Mon Sep 17 00:00:00 2001 From: Eashan Hatti Date: Sun, 3 Aug 2025 20:59:27 -0700 Subject: [PATCH 09/10] Fix read --- src/transport/stdio.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transport/stdio.rs b/src/transport/stdio.rs index 2fab6b8..145659f 100644 --- a/src/transport/stdio.rs +++ b/src/transport/stdio.rs @@ -59,7 +59,10 @@ where } let message = match serde_json::from_str::(trimmed) { Ok(m) => Ok(m), - Err(err) => Err(Error::Serialization(err.to_string())), + Err(err) => { + tracing::warn!(?err, ?trimmed, "Message read error"); + continue + }, }; let _ = sender_clone.send(message); From a5d2051d1ce3670dd0bb4652b92580b6e8c4be5b Mon Sep 17 00:00:00 2001 From: Julien Camus Date: Fri, 7 Aug 2026 10:01:33 +0200 Subject: [PATCH 10/10] feat(client): make the initialize handshake timeout configurable request() applies the same 30s ceiling to every JSON-RPC round trip, including initialize. That's fine for steady-state calls (list_tools, call_tool) against a server that's already running, but it also caps how long a command is given to start in the first place -- and for commands like `npx`/`uvx`, "start" can include downloading and installing a package on first run. A consumer with a slow registry or a large package has no way to give that one-time cost more room without also loosening detection of a genuinely hung steady-state call. Split the two: `request_with_timeout` takes an explicit timeout and `request` keeps calling it with the existing 30s default, so no caller of `request`, `list_tools`, or `call_tool` sees a behavior change. `initialize_with_timeout` and `ClientBuilder::handshake_timeout` do the same for the handshake specifically -- `spawn_and_initialize` uses the override when set, or the unchanged 30s default otherwise. --- src/client/builder.rs | 33 +++++++++++++++++++++-- src/client/mod.rs | 62 ++++++++++++++++++++++++++++++++++++------- src/client/test.rs | 48 +++++++++++++++++++++++++++++++++ 3 files changed, 132 insertions(+), 11 deletions(-) diff --git a/src/client/builder.rs b/src/client/builder.rs index f93d8d3..ab77489 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -6,6 +6,7 @@ use std::collections::HashMap; use std::path::PathBuf; use std::process::Stdio; use std::sync::Arc; +use std::time::Duration; use tempfile::NamedTempFile; use tokio::process::Command; @@ -24,6 +25,10 @@ pub struct ClientBuilder { capabilities: Option, /// Environment variables for the subprocess. env: HashMap, + /// How long `spawn_and_initialize` waits for the `initialize` response. Defaults to 30 + /// seconds (see [`Client::initialize`]) when not set — override for commands whose first + /// run needs to install something (`npx`, `uvx`) before they can speak the protocol. + handshake_timeout: Option, } impl ClientBuilder { @@ -36,9 +41,20 @@ impl ClientBuilder { implementation: None, capabilities: None, env: HashMap::new(), + handshake_timeout: None, } } + /// Overrides how long `spawn_and_initialize` waits for the `initialize` response. + /// + /// Only the handshake is affected — requests made after the server is up + /// (`list_tools`, `call_tool`) keep their own independent timeout. + pub fn handshake_timeout(mut self, timeout: Duration) -> Self { + tracing::trace!(?timeout, "Setting handshake timeout for ClientBuilder"); + self.handshake_timeout = Some(timeout); + self + } + pub fn arg(mut self, arg: &str) -> Self { tracing::trace!(%arg, "Adding argument to ClientBuilder"); self.args.push(arg.to_string()); @@ -181,10 +197,23 @@ impl ClientBuilder { /// /// Returns an error if the command cannot be spawned, or if initialization fails. pub async fn spawn_and_initialize(self) -> Result { + let handshake_timeout = self.handshake_timeout; let (mut client, implementation, capabilities) = self.spawn().await?; - tracing::debug!(?implementation, ?capabilities, "Initializing client"); - client.initialize(implementation, capabilities).await?; + tracing::debug!( + ?implementation, + ?capabilities, + ?handshake_timeout, + "Initializing client" + ); + match handshake_timeout { + Some(timeout) => { + client + .initialize_with_timeout(implementation, capabilities, timeout) + .await? + } + None => client.initialize(implementation, capabilities).await?, + }; tracing::info!("MCP client successfully spawned and initialized"); Ok(client) diff --git a/src/client/mod.rs b/src/client/mod.rs index f613a88..2bc7761 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -106,17 +106,38 @@ impl Client { client } + /// Initializes the client, waiting up to 30 seconds for the server's response. + /// + /// Use [`Client::initialize_with_timeout`] for servers whose startup command needs + /// longer than that — e.g. `npx`/`uvx` installing a package on first run. + pub async fn initialize( + &mut self, + implementation: Implementation, + capabilities: ClientCapabilities, + ) -> Result { + self.initialize_with_timeout(implementation, capabilities, Duration::from_secs(30)) + .await + } + /// Initializes the client by sending an "initialize" request containing: /// - client implementation info /// - client capabilities /// - protocol version /// + /// `timeout` bounds only this handshake — it does not change the 30-second timeout + /// applied to requests made after initialization (`list_tools`, `call_tool`, etc. via + /// [`Client::request`]). Startup cost (installing a package, cold-starting a runtime) + /// is a one-time cost that can legitimately run longer than a single steady-state call + /// should be allowed to hang for; keeping the two independent means a slow install + /// doesn't have to buy every later call room to hang too. + /// /// On success, updates the client's `server_capabilities` field and sends an /// `initialized` notification to the server. - pub async fn initialize( + pub async fn initialize_with_timeout( &mut self, implementation: Implementation, capabilities: ClientCapabilities, + timeout: Duration, ) -> Result { // Set connection state to Connected { @@ -127,7 +148,7 @@ impl Client { *state = ConnectionState::Connected; } - tracing::debug!(?implementation, "Initializing MCP client"); + tracing::debug!(?implementation, ?timeout, "Initializing MCP client"); let params = serde_json::json!({ "clientInfo": implementation, @@ -135,7 +156,9 @@ impl Client { "protocolVersion": crate::LATEST_PROTOCOL_VERSION, }); - let response = self.request("initialize", Some(params)).await?; + let response = self + .request_with_timeout("initialize", Some(params), timeout) + .await?; let init_result: InitializeResult = serde_json::from_value(response)?; tracing::debug!(?init_result, "Received initialization response"); @@ -171,6 +194,23 @@ impl Client { &mut self, method: &str, params: Option, + ) -> Result { + self.request_with_timeout(method, params, Duration::from_secs(30)) + .await + } + + /// Sends a request to the server with the given method and optional parameters, + /// then waits up to `timeout` for a matching response. + /// + /// # Errors + /// + /// Returns an error if the transport fails, the server returns an error, + /// or no response is received within `timeout`. + pub async fn request_with_timeout( + &mut self, + method: &str, + params: Option, + timeout: Duration, ) -> Result { // Increment request ID let mut counter = self.request_counter.write().await; @@ -178,14 +218,18 @@ impl Client { let id = RequestId::Number(*counter); let request = Request::new(method, params, id.clone()); - tracing::debug!(?request, "Sending MCP request"); + tracing::debug!(?request, ?timeout, "Sending MCP request"); // Send request self.transport.send(Message::Request(request)).await?; - // Wait for a matching response (by request ID) or a 30s timeout + // Wait for a matching response (by request ID) or the timeout let mut rx = self.response_receiver.lock().await; + // Poll process liveness in increments no larger than 300ms, so a process that exits + // early is still caught promptly even when `timeout` is large. + let liveness_checks = (timeout.as_millis() / 300).max(1) as u32; + tokio::select! { // Branch 1: Handle the message receiving logic result = async { @@ -227,9 +271,9 @@ impl Client { )) } => result, - // Branch 2: Periodically check if the process is still alive, or timeout after 30s + // Branch 2: Periodically check if the process is still alive, or time out result = async { - for _ in 1..=100 { + for _ in 1..=liveness_checks { tokio::time::sleep(Duration::from_millis(300)).await; if let Some(process) = &mut self.subprocess { @@ -245,9 +289,9 @@ impl Client { } } - tracing::error!("Request to '{}' timed out after 30 seconds", method); + tracing::error!("Request to '{}' timed out after {:?}", method, timeout); Err(Error::Other(format!( - "Request to '{method}' timed out after 30 seconds" + "Request to '{method}' timed out after {timeout:?}" ))) } => result, } diff --git a/src/client/test.rs b/src/client/test.rs index 0c113b4..db900a3 100644 --- a/src/client/test.rs +++ b/src/client/test.rs @@ -274,3 +274,51 @@ async fn test_set_log_level() -> Result<(), Error> { } Ok(()) } + +/// `handshake_timeout` bounds only the `initialize` wait, and is independent of the +/// steady-state per-request timeout: `sleep` never speaks the protocol, so the process stays +/// alive and never exits, which means the only way this returns in well under 30 seconds is +/// if the override is actually being used instead of the default. +#[tokio::test] +async fn test_handshake_timeout_overrides_the_default() { + let start = std::time::Instant::now(); + + let result = ClientBuilder::new("sleep") + .arg("100") + .handshake_timeout(std::time::Duration::from_millis(500)) + .spawn_and_initialize() + .await; + + let elapsed = start.elapsed(); + assert!(result.is_err(), "sleep never responds, so this must fail"); + assert!( + elapsed < std::time::Duration::from_secs(5), + "expected the 500ms override to apply, took {elapsed:?}" + ); +} + +/// Without an explicit override, `request()` keeps its historical 30-second default — +/// callers that never opt into a custom timeout must see no behavior change. `sleep` never +/// responds and never exits within our 2-second probe window, so a plain `request()` call +/// should still be pending — unlike the 500ms-override test above, which resolves quickly. +#[tokio::test] +async fn test_request_default_timeout_is_unchanged_by_handshake_timeout() { + let (mut client, _implementation, _capabilities) = ClientBuilder::new("sleep") + .arg("100") + .spawn() + .await + .expect("spawning sleep should not fail"); + + let start = std::time::Instant::now(); + let result = tokio::time::timeout( + std::time::Duration::from_secs(2), + client.request("ping", None), + ) + .await; + + assert!( + result.is_err(), + "request() should still be waiting well before the 30s default, elapsed {:?}", + start.elapsed() + ); +}