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.lock b/Cargo.lock index 18f48d4..22322bf 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,33 @@ 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 = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] [[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 = "cc" +version = "1.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be714c154be609ec7f5dad223a33bf1482fff90472de28f7362806e6d4832b8c" +dependencies = [ + "shlex", +] [[package]] name = "cfg-if" @@ -89,6 +107,68 @@ 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" @@ -96,7 +176,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.59.0", ] [[package]] @@ -105,6 +185,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 +304,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,23 +343,185 @@ version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +[[package]] +name = "http" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4a85d31aea989eead29a3aaf9e1115a180df8282431156e533de47660892565" +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.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.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe7db12097d22ec582439daf8618b8fdd1a7bef6270e9af3b1ebcd30893cf413" + +[[package]] +name = "litemap" +version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d26c52dbd32dccf2d10cac7725f8eae5296885fb5703b261f7d0a0739ec807ab" +checksum = "23fb14cb19457329c82206317a5663005a4d404783dc74f4252769b0d5f42856" [[package]] name = "lock_api" @@ -240,6 +533,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 +546,17 @@ dependencies = [ "async-stream", "async-trait", "futures", + "futures-util", + "http", + "rand 0.8.5", "serde", "serde_json", "tempfile", "thiserror", "tokio", + "tokio-tungstenite", "tracing", + "url", ] [[package]] @@ -263,9 +567,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 +582,84 @@ checksum = "2886843bf800fba2e3377cff24abf6379b4c4d5c6681eaf9ea5b0d15090450bd" dependencies = [ "libc", "wasi 0.11.0+wasi-snapshot-preview1", - "windows-sys", + "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.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 = "d75b0bedcc4fe52caa0e03d9f1151a323e4aa5e2d78ba3580400cd3c9e2bc4bc" + +[[package]] +name = "openssl" +version = "0.10.71" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1261fe7e33c73b354eab43b1273a57c8f967d0391e80353e51f764ac02cf6775" +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" @@ -319,11 +684,17 @@ 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.15" +version = "0.2.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "915a1e146535de9163f3987b8944ed8cf49a18bb0056bcebcdcece385cece4ff" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" [[package]] name = "pin-utils" @@ -331,29 +702,104 @@ 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.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 = "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 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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "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]] +name = "rand_core" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38" +dependencies = [ + "getrandom 0.3.1", +] + [[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 +812,31 @@ 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 = "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" @@ -389,20 +844,43 @@ 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.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 +889,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", @@ -421,6 +899,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" @@ -441,9 +936,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,59 +947,85 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c970269d99b64e60ec3bd6ad27270092a5394c4e309314b18ae3fe575695fbe8" dependencies = [ "libc", - "windows-sys", + "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.90" +version = "2.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "919d3b74a5dd0ccd15aeb8f93e7006bd9e14c295087c9896a110f490752bcf31" +checksum = "b09a44accad81e1ba1cd74a32461ba89dee89095ba17b32f5d03683b1b1fc2a0" dependencies = [ "proc-macro2", "quote", "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.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", + "getrandom 0.3.1", "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", "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.42.0" +version = "1.44.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cec9b21b0450273377fc97bd4c33a8acffc8c996c987a7c5b319a0083707551" +checksum = "f382da615b842244d4b8738c82ed1275e6c5dd90c459a30941cd07080b06c91a" dependencies = [ "backtrace", "bytes", @@ -515,20 +1036,44 @@ 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", "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.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a9daff607c6d2bf6c16fd681ccb7eecc83e4e2cdc1ca067ffaadfca5de7f084" +dependencies = [ + "futures-util", + "log", + "native-tls", + "tokio", + "tokio-native-tls", + "tungstenite", +] + [[package]] name = "tracing" version = "0.1.41" @@ -560,11 +1105,76 @@ dependencies = [ "once_cell", ] +[[package]] +name = "tungstenite" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4793cb5e56680ecbb1d843515b23b6de9a75eb04b66643e256a396d43be33c13" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "native-tls", + "rand 0.9.0", + "sha1", + "thiserror", + "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.14" +version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "adb9e6ca4f869e1180728b7950e35922a7fc6397f7b641499e8f3ef06e50dc83" +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" @@ -590,6 +1200,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" @@ -662,3 +1281,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 609052c..5f4d5cd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,18 +4,26 @@ 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" +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/builder.rs b/src/client/builder.rs index 6beaf44..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()); @@ -90,8 +106,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 +138,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)); @@ -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 4df36d4..2bc7761 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,21 +1,20 @@ +use futures::Stream; use futures::StreamExt; -use serde_json::Value; -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::{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,27 +35,47 @@ 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 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 { /// 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, + 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. @@ -87,19 +106,49 @@ 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 { - tracing::info!(?implementation, "Initializing MCP client"); + // 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::debug!(?implementation, ?timeout, "Initializing MCP client"); let params = serde_json::json!({ "clientInfo": implementation, @@ -107,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"); @@ -119,7 +170,16 @@ impl Client { tracing::debug!("Sending initialized notification"); self.notify("notifications/initialized", None).await?; - tracing::info!("MCP client initialization complete"); + // Update initialization state + { + let mut initialized = self.initialized.write().await; + *initialized = true; + + let mut state = self.connection_state.write().await; + *state = ConnectionState::Initialized; + } + + tracing::debug!("MCP client initialization complete"); Ok(init_result) } @@ -134,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; @@ -141,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 { @@ -189,12 +270,12 @@ impl Client { "Connection closed while waiting for response", )) } => 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 { match process.try_wait() { Ok(None) => continue, @@ -207,10 +288,10 @@ 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, } @@ -237,34 +318,76 @@ 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 - } + /// 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(()); + } + *state = ConnectionState::ShuttingDown; + } - 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; + 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 } - 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; + 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::debug!("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"); @@ -350,20 +473,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 +499,53 @@ 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}"); + // 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; } }); } -} \ No newline at end of file +} + +// 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/client/test.rs b/src/client/test.rs index f96b9f1..db900a3 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( @@ -277,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() + ); +} 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/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..02ff300 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; @@ -159,3 +149,4 @@ pub trait Transport: Send + Sync + 'static { } pub mod stdio; +pub mod websocket; diff --git a/src/transport/stdio.rs b/src/transport/stdio.rs index 48acccb..145659f 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}, }; @@ -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); diff --git a/src/transport/websocket.rs b/src/transport/websocket.rs new file mode 100644 index 0000000..43b7264 --- /dev/null +++ b/src/transport/websocket.rs @@ -0,0 +1,394 @@ +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::{ + 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; + +/// 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. + writer: Arc>, WsMessage>>>, + /// 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>, + /// Configuration for ping behavior + ping_config: PingConfig, +} + +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 { + 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(); + + // 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)))?; + + // 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)))?; + + // 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); + + // 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(); + let sender_clone = sender.clone(); + tokio::spawn(async move { + tracing::debug!("Starting WebSocket reader task"); + + while let Some(result) = reader.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 + 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 handled by the ping task + } + 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 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 { + Self::setup_ping_task( + writer_for_ping, + ping_stop_clone, + ping_config_clone, + ping_failures_clone, + ) + .await; + }); + + Ok(WebSocketTransport { + writer, + receiver, + _sender: sender, + ping_stop, + ping_config, + }) + } + + /// 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 + } + + /// 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>, WsMessage>>>, + ping_stop: Arc>, + ping_config: PingConfig, + ping_failures: Arc>, + ) { + let ping_interval = Duration::from_secs(ping_config.interval_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; + + // 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; + } + } + } + + tracing::debug!("WebSocket ping task terminated"); + } +} + +#[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.into())) + .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 with a proper close frame + let mut writer = self.writer.lock().await; + match writer.close().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(()) + } + } + } +}