diff --git a/.release-plz.toml b/.release-plz.toml index 21ce066..8d368cf 100644 --- a/.release-plz.toml +++ b/.release-plz.toml @@ -1,6 +1,53 @@ [workspace] git_release_type = "auto" +# Only release when merging the release PR (avoids race conditions with squash-merge) +release_always = false + +# Keep dependencies updated in Cargo.lock +dependencies_update = true + +# Tag release PRs for visibility +pr_labels = ["release"] + [[package]] name = "agent-client-protocol" git_tag_name = "v{{ version }}" + +[[package]] +name = "agent-client-protocol-conductor" +git_tag_name = "agent-client-protocol-conductor-v{{ version }}" + +[[package]] +name = "agent-client-protocol-cookbook" +git_tag_name = "agent-client-protocol-cookbook-v{{ version }}" + +[[package]] +name = "agent-client-protocol-core" +git_tag_name = "agent-client-protocol-core-v{{ version }}" + +[[package]] +name = "agent-client-protocol-derive" +git_tag_name = "agent-client-protocol-derive-v{{ version }}" + +[[package]] +name = "agent-client-protocol-rmcp" +git_tag_name = "agent-client-protocol-rmcp-v{{ version }}" + +[[package]] +name = "agent-client-protocol-test" +git_tag_name = "agent-client-protocol-test-v{{ version }}" +# Don't publish test utilities to crates.io +publish = false + +[[package]] +name = "agent-client-protocol-tokio" +git_tag_name = "agent-client-protocol-tokio-v{{ version }}" + +[[package]] +name = "agent-client-protocol-trace-viewer" +git_tag_name = "agent-client-protocol-trace-viewer-v{{ version }}" + +[[package]] +name = "agent-client-protocol-yopo" +git_tag_name = "agent-client-protocol-yopo-v{{ version }}" diff --git a/Cargo.lock b/Cargo.lock index a0eaa72..85dcffe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -24,6 +24,111 @@ dependencies = [ "tokio-util", ] +[[package]] +name = "agent-client-protocol-conductor" +version = "0.1.0" +dependencies = [ + "agent-client-protocol-core", + "agent-client-protocol-schema", + "agent-client-protocol-test", + "agent-client-protocol-tokio", + "agent-client-protocol-trace-viewer", + "agent-client-protocol-yopo", + "anyhow", + "async-stream", + "axum", + "chrono", + "clap", + "expect-test", + "futures", + "futures-concurrency", + "hyper", + "regex", + "rmcp", + "rustc-hash", + "schemars", + "serde", + "serde_json", + "shell-words", + "strip-ansi-escapes", + "thiserror", + "tokio", + "tokio-util", + "tower", + "tracing", + "tracing-appender", + "tracing-subscriber", + "uuid", +] + +[[package]] +name = "agent-client-protocol-cookbook" +version = "0.1.0" +dependencies = [ + "agent-client-protocol-conductor", + "agent-client-protocol-core", + "agent-client-protocol-rmcp", + "agent-client-protocol-tokio", + "rmcp", + "schemars", + "serde", + "serde_json", + "tokio", + "tracing", +] + +[[package]] +name = "agent-client-protocol-core" +version = "0.1.0" +dependencies = [ + "agent-client-protocol-derive", + "agent-client-protocol-schema", + "agent-client-protocol-test", + "anyhow", + "boxfnonce", + "clap", + "expect-test", + "futures", + "futures-concurrency", + "jsonrpcmsg", + "rmcp", + "rustc-hash", + "schemars", + "serde", + "serde_json", + "shell-words", + "thiserror", + "tokio", + "tokio-util", + "tracing", + "uuid", +] + +[[package]] +name = "agent-client-protocol-derive" +version = "0.1.0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "agent-client-protocol-rmcp" +version = "0.1.0" +dependencies = [ + "agent-client-protocol-core", + "futures", + "futures-concurrency", + "rmcp", + "schemars", + "serde", + "tokio", + "tokio-util", + "tracing", + "tracing-subscriber", +] + [[package]] name = "agent-client-protocol-schema" version = "0.11.4" @@ -38,6 +143,68 @@ dependencies = [ "strum", ] +[[package]] +name = "agent-client-protocol-test" +version = "0.1.0" +dependencies = [ + "agent-client-protocol-core", + "agent-client-protocol-tokio", + "agent-client-protocol-yopo", + "anyhow", + "futures", + "rmcp", + "schemars", + "serde", + "serde_json", + "tokio", + "tokio-util", + "tracing", + "tracing-subscriber", + "uuid", +] + +[[package]] +name = "agent-client-protocol-tokio" +version = "0.1.0" +dependencies = [ + "agent-client-protocol-core", + "agent-client-protocol-test", + "expect-test", + "futures", + "serde", + "serde_json", + "shell-words", + "tokio", + "tokio-util", +] + +[[package]] +name = "agent-client-protocol-trace-viewer" +version = "0.1.0" +dependencies = [ + "anyhow", + "axum", + "clap", + "open", + "serde", + "serde_json", + "tokio", + "tower", + "tower-http", +] + +[[package]] +name = "agent-client-protocol-yopo" +version = "0.1.0" +dependencies = [ + "agent-client-protocol-core", + "agent-client-protocol-tokio", + "clap", + "tokio", + "tracing", + "tracing-subscriber", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -47,6 +214,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anstream" version = "1.0.0" @@ -115,6 +291,28 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -132,18 +330,104 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "axum" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8" +dependencies = [ + "axum-core", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + [[package]] name = "bitflags" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +[[package]] +name = "boxfnonce" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5988cb1d626264ac94100be357308f29ff7cbdd3b36bda27f450a4ee3f713426" + +[[package]] +name = "bumpalo" +version = "3.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" + [[package]] name = "bytes" version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +[[package]] +name = "cc" +version = "1.2.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1e928d4b69e3077709075a938a05ffbedfa53a84c8f766efbf8220bb1ff60e1" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -156,6 +440,60 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chrono" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "clap" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b193af5b67834b676abd72466a96c1024e6a6ad978a1f484bd90b85c94041351" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1110bd8a634a1ab8cb04345d8d878267d57c3cf1b38d91b71af6686408bbca6a" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "clipboard-win" version = "5.4.1" @@ -189,12 +527,70 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "darling" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25ae13da2f202d56bd7f91c25fba009e7717a1e4a1cc98a76d844b65ae912e9d" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9865a50f7c335f53564bb694ef660825eb8610e0a53d3e11bf1b0d3df31e03b0" +dependencies = [ + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" +dependencies = [ + "darling_core", + "quote", + "syn", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "powerfmt", +] + [[package]] name = "derive_more" version = "2.1.1" @@ -224,6 +620,23 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" +[[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 = "dissimilar" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aeda16ab4059c5fd2a83f2b9c9e9c981327b18aa8e3b313f7e6563799d4f093e" + [[package]] name = "dyn-clone" version = "1.0.20" @@ -259,6 +672,12 @@ dependencies = [ "log", ] +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + [[package]] name = "errno" version = "0.3.14" @@ -296,6 +715,16 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "expect-test" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63af43ff4431e848fb47472a920f14fa71c24de13255a5692e93d4e90302acb0" +dependencies = [ + "dissimilar", + "once_cell", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -313,6 +742,33 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fixedbitset" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d674e81391d1e1ab681a28d99df07927c6d4aa5b027d7da16ba32d1d21ecd99" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + [[package]] name = "futures" version = "0.3.32" @@ -338,6 +794,19 @@ dependencies = [ "futures-sink", ] +[[package]] +name = "futures-concurrency" +version = "7.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "175cd8cca9e1d45b87f18ffa75088f2099e3c4fe5e2f83e42de112560bea8ea6" +dependencies = [ + "fixedbitset", + "futures-core", + "futures-lite", + "pin-project", + "smallvec", +] + [[package]] name = "futures-core" version = "0.3.32" @@ -361,6 +830,19 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" +[[package]] +name = "futures-lite" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad" +dependencies = [ + "fastrand", + "futures-core", + "futures-io", + "parking", + "pin-project-lite", +] + [[package]] name = "futures-macro" version = "0.3.32" @@ -401,6 +883,34 @@ dependencies = [ "slab", ] +[[package]] +name = "getrandom" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", + "wasip3", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" + [[package]] name = "heck" version = "0.5.0" @@ -416,6 +926,286 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "http" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "http-range-header" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9171a2ea8a68358193d15dd5d70c1c10a2afc3e7e4c5bc92bc9f025cebd7359c" + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6299f016b246a94207e63da54dbe807655bf9e00044f73ded42c3ac5305fbcca" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "icu_collections" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43" +dependencies = [ + "displaydoc", + "potential_utf", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" + +[[package]] +name = "icu_properties" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af" + +[[package]] +name = "icu_provider" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" +dependencies = [ + "equivalent", + "hashbrown 0.16.1", + "serde", + "serde_core", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "iri-string" +version = "0.7.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25e659a4bb38e810ebc252e53b5814ff908a8c58c2a9ce2fae1bbec24cbf4e20" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "is-docker" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "928bae27f42bc99b60d9ac7334e3a21d10ad8f1835a4e12ec3ec0464765ed1b3" +dependencies = [ + "once_cell", +] + +[[package]] +name = "is-wsl" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "173609498df190136aa7dea1a91db051746d339e18476eed5ca40521f02d7aa5" +dependencies = [ + "is-docker", + "once_cell", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -452,6 +1242,40 @@ dependencies = [ "syn", ] +[[package]] +name = "js-sys" +version = "0.3.94" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e04e2ef80ce82e13552136fabeef8a5ed1f985a96805761cbb9a2c34e7664d9" +dependencies = [ + "cfg-if", + "futures-util", + "once_cell", + "wasm-bindgen", +] + +[[package]] +name = "jsonrpcmsg" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d833a15225c779251e13929203518c2ff26e2fe0f322d584b213f4f4dad37bd" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + [[package]] name = "libc" version = "0.2.183" @@ -464,6 +1288,12 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" +[[package]] +name = "litemap" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" + [[package]] name = "lock_api" version = "0.4.14" @@ -479,12 +1309,43 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + [[package]] name = "memchr" version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "mime_guess" +version = "2.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7c44f8e672c00fe5308fa235f821cb4198414e1c77935c1ab6948d3fd78550e" +dependencies = [ + "mime", + "unicase", +] + [[package]] name = "mio" version = "1.2.0" @@ -517,12 +1378,65 @@ dependencies = [ "libc", ] +[[package]] +name = "nix" +version = "0.31.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "num-conv" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + [[package]] name = "once_cell_polyfill" version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" +[[package]] +name = "open" +version = "5.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43bb73a7fa3799b198970490a51174027ba0d4ec504b03cd08caf513d40024bc" +dependencies = [ + "is-wsl", + "libc", + "pathdiff", +] + [[package]] name = "parking" version = "2.2.1" @@ -553,8 +1467,46 @@ dependencies = [ ] [[package]] -name = "pin-project-lite" -version = "0.2.17" +name = "pastey" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b867cad97c0791bbd3aaa6472142568c6c9e8f71937e98379f584cfb0cf35bec" + +[[package]] +name = "pathdiff" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df94ce210e5bc13cb6651479fa48d14f601d9858cfe0467f43ae157023b938d3" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" @@ -584,6 +1536,21 @@ dependencies = [ "portable-atomic", ] +[[package]] +name = "potential_utf" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77" +dependencies = [ + "zerovec", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "pretty_assertions" version = "1.4.1" @@ -594,6 +1561,16 @@ dependencies = [ "yansi", ] +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -603,6 +1580,20 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "process-wrap" +version = "9.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e842efad9119158434d193c6682e2ebee4b44d6ad801d7b349623b3f57cdf55" +dependencies = [ + "futures", + "indexmap", + "nix 0.31.2", + "tokio", + "tracing", + "windows", +] + [[package]] name = "quote" version = "1.0.45" @@ -612,6 +1603,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "radix_trie" version = "0.2.1" @@ -680,6 +1677,86 @@ version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +[[package]] +name = "reqwest" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab3f43e3283ab1488b624b44b0e988d0acea0b3214e694730a055cb6b2efa801" +dependencies = [ + "base64", + "bytes", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "serde", + "serde_json", + "sync_wrapper", + "tokio", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + +[[package]] +name = "rmcp" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2231b2c085b371c01bc90c0e6c1cab8834711b6394533375bdbf870b0166d419" +dependencies = [ + "async-trait", + "base64", + "chrono", + "futures", + "http", + "pastey", + "pin-project-lite", + "process-wrap", + "reqwest", + "rmcp-macros", + "schemars", + "serde", + "serde_json", + "sse-stream", + "thiserror", + "tokio", + "tokio-stream", + "tokio-util", + "tracing", +] + +[[package]] +name = "rmcp-macros" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36ea0e100fadf81be85d7ff70f86cd805c7572601d4ab2946207f36540854b43" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "serde_json", + "syn", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + [[package]] name = "rustc_version" version = "0.4.1" @@ -702,6 +1779,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + [[package]] name = "rustyline" version = "17.0.2" @@ -716,7 +1799,7 @@ dependencies = [ "libc", "log", "memchr", - "nix", + "nix 0.30.1", "radix_trie", "unicode-segmentation", "unicode-width", @@ -724,12 +1807,19 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "schemars" version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" dependencies = [ + "chrono", "dyn-clone", "ref-cast", "schemars_derive", @@ -756,151 +1846,774 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" [[package]] -name = "semver" -version = "1.0.27" +name = "semver" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shell-words" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc6fe69c597f9c37bfeeeeeb33da3530379845f10be461a66d16d03eca2ded77" + +[[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.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + +[[package]] +name = "socket2" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "sse-stream" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb4dc4d33c68ec1f27d386b5610a351922656e1fdf5c05bbaad930cd1519479a" +dependencies = [ + "bytes", + "futures-util", + "http-body", + "http-body-util", + "pin-project-lite", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "strip-ansi-escapes" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a8f8038e7e7969abb3f1b7c2a811225e9296da208539e0f79c5251d6cac0025" +dependencies = [ + "vte", +] + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "time" +version = "0.3.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" + +[[package]] +name = "time-macros" +version = "0.2.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215" +dependencies = [ + "num-conv", + "time-core", +] + +[[package]] +name = "tinystr" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tokio" +version = "1.50.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +dependencies = [ + "bytes", + "libc", + "mio", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-io", + "futures-sink", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" +dependencies = [ + "bitflags", + "bytes", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "http-range-header", + "httpdate", + "iri-string", + "mime", + "mime_guess", + "percent-encoding", + "pin-project-lite", + "tokio", + "tokio-util", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-appender" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "786d480bce6247ab75f005b14ae1624ad978d3029d9113f0a22fa1ac773faeaf" +dependencies = [ + "crossbeam-channel", + "thiserror", + "time", + "tracing-subscriber", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-serde" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704b1aeb7be0d0a84fc9828cae51dab5970fee5088f83d1dd7ee6f6246fc6ff1" +dependencies = [ + "serde", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "serde", + "serde_json", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", + "tracing-serde", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "unicase" +version = "2.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-segmentation" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "uuid" +version = "1.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ac8b6f42ead25368cf5b098aeb3dc8a1a2c05a3eee8a9a1a68c640edbfc79d9" +dependencies = [ + "getrandom", + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "vte" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "231fdcd7ef3037e8330d8e17e61011a2c244126acc0a982f4040ac3f9f0bc077" +dependencies = [ + "memchr", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.2+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" +checksum = "0551fc1bb415591e3372d0bc4780db7e587d84e2a7e79da121051c5c4b89d0b0" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] [[package]] -name = "serde" -version = "1.0.228" +name = "wasm-bindgen-futures" +version = "0.4.67" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "03623de6905b7206edd0a75f69f747f134b7f0a2323392d664448bf2d3c5d87e" dependencies = [ - "serde_core", - "serde_derive", + "js-sys", + "wasm-bindgen", ] [[package]] -name = "serde_core" -version = "1.0.228" +name = "wasm-bindgen-macro" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "7fbdf9a35adf44786aecd5ff89b4563a90325f9da0923236f6104e603c7e86be" dependencies = [ - "serde_derive", + "quote", + "wasm-bindgen-macro-support", ] [[package]] -name = "serde_derive" -version = "1.0.228" +name = "wasm-bindgen-macro-support" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "dca9693ef2bab6d4e6707234500350d8dad079eb508dca05530c85dc3a529ff2" dependencies = [ + "bumpalo", "proc-macro2", "quote", "syn", + "wasm-bindgen-shared", ] [[package]] -name = "serde_derive_internals" -version = "0.29.1" +name = "wasm-bindgen-shared" +version = "0.2.117" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +checksum = "39129a682a6d2d841b6c429d0c51e5cb0ed1a03829d8b3d1e69a011e62cb3d3b" dependencies = [ - "proc-macro2", - "quote", - "syn", + "unicode-ident", ] [[package]] -name = "serde_json" -version = "1.0.149" +name = "wasm-encoder" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" dependencies = [ - "itoa", - "memchr", - "serde", - "serde_core", - "zmij", + "leb128fmt", + "wasmparser", ] [[package]] -name = "signal-hook-registry" -version = "1.4.8" +name = "wasm-metadata" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" dependencies = [ - "errno", - "libc", + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser", ] [[package]] -name = "slab" -version = "0.4.12" +name = "wasm-streams" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] [[package]] -name = "smallvec" -version = "1.15.1" +name = "wasmparser" +version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", +] [[package]] -name = "socket2" -version = "0.6.3" +name = "web-sys" +version = "0.3.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" +checksum = "cd70027e39b12f0849461e08ffc50b9cd7688d942c1c8e3c7b22273236b4dd0a" dependencies = [ - "libc", - "windows-sys 0.61.2", + "js-sys", + "wasm-bindgen", ] [[package]] -name = "strum" -version = "0.28.0" +name = "windows" +version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" dependencies = [ - "strum_macros", + "windows-collections", + "windows-core", + "windows-future", + "windows-numerics", ] [[package]] -name = "strum_macros" -version = "0.28.0" +name = "windows-collections" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab85eea0270ee17587ed4156089e10b9e6880ee688791d45a905f5b1ca36f664" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn", + "windows-core", ] [[package]] -name = "syn" -version = "2.0.117" +name = "windows-core" +version = "0.62.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ - "proc-macro2", - "quote", - "unicode-ident", + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", ] [[package]] -name = "tokio" -version = "1.50.0" +name = "windows-future" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" dependencies = [ - "bytes", - "libc", - "mio", - "parking_lot", - "pin-project-lite", - "signal-hook-registry", - "socket2", - "tokio-macros", - "windows-sys 0.61.2", + "windows-core", + "windows-link", + "windows-threading", ] [[package]] -name = "tokio-macros" -version = "2.6.1" +name = "windows-implement" +version = "0.60.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", @@ -908,60 +2621,49 @@ dependencies = [ ] [[package]] -name = "tokio-util" -version = "0.7.18" +name = "windows-interface" +version = "0.59.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ - "bytes", - "futures-core", - "futures-io", - "futures-sink", - "pin-project-lite", - "tokio", + "proc-macro2", + "quote", + "syn", ] [[package]] -name = "unicode-ident" -version = "1.0.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" - -[[package]] -name = "unicode-segmentation" -version = "1.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" - -[[package]] -name = "unicode-width" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" - -[[package]] -name = "unicode-xid" -version = "0.2.6" +name = "windows-link" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" [[package]] -name = "utf8parse" -version = "0.2.2" +name = "windows-numerics" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core", + "windows-link", +] [[package]] -name = "wasi" -version = "0.11.1+wasi-snapshot-preview1" +name = "windows-result" +version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] [[package]] -name = "windows-link" -version = "0.2.1" +name = "windows-strings" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] [[package]] name = "windows-sys" @@ -1023,6 +2725,15 @@ dependencies = [ "windows_x86_64_msvc 0.53.1", ] +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" @@ -1119,12 +2830,183 @@ version = "0.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +[[package]] +name = "wit-bindgen" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap", + "prettyplease", + "syn", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + +[[package]] +name = "writeable" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" + [[package]] name = "yansi" version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" +[[package]] +name = "yoke" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[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 = "zerotrie" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a59c17a5562d507e4b54960e8569ebee33bee890c70aa3fe7b97e85a9fd7851" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c28719294829477f525be0186d13efa9a3c602f7ec202ca9e353d310fb9a002" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zmij" version = "1.0.21" diff --git a/Cargo.toml b/Cargo.toml index b935f36..7718d3d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,15 @@ [workspace] members = [ "src/agent-client-protocol", + "src/agent-client-protocol-conductor", + "src/agent-client-protocol-cookbook", + "src/agent-client-protocol-core", + "src/agent-client-protocol-derive", + "src/agent-client-protocol-rmcp", + "src/agent-client-protocol-test", + "src/agent-client-protocol-tokio", + "src/agent-client-protocol-trace-viewer", + "src/yopo", ] resolver = "3" @@ -12,6 +21,16 @@ repository = "https://github.com/agentclientprotocol/rust-sdk" homepage = "https://github.com/agentclientprotocol/rust-sdk" [workspace.dependencies] +# Internal crates +agent-client-protocol-conductor = { path = "src/agent-client-protocol-conductor", version = "0.1.0" } +agent-client-protocol-core = { path = "src/agent-client-protocol-core", version = "0.1.0" } +agent-client-protocol-derive = { path = "src/agent-client-protocol-derive", version = "0.1.0" } +agent-client-protocol-rmcp = { path = "src/agent-client-protocol-rmcp", version = "0.1.0" } +agent-client-protocol-test = { path = "src/agent-client-protocol-test", version = "0.1.0" } +agent-client-protocol-tokio = { path = "src/agent-client-protocol-tokio", version = "0.1.0" } +agent-client-protocol-trace-viewer = { path = "src/agent-client-protocol-trace-viewer", version = "0.1.0" } +yopo = { package = "agent-client-protocol-yopo", path = "src/yopo", version = "0.1.0" } + # Core async runtime tokio = { version = "1.48", features = ["full"] } tokio-util = { version = "0.7", features = ["compat"] } @@ -39,24 +58,42 @@ tracing-appender = "0.2" tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] } # MCP SDK -rmcp = { version = "0.8.3", features = ["server", "transport-io", "schemars"] } +rmcp = { version = "1.2.0", features = ["server", "transport-io", "schemars"] } # CLI parsing clap = { version = "4.5", features = ["derive"] } +# HTTP +axum = "0.8" +hyper = "1.0" +tower = "0.5" +tower-http = { version = "0.6", features = ["fs"] } + # Other dependencies async-broadcast = "0.7" +async-stream = "0.3.6" async-trait = "0.1" boxfnonce = "0.1.1" chrono = "0.4" derive_more = { version = "2", features = ["from"] } futures = "0.3.31" +futures-concurrency = "7.6.3" fxhash = "0.2.1" jsonrpcmsg = "0.1.2" +open = "5" +rustc-hash = "2.1.1" +shell-words = "1.1" +strip-ansi-escapes = "0.2" + +# Proc macros +proc-macro2 = "1" +quote = "1" +syn = { version = "2", features = ["full", "extra-traits"] } # Testing expect-test = "1.5" +regex = "1.12.2" futures-util = { version = "0.3", features = ["io"] } piper = "0.2" pretty_assertions = "1" @@ -82,4 +119,5 @@ needless_pass_by_value = "allow" similar_names = "allow" struct_field_names = "allow" too_many_lines = "allow" +type_complexity = "allow" wildcard_imports = "allow" diff --git a/src/agent-client-protocol-conductor/CHANGELOG.md b/src/agent-client-protocol-conductor/CHANGELOG.md new file mode 100644 index 0000000..9e0e2bb --- /dev/null +++ b/src/agent-client-protocol-conductor/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-conductor/Cargo.toml b/src/agent-client-protocol-conductor/Cargo.toml new file mode 100644 index 0000000..d3da12d --- /dev/null +++ b/src/agent-client-protocol-conductor/Cargo.toml @@ -0,0 +1,57 @@ +[package] +name = "agent-client-protocol-conductor" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "Conductor for orchestrating Agent Client Protocol proxy chains" +keywords = ["acp", "agent", "conductor", "ai"] +categories = ["development-tools"] + +[[bin]] +name = "agent-client-protocol-conductor" +path = "src/main.rs" + +[features] +test-support = [] + +[dependencies] +agent-client-protocol-core.workspace = true +agent-client-protocol-schema.workspace = true +agent-client-protocol-tokio.workspace = true +agent-client-protocol-trace-viewer.workspace = true +anyhow.workspace = true +async-stream.workspace = true +axum.workspace = true +chrono.workspace = true +clap.workspace = true +futures.workspace = true +futures-concurrency.workspace = true +hyper.workspace = true +rustc-hash.workspace = true +serde.workspace = true +serde_json.workspace = true +shell-words.workspace = true +strip-ansi-escapes.workspace = true +thiserror.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tower.workspace = true +tracing.workspace = true +tracing-appender.workspace = true +tracing-subscriber.workspace = true +uuid.workspace = true + +[dev-dependencies] +agent-client-protocol-test.workspace = true +agent-client-protocol-tokio.workspace = true +yopo.workspace = true +expect-test.workspace = true +regex.workspace = true +rmcp = { workspace = true, features = ["client", "server", "transport-io", "transport-child-process"] } +schemars.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-conductor/README.md b/src/agent-client-protocol-conductor/README.md new file mode 100644 index 0000000..8b3e62c --- /dev/null +++ b/src/agent-client-protocol-conductor/README.md @@ -0,0 +1,73 @@ +# agent-client-protocol-conductor + +Binary for orchestrating [ACP](https://agentclientprotocol.com/) proxy chains. + +## What is the conductor? + +The conductor is a tool that manages proxy chains — it spawns proxy components and the base agent, then routes messages between them. From the editor's perspective, the conductor appears as a single ACP agent. + +``` +Editor ← stdio → Conductor → Proxy 1 → Proxy 2 → Agent +``` + +## Usage + +### Agent Mode + +Orchestrate a chain of proxies in front of an agent: + +```bash +# Chain format: proxy1 proxy2 ... agent +agent-client-protocol-conductor agent "python proxy1.py" "python proxy2.py" "python base-agent.py" +``` + +The conductor: + +1. Spawns each component as a subprocess +2. Connects them in a chain +3. Presents as a single agent on stdin/stdout +4. Manages the lifecycle of all processes + +### MCP Bridge Mode + +Connect stdio to a TCP-based MCP server: + +```bash +# Bridge stdio to MCP server on localhost:8080 +agent-client-protocol-conductor mcp 8080 +``` + +This allows stdio-based tools to communicate with TCP MCP servers. + +## How It Works + +**Component Communication:** + +- Editor talks to conductor via stdio +- Conductor uses `_proxy/successor/*` protocol extensions to route messages +- Each proxy can intercept, transform, or forward messages +- Final agent receives standard ACP messages + +**Process Management:** + +- All components are spawned as child processes +- When conductor exits, all children are terminated +- Errors in any component bring down the entire chain + +## Building + +```bash +cargo build --release -p agent-client-protocol-conductor +``` + +Binary will be at `target/release/agent-client-protocol-conductor`. + +## Related Crates + +- **[agent-client-protocol-core](../agent-client-protocol-core/)** — Core ACP protocol types and traits +- **[agent-client-protocol-tokio](../agent-client-protocol-tokio/)** — Tokio utilities for process spawning +- **[agent-client-protocol-trace-viewer](../agent-client-protocol-trace-viewer/)** — Interactive trace visualization + +## License + +Apache-2.0 diff --git a/src/agent-client-protocol-conductor/src/conductor.rs b/src/agent-client-protocol-conductor/src/conductor.rs new file mode 100644 index 0000000..9ca1641 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/conductor.rs @@ -0,0 +1,1619 @@ +//! # Conductor: SACP Proxy Chain Orchestrator +//! +//! This module implements the Conductor conductor, which orchestrates a chain of +//! proxy components that sit between an editor and an agent, transforming the +//! Agent-Client Protocol (ACP) stream bidirectionally. +//! +//! ## Architecture Overview +//! +//! The conductor builds and manages a chain of components: +//! +//! ```text +//! Editor <-ACP-> [Component 0] <-ACP-> [Component 1] <-ACP-> ... <-ACP-> Agent +//! ``` +//! +//! Each component receives ACP messages, can transform them, and forwards them +//! to the next component in the chain. The conductor: +//! +//! 1. Spawns each component as a subprocess +//! 2. Establishes bidirectional JSON-RPC connections with each component +//! 3. Routes messages between editor, components, and agent +//! 4. Distinguishes proxy vs agent components via distinct request types +//! +//! ## Recursive Chain Building +//! +//! The chain is built recursively through the `_proxy/successor/*` protocol: +//! +//! 1. Editor connects to Component 0 via the conductor +//! 2. When Component 0 wants to communicate with its successor, it sends +//! requests/notifications with method prefix `_proxy/successor/` +//! 3. The conductor intercepts these messages, strips the prefix, and forwards +//! to Component 1 +//! 4. Component 1 does the same for Component 2, and so on +//! 5. The last component talks directly to the agent (no `_proxy/successor/` prefix) +//! +//! This allows each component to be written as if it's talking to a single successor, +//! without knowing about the full chain. +//! +//! ## Proxy vs Agent Initialization +//! +//! Components discover whether they're a proxy or agent via the initialization request they receive: +//! +//! - **Proxy components**: Receive `InitializeProxyRequest` (`_proxy/initialize` method) +//! - **Agent component**: Receives standard `InitializeRequest` (`initialize` method) +//! +//! The conductor sends `InitializeProxyRequest` to all proxy components in the chain, +//! and `InitializeRequest` only to the final agent component. This allows proxies to +//! know they should forward messages to a successor, while agents know they are the +//! terminal component +//! +//! ## Message Routing +//! +//! The conductor runs an event loop processing messages from: +//! +//! - **Editor to first component**: Standard ACP messages +//! - **Component to successor**: Via `_proxy/successor/*` prefix +//! - **Component responses**: Via futures channels back to requesters +//! +//! The message flow ensures bidirectional communication while maintaining the +//! abstraction that each component only knows about its immediate successor. +//! +//! ## Lazy Component Initialization +//! +//! Components are instantiated lazily when the first `initialize` request is received +//! from the editor. This enables dynamic proxy chain construction based on client capabilities. +//! +//! ### Simple Usage +//! +//! Pass a Vec of components that implement `Component`: +//! +//! ```ignore +//! let conductor = Conductor::new( +//! "my-conductor", +//! vec![proxy1, proxy2, agent], +//! None, +//! ); +//! ``` +//! +//! All components are spawned in order when the editor sends the first `initialize` request. +//! +//! ### Dynamic Component Selection +//! +//! Pass a closure to examine the `InitializeRequest` and dynamically construct the chain: +//! +//! ```ignore +//! let conductor = Conductor::new( +//! "my-conductor", +//! |cx, conductor_tx, init_req| async move { +//! // Examine capabilities +//! let needs_auth = has_auth_capability(&init_req); +//! +//! let mut components = Vec::new(); +//! if needs_auth { +//! components.push(spawn_auth_proxy(&cx, &conductor_tx)?); +//! } +//! components.push(spawn_agent(&cx, &conductor_tx)?); +//! +//! // Return (potentially modified) request and component list +//! Ok((init_req, components)) +//! }, +//! None, +//! ); +//! ``` +//! +//! The closure receives: +//! - `cx: &ConnectionTo` - Connection context for spawning components +//! - `conductor_tx: &mpsc::Sender` - Channel for message routing +//! - `init_req: InitializeRequest` - The Initialize request from the editor +//! +//! And returns: +//! - Modified `InitializeRequest` to forward downstream +//! - `Vec` of spawned components + +use std::{collections::HashMap, sync::Arc}; + +use agent_client_protocol_core::{ + Agent, BoxFuture, Client, Conductor, ConnectTo, Dispatch, DynConnectTo, Error, JsonRpcMessage, + Proxy, Role, RunWithConnectionTo, role::HasPeer, util::MatchDispatch, +}; +use agent_client_protocol_core::{ + Builder, ConnectionTo, JsonRpcNotification, JsonRpcRequest, SentRequest, UntypedMessage, +}; +use agent_client_protocol_core::{ + HandleDispatchFrom, + schema::{InitializeProxyRequest, InitializeRequest, NewSessionRequest}, + util::MatchDispatchFrom, +}; +use agent_client_protocol_core::{ + Handled, + schema::{ + McpConnectRequest, McpConnectResponse, McpDisconnectNotification, McpOverAcpMessage, + SuccessorMessage, + }, +}; +use futures::{ + SinkExt, StreamExt, + channel::mpsc::{self}, +}; +use tracing::{debug, info}; + +use crate::conductor::mcp_bridge::{ + McpBridgeConnection, McpBridgeConnectionActor, McpBridgeListeners, +}; + +mod mcp_bridge; + +/// The conductor manages the proxy chain lifecycle and message routing. +/// +/// It maintains connections to all components in the chain and routes messages +/// bidirectionally between the editor, components, and agent. +/// +#[derive(Debug)] +pub struct ConductorImpl { + host: Host, + name: String, + instantiator: Host::Instantiator, + mcp_bridge_mode: crate::McpBridgeMode, + trace_writer: Option, +} + +impl ConductorImpl { + pub fn new( + host: Host, + name: impl ToString, + instantiator: Host::Instantiator, + mcp_bridge_mode: crate::McpBridgeMode, + ) -> Self { + ConductorImpl { + name: name.to_string(), + host, + instantiator, + mcp_bridge_mode, + trace_writer: None, + } + } +} + +impl ConductorImpl { + /// Create a conductor in agent mode (the last component is an agent). + pub fn new_agent( + name: impl ToString, + instantiator: impl InstantiateProxiesAndAgent + 'static, + mcp_bridge_mode: crate::McpBridgeMode, + ) -> Self { + ConductorImpl::new(Agent, name, Box::new(instantiator), mcp_bridge_mode) + } +} + +impl ConductorImpl { + /// Create a conductor in proxy mode (forwards to another conductor). + pub fn new_proxy( + name: impl ToString, + instantiator: impl InstantiateProxies + 'static, + mcp_bridge_mode: crate::McpBridgeMode, + ) -> Self { + ConductorImpl::new(Proxy, name, Box::new(instantiator), mcp_bridge_mode) + } +} + +impl ConductorImpl { + /// Enable trace logging to a custom destination. + /// + /// Use `agent-client-protocol-trace-viewer` to view the trace as an interactive sequence diagram. + #[must_use] + pub fn trace_to(mut self, dest: impl crate::trace::WriteEvent) -> Self { + self.trace_writer = Some(crate::trace::TraceWriter::new(dest)); + self + } + + /// Enable trace logging to a file path. + /// + /// Events will be written as newline-delimited JSON (`.jsons` format). + /// Use `agent-client-protocol-trace-viewer` to view the trace as an interactive sequence diagram. + pub fn trace_to_path(mut self, path: impl AsRef) -> std::io::Result { + self.trace_writer = Some(crate::trace::TraceWriter::from_path(path)?); + Ok(self) + } + + /// Enable trace logging with an existing TraceWriter. + #[must_use] + pub fn with_trace_writer(mut self, writer: crate::trace::TraceWriter) -> Self { + self.trace_writer = Some(writer); + self + } + + /// Run the conductor with a transport. + pub async fn run( + self, + transport: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + let (conductor_tx, conductor_rx) = mpsc::channel(128 /* chosen arbitrarily */); + + // Set up tracing if enabled - spawn writer task and get handle + let trace_handle; + let trace_future: BoxFuture<'static, Result<(), agent_client_protocol_core::Error>>; + if let Some((h, f)) = self.trace_writer.map(super::trace::TraceWriter::spawn) { + trace_handle = Some(h); + trace_future = Box::pin(f); + } else { + trace_handle = None; + trace_future = Box::pin(std::future::ready(Ok(()))); + } + + let responder = ConductorResponder { + conductor_rx, + conductor_tx: conductor_tx.clone(), + instantiator: Some(self.instantiator), + bridge_listeners: McpBridgeListeners::default(), + bridge_connections: HashMap::default(), + mcp_bridge_mode: self.mcp_bridge_mode, + proxies: Vec::default(), + successor: Arc::new(agent_client_protocol_core::util::internal_error( + "successor not initialized", + )), + trace_handle, + host: self.host.clone(), + }; + + Builder::new_with( + self.host.clone(), + ConductorMessageHandler { + conductor_tx, + host: self.host.clone(), + }, + ) + .name(self.name) + .with_responder(responder) + .with_spawned(|_cx| trace_future) + .connect_to(transport) + .await + } + + async fn incoming_message_from_client( + conductor_tx: &mut mpsc::Sender, + message: Dispatch, + ) -> Result<(), agent_client_protocol_core::Error> { + conductor_tx + .send(ConductorMessage::LeftToRight { + target_component_index: 0, + message, + }) + .await + .map_err(agent_client_protocol_core::util::internal_error) + } + + async fn incoming_message_from_agent( + conductor_tx: &mut mpsc::Sender, + message: Dispatch, + ) -> Result<(), agent_client_protocol_core::Error> { + conductor_tx + .send(ConductorMessage::RightToLeft { + source_component_index: SourceComponentIndex::Successor, + message, + }) + .await + .map_err(agent_client_protocol_core::util::internal_error) + } +} + +impl ConnectTo for ConductorImpl { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + self.run(client).await + } +} + +struct ConductorMessageHandler { + conductor_tx: mpsc::Sender, + host: Host, +} + +impl HandleDispatchFrom + for ConductorMessageHandler +{ + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + connection: agent_client_protocol_core::ConnectionTo, + ) -> Result, agent_client_protocol_core::Error> + { + self.host + .handle_dispatch(message, connection, &mut self.conductor_tx) + .await + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + "ConductorMessageHandler" + } +} + +/// The conductor manages the proxy chain lifecycle and message routing. +/// +/// It maintains connections to all components in the chain and routes messages +/// bidirectionally between the editor, components, and agent. +/// +pub struct ConductorResponder +where + Host: ConductorHostRole, +{ + conductor_rx: mpsc::Receiver, + + conductor_tx: mpsc::Sender, + + /// Manages the TCP listeners for MCP connections that will be proxied over ACP. + bridge_listeners: McpBridgeListeners, + + /// Manages active connections to MCP clients. + bridge_connections: HashMap, + + /// The instantiator for lazy initialization. + /// Set to None after components are instantiated. + instantiator: Option, + + /// The chain of proxies before the agent (if any). + /// + /// Populated lazily when the first Initialize request is received. + proxies: Vec>, + + /// If the conductor is operating in agent mode, this will direct messages to the agent. + /// If the conductor is operating in proxy mode, this will direct messages to the successor. + /// Populated lazily when the first Initialize request is received; the initial value just returns errors. + successor: Arc>, + + /// Mode for the MCP bridge (determines how to spawn bridge processes). + mcp_bridge_mode: crate::McpBridgeMode, + + /// Optional trace handle for sequence diagram visualization. + trace_handle: Option, + + /// Defines what sort of link we have + host: Host, +} + +impl std::fmt::Debug for ConductorResponder +where + Host: ConductorHostRole, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ConductorResponder") + .field("conductor_rx", &self.conductor_rx) + .field("conductor_tx", &self.conductor_tx) + .field("bridge_listeners", &self.bridge_listeners) + .field("bridge_connections", &self.bridge_connections) + .field("proxies", &self.proxies) + .field("mcp_bridge_mode", &self.mcp_bridge_mode) + .field("trace_handle", &self.trace_handle) + .field("host", &self.host) + .finish_non_exhaustive() + } +} + +impl RunWithConnectionTo for ConductorResponder +where + Host: ConductorHostRole, +{ + async fn run_with_connection_to( + mut self, + connection: ConnectionTo, + ) -> Result<(), agent_client_protocol_core::Error> { + // Components are now spawned lazily in forward_initialize_request + // when the first Initialize request is received. + + // This is the "central actor" of the conductor. Most other things forward messages + // via `conductor_tx` into this loop. This lets us serialize the conductor's activity. + while let Some(message) = self.conductor_rx.next().await { + self.handle_conductor_message(connection.clone(), message) + .await?; + } + Ok(()) + } +} + +impl ConductorResponder +where + Host: ConductorHostRole, +{ + /// Recursively spawns components and builds the proxy chain. + /// + /// This function implements the recursive chain building pattern: + /// 1. Pop the next component from the `providers` list + /// 2. Create the component (either spawn subprocess or use mock) + /// 3. Set up JSON-RPC connection and message handlers + /// 4. Recursively call itself to spawn the next component + /// 5. When no components remain, start the message routing loop via `serve()` + /// + /// Central message handling logic for the conductor. + /// The conductor routes all [`ConductorMessage`] messages through to this function. + /// Each message corresponds to a request or notification from one component to another. + /// The conductor ferries messages from one place to another, sometimes making modifications along the way. + /// Note that *responses to requests* are sent *directly* without going through this loop. + /// + /// The names we use are + /// + /// * The *client* is the originator of all ACP traffic, typically an editor or GUI. + /// * Then there is a sequence of *components* consisting of: + /// * Zero or more *proxies*, which receive messages and forward them to the next component in the chain. + /// * And finally the *agent*, which is the final component in the chain and handles the actual work. + /// + /// For the most part, we just pass messages through the chain without modification, but there are a few exceptions: + /// + /// * We send `InitializeProxyRequest` to proxy components and `InitializeRequest` to the agent component. + /// * We modify "session/new" requests that use `acp:...` as the URL for an MCP server to redirect + /// through a stdio server that runs on localhost and bridges messages. + async fn handle_conductor_message( + &mut self, + client: ConnectionTo, + message: ConductorMessage, + ) -> Result<(), agent_client_protocol_core::Error> { + tracing::debug!(?message, "handle_conductor_message"); + + match message { + ConductorMessage::LeftToRight { + target_component_index, + message, + } => { + // Tracing happens inside forward_client_to_agent_message, after initialization, + // so that component_name() has access to the populated proxies list. + self.forward_client_to_agent_message(target_component_index, message, client) + .await + } + + ConductorMessage::RightToLeft { + source_component_index, + message, + } => { + tracing::debug!( + ?source_component_index, + message_method = ?message.method(), + "Conductor: AgentToClient received" + ); + self.send_message_to_predecessor_of(client, source_component_index, message) + } + + // New MCP connection request. Send it back along the chain to get a connection id. + // When the connection id arrives, send a message back into this conductor loop with + // the connection id and the (as yet unspawned) actor. + ConductorMessage::McpConnectionReceived { + acp_url, + connection, + actor, + } => { + // MCP connection requests always come from the agent + // (we must be in agent mode, in fact), so send the MCP request + // to the final proxy. + self.send_request_to_predecessor_of( + client, + self.proxies.len(), + McpConnectRequest { + acp_url, + meta: None, + }, + ) + .on_receiving_result({ + let mut conductor_tx = self.conductor_tx.clone(); + async move |result| { + match result { + Ok(response) => conductor_tx + .send(ConductorMessage::McpConnectionEstablished { + response, + actor, + connection, + }) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error()), + Err(_) => { + // Error occurred, just drop the connection. + Ok(()) + } + } + } + }) + } + + // MCP connection successfully established. Spawn the actor + // and insert the connection into our map for future reference. + ConductorMessage::McpConnectionEstablished { + response: McpConnectResponse { connection_id, .. }, + actor, + connection, + } => { + self.bridge_connections + .insert(connection_id.clone(), connection); + client.spawn(actor.run(connection_id)) + } + + // Message meant for the MCP client received. Forward it to the appropriate actor's mailbox. + ConductorMessage::McpClientToMcpServer { + connection_id, + message, + } => { + let wrapped = message.map( + |request, responder| { + ( + McpOverAcpMessage { + connection_id: connection_id.clone(), + message: request, + meta: None, + }, + responder, + ) + }, + |notification| McpOverAcpMessage { + connection_id: connection_id.clone(), + message: notification, + meta: None, + }, + ); + + // We only get MCP-over-ACP requests when we are in bridging MCP for the final agent, + // so send them to the final proxy. + self.send_message_to_predecessor_of( + client, + SourceComponentIndex::Successor, + wrapped, + ) + } + + // MCP client disconnected. Remove it from our map and send the + // notification backwards along the chain. + ConductorMessage::McpConnectionDisconnected { notification } => { + // We only get MCP-over-ACP requests when we are in bridging MCP for the final agent. + + self.bridge_connections.remove(¬ification.connection_id); + self.send_notification_to_predecessor_of(client, self.proxies.len(), notification) + } + } + } + + /// Send a message (request or notification) to the predecessor of the given component. + /// + /// This is a bit subtle because the relationship of the conductor + /// is different depending on who will be receiving the message: + /// * If the message is going to the conductor's client, then no changes + /// are needed, as the conductor is sending an agent-to-client message and + /// the conductor is acting as the agent. + /// * If the message is going to a proxy component, then we have to wrap + /// it in a "from successor" wrapper, because the conductor is the + /// proxy's client. + fn send_message_to_predecessor_of( + &mut self, + client: ConnectionTo, + source_component_index: SourceComponentIndex, + message: Dispatch, + ) -> Result<(), agent_client_protocol_core::Error> + where + Req::Response: Send, + { + let source_component_index = match source_component_index { + SourceComponentIndex::Successor => self.proxies.len(), + SourceComponentIndex::Proxy(index) => index, + }; + + match message { + Dispatch::Request(request, responder) => self + .send_request_to_predecessor_of(client, source_component_index, request) + .forward_response_to(responder), + Dispatch::Notification(notification) => self.send_notification_to_predecessor_of( + client, + source_component_index, + notification, + ), + Dispatch::Response(result, router) => router.respond_with_result(result), + } + } + + fn send_request_to_predecessor_of( + &mut self, + client_connection: ConnectionTo, + source_component_index: usize, + request: Req, + ) -> SentRequest { + if source_component_index == 0 { + client_connection.send_request_to(Client, request) + } else { + self.proxies[source_component_index - 1].send_request(SuccessorMessage { + message: request, + meta: None, + }) + } + } + + /// Send a notification to the predecessor of the given component. + /// + /// This is a bit subtle because the relationship of the conductor + /// is different depending on who will be receiving the message: + /// * If the notification is going to the conductor's client, then no changes + /// are needed, as the conductor is sending an agent-to-client message and + /// the conductor is acting as the agent. + /// * If the notification is going to a proxy component, then we have to wrap + /// it in a "from successor" wrapper, because the conductor is the + /// proxy's client. + fn send_notification_to_predecessor_of( + &mut self, + client: ConnectionTo, + source_component_index: usize, + notification: N, + ) -> Result<(), agent_client_protocol_core::Error> { + tracing::debug!( + source_component_index, + proxies_len = self.proxies.len(), + "send_notification_to_predecessor_of" + ); + if source_component_index == 0 { + tracing::debug!("Sending notification directly to client"); + client.send_notification_to(Client, notification) + } else { + tracing::debug!( + target_proxy = source_component_index - 1, + "Sending notification wrapped as SuccessorMessage to proxy" + ); + self.proxies[source_component_index - 1].send_notification(SuccessorMessage { + message: notification, + meta: None, + }) + } + } + + /// Send a message (request or notification) from 'left to right'. + /// Left-to-right means from the client or an intermediate proxy to the component + /// at `target_component_index` (could be a proxy or the agent). + /// Makes changes to select messages along the way (e.g., `initialize` and `session/new`). + async fn forward_client_to_agent_message( + &mut self, + target_component_index: usize, + message: Dispatch, + client: ConnectionTo, + ) -> Result<(), agent_client_protocol_core::Error> { + tracing::trace!( + target_component_index, + ?message, + "forward_client_to_agent_message" + ); + + // Ensure components are initialized before processing any message. + let message = self.ensure_initialized(client.clone(), message).await?; + + // In proxy mode, if the target is beyond our component chain, + // forward to the conductor's own successor (via client connection) + if target_component_index < self.proxies.len() { + self.forward_message_from_client_to_proxy(target_component_index, message) + .await + } else { + assert_eq!(target_component_index, self.proxies.len()); + + debug!( + target_component_index, + proxies_count = self.proxies.len(), + "Proxy mode: forwarding successor message to conductor's successor" + ); + let successor = self.successor.clone(); + successor.send_message(message, client, self).await + } + } + + /// Ensures components are initialized before processing messages. + /// + /// If components haven't been initialized yet, this expects the first message + /// to be an `initialize` request and uses it to spawn the component chain. + /// + /// Returns: + /// - `Ok(Some(message))` - Components are initialized, continue processing this message + /// - `Ok(None)` - An error response was sent, caller should return early + /// - `Err(_)` - A fatal error occurred + async fn ensure_initialized( + &mut self, + client: ConnectionTo, + message: Dispatch, + ) -> Result { + // Already initialized - pass through + let Some(instantiator) = self.instantiator.take() else { + return Ok(message); + }; + + let host = self.host.clone(); + let message = host.initialize(message, client, instantiator, self).await?; + Ok(message) + } + + /// Wrap a proxy component with tracing if tracing is enabled. + /// + /// Returns the component unchanged if tracing is disabled. + fn trace_proxy( + &self, + proxy_index: ComponentIndex, + successor_index: ComponentIndex, + component: impl ConnectTo, + ) -> DynConnectTo { + match &self.trace_handle { + Some(trace_handle) => { + trace_handle.bridge_component(proxy_index, successor_index, component) + } + None => DynConnectTo::new(component), + } + } + + /// Spawn proxy components and add them to the proxies list. + fn spawn_proxies( + &mut self, + client: ConnectionTo, + proxy_components: Vec>, + ) -> Result<(), agent_client_protocol_core::Error> { + assert!(self.proxies.is_empty()); + + let num_proxies = proxy_components.len(); + info!(proxy_count = num_proxies, "spawn_proxies"); + + // Special case: if there are no user-defined proxies + // but tracing is enabled, we make a dummy proxy that just + // passes through messages but which can trigger the + // tracing events. + if self.trace_handle.is_some() && num_proxies == 0 { + self.connect_to_proxy( + &client, + 0, + ComponentIndex::Client, + ComponentIndex::Agent, + Proxy.builder(), + )?; + } else { + // Spawn each proxy component + for (component_index, dyn_component) in proxy_components.into_iter().enumerate() { + debug!(component_index, "spawning proxy"); + + self.connect_to_proxy( + &client, + component_index, + ComponentIndex::Proxy(component_index), + ComponentIndex::successor_of(component_index, num_proxies), + dyn_component, + )?; + } + } + + info!(proxy_count = self.proxies.len(), "Proxies spawned"); + + Ok(()) + } + + /// Create a connection to the proxy with index `component_index` implemented in `component`. + /// + /// If tracing is enabled, the proxy's index is `trace_proxy_index` and its successor is `trace_successor_index`. + fn connect_to_proxy( + &mut self, + client: &ConnectionTo, + component_index: usize, + trace_proxy_index: ComponentIndex, + trace_successor_index: ComponentIndex, + component: impl ConnectTo, + ) -> Result<(), Error> { + let connection_builder = self.connection_to_proxy(component_index); + let connect_component = + self.trace_proxy(trace_proxy_index, trace_successor_index, component); + let proxy_connection = client.spawn_connection(connection_builder, connect_component)?; + self.proxies.push(proxy_connection); + Ok(()) + } + + /// Create the conductor's connection to the proxy with index `component_index`. + /// + /// Outgoing messages received from the proxy are sent to `self.conductor_tx` as either + /// left-to-right or right-to-left messages depending on whether they are wrapped + /// in `SuccessorMessage`. + fn connection_to_proxy( + &mut self, + component_index: usize, + ) -> Builder + 'static> { + type SuccessorDispatch = Dispatch; + let mut conductor_tx = self.conductor_tx.clone(); + Conductor + .builder() + .name(format!("conductor-to-component({component_index})")) + // Intercept messages sent by the proxy. + .on_receive_dispatch( + async move |dispatch: Dispatch, _connection| { + MatchDispatch::new(dispatch) + .if_message(async |dispatch: SuccessorDispatch| { + // ------------------ + // SuccessorMessages sent by the proxy go to its successor. + // + // Subtle point: + // + // `ConductorToProxy` has only a single peer, `Agent`. This means that we see + // "successor messages" in their "desugared form". So when we intercept an *outgoing* + // message that matches `SuccessorMessage`, it could be one of three things + // + // - A request being sent by the proxy to its successor (hence going left->right) + // - A notification being sent by the proxy to its successor (hence going left->right) + // - A response to a request sent to the proxy *by* its successor. Here, the *request* + // was going right->left, but the *response* (the message we are processing now) + // is going left->right. + // + // So, in all cases, we forward as a left->right message. + + conductor_tx + .send(ConductorMessage::LeftToRight { + target_component_index: component_index + 1, + message: dispatch.map(|r, cx| (r.message, cx), |n| n.message), + }) + .await + .map_err(agent_client_protocol_core::util::internal_error) + }) + .await + .otherwise(async |dispatch| { + // Other messagrs send by the proxy go its predecessor. + // As in the previous handler: + // + // Messages here are seen in their "desugared form", so we are seeing + // one of three things + // + // - A request being sent by the proxy to its predecessor (hence going right->left) + // - A notification being sent by the proxy to its predecessor (hence going right->left) + // - A response to a request sent to the proxy *by* its predecessor. Here, the *request* + // was going left->right, but the *response* (the message we are processing now) + // is going right->left. + // + // So, in all cases, we forward as a right->left message. + + let message = ConductorMessage::RightToLeft { + source_component_index: SourceComponentIndex::Proxy( + component_index, + ), + message: dispatch, + }; + conductor_tx + .send(message) + .await + .map_err(agent_client_protocol_core::util::internal_error) + }) + .await + }, + agent_client_protocol_core::on_receive_dispatch!(), + ) + } + + async fn forward_message_from_client_to_proxy( + &mut self, + target_component_index: usize, + message: Dispatch, + ) -> Result<(), agent_client_protocol_core::Error> { + tracing::debug!(?message, "forward_message_to_proxy"); + + MatchDispatch::new(message) + .if_request(async |_request: InitializeProxyRequest, responder| { + responder.respond_with_error( + agent_client_protocol_core::Error::invalid_request() + .data("initialize/proxy requests are only sent by the conductor"), + ) + }) + .await + .if_request(async |request: InitializeRequest, responder| { + // The pattern for `Initialize` messages is a bit subtle. + // Proxy receive incoming `Initialize` messages as if they + // were a client. The conductor (us) intercepts these and + // converts them to an `InitializeProxyRequest`. + // + // The proxy will then initialize itself and forward an `Initialize` + // request to its successor. + self.proxies[target_component_index] + .send_request(InitializeProxyRequest::from(request)) + .on_receiving_result(async move |result| { + tracing::debug!(?result, "got initialize_proxy response from proxy"); + responder.respond_with_result(result) + }) + }) + .await + .otherwise(async |message| { + // Otherwise, just send the message along "as is". + self.proxies[target_component_index].send_proxied_message(message) + }) + .await + } + + /// Invoked when sending a message from the conductor to the agent that it manages. + /// This is called by `self.successor`'s [`ConductorSuccessor::send_message`] + /// method when `Link = ConductorToClient` (i.e., the conductor is not itself + /// running as a proxy). + async fn forward_message_to_agent( + &mut self, + client_connection: ConnectionTo, + message: Dispatch, + agent_connection: ConnectionTo, + ) -> Result<(), Error> { + MatchDispatch::new(message) + .if_request(async |_request: InitializeProxyRequest, responder| { + responder.respond_with_error( + agent_client_protocol_core::Error::invalid_request() + .data("initialize/proxy requests are only sent by the conductor"), + ) + }) + .await + .if_request(async |mut request: NewSessionRequest, responder| { + // When forwarding "session/new" to the agent, + // we adjust MCP servers to manage "acp:" URLs. + for mcp_server in &mut request.mcp_servers { + self.bridge_listeners + .transform_mcp_server( + client_connection.clone(), + mcp_server, + &self.conductor_tx, + &self.mcp_bridge_mode, + ) + .await?; + } + + agent_connection + .send_request(request) + .forward_response_to(responder) + }) + .await + .if_request( + async |request: McpOverAcpMessage, responder| { + let McpOverAcpMessage { + connection_id, + message: mcp_request, + .. + } = request; + self.bridge_connections + .get_mut(&connection_id) + .ok_or_else(|| { + agent_client_protocol_core::util::internal_error(format!( + "unknown connection id: {connection_id}" + )) + })? + .send(Dispatch::Request(mcp_request, responder)) + .await + }, + ) + .await + .if_notification(async |notification: McpOverAcpMessage| { + let McpOverAcpMessage { + connection_id, + message: mcp_notification, + .. + } = notification; + self.bridge_connections + .get_mut(&connection_id) + .ok_or_else(|| { + agent_client_protocol_core::util::internal_error(format!( + "unknown connection id: {connection_id}" + )) + })? + .send(Dispatch::Notification(mcp_notification)) + .await + }) + .await + .otherwise(async |message| { + // Otherwise, just send the message along "as is". + agent_connection.send_proxied_message_to(Agent, message) + }) + .await + } +} + +/// Identifies a component in the conductor's chain for tracing purposes. +/// +/// Used to track message sources and destinations through the proxy chain. +#[derive(Debug, Clone, Copy)] +pub enum ComponentIndex { + /// The client (editor) at the start of the chain. + Client, + + /// A proxy component at the given index. + Proxy(usize), + + /// The successor (agent in agent mode, outer conductor in proxy mode). + Agent, +} + +impl ComponentIndex { + /// Return the index for the predecessor of `proxy_index`, which might be `Client`. + #[must_use] + pub fn predecessor_of(proxy_index: usize) -> Self { + match proxy_index.checked_sub(1) { + Some(p_i) => ComponentIndex::Proxy(p_i), + None => ComponentIndex::Client, + } + } + + /// Return the index for the predecessor of `proxy_index`, which might be `Client`. + #[must_use] + pub fn successor_of(proxy_index: usize, num_proxies: usize) -> Self { + if proxy_index == num_proxies { + ComponentIndex::Agent + } else { + ComponentIndex::Proxy(proxy_index + 1) + } + } +} + +/// Identifies the source of an agent-to-client message. +/// +/// This enum handles the fact that the conductor may receive messages from two different sources: +/// 1. From one of its managed components (identified by index) +/// 2. From the conductor's own successor in a larger proxy chain (when in proxy mode) +#[derive(Debug, Clone, Copy)] +pub enum SourceComponentIndex { + /// Message from a specific component at the given index in the managed chain. + Proxy(usize), + + /// Message from the conductor's agent or successor. + Successor, +} + +/// Trait for lazy proxy instantiation (proxy mode). +/// +/// Used by conductors in proxy mode (`ConductorToConductor`) where all components +/// are proxies that forward to an outer conductor. +pub trait InstantiateProxies: Send { + /// Instantiate proxy components based on the Initialize request. + /// + /// Returns proxy components typed as `DynConnectTo` since proxies + /// communicate with the conductor. + fn instantiate_proxies( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + (InitializeRequest, Vec>), + agent_client_protocol_core::Error, + >, + >; +} + +/// Simple implementation: provide all proxy components unconditionally. +/// +/// Requires `T: ConnectTo`. +impl InstantiateProxies for Vec +where + T: ConnectTo + 'static, +{ + fn instantiate_proxies( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + (InitializeRequest, Vec>), + agent_client_protocol_core::Error, + >, + > { + Box::pin(async move { + let components: Vec> = + (*self).into_iter().map(|c| DynConnectTo::new(c)).collect(); + Ok((req, components)) + }) + } +} + +/// Dynamic implementation: closure receives the Initialize request and returns proxies. +impl InstantiateProxies for F +where + F: FnOnce(InitializeRequest) -> Fut + Send + 'static, + Fut: std::future::Future< + Output = Result< + (InitializeRequest, Vec>), + agent_client_protocol_core::Error, + >, + > + Send + + 'static, +{ + fn instantiate_proxies( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + (InitializeRequest, Vec>), + agent_client_protocol_core::Error, + >, + > { + Box::pin(async move { (*self)(req).await }) + } +} + +/// Trait for lazy proxy and agent instantiation (agent mode). +/// +/// Used by conductors in agent mode (`ConductorToClient`) where there are +/// zero or more proxies followed by an agent component. +pub trait InstantiateProxiesAndAgent: Send { + /// Instantiate proxy and agent components based on the Initialize request. + /// + /// Returns the (possibly modified) request, a vector of proxy components + /// (typed as `DynConnectTo`), and the agent component + /// (typed as `DynConnectTo`). + fn instantiate_proxies_and_agent( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + ( + InitializeRequest, + Vec>, + DynConnectTo, + ), + agent_client_protocol_core::Error, + >, + >; +} + +/// Wrapper to convert a single agent component (no proxies) into InstantiateProxiesAndAgent. +#[derive(Debug)] +pub struct AgentOnly(pub A); + +impl + 'static> InstantiateProxiesAndAgent for AgentOnly { + fn instantiate_proxies_and_agent( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + ( + InitializeRequest, + Vec>, + DynConnectTo, + ), + agent_client_protocol_core::Error, + >, + > { + Box::pin(async move { Ok((req, Vec::new(), DynConnectTo::new(self.0))) }) + } +} + +/// Builder for creating proxies and agent components. +/// +/// # Example +/// ```ignore +/// ProxiesAndAgent::new(ElizaAgent::new()) +/// .proxy(LoggingProxy::new()) +/// .proxy(AuthProxy::new()) +/// ``` +#[derive(Debug)] +pub struct ProxiesAndAgent { + proxies: Vec>, + agent: DynConnectTo, +} + +impl ProxiesAndAgent { + /// Create a new builder with the given agent component. + pub fn new(agent: impl ConnectTo + 'static) -> Self { + Self { + proxies: vec![], + agent: DynConnectTo::new(agent), + } + } + + /// Add a single proxy component. + #[must_use] + pub fn proxy(mut self, proxy: impl ConnectTo + 'static) -> Self { + self.proxies.push(DynConnectTo::new(proxy)); + self + } + + /// Add multiple proxy components. + #[must_use] + pub fn proxies(mut self, proxies: I) -> Self + where + P: ConnectTo + 'static, + I: IntoIterator, + { + self.proxies + .extend(proxies.into_iter().map(DynConnectTo::new)); + self + } +} + +impl InstantiateProxiesAndAgent for ProxiesAndAgent { + fn instantiate_proxies_and_agent( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + ( + InitializeRequest, + Vec>, + DynConnectTo, + ), + agent_client_protocol_core::Error, + >, + > { + Box::pin(async move { Ok((req, self.proxies, self.agent)) }) + } +} + +/// Dynamic implementation: closure receives the Initialize request and returns proxies + agent. +impl InstantiateProxiesAndAgent for F +where + F: FnOnce(InitializeRequest) -> Fut + Send + 'static, + Fut: std::future::Future< + Output = Result< + ( + InitializeRequest, + Vec>, + DynConnectTo, + ), + agent_client_protocol_core::Error, + >, + > + Send + + 'static, +{ + fn instantiate_proxies_and_agent( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + ( + InitializeRequest, + Vec>, + DynConnectTo, + ), + agent_client_protocol_core::Error, + >, + > { + Box::pin(async move { (*self)(req).await }) + } +} + +/// Messages sent to the conductor's main event loop for routing. +/// +/// These messages enable the conductor to route communication between: +/// - The editor and the first component +/// - Components and their successors in the chain +/// - Components and their clients (editor or predecessor) +/// +/// All spawned tasks send messages via this enum through a shared channel, +/// allowing centralized routing logic in the `serve()` loop. +#[derive(Debug)] +pub enum ConductorMessage { + /// If this message is a request or notification, then it is going "left-to-right" + /// (e.g., a component making a request of its successor). + /// + /// If this message is a response, then it is going right-to-left + /// (i.e., the successor answering a request made by its predecessor). + LeftToRight { + target_component_index: usize, + message: Dispatch, + }, + + /// If this message is a request or notification, then it is going "right-to-left" + /// (e.g., a component making a request of its predecessor). + /// + /// If this message is a response, then it is going "left-to-right" + /// (i.e., the predecessor answering a request made by its successor). + RightToLeft { + source_component_index: SourceComponentIndex, + message: Dispatch, + }, + + /// A pending MCP bridge connection request request. + /// The request must be sent back over ACP to receive the connection-id. + /// Once the connection-id is received, the actor must be spawned. + McpConnectionReceived { + /// The acp:$UUID URL identifying this bridge + acp_url: String, + + /// The actor that should be spawned once the connection-id is available. + actor: McpBridgeConnectionActor, + + /// The connection to the bridge + connection: McpBridgeConnection, + }, + + /// A pending MCP bridge connection request request. + /// The request must be sent back over ACP to receive the connection-id. + /// Once the connection-id is received, the actor must be spawned. + McpConnectionEstablished { + response: McpConnectResponse, + + /// The actor that should be spawned once the connection-id is available. + actor: McpBridgeConnectionActor, + + /// The connection to the bridge + connection: McpBridgeConnection, + }, + + /// MCP message (request or notification) received from a bridge that needs to be routed to the final proxy. + /// + /// Sent when the bridge receives an MCP tool call from the agent and forwards it + /// to the conductor via TCP. The conductor routes this to the appropriate proxy component. + McpClientToMcpServer { + connection_id: String, + message: Dispatch, + }, + + /// Message sent when MCP client disconnects + McpConnectionDisconnected { + notification: McpDisconnectNotification, + }, +} + +/// Trait implemented for the two links the conductor can use: +/// +/// * ConductorToClient -- conductor is acting as an agent, so when its last proxy sends to its successor, the conductor sends that message to its agent component +/// * ConductorToConductor -- conductor is acting as a proxy, so when its last proxy sends to its successor, the (inner) conductor sends that message to its successor, via the outer conductor +pub trait ConductorHostRole: Role> { + /// The type used to instantiate components for this link type. + type Instantiator: Send; + + /// Handle initialization: parse the init request, instantiate components, and spawn them. + /// + /// Takes ownership of the instantiator and returns the (possibly modified) init request + /// wrapped in a Dispatch for forwarding. + fn initialize( + &self, + message: Dispatch, + connection: ConnectionTo, + instantiator: Self::Instantiator, + responder: &mut ConductorResponder, + ) -> impl Future> + Send; + + /// Handle an incoming message from the client or conductor, depending on `Self` + fn handle_dispatch( + &self, + message: Dispatch, + connection: ConnectionTo, + conductor_tx: &mut mpsc::Sender, + ) -> impl Future, agent_client_protocol_core::Error>> + Send; +} + +/// Conductor acting as an agent +impl ConductorHostRole for Agent { + type Instantiator = Box; + + async fn initialize( + &self, + message: Dispatch, + client_connection: ConnectionTo, + instantiator: Self::Instantiator, + responder: &mut ConductorResponder, + ) -> Result { + let invalid_request = || Error::invalid_request().data("expected `initialize` request"); + + // Not yet initialized - expect an initialize request. + // Error if we get anything else. + let Dispatch::Request(request, init_responder) = message else { + message.respond_with_error(invalid_request(), client_connection.clone())?; + return Err(invalid_request()); + }; + if !InitializeRequest::matches_method(request.method()) { + init_responder.respond_with_error(invalid_request())?; + return Err(invalid_request()); + } + + let init_request = + match InitializeRequest::parse_message(request.method(), request.params()) { + Ok(r) => r, + Err(error) => { + init_responder.respond_with_error(error)?; + return Err(invalid_request()); + } + }; + + // Instantiate proxies and agent + let (modified_req, proxy_components, agent_component) = instantiator + .instantiate_proxies_and_agent(init_request) + .await?; + + // Spawn the agent component + debug!(?agent_component, "spawning agent"); + + let connection_to_agent = client_connection.spawn_connection( + Client + .builder() + .name("conductor-to-agent") + // Intercept agent-to-client messages from the agent. + .on_receive_dispatch( + { + let mut conductor_tx = responder.conductor_tx.clone(); + async move |dispatch: Dispatch, _cx| { + conductor_tx + .send(ConductorMessage::RightToLeft { + source_component_index: SourceComponentIndex::Successor, + message: dispatch, + }) + .await + .map_err(agent_client_protocol_core::util::internal_error) + } + }, + agent_client_protocol_core::on_receive_dispatch!(), + ), + agent_component, + )?; + responder.successor = Arc::new(connection_to_agent); + + // Spawn the proxy components + responder.spawn_proxies(client_connection.clone(), proxy_components)?; + + Ok(Dispatch::Request( + modified_req.to_untyped_message()?, + init_responder, + )) + } + + async fn handle_dispatch( + &self, + message: Dispatch, + client_connection: ConnectionTo, + conductor_tx: &mut mpsc::Sender, + ) -> Result, agent_client_protocol_core::Error> { + tracing::debug!( + method = ?message.method(), + "ConductorToClient::handle_dispatch" + ); + MatchDispatchFrom::new(message, &client_connection) + // Any incoming messages from the client are client-to-agent messages targeting the first component. + .if_message_from(Client, async move |message: Dispatch| { + tracing::debug!( + method = ?message.method(), + "ConductorToClient::handle_dispatch - matched Client" + ); + ConductorImpl::::incoming_message_from_client(conductor_tx, message).await + }) + .await + .done() + } +} + +/// Conductor acting as a proxy +impl ConductorHostRole for Proxy { + type Instantiator = Box; + + async fn initialize( + &self, + message: Dispatch, + client_connection: ConnectionTo, + instantiator: Self::Instantiator, + responder: &mut ConductorResponder, + ) -> Result { + let invalid_request = || Error::invalid_request().data("expected `initialize` request"); + + // Not yet initialized - expect an InitializeProxy request. + // Error if we get anything else. + let Dispatch::Request(request, init_responder) = message else { + message.respond_with_error(invalid_request(), client_connection.clone())?; + return Err(invalid_request()); + }; + if !InitializeProxyRequest::matches_method(request.method()) { + init_responder.respond_with_error(invalid_request())?; + return Err(invalid_request()); + } + + let InitializeProxyRequest { initialize } = + match InitializeProxyRequest::parse_message(request.method(), request.params()) { + Ok(r) => r, + Err(error) => { + init_responder.respond_with_error(error)?; + return Err(invalid_request()); + } + }; + + tracing::debug!("ensure_initialized: InitializeProxyRequest (proxy mode)"); + + // Instantiate proxies (no agent in proxy mode) + let (modified_req, proxy_components) = instantiator.instantiate_proxies(initialize).await?; + + // In proxy mode, our successor is the outer conductor (via our client connection) + responder.successor = Arc::new(GrandSuccessor); + + // Spawn the proxy components + responder.spawn_proxies(client_connection.clone(), proxy_components)?; + + Ok(Dispatch::Request( + modified_req.to_untyped_message()?, + init_responder, + )) + } + + async fn handle_dispatch( + &self, + message: Dispatch, + client_connection: ConnectionTo, + conductor_tx: &mut mpsc::Sender, + ) -> Result, agent_client_protocol_core::Error> { + tracing::debug!( + method = ?message.method(), + ?message, + "ConductorToConductor::handle_dispatch" + ); + MatchDispatchFrom::new(message, &client_connection) + .if_message_from(Agent, { + // Messages from our successor arrive already unwrapped + // (RemoteRoleStyle::Successor strips the SuccessorMessage envelope). + async |message: Dispatch| { + tracing::debug!( + method = ?message.method(), + "ConductorToConductor::handle_dispatch - matched Agent" + ); + let mut conductor_tx = conductor_tx.clone(); + ConductorImpl::::incoming_message_from_agent(&mut conductor_tx, message) + .await + } + }) + .await + // Any incoming messages from the client are client-to-agent messages targeting the first component. + .if_message_from(Client, async |message: Dispatch| { + tracing::debug!( + method = ?message.method(), + "ConductorToConductor::handle_dispatch - matched Client" + ); + let mut conductor_tx = conductor_tx.clone(); + ConductorImpl::::incoming_message_from_client(&mut conductor_tx, message) + .await + }) + .await + .done() + } +} + +pub trait ConductorSuccessor: Send + Sync + 'static { + /// Send a message to the successor. + fn send_message<'a>( + &self, + message: Dispatch, + connection_to_conductor: ConnectionTo, + responder: &'a mut ConductorResponder, + ) -> BoxFuture<'a, Result<(), agent_client_protocol_core::Error>>; +} + +impl ConductorSuccessor for agent_client_protocol_core::Error { + fn send_message<'a>( + &self, + #[expect(unused_variables)] message: Dispatch, + #[expect(unused_variables)] connection_to_conductor: ConnectionTo, + #[expect(unused_variables)] responder: &'a mut ConductorResponder, + ) -> BoxFuture<'a, Result<(), agent_client_protocol_core::Error>> { + let error = self.clone(); + Box::pin(std::future::ready(Err(error))) + } +} + +/// A dummy type handling messages sent to the conductor's +/// successor when it is acting as a proxy. +struct GrandSuccessor; + +/// When the conductor is acting as an proxy, messages sent by +/// the last proxy go to the conductor's successor. +/// +/// ```text +/// client --> Conductor -----------------------------> GrandSuccessor +/// | | +/// +-> Proxy[0] -> ... -> Proxy[n-1] -+ +/// ``` +impl ConductorSuccessor for GrandSuccessor { + fn send_message<'a>( + &self, + message: Dispatch, + connection: ConnectionTo, + _responder: &'a mut ConductorResponder, + ) -> BoxFuture<'a, Result<(), agent_client_protocol_core::Error>> { + Box::pin(async move { + debug!("Proxy mode: forwarding successor message to conductor's successor"); + connection.send_proxied_message_to(Agent, message) + }) + } +} + +/// When the conductor is acting as an agent, messages sent by +/// the last proxy to its successor go to the internal agent +/// (`self`). +impl ConductorSuccessor for ConnectionTo { + fn send_message<'a>( + &self, + message: Dispatch, + connection: ConnectionTo, + responder: &'a mut ConductorResponder, + ) -> BoxFuture<'a, Result<(), agent_client_protocol_core::Error>> { + let connection_to_agent = self.clone(); + Box::pin(async move { + debug!("Proxy mode: forwarding successor message to conductor's successor"); + responder + .forward_message_to_agent(connection, message, connection_to_agent) + .await + }) + } +} diff --git a/src/agent-client-protocol-conductor/src/conductor/mcp_bridge.rs b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge.rs new file mode 100644 index 0000000..4e9dae5 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge.rs @@ -0,0 +1,173 @@ +pub mod actor; +pub mod http; +pub mod stdio; + +use std::collections::HashMap; +use std::path::PathBuf; + +use agent_client_protocol_core::schema::{McpServer, McpServerHttp, McpServerStdio}; +use agent_client_protocol_core::{ConnectionTo, Dispatch, Role}; +use futures::{SinkExt, channel::mpsc}; +use tokio::net::TcpListener; +use tracing::info; + +pub use self::actor::McpBridgeConnectionActor; +use crate::conductor::ConductorMessage; + +/// Maintains bridges for MCP message routing. +#[derive(Default, Debug)] +pub struct McpBridgeListeners { + /// Mapping of acp:$UUID URLs to TCP bridge information for MCP message routing + listeners: HashMap, +} + +/// Information about an MCP bridge that is listening for connections from MCP clients. +#[derive(Clone, Debug)] +pub(super) struct McpBridgeListener { + /// The replacement MCP server + pub server: McpServer, +} + +/// Connection handle for sending messages to an MCP client. +#[derive(Clone, Debug)] +pub struct McpBridgeConnection { + /// Channel to send messages from MCP server (ACP proxy) to the MCP client (ACP agent). + to_mcp_client_tx: mpsc::Sender, +} + +impl McpBridgeConnection { + pub fn new(to_mcp_client_tx: mpsc::Sender) -> Self { + Self { to_mcp_client_tx } + } + + pub async fn send( + &mut self, + message: Dispatch, + ) -> Result<(), agent_client_protocol_core::Error> { + self.to_mcp_client_tx + .send(message) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + } +} + +impl McpBridgeListeners { + /// Transforms MCP servers with `acp:$UUID` URLs for agents that need bridging. + /// + /// For each MCP server with an `acp:` URL: + /// 1. Spawns a TCP listener on an ephemeral port + /// 2. Stores the mapping for message routing + /// 3. Transforms the server to use either stdio or HTTP transport depending on bridge mode + /// + /// Other MCP servers are left unchanged. + pub async fn transform_mcp_server( + &mut self, + connection: ConnectionTo, + mcp_server: &mut McpServer, + conductor_tx: &mpsc::Sender, + mcp_bridge_mode: &crate::McpBridgeMode, + ) -> Result<(), agent_client_protocol_core::Error> { + use agent_client_protocol_core::schema::McpServer; + + let McpServer::Http(http) = mcp_server else { + return Ok(()); + }; + + if !http.url.starts_with("acp:") { + return Ok(()); + } + + if !http.headers.is_empty() { + return Err(agent_client_protocol_core::Error::internal_error()); + } + + let name = &http.name; + let url = &http.url; + + info!( + server_name = name, + acp_url = url, + "Detected MCP server with ACP transport, spawning TCP bridge" + ); + + // Create oneshot channel for session_id delivery + let transformed = self + .spawn_bridge(connection, name, url, conductor_tx, mcp_bridge_mode) + .await?; + *mcp_server = transformed; + Ok(()) + } + + /// Spawn a bridge listener (HTTP or stdio) for an MCP server with ACP transport + async fn spawn_bridge( + &mut self, + connection: ConnectionTo, + server_name: &str, + acp_url: &str, + conductor_tx: &mpsc::Sender, + mcp_bridge_mode: &crate::McpBridgeMode, + ) -> anyhow::Result { + // If there is already a listener for the ACP URL, return its server + if let Some(listener) = self.listeners.get(acp_url) { + return Ok(listener.server.clone()); + } + + // Bind to ephemeral port + let tcp_listener = TcpListener::bind("127.0.0.1:0").await?; + let tcp_port = tcp_listener.local_addr()?.port(); + + info!(acp_url = acp_url, tcp_port, "Bound listener for MCP bridge"); + + let new_server = match mcp_bridge_mode { + crate::McpBridgeMode::Stdio { conductor_command } => McpServer::Stdio( + McpServerStdio::new( + server_name.to_string(), + PathBuf::from(&conductor_command[0]), + ) + .args( + conductor_command[1..] + .iter() + .cloned() + .chain(vec!["mcp".to_string(), format!("{tcp_port}")]) + .collect::>(), + ), + ), + + crate::McpBridgeMode::Http => McpServer::Http(McpServerHttp::new( + server_name.to_string(), + format!("http://localhost:{tcp_port}"), + )), + }; + + // remember for later + self.listeners.insert( + acp_url.to_string(), + McpBridgeListener { + server: new_server.clone(), + }, + ); + + connection.spawn({ + let acp_url = acp_url.to_string(); + let conductor_tx = conductor_tx.clone(); + let mcp_bridge_mode = mcp_bridge_mode.clone(); + async move { + info!( + acp_url = acp_url, + tcp_port, "now accepting bridge connections" + ); + + match mcp_bridge_mode { + crate::McpBridgeMode::Stdio { + conductor_command: _, + } => stdio::run_tcp_listener(tcp_listener, acp_url, conductor_tx).await, + crate::McpBridgeMode::Http => { + http::run_http_listener(tcp_listener, acp_url, conductor_tx).await + } + } + } + })?; + + Ok(new_server) + } +} diff --git a/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/actor.rs b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/actor.rs new file mode 100644 index 0000000..b1a4c71 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/actor.rs @@ -0,0 +1,88 @@ +use agent_client_protocol_core::{ + ConnectTo, Dispatch, DynConnectTo, role::mcp, schema::McpDisconnectNotification, +}; +use futures::{SinkExt as _, StreamExt as _, channel::mpsc}; +use tracing::info; + +use crate::conductor::ConductorMessage; + +/// Trait for actors that handle MCP bridge connections. +/// +/// Implementations bridge between MCP clients and the conductor's ACP message flow. +#[derive(Debug)] +pub struct McpBridgeConnectionActor { + /// How to connect to the MCP server + transport: DynConnectTo, + + /// Sender for messages to the conductor + conductor_tx: mpsc::Sender, + + /// Receiver for messages from the conductor to the MCP client + to_mcp_client_rx: mpsc::Receiver, +} + +impl McpBridgeConnectionActor { + pub fn new( + component: impl ConnectTo, + conductor_tx: mpsc::Sender, + to_mcp_client_rx: mpsc::Receiver, + ) -> Self { + Self { + transport: DynConnectTo::new(component), + conductor_tx, + to_mcp_client_rx, + } + } + + pub async fn run(self, connection_id: String) -> Result<(), agent_client_protocol_core::Error> { + info!(connection_id, "MCP bridge connected"); + + let McpBridgeConnectionActor { + transport, + mut conductor_tx, + to_mcp_client_rx, + } = self; + + let client = mcp::Client + .builder() + .name(format!("mpc-client-to-conductor({connection_id})")) + // When we receive a message from the MCP client, forward it to the conductor + .on_receive_dispatch( + { + let mut conductor_tx = conductor_tx.clone(); + let connection_id = connection_id.clone(); + async move |message: agent_client_protocol_core::Dispatch, _cx| { + conductor_tx + .send(ConductorMessage::McpClientToMcpServer { + connection_id: connection_id.clone(), + message, + }) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + } + }, + agent_client_protocol_core::on_receive_dispatch!(), + ) + // When we receive messages from the conductor, forward them to the MCP client + .connect_with(transport, async move |mcp_connection_to_client| { + let mut to_mcp_client_rx = to_mcp_client_rx; + while let Some(message) = to_mcp_client_rx.next().await { + mcp_connection_to_client.send_proxied_message(message)?; + } + Ok(()) + }); + let result = Box::pin(client).await; + + conductor_tx + .send(ConductorMessage::McpConnectionDisconnected { + notification: McpDisconnectNotification { + connection_id, + meta: None, + }, + }) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error())?; + + result + } +} diff --git a/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/http.rs b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/http.rs new file mode 100644 index 0000000..2b90998 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/http.rs @@ -0,0 +1,581 @@ +use agent_client_protocol_core::{BoxFuture, Channel, ConnectTo, jsonrpcmsg::Message, role::mcp}; +use axum::{ + Router, + extract::State, + http::StatusCode, + response::{IntoResponse, Response, Sse}, + routing::post, +}; +use futures::{SinkExt, StreamExt as _, channel::mpsc, future::Either, stream::Stream}; +use futures_concurrency::future::FutureExt as _; +use futures_concurrency::stream::StreamExt as _; +use rustc_hash::FxHashMap; +use std::{ + collections::{HashMap, VecDeque}, + pin::pin, + sync::Arc, +}; +use tokio::net::TcpListener; + +use crate::conductor::{ + ConductorMessage, + mcp_bridge::{McpBridgeConnection, McpBridgeConnectionActor}, +}; + +/// Runs an HTTP listener for MCP bridge connections +pub async fn run_http_listener( + tcp_listener: TcpListener, + acp_url: String, + mut conductor_tx: mpsc::Sender, +) -> Result<(), agent_client_protocol_core::Error> { + let (to_mcp_client_tx, to_mcp_client_rx) = mpsc::channel(128); + + // When we send this message to the conductor, + // it is going to go through a step or two and eventually + // spawn the McpBridgeConnectionActor, which will ferry MCP requests + // back and forth. + conductor_tx + .send(ConductorMessage::McpConnectionReceived { + acp_url, + actor: McpBridgeConnectionActor::new( + HttpMcpBridge::new(tcp_listener), + conductor_tx.clone(), + to_mcp_client_rx, + ), + connection: McpBridgeConnection::new(to_mcp_client_tx), + }) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error())?; + + Ok(()) +} + +/// A component that receives HTTP requests/responses using the HTTP transport defined by the MCP protocol. +struct HttpMcpBridge { + listener: tokio::net::TcpListener, +} + +impl HttpMcpBridge { + /// Creates a new HTTP-MCP bridge from an existing TCP listener. + fn new(listener: tokio::net::TcpListener) -> Self { + Self { listener } + } +} + +impl ConnectTo for HttpMcpBridge { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + let (channel, serve_self) = self.into_channel_and_future(); + match futures::future::select(pin!(client.connect_to(channel)), serve_self).await { + Either::Left((result, _)) | Either::Right((result, _)) => result, + } + } + + fn into_channel_and_future( + self, + ) -> ( + Channel, + BoxFuture<'static, Result<(), agent_client_protocol_core::Error>>, + ) + where + Self: Sized, + { + let (channel_a, channel_b) = Channel::duplex(); + (channel_a, Box::pin(run(self.listener, channel_b))) + } +} + +/// Error type that we use to respond to malformed HTTP requests. +#[derive(Debug, thiserror::Error)] +#[error(transparent)] +struct HttpError(#[from] agent_client_protocol_core::Error); + +impl From for HttpError { + fn from(error: axum::Error) -> Self { + HttpError(agent_client_protocol_core::util::internal_error(error)) + } +} + +impl IntoResponse for HttpError { + fn into_response(self) -> Response { + let message = format!("Error: {}", self.0); + (StatusCode::INTERNAL_SERVER_ERROR, message).into_response() + } +} + +/// Run a webserver listening on `listener` for HTTP requests at `/` +/// and communicating those requests over `channel` to the JSON-RPC server. +async fn run( + listener: TcpListener, + channel: Channel, +) -> Result<(), agent_client_protocol_core::Error> { + let (registration_tx, registration_rx) = mpsc::unbounded(); + + let state = BridgeState { registration_tx }; + + // The way that the MCP protocol works is a bit "special". + // + // Clients *POST* messages to `/`. Those are submitted to the MCP server. + // If the message is a REQUEST, then the client waits until it gets a reply. + // It expects the server to close the connection after responding. + // + // Clients can also issue a *GET* request. This will result in a stream of messages. + // + // Non-reply messages can be send to any open stream (POST, GET, etc) but must be sent to + // exactly one. + // + // There are provisions for "resuming" from a blocked point by tagging each message in the SSE stream + // with an id, but we are not implementing that because I am lazy. + async { + let app = Router::new() + .route("/", post(handle_post).get(handle_get)) + .with_state(Arc::new(state)); + + axum::serve(listener, app) + .await + .map_err(agent_client_protocol_core::util::internal_error) + } + .race(RunningServer::new().run(channel, registration_rx)) + .await +} + +/// The state we pass to our POST/GET handlers. +struct BridgeState { + /// Where to send registration messages. + registration_tx: mpsc::UnboundedSender, +} + +/// Messages from HTTP handlers to the bridge server. +#[derive(Debug)] +enum HttpMessage { + /// A JSON-RPC request (has an id, expects a response via the channel) + Request { + http_request_id: uuid::Uuid, + request: agent_client_protocol_core::jsonrpcmsg::Request, + response_tx: mpsc::UnboundedSender, + }, + /// A JSON-RPC notification (no id, no response expected) + Notification { + http_request_id: uuid::Uuid, + request: agent_client_protocol_core::jsonrpcmsg::Request, + }, + /// A JSON-RPC response from the client + Response { + http_request_id: uuid::Uuid, + response: agent_client_protocol_core::jsonrpcmsg::Response, + }, + /// A GET request to open an SSE stream for server-initiated messages + Get { + http_request_id: uuid::Uuid, + response_tx: mpsc::UnboundedSender, + }, +} + +/// Clone of `agent_client_protocol_core::jsonrpcmsg::Id` since for unfathomable reasons that does not impl Hash +#[derive(Eq, PartialEq, PartialOrd, Ord, Hash, Debug, Clone)] +enum JsonRpcId { + /// String identifier + String(String), + /// Numeric identifier + Number(u64), + /// Null identifier (for notifications) + Null, +} + +impl From for JsonRpcId { + fn from(id: agent_client_protocol_core::jsonrpcmsg::Id) -> Self { + match id { + agent_client_protocol_core::jsonrpcmsg::Id::String(s) => JsonRpcId::String(s), + agent_client_protocol_core::jsonrpcmsg::Id::Number(n) => JsonRpcId::Number(n), + agent_client_protocol_core::jsonrpcmsg::Id::Null => JsonRpcId::Null, + } + } +} + +struct RunningServer { + waiting_sessions: FxHashMap, + general_sessions: Vec, + message_deque: VecDeque, +} + +impl RunningServer { + fn new() -> Self { + RunningServer { + waiting_sessions: HashMap::default(), + general_sessions: Vec::default(), + message_deque: VecDeque::with_capacity(32), + } + } + + /// The main loop: listen for incoming HTTP messages and outgoing JSON-RPC messages. + /// + /// # Parameters + /// + /// * `channel`: The channel to use for sending/receiving JSON-RPC messages. + /// * `http_rx`: The receiver for messages from HTTP handlers. + async fn run( + mut self, + mut channel: Channel, + http_rx: mpsc::UnboundedReceiver, + ) -> Result<(), agent_client_protocol_core::Error> { + #[derive(Debug)] + enum MultiplexMessage { + FromHttpToChannel(HttpMessage), + FromChannelToHttp( + Result< + agent_client_protocol_core::jsonrpcmsg::Message, + agent_client_protocol_core::Error, + >, + ), + } + + let mut merged_stream = http_rx + .map(MultiplexMessage::FromHttpToChannel) + .merge(channel.rx.map(MultiplexMessage::FromChannelToHttp)); + + while let Some(message) = merged_stream.next().await { + tracing::trace!(?message, "received message"); + + match message { + MultiplexMessage::FromHttpToChannel(http_message) => { + self.handle_http_message(http_message, &mut channel.tx)?; + } + + MultiplexMessage::FromChannelToHttp(message) => { + let message = message.unwrap_or_else(|err| { + agent_client_protocol_core::jsonrpcmsg::Message::Response( + agent_client_protocol_core::jsonrpcmsg::Response::error( + agent_client_protocol_core::util::into_jsonrpc_error(err), + None, + ), + ) + }); + tracing::debug!( + queue_len = self.message_deque.len() + 1, + ?message, + "enqueuing outgoing message" + ); + self.message_deque.push_back(message); + } + } + + self.drain_jsonrpc_messages(); + } + + tracing::trace!("http connection terminating"); + + Ok(()) + } + + /// Handle an incoming HTTP message (request, notification, response, or GET). + fn handle_http_message( + &mut self, + message: HttpMessage, + channel_tx: &mut mpsc::UnboundedSender< + Result< + agent_client_protocol_core::jsonrpcmsg::Message, + agent_client_protocol_core::Error, + >, + >, + ) -> Result<(), agent_client_protocol_core::Error> { + match message { + HttpMessage::Request { + http_request_id, + request, + response_tx, + } => { + tracing::debug!(%http_request_id, ?request, "handling request"); + let request_id = request.id.clone().map(JsonRpcId::from); + + // Send to the JSON-RPC server + channel_tx + .unbounded_send(Ok(Message::Request(request))) + .map_err(agent_client_protocol_core::util::internal_error)?; + + // Register to receive the response + let session = RegisteredSession::new(response_tx); + if let Some(id) = request_id { + tracing::debug!(%http_request_id, session_id = %session.id, ?id, "registering waiting session"); + self.waiting_sessions.insert(id, session); + } else { + // Request without id - treat like a general session + tracing::debug!(%http_request_id, session_id = %session.id, "registering general session (request without id)"); + self.general_sessions.push(session); + } + } + + HttpMessage::Notification { + http_request_id, + request, + } => { + tracing::debug!(%http_request_id, ?request, "handling notification"); + // Just forward to the server, no response tracking needed + channel_tx + .unbounded_send(Ok(Message::Request(request))) + .map_err(agent_client_protocol_core::util::internal_error)?; + } + + HttpMessage::Response { + http_request_id, + response, + } => { + tracing::debug!(%http_request_id, ?response, "handling response"); + // Forward to the server + channel_tx + .unbounded_send(Ok(Message::Response(response))) + .map_err(agent_client_protocol_core::util::internal_error)?; + } + + HttpMessage::Get { + http_request_id, + response_tx, + } => { + let session = RegisteredSession::new(response_tx); + tracing::debug!( + %http_request_id, + session_id = %session.id, + queued_messages = self.message_deque.len(), + "handling GET (opening SSE stream)" + ); + // Register as a general session to receive server-initiated messages + self.general_sessions.push(session); + } + } + + // Purge closed sessions for good hygiene + self.purge_closed_sessions(); + + Ok(()) + } + + /// Remove messages from the queue and send them. + /// Stop if we cannot find places to send them. + fn drain_jsonrpc_messages(&mut self) { + if !self.message_deque.is_empty() { + tracing::debug!( + queue_len = self.message_deque.len(), + general_sessions = self.general_sessions.len(), + waiting_sessions = self.waiting_sessions.len(), + "draining message queue" + ); + } + + while let Some(message) = self.message_deque.pop_front() { + match self.try_dispatch_jsonrpc_message(message) { + None => { + tracing::debug!( + remaining = self.message_deque.len(), + "message dispatched successfully" + ); + } + + Some(message) => { + tracing::debug!( + remaining = self.message_deque.len() + 1, + "no available session, re-enqueuing message" + ); + self.message_deque.push_front(message); + break; + } + } + } + } + + /// Invoked when there is an outgoing JSON-RPC message to send. + /// Tries to find a suitable place to send it. + /// If it succeeds, returns `None`. + /// If there is no place to send it, returns `Some(message)`. + fn try_dispatch_jsonrpc_message( + &mut self, + mut message: agent_client_protocol_core::jsonrpcmsg::Message, + ) -> Option { + // Extract the id of the message we are replying to, if any + let message_id = match &message { + Message::Response(response) => response.id.as_ref().map(|v| v.clone().into()), + Message::Request(_) => None, + }; + + tracing::debug!(?message_id, "attempting to dispatch JSON-RPC message"); + + // If there is a specific id, try to send the message to that sender. + // This also removes them from the list of waiting sessions. + if let Some(ref message_id) = message_id + && let Some(session) = self.waiting_sessions.remove(message_id) + { + tracing::debug!(session_id = %session.id, "found waiting session, attempting send"); + + match session.outgoing_tx.unbounded_send(message) { + // Successfully sent the message, return + Ok(()) => { + tracing::debug!(session_id = %session.id, "sent to waiting session"); + return None; + } + + // If the sender died, just recover the message and send it to anyone. + Err(m) => { + tracing::debug!(session_id = %session.id, "waiting session disconnected"); + // If that sender is dead, remove them from the list + // and recover the message. + assert!(m.is_disconnected()); + message = m.into_inner(); + } + } + } + + // Try to find *somewhere* to send the message + self.purge_closed_sessions(); + tracing::debug!( + general_sessions = self.general_sessions.len(), + waiting_sessions = self.waiting_sessions.len(), + "trying to find any active session" + ); + let all_sessions = self + .general_sessions + .iter_mut() + .chain(self.waiting_sessions.values_mut()); + for session in all_sessions { + tracing::trace!(session_id = %session.id, "trying session"); + match session.outgoing_tx.unbounded_send(message) { + Ok(()) => { + tracing::debug!(session_id = %session.id, "sent to session"); + return None; + } + + Err(m) => { + tracing::debug!(session_id = %session.id, "session disconnected, trying next"); + assert!(m.is_disconnected()); + message = m.into_inner(); + } + } + } + + // If we don't find anywhere to send the message, return it. + Some(message) + } + + /// Purge sessions from the bridge state where the receiver is closed. + /// This happens when the HTTP client disconnects. + fn purge_closed_sessions(&mut self) { + self.general_sessions + .retain(|session| !session.outgoing_tx.is_closed()); + self.waiting_sessions + .retain(|_, session| !session.outgoing_tx.is_closed()); + } +} + +struct RegisteredSession { + id: uuid::Uuid, + outgoing_tx: mpsc::UnboundedSender, +} + +impl RegisteredSession { + fn new( + outgoing_tx: mpsc::UnboundedSender, + ) -> Self { + Self { + id: uuid::Uuid::new_v4(), + outgoing_tx, + } + } +} + +/// Accept a POST request carrying a JSON-RPC message from an MCP client. +/// For requests (messages with id), we return an SSE stream. +/// For notifications/responses (messages without id), we return 202 Accepted. +async fn handle_post( + State(state): State>, + body: String, +) -> Result { + let http_request_id = uuid::Uuid::new_v4(); + + // Parse incoming JSON-RPC message + let message: agent_client_protocol_core::jsonrpcmsg::Message = + serde_json::from_str(&body).map_err(agent_client_protocol_core::util::parse_error)?; + + match message { + Message::Request(request) if request.id.is_some() => { + tracing::debug!(%http_request_id, method = %request.method, "POST request received"); + // Request with id - return SSE stream for response + let (tx, mut rx) = mpsc::unbounded(); + state + .registration_tx + .unbounded_send(HttpMessage::Request { + http_request_id, + request, + response_tx: tx, + }) + .map_err(agent_client_protocol_core::util::internal_error)?; + + let stream = async_stream::stream! { + while let Some(message) = rx.next().await { + tracing::debug!(%http_request_id, "sending SSE event"); + match axum::response::sse::Event::default().json_data(message) { + Ok(v) => yield Ok(v), + Err(e) => yield Err(HttpError::from(e)), + } + } + tracing::debug!(%http_request_id, "SSE stream completed"); + }; + Ok(Sse::new(stream).into_response()) + } + + Message::Request(request) => { + tracing::debug!(%http_request_id, method = %request.method, "POST notification received"); + // Request without id is a notification + state + .registration_tx + .unbounded_send(HttpMessage::Notification { + http_request_id, + request, + }) + .map_err(agent_client_protocol_core::util::internal_error)?; + Ok(StatusCode::ACCEPTED.into_response()) + } + + Message::Response(response) => { + tracing::debug!(%http_request_id, "POST response received"); + // Response from client (rare, but possible in MCP) + state + .registration_tx + .unbounded_send(HttpMessage::Response { + http_request_id, + response, + }) + .map_err(agent_client_protocol_core::util::internal_error)?; + Ok(StatusCode::ACCEPTED.into_response()) + } + } +} + +/// Accept a GET request from an MCP client. +/// Opens an SSE stream for server-initiated messages. +async fn handle_get( + State(state): State>, +) -> Result>>, HttpError> { + let http_request_id = uuid::Uuid::new_v4(); + tracing::debug!(%http_request_id, "GET request received"); + + let (tx, mut rx) = mpsc::unbounded(); + state + .registration_tx + .unbounded_send(HttpMessage::Get { + http_request_id, + response_tx: tx, + }) + .map_err(agent_client_protocol_core::util::internal_error)?; + + let stream = async_stream::stream! { + while let Some(message) = rx.next().await { + tracing::debug!(%http_request_id, "sending SSE event"); + match axum::response::sse::Event::default().json_data(message) { + Ok(v) => yield Ok(v), + Err(e) => yield Err(HttpError::from(e)), + } + } + tracing::debug!(%http_request_id, "SSE stream completed"); + }; + + Ok(Sse::new(stream)) +} diff --git a/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/stdio.rs b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/stdio.rs new file mode 100644 index 0000000..3b15c64 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/conductor/mcp_bridge/stdio.rs @@ -0,0 +1,54 @@ +use agent_client_protocol_core::Dispatch; +use futures::{SinkExt, channel::mpsc}; +use tokio::net::{TcpListener, TcpStream}; +use tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _}; + +use crate::conductor::ConductorMessage; + +use super::{McpBridgeConnection, McpBridgeConnectionActor}; + +/// Runs the stdio bridge TCP listener, accepting connections and creating bridge actors for each. +/// +/// Loops indefinitely, accepting incoming TCP connections and spawning an MCP bridge actor +/// for each connection to handle bidirectional message forwarding between the MCP client +/// and the conductor. +pub async fn run_tcp_listener( + tcp_listener: TcpListener, + acp_url: String, + mut conductor_tx: mpsc::Sender, +) -> Result<(), agent_client_protocol_core::Error> { + // Accept connections + loop { + let (stream, _addr) = tcp_listener + .accept() + .await + .map_err(agent_client_protocol_core::Error::into_internal_error)?; + + let (to_mcp_client_tx, to_mcp_client_rx) = mpsc::channel(128); + + conductor_tx + .send(ConductorMessage::McpConnectionReceived { + acp_url: acp_url.clone(), + actor: make_stdio_actor(stream, conductor_tx.clone(), to_mcp_client_rx), + connection: McpBridgeConnection::new(to_mcp_client_tx), + }) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error())?; + } +} + +fn make_stdio_actor( + stream: TcpStream, + conductor_tx: mpsc::Sender, + to_mcp_client_rx: mpsc::Receiver, +) -> McpBridgeConnectionActor { + let (read_half, write_half) = stream.into_split(); + + // Establish bidirectional JSON-RPC connection + // The bridge will send MCP requests (tools/call, etc.) to the conductor + // The conductor can also send responses back + let transport = + agent_client_protocol_core::ByteStreams::new(write_half.compat_write(), read_half.compat()); + + McpBridgeConnectionActor::new(transport, conductor_tx, to_mcp_client_rx) +} diff --git a/src/agent-client-protocol-conductor/src/debug_logger.rs b/src/agent-client-protocol-conductor/src/debug_logger.rs new file mode 100644 index 0000000..39bf836 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/debug_logger.rs @@ -0,0 +1,179 @@ +//! Debug logging for conductor + +use chrono::Local; +use std::io::Write; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::Instant; +use tokio::fs::OpenOptions; +use tokio::io::AsyncWriteExt; +use tokio::sync::Mutex; + +/// A debug logger that writes lines to a timestamped log file +pub struct DebugLogger { + writer: Arc>>, + start_time: Instant, +} + +impl DebugLogger { + /// Create a new debug logger with a timestamped log file + pub async fn new( + debug_dir: Option, + component_commands: &[String], + ) -> Result { + // Create log directory + let log_dir = debug_dir.unwrap_or_else(|| PathBuf::from(".")); + tokio::fs::create_dir_all(&log_dir).await?; + + // Create timestamped log file + let timestamp = Local::now().format("%Y%m%d-%H%M%S"); + let log_file = log_dir.join(format!("{timestamp}.log")); + + let file = OpenOptions::new() + .create(true) + .append(true) + .open(&log_file) + .await?; + + let mut writer = tokio::io::BufWriter::new(file); + + // Write header + writer.write_all(b"=== Conductor Debug Log ===\n").await?; + writer + .write_all(format!("Started: {}\n", Local::now().to_rfc3339()).as_bytes()) + .await?; + writer.write_all(b"Components:\n").await?; + for (i, cmd) in component_commands.iter().enumerate() { + writer + .write_all(format!(" {i}: {cmd}\n").as_bytes()) + .await?; + } + writer.write_all(b"========================\n").await?; + writer.flush().await?; + + eprintln!("Debug logging to: {}", log_file.display()); + + Ok(Self { + writer: Arc::new(Mutex::new(writer)), + start_time: Instant::now(), + }) + } + + /// Get the elapsed time since the logger started, in milliseconds + fn elapsed_ms(&self) -> u128 { + self.start_time.elapsed().as_millis() + } + + /// Create a callback for a specific component + pub fn create_callback( + &self, + component_label: String, + ) -> impl Fn(&str, agent_client_protocol_tokio::LineDirection) + Send + Sync + 'static { + let writer = self.writer.clone(); + let start_time = self.start_time; + move |line: &str, direction: agent_client_protocol_tokio::LineDirection| { + let writer = writer.clone(); + let component_label = component_label.clone(); + let line = line.to_string(); + let elapsed_ms = start_time.elapsed().as_millis(); + tokio::spawn(async move { + let arrow = match direction { + agent_client_protocol_tokio::LineDirection::Stdin => "→", + agent_client_protocol_tokio::LineDirection::Stdout => "←", + agent_client_protocol_tokio::LineDirection::Stderr => "!", + }; + + // Strip ANSI escape codes from stderr to keep logs clean + let cleaned_line = if matches!( + direction, + agent_client_protocol_tokio::LineDirection::Stderr + ) { + let bytes = strip_ansi_escapes::strip(&line); + String::from_utf8_lossy(&bytes).to_string() + } else { + line + }; + + let log_line = + format!("{component_label} {arrow} +{elapsed_ms}ms {cleaned_line}\n"); + let mut writer = writer.lock().await; + drop(writer.write_all(log_line.as_bytes()).await); + drop(writer.flush().await); + }); + } + } + + /// Write a tracing log line to the debug file + /// This is synchronous and blocks, suitable for use with tracing's MakeWriter + pub fn write_tracing_log(&self, line: &str) { + let writer = self.writer.clone(); + let line = line.to_string(); + let elapsed_ms = self.elapsed_ms(); + tokio::spawn(async move { + // Strip ANSI escape codes to keep logs clean + let bytes = strip_ansi_escapes::strip(&line); + let cleaned_line = String::from_utf8_lossy(&bytes); + + let log_line = format!("C ! +{}ms {}\n", elapsed_ms, cleaned_line.trim_end()); + let mut writer = writer.lock().await; + drop(writer.write_all(log_line.as_bytes()).await); + drop(writer.flush().await); + }); + } + + /// Create a writer for tracing logs + pub fn create_tracing_writer(&self) -> DebugLogWriter { + DebugLogWriter { + logger: self.clone(), + buffer: Vec::new(), + } + } +} + +impl Clone for DebugLogger { + fn clone(&self) -> Self { + Self { + writer: self.writer.clone(), + start_time: self.start_time, + } + } +} + +/// A writer that sends tracing logs to the debug logger +pub struct DebugLogWriter { + logger: DebugLogger, + buffer: Vec, +} + +impl Write for DebugLogWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.buffer.extend_from_slice(buf); + + // Write complete lines + while let Some(newline_pos) = self.buffer.iter().position(|&b| b == b'\n') { + let line = self.buffer.drain(..=newline_pos).collect::>(); + let line_str = String::from_utf8_lossy(&line); + self.logger.write_tracing_log(&line_str); + } + + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + if !self.buffer.is_empty() { + let line = self.buffer.drain(..).collect::>(); + let line_str = String::from_utf8_lossy(&line); + self.logger.write_tracing_log(&line_str); + } + Ok(()) + } +} + +impl Clone for DebugLogWriter { + fn clone(&self) -> Self { + Self { + logger: self.logger.clone(), + buffer: Vec::new(), + } + } +} diff --git a/src/agent-client-protocol-conductor/src/lib.rs b/src/agent-client-protocol-conductor/src/lib.rs new file mode 100644 index 0000000..5c8f0c5 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/lib.rs @@ -0,0 +1,403 @@ +//! # agent-client-protocol-conductor +//! +//! Binary for orchestrating ACP proxy chains. +//! +//! ## What is the conductor? +//! +//! The conductor is a tool that manages proxy chains - it spawns proxy components and the base agent, +//! then routes messages between them. From the editor's perspective, the conductor appears as a single ACP agent. +//! +//! ```text +//! Editor ← stdio → Conductor → Proxy 1 → Proxy 2 → Agent +//! ``` +//! +//! ## Usage +//! +//! ### Agent Mode +//! +//! Orchestrate a chain of proxies in front of an agent: +//! +//! ```bash +//! # Chain format: proxy1 proxy2 ... agent +//! agent-client-protocol-conductor agent "python proxy1.py" "python proxy2.py" "python base-agent.py" +//! ``` +//! +//! The conductor: +//! 1. Spawns each component as a subprocess +//! 2. Connects them in a chain +//! 3. Presents as a single agent on stdin/stdout +//! 4. Manages the lifecycle of all processes +//! +//! ### MCP Bridge Mode +//! +//! Connect stdio to a TCP-based MCP server: +//! +//! ```bash +//! # Bridge stdio to MCP server on localhost:8080 +//! agent-client-protocol-conductor mcp 8080 +//! ``` +//! +//! This allows stdio-based tools to communicate with TCP MCP servers. +//! +//! ## How It Works +//! +//! **Component Communication:** +//! - Editor talks to conductor via stdio +//! - Conductor uses `_proxy/successor/*` protocol extensions to route messages +//! - Each proxy can intercept, transform, or forward messages +//! - Final agent receives standard ACP messages +//! +//! **Process Management:** +//! - All components are spawned as child processes +//! - When conductor exits, all children are terminated +//! - Errors in any component bring down the entire chain +//! +//! ## Example Use Case +//! +//! Add Sparkle embodiment + custom tools to any agent: +//! +//! ```bash +//! agent-client-protocol-conductor agent \ +//! "sparkle-acp-proxy" \ +//! "my-custom-tools-proxy" \ +//! "claude-agent" +//! ``` +//! +//! This creates a stack where: +//! 1. Sparkle proxy injects MCP servers and prepends embodiment +//! 2. Custom tools proxy adds domain-specific functionality +//! 3. Base agent handles the actual AI responses +//! +//! ## Related Crates +//! +//! - **[agent-client-protocol-proxy](https://crates.io/crates/agent-client-protocol-proxy)** - Framework for building proxy components +//! - **[agent-client-protocol-core](https://crates.io/crates/agent-client-protocol-core)** - Core ACP SDK +//! - **[agent-client-protocol-tokio](https://crates.io/crates/agent-client-protocol-tokio)** - Tokio utilities for process spawning + +use std::path::PathBuf; +use std::str::FromStr; + +/// Core conductor logic for orchestrating proxy chains +mod conductor; +/// Debug logging for conductor +mod debug_logger; +/// MCP bridge functionality for TCP-based MCP servers +mod mcp_bridge; +mod snoop; +/// Trace event types for sequence diagram viewer +pub mod trace; + +pub use self::conductor::*; + +use clap::{Parser, Subcommand}; + +use agent_client_protocol_core::{Client, Conductor, DynConnectTo, schema::InitializeRequest}; +use agent_client_protocol_tokio::{AcpAgent, Stdio}; +use tracing::Instrument; +use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; + +/// Wrapper for command-line component lists that can serve as either +/// proxies-only (for proxy mode) or proxies+agent (for agent mode). +/// +/// This exists because `AcpAgent` implements `Component` for all `L`, +/// so a `Vec` can be used as either a list of proxies or as +/// proxies + final agent depending on the conductor mode. +#[derive(Debug)] +pub struct CommandLineComponents(pub Vec); + +impl InstantiateProxies for CommandLineComponents { + fn instantiate_proxies( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + (InitializeRequest, Vec>), + agent_client_protocol_core::Error, + >, + > { + Box::pin(async move { + let proxies = self.0.into_iter().map(DynConnectTo::new).collect(); + Ok((req, proxies)) + }) + } +} + +impl InstantiateProxiesAndAgent for CommandLineComponents { + fn instantiate_proxies_and_agent( + self: Box, + req: InitializeRequest, + ) -> futures::future::BoxFuture< + 'static, + Result< + ( + InitializeRequest, + Vec>, + DynConnectTo, + ), + agent_client_protocol_core::Error, + >, + > { + Box::pin(async move { + let mut iter = self.0.into_iter().peekable(); + let mut proxies: Vec> = Vec::new(); + + // All but the last element are proxies + while let Some(component) = iter.next() { + if iter.peek().is_some() { + proxies.push(DynConnectTo::new(component)); + } else { + // Last element is the agent + let agent = DynConnectTo::new(component); + return Ok((req, proxies, agent)); + } + } + + Err(agent_client_protocol_core::util::internal_error( + "no agent component in list", + )) + }) + } +} + +/// Wrapper to implement WriteEvent for TraceHandle. +struct TraceHandleWriter(agent_client_protocol_trace_viewer::TraceHandle); + +impl trace::WriteEvent for TraceHandleWriter { + fn write_event(&mut self, event: &trace::TraceEvent) -> std::io::Result<()> { + let value = serde_json::to_value(event).map_err(std::io::Error::other)?; + self.0.push(value); + Ok(()) + } +} + +/// Mode for the MCP bridge. +#[derive(Debug, Clone, Default)] +pub enum McpBridgeMode { + /// Use stdio-based MCP bridge with a conductor subprocess. + Stdio { + /// Command and args to spawn conductor MCP bridge processes. + /// E.g., vec!["conductor"] or vec!["cargo", "run", "-p", "conductor", "--"] + conductor_command: Vec, + }, + + /// Use HTTP-based MCP bridge + #[default] + Http, +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +pub struct ConductorArgs { + /// Enable debug logging of all stdin/stdout/stderr from components + #[arg(long)] + pub debug: bool, + + /// Directory for debug log files (defaults to current directory) + #[arg(long)] + pub debug_dir: Option, + + /// Set log level (e.g., "trace", "debug", "info", "warn", "error", or module-specific like "conductor=debug") + /// Only applies when --debug is enabled + #[arg(long)] + pub log: Option, + + /// Path to write trace events for sequence diagram visualization. + /// Events are written as newline-delimited JSON (.jsons format). + #[arg(long)] + pub trace: Option, + + /// Serve trace viewer in browser with live updates. + /// Can be used alone (in-memory) or with --trace (file-backed). + #[arg(long)] + pub serve: bool, + + #[command(subcommand)] + pub command: ConductorCommand, +} + +#[derive(Subcommand, Debug)] +pub enum ConductorCommand { + /// Run as agent orchestrator managing a proxy chain + Agent { + /// Name of the agent + #[arg(short, long, default_value = "conductor")] + name: String, + + /// List of commands to chain together; the final command must be the agent. + components: Vec, + }, + + /// Run as a proxy orchestrating a proxy chain + Proxy { + /// Name of the proxy + #[arg(short, long, default_value = "conductor")] + name: String, + + /// List of proxy commands to chain together + proxies: Vec, + }, + + /// Run as MCP bridge connecting stdio to TCP + Mcp { + /// TCP port to connect to on localhost + port: u16, + }, +} + +impl ConductorArgs { + /// Main entry point that sets up tracing and runs the conductor + pub async fn main(self) -> anyhow::Result<()> { + let pid = std::process::id(); + let cwd = std::env::current_dir() + .map_or_else(|_| "".to_string(), |p| p.display().to_string()); + + // Only set up tracing if --debug is enabled + let debug_logger = if self.debug { + // Extract proxy list to create the debug logger + let components = match &self.command { + ConductorCommand::Agent { components, .. } => components.clone(), + ConductorCommand::Proxy { proxies, .. } => proxies.clone(), + ConductorCommand::Mcp { .. } => Vec::new(), + }; + + // Create debug logger + Some( + debug_logger::DebugLogger::new(self.debug_dir.clone(), &components) + .await + .map_err(|e| anyhow::anyhow!("Failed to create debug logger: {e}"))?, + ) + } else { + None + }; + + if let Some(debug_logger) = &debug_logger { + // Set up log level from --log flag, defaulting to "info" + let log_level = self.log.as_deref().unwrap_or("info"); + + // Set up tracing to write to the debug file with "C !" prefix + let tracing_writer = debug_logger.create_tracing_writer(); + tracing_subscriber::registry() + .with(EnvFilter::new(log_level)) + .with( + tracing_subscriber::fmt::layer() + .with_target(true) + .with_writer(move || tracing_writer.clone()), + ) + .init(); + + tracing::info!(pid = %pid, cwd = %cwd, level = %log_level, "Conductor starting with debug logging"); + } + + // Set up tracing based on --trace and --serve flags + let (trace_writer, _viewer_server) = match (&self.trace, self.serve) { + // --trace only: write to file + (Some(trace_path), false) => { + let writer = trace::TraceWriter::from_path(trace_path) + .map_err(|e| anyhow::anyhow!("Failed to create trace writer: {e}"))?; + (Some(writer), None) + } + // --serve only: in-memory with viewer + (None, true) => { + let (handle, server) = agent_client_protocol_trace_viewer::serve_memory( + agent_client_protocol_trace_viewer::TraceViewerConfig::default(), + )?; + let writer = trace::TraceWriter::new(TraceHandleWriter(handle)); + (Some(writer), Some(tokio::spawn(server))) + } + // --trace --serve: write to file and serve it + (Some(trace_path), true) => { + let writer = trace::TraceWriter::from_path(trace_path) + .map_err(|e| anyhow::anyhow!("Failed to create trace writer: {e}"))?; + let server = agent_client_protocol_trace_viewer::serve_file( + trace_path.clone(), + agent_client_protocol_trace_viewer::TraceViewerConfig::default(), + ); + (Some(writer), Some(tokio::spawn(server))) + } + // Neither: no tracing + (None, false) => (None, None), + }; + + Box::pin( + self.run(debug_logger.as_ref(), trace_writer) + .instrument(tracing::info_span!("conductor", pid = %pid, cwd = %cwd)), + ) + .await + .map_err(|err| anyhow::anyhow!("{err}")) + } + + async fn run( + self, + debug_logger: Option<&debug_logger::DebugLogger>, + trace_writer: Option, + ) -> Result<(), agent_client_protocol_core::Error> { + match self.command { + ConductorCommand::Agent { name, components } => { + Box::pin(initialize_conductor( + debug_logger, + trace_writer, + name, + components, + ConductorImpl::new_agent, + )) + .await + } + ConductorCommand::Proxy { name, proxies } => { + Box::pin(initialize_conductor( + debug_logger, + trace_writer, + name, + proxies, + ConductorImpl::new_proxy, + )) + .await + } + ConductorCommand::Mcp { port } => mcp_bridge::run_mcp_bridge(port).await, + } + } +} + +async fn initialize_conductor( + debug_logger: Option<&debug_logger::DebugLogger>, + trace_writer: Option, + name: String, + components: Vec, + new_conductor: impl FnOnce( + String, + CommandLineComponents, + crate::McpBridgeMode, + ) -> ConductorImpl, +) -> Result<(), agent_client_protocol_core::Error> { + // Parse agents and optionally wrap with debug callbacks + let providers: Vec = components + .into_iter() + .enumerate() + .map(|(i, s)| { + let mut agent = AcpAgent::from_str(&s)?; + if let Some(logger) = debug_logger { + agent = agent.with_debug(logger.create_callback(i.to_string())); + } + Ok(agent) + }) + .collect::, agent_client_protocol_core::Error>>()?; + + // Create Stdio component with optional debug logging + let stdio = if let Some(logger) = debug_logger { + Stdio::new().with_debug(logger.create_callback("C".to_string())) + } else { + Stdio::new() + }; + + // Create conductor with optional trace writer + let mut conductor = new_conductor( + name, + CommandLineComponents(providers), + McpBridgeMode::default(), + ); + if let Some(writer) = trace_writer { + conductor = conductor.with_trace_writer(writer); + } + + conductor.run(stdio).await +} diff --git a/src/agent-client-protocol-conductor/src/main.rs b/src/agent-client-protocol-conductor/src/main.rs new file mode 100644 index 0000000..c02cede --- /dev/null +++ b/src/agent-client-protocol-conductor/src/main.rs @@ -0,0 +1,7 @@ +use agent_client_protocol_conductor::ConductorArgs; +use clap::Parser; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + ConductorArgs::parse().main().await +} diff --git a/src/agent-client-protocol-conductor/src/mcp_bridge.rs b/src/agent-client-protocol-conductor/src/mcp_bridge.rs new file mode 100644 index 0000000..d6b3ca0 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/mcp_bridge.rs @@ -0,0 +1,122 @@ +//! MCP Bridge: Bridges MCP JSON-RPC over stdio to TCP connection +//! +//! This module implements `conductor mcp $port` mode, which acts as an MCP server +//! over stdio but forwards all messages to/from a TCP connection on localhost:$port. +//! +//! The main conductor (in agent mode) listens on the TCP port and translates between +//! TCP (raw JSON-RPC) and ACP `_mcp/*` extension messages. + +use anyhow::Context; +use serde_json::Value; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::net::TcpStream; + +/// Run the MCP bridge: stdio ↔ TCP +/// +/// Reads MCP JSON-RPC messages from stdin, forwards to TCP connection. +/// Reads responses from TCP, writes to stdout. +pub async fn run_mcp_bridge(port: u16) -> Result<(), agent_client_protocol_core::Error> { + tracing::info!("MCP bridge starting, connecting to localhost:{}", port); + + // Connect to the main conductor via TCP + let stream = connect_with_retry(port).await?; + let (tcp_read, mut tcp_write) = stream.into_split(); + + // Set up stdio + let stdin = tokio::io::stdin(); + let stdout = tokio::io::stdout(); + let mut stdin_reader = BufReader::new(stdin); + let mut stdout_writer = stdout; + let mut tcp_reader = BufReader::new(tcp_read); + + // Prepare line buffers + let mut stdin_line = String::new(); + let mut tcp_line = String::new(); + + tracing::info!("MCP bridge connected, starting message loop"); + + loop { + tokio::select! { + // Read from stdin → send to TCP + result = stdin_reader.read_line(&mut stdin_line) => { + let n = result.context("Failed to read from stdin")?; + + if n == 0 { + tracing::info!("Stdin closed, shutting down bridge"); + break; + } + + // Parse to validate JSON + drop(serde_json::from_str::(stdin_line.trim()) + .context("Invalid JSON from stdin")?); + + tracing::debug!("Bridge: stdin → TCP: {}", stdin_line.trim()); + + // Forward to TCP + tcp_write.write_all(stdin_line.as_bytes()).await + .context("Failed to write to TCP")?; + tcp_write.flush().await + .context("Failed to flush TCP")?; + + stdin_line.clear(); + } + + // Read from TCP → send to stdout + result = tcp_reader.read_line(&mut tcp_line) => { + let n = result.context("Failed to read from TCP")?; + + if n == 0 { + tracing::info!("TCP connection closed, shutting down bridge"); + break; + } + + // Parse to validate JSON + drop(serde_json::from_str::(tcp_line.trim()) + .context("Invalid JSON from TCP")?); + + tracing::debug!("Bridge: TCP → stdout: {}", tcp_line.trim()); + + // Forward to stdout + stdout_writer.write_all(tcp_line.as_bytes()).await + .context("Failed to write to stdout")?; + stdout_writer.flush().await + .context("Failed to flush stdout")?; + + tcp_line.clear(); + } + } + } + + tracing::info!("MCP bridge shutting down"); + Ok(()) +} + +/// Connect to TCP port with retry logic +async fn connect_with_retry(port: u16) -> Result { + let max_retries = 10; + let mut retry_delay_ms = 50; + + for attempt in 1..=max_retries { + match TcpStream::connect(format!("127.0.0.1:{port}")).await { + Ok(stream) => { + tracing::info!("Connected to localhost:{} on attempt {}", port, attempt); + return Ok(stream); + } + Err(e) if attempt < max_retries => { + tracing::debug!( + "Connection attempt {} failed: {}, retrying in {}ms", + attempt, + e, + retry_delay_ms + ); + tokio::time::sleep(tokio::time::Duration::from_millis(retry_delay_ms)).await; + retry_delay_ms = (retry_delay_ms * 2).min(1000); // Exponential backoff, max 1s + } + Err(e) => { + return Err(agent_client_protocol_core::Error::into_internal_error(e)); + } + } + } + + unreachable!() +} diff --git a/src/agent-client-protocol-conductor/src/snoop.rs b/src/agent-client-protocol-conductor/src/snoop.rs new file mode 100644 index 0000000..455fd32 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/snoop.rs @@ -0,0 +1,91 @@ +use agent_client_protocol_core::{Channel, ConnectTo, DynConnectTo, Role, jsonrpcmsg}; +use futures::StreamExt; +use futures_concurrency::future::TryJoin; + +pub struct SnooperComponent { + base_component: DynConnectTo, + incoming_message: Box< + dyn FnMut(&jsonrpcmsg::Message) -> Result<(), agent_client_protocol_core::Error> + + Send + + Sync, + >, + outgoing_message: Box< + dyn FnMut(&jsonrpcmsg::Message) -> Result<(), agent_client_protocol_core::Error> + + Send + + Sync, + >, +} + +impl SnooperComponent { + pub fn new( + base_component: impl ConnectTo, + incoming_message: impl FnMut( + &jsonrpcmsg::Message, + ) -> Result<(), agent_client_protocol_core::Error> + + Send + + Sync + + 'static, + outgoing_message: impl FnMut( + &jsonrpcmsg::Message, + ) -> Result<(), agent_client_protocol_core::Error> + + Send + + Sync + + 'static, + ) -> Self { + Self { + base_component: DynConnectTo::new(base_component), + incoming_message: Box::new(incoming_message), + outgoing_message: Box::new(outgoing_message), + } + } +} + +impl ConnectTo for SnooperComponent { + async fn connect_to( + mut self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + let (client_a, mut client_b) = Channel::duplex(); + + let client_future = client.connect_to(client_a); + + let (mut base_channel, base_future) = self.base_component.into_channel_and_future(); + + // Read messages send by `client`. These are 'incoming' to our wrapped + // component. + let snoop_incoming = async { + while let Some(msg) = client_b.rx.next().await { + if let Ok(msg) = &msg { + (self.incoming_message)(msg)?; + } + + base_channel + .tx + .unbounded_send(msg) + .map_err(agent_client_protocol_core::util::internal_error)?; + } + Ok(()) + }; + + // Read messages send by `base`. These are 'outgoing' from our wrapped + // component. + let snoop_outgoing = async { + while let Some(msg) = base_channel.rx.next().await { + if let Ok(msg) = &msg { + (self.outgoing_message)(msg)?; + } + + client_b + .tx + .unbounded_send(msg) + .map_err(agent_client_protocol_core::util::internal_error)?; + } + Ok(()) + }; + + (client_future, base_future, snoop_incoming, snoop_outgoing) + .try_join() + .await?; + Ok(()) + } +} diff --git a/src/agent-client-protocol-conductor/src/trace.rs b/src/agent-client-protocol-conductor/src/trace.rs new file mode 100644 index 0000000..2573ca9 --- /dev/null +++ b/src/agent-client-protocol-conductor/src/trace.rs @@ -0,0 +1,545 @@ +//! Trace event types for the sequence diagram viewer. +//! +//! Events are serialized as newline-delimited JSON (`.jsons` files). +//! The viewer loads these files to render interactive sequence diagrams. + +use std::collections::HashMap; +use std::fs::OpenOptions; +use std::io::{BufWriter, Write}; +use std::path::Path; +use std::time::Instant; + +use agent_client_protocol_core::schema::{McpOverAcpMessage, SuccessorMessage}; +use agent_client_protocol_core::{DynConnectTo, JsonRpcMessage, Role, UntypedMessage, jsonrpcmsg}; +use rustc_hash::FxHashMap; +use serde::{Deserialize, Serialize}; + +use crate::ComponentIndex; +use crate::snoop::SnooperComponent; + +/// A trace event representing message flow between components. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +#[non_exhaustive] +pub enum TraceEvent { + /// A JSON-RPC request from one component to another. + Request(RequestEvent), + + /// A JSON-RPC response to a prior request. + Response(ResponseEvent), + + /// A JSON-RPC notification (no response expected). + Notification(NotificationEvent), +} + +/// Protocol type for messages. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +#[non_exhaustive] +pub enum Protocol { + /// Standard ACP protocol messages. + Acp, + /// MCP-over-ACP messages (agent calling proxy's MCP server). + Mcp, +} + +/// A JSON-RPC request from one component to another. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct RequestEvent { + /// Monotonic timestamp (seconds since trace start). + pub ts: f64, + + /// Protocol: ACP or MCP. + pub protocol: Protocol, + + /// Source component (e.g., "client", "proxy:0", "proxy:1", "agent"). + pub from: String, + + /// Destination component. + pub to: String, + + /// JSON-RPC request ID (for correlating with response). + pub id: serde_json::Value, + + /// JSON-RPC method name. + pub method: String, + + /// ACP session ID, if known (null before session/new completes). + #[serde(skip_serializing_if = "Option::is_none")] + pub session: Option, + + /// Full request params. + pub params: serde_json::Value, +} + +/// A JSON-RPC response to a prior request. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct ResponseEvent { + /// Monotonic timestamp (seconds since trace start). + pub ts: f64, + + /// Source component (who sent the response). + pub from: String, + + /// Destination component (who receives the response). + pub to: String, + + /// JSON-RPC request ID this responds to. + pub id: serde_json::Value, + + /// True if this is an error response. + pub is_error: bool, + + /// Response result or error object. + pub payload: serde_json::Value, +} + +/// A JSON-RPC notification (no response expected). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] +pub struct NotificationEvent { + /// Monotonic timestamp (seconds since trace start). + pub ts: f64, + + /// Protocol: ACP or MCP. + pub protocol: Protocol, + + /// Source component. + pub from: String, + + /// Destination component. + pub to: String, + + /// JSON-RPC method name. + pub method: String, + + /// ACP session ID, if known. + #[serde(skip_serializing_if = "Option::is_none")] + pub session: Option, + + /// Full notification params. + pub params: serde_json::Value, +} + +/// Trait for destinations that can receive trace events. +pub trait WriteEvent: Send + 'static { + /// Write a trace event to the destination. + fn write_event(&mut self, event: &TraceEvent) -> std::io::Result<()>; +} + +/// Writes trace events as newline-delimited JSON to a `Write` impl. +pub(crate) struct EventWriter { + writer: W, +} + +impl EventWriter { + pub fn new(writer: W) -> Self { + Self { writer } + } +} + +impl WriteEvent for EventWriter { + fn write_event(&mut self, event: &TraceEvent) -> std::io::Result<()> { + serde_json::to_writer(&mut self.writer, event).map_err(std::io::Error::other)?; + self.writer.write_all(b"\n")?; + self.writer.flush() + } +} + +/// Impl for UnboundedSender - sends events to a channel (useful for testing). +impl WriteEvent for futures::channel::mpsc::UnboundedSender { + fn write_event(&mut self, event: &TraceEvent) -> std::io::Result<()> { + self.unbounded_send(event.clone()) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::BrokenPipe, e)) + } +} + +/// Writer for trace events. +pub struct TraceWriter { + dest: Box, + start_time: Instant, + + /// When we trace a request, we store its id along with the + /// details here. When we see responses, we try to match them up. + request_details: FxHashMap, +} + +impl std::fmt::Debug for TraceWriter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TraceWriter") + .field("start_time", &self.start_time) + .finish_non_exhaustive() + } +} + +struct RequestDetails { + #[expect(dead_code)] + protocol: Protocol, + + #[expect(dead_code)] + method: String, + + request_from: ComponentIndex, + request_to: ComponentIndex, +} + +impl TraceWriter { + /// Create a new trace writer from any WriteEvent destination. + pub fn new(dest: D) -> Self { + Self { + dest: Box::new(dest), + start_time: Instant::now(), + request_details: HashMap::default(), + } + } + + /// Create a new trace writer that writes to a file path. + pub fn from_path(path: impl AsRef) -> std::io::Result { + let file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(path.as_ref())?; + Ok(Self::new(EventWriter::new(BufWriter::new(file)))) + } + + /// Get the elapsed time since trace start, in seconds. + fn elapsed(&self) -> f64 { + self.start_time.elapsed().as_secs_f64() + } + + /// Write a trace event. + fn write_event(&mut self, event: &TraceEvent) { + // Ignore errors - tracing should not break the conductor + drop(self.dest.write_event(event)); + } + + /// Write a request event. + #[expect(clippy::too_many_arguments)] + fn request( + &mut self, + protocol: Protocol, + from: ComponentIndex, + to: ComponentIndex, + id: serde_json::Value, + method: String, + session: Option, + params: serde_json::Value, + ) { + self.request_details.insert( + id.clone(), + RequestDetails { + protocol, + method: method.clone(), + request_from: from, + request_to: to, + }, + ); + self.write_event(&TraceEvent::Request(RequestEvent { + ts: self.elapsed(), + protocol, + from: format!("{from:?}"), + to: format!("{to:?}"), + id, + method, + session, + params, + })); + } + + /// Write a response event. + fn response( + &mut self, + from: ComponentIndex, + to: ComponentIndex, + id: serde_json::Value, + is_error: bool, + payload: serde_json::Value, + ) { + self.write_event(&TraceEvent::Response(ResponseEvent { + ts: self.elapsed(), + from: format!("{from:?}"), + to: format!("{to:?}"), + id, + is_error, + payload, + })); + } + + /// Write a notification event. + fn notification( + &mut self, + protocol: Protocol, + from: ComponentIndex, + to: ComponentIndex, + method: impl Into, + session: Option, + params: serde_json::Value, + ) { + self.write_event(&TraceEvent::Notification(NotificationEvent { + ts: self.elapsed(), + protocol, + from: format!("{from:?}"), + to: format!("{to:?}"), + method: method.into(), + session, + params, + })); + } + + /// Trace a raw JSON-RPC message being sent from one component to another. + fn trace_message(&mut self, traced_message: TracedMessage) { + let TracedMessage { + component_index, + successor_index, + incoming, + message, + } = traced_message; + + // We get every message going into or out of a proxy. This includes + // a fair number of duplicates: for example, if proxy P0 sends to P1, + // we'll get it as an *outgoing* message from P0 and an *incoming* message to P1. + // So we want to keep just one copy. + // + // We retain: + // + // * Incoming requests/notifications targeting a PROXY. + // * Incoming requests/notifications targeting the AGENT. + + match message { + jsonrpcmsg::Message::Request(req) => { + let MessageInfo { + successor, + id, + protocol, + method, + params, + } = MessageInfo::from_req(req); + + let (from, to) = match (successor, incoming, component_index, successor_index) { + // An incoming request/notification to a proxy from its predecessor. + (Successor(false), Incoming(true), ComponentIndex::Proxy(proxy_index), _) => ( + ComponentIndex::predecessor_of(proxy_index), + ComponentIndex::Proxy(proxy_index), + ), + + // An incoming request/notification to any component from its successor. + // + // This includes incoming messages to the client in the case where we have no proxies. + (Successor(true), Incoming(true), component_index, successor_index) => { + (successor_index, component_index) + } + + // An outgoing request/notification from a component to its successor + // *and* its successor is not a proxy. + // + // (If its successor is a proxy, we ignore it, because we'll also see the + // message in "incoming" form). + (Successor(true), Incoming(false), component_index, ComponentIndex::Agent) => { + (component_index, ComponentIndex::Agent) + } + + _ => return, + }; + + match id { + Some(id) => { + self.request(protocol, from, to, id_to_json(&id), method, None, params); + } + None => { + self.notification(protocol, from, to, method, None, params); + } + } + } + jsonrpcmsg::Message::Response(resp) => { + // Lookup the response by its id. + // All of the messages we are intercepting go to our proxies, + // and we always assign them globally unique + if let Some(id) = resp.id { + let id = id_to_json(&id); + if let Some(RequestDetails { + protocol: _, + method: _, + request_from, + request_to, + }) = self.request_details.remove(&id) + { + let (is_error, payload) = match (&resp.result, &resp.error) { + (Some(result), _) => (false, result.clone()), + (_, Some(error)) => { + (true, serde_json::to_value(error).unwrap_or_default()) + } + (None, None) => (false, serde_json::Value::Null), + }; + self.response(request_to, request_from, id, is_error, payload); + } + } + } + } + } + + /// Spawn a trace writer task. + /// + /// Returns a `TraceHandle` that can be cloned and used from multiple tasks, + /// and a future that should be spawned (e.g., via `with_spawned`). + pub(crate) fn spawn( + mut self: TraceWriter, + ) -> ( + TraceHandle, + impl std::future::Future>, + ) { + use futures::StreamExt; + + let (tx, mut rx) = futures::channel::mpsc::unbounded(); + + let future = async move { + while let Some(event) = rx.next().await { + self.trace_message(event); + } + Ok(()) + }; + + (TraceHandle { tx }, future) + } +} + +/// A cloneable handle for sending trace events to the trace writer task. +/// +/// Create with [`spawn_trace_writer`], then clone and pass to bridges. +#[derive(Clone, Debug)] +pub(crate) struct TraceHandle { + tx: futures::channel::mpsc::UnboundedSender, +} + +impl TraceHandle { + /// Trace a raw JSON-RPC message being sent from one component to another. + fn trace_message( + &self, + component_index: ComponentIndex, + successor_index: ComponentIndex, + incoming: Incoming, + message: &jsonrpcmsg::Message, + ) -> Result<(), agent_client_protocol_core::Error> { + self.tx + .unbounded_send(TracedMessage { + component_index, + successor_index, + incoming, + message: message.clone(), + }) + .map_err(agent_client_protocol_core::util::internal_error) + } + + /// Create a tracing bridge that wraps a proxy component. + /// + /// Spawns a bridge task that forwards messages between the channel and the component + /// while tracing them. Returns the wrapped component. + /// + /// Tracing strategy: + /// - **Left→Right (incoming)**: Trace requests/notifications, skip responses + /// - **Right→Left (outgoing)**: Trace responses, and if `trace_outgoing_requests` is true, + /// also trace requests/notifications (needed for edge bridges at conductor boundaries) + /// + /// - `cx`: Connection context for spawning the bridge task + /// - `left_name`: Logical name of the component on the "left" side (e.g., "client", "proxy:0") + /// - `right_name`: Logical name of the component on the "right" side (e.g., "proxy:0", "agent") + /// - `component`: The component to wrap + pub fn bridge_component( + &self, + proxy_index: ComponentIndex, + successor_index: ComponentIndex, + proxy: impl agent_client_protocol_core::ConnectTo, + ) -> DynConnectTo { + DynConnectTo::new(SnooperComponent::new( + proxy, + { + let trace_handle = self.clone(); + move |msg| { + trace_handle.trace_message(proxy_index, successor_index, Incoming(true), msg) + } + }, + { + let trace_handle = self.clone(); + move |msg| { + trace_handle.trace_message(proxy_index, successor_index, Incoming(false), msg) + } + }, + )) + } +} + +/// Convert a jsonrpcmsg::Id to serde_json::Value. +fn id_to_json(id: &jsonrpcmsg::Id) -> serde_json::Value { + match id { + jsonrpcmsg::Id::String(s) => serde_json::Value::String(s.clone()), + jsonrpcmsg::Id::Number(n) => serde_json::Value::Number((*n).into()), + jsonrpcmsg::Id::Null => serde_json::Value::Null, + } +} + +/// A message observed going over a channel connected to `left` and `right`. +/// This could be a successor message, a mcp-over-acp message, etc. +#[derive(Debug)] +struct TracedMessage { + component_index: ComponentIndex, + successor_index: ComponentIndex, + incoming: Incoming, + message: jsonrpcmsg::Message, +} + +/// Fully interpreted message info. +#[derive(Debug)] +struct MessageInfo { + successor: Successor, + id: Option, + protocol: Protocol, + method: String, + params: serde_json::Value, +} + +#[derive(Copy, Clone, Debug)] +struct Successor(bool); + +#[derive(Copy, Clone, Debug)] +struct Incoming(bool); + +impl MessageInfo { + /// Extract logical message info from method and params. + /// + /// This unwraps protocol wrappers to get the "real" message: + /// - `_proxy/successor` messages are unwrapped to get the inner message + /// - `_proxy/initialize` messages are unwrapped to get `initialize` + /// - `_mcp/message` messages are detected and marked as MCP protocol + /// + /// Returns (protocol, method, params). + fn from_req(req: jsonrpcmsg::Request) -> Self { + let untyped = UntypedMessage::parse_message(&req.method, &req.params) + .expect("untyped message is infallible"); + Self::from_untyped(Successor(false), req.id, Protocol::Acp, untyped) + } + + fn from_untyped( + successor: Successor, + id: Option, + protocol: Protocol, + untyped: UntypedMessage, + ) -> Self { + if let Ok(m) = SuccessorMessage::parse_message(&untyped.method, &untyped.params) { + return Self::from_untyped(Successor(true), id, protocol, m.message); + } + + if let Ok(m) = McpOverAcpMessage::parse_message(&untyped.method, &untyped.params) { + return Self::from_untyped(successor, id, Protocol::Mcp, m.message); + } + + Self { + successor, + id, + protocol, + method: untyped.method, + params: untyped.params, + } + } +} diff --git a/src/agent-client-protocol-conductor/tests/arrow_proxy_eliza.rs b/src/agent-client-protocol-conductor/tests/arrow_proxy_eliza.rs new file mode 100644 index 0000000..d2231a4 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/arrow_proxy_eliza.rs @@ -0,0 +1,78 @@ +//! Integration test for conductor with arrow proxy and test agent. +//! +//! This test verifies that: +//! 1. Conductor can orchestrate a proxy chain with arrow proxy + test agent +//! 2. Session updates from test agent get the '>' prefix from arrow proxy +//! 3. The full proxy chain works end-to-end +//! +//! Run `just prep-tests` before running this test. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_test::test_binaries::{arrow_proxy_example, testy}; +use agent_client_protocol_test::testy::TestyCommand; +use agent_client_protocol_tokio::AcpAgent; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +#[tokio::test] +async fn test_conductor_with_arrow_proxy_and_test_agent() +-> Result<(), agent_client_protocol_core::Error> { + // Create the component chain: arrow_proxy -> test_agent + // Uses pre-built binaries to avoid cargo run races during `cargo test --all` + let arrow_proxy_agent = + AcpAgent::from_args([arrow_proxy_example().to_string_lossy().to_string()])?; + let test_agent = testy(); + + // Create duplex streams for editor <-> conductor communication + let (editor_write, conductor_read) = duplex(8192); + let (conductor_write, editor_read) = duplex(8192); + + // Spawn the conductor + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(test_agent).proxy(arrow_proxy_agent), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Wait for editor to complete and get the result + let result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + let result = Box::pin(yopo::prompt( + agent_client_protocol_core::ByteStreams::new( + editor_write.compat_write(), + editor_read.compat(), + ), + TestyCommand::Greet.to_prompt(), + )) + .await?; + + tracing::debug!(?result, "Received response from arrow proxy chain"); + + assert!( + result.starts_with('>'), + "Expected response to start with '>' from arrow proxy, got: {result}" + ); + + Ok::(result) + }) + .await + .expect("Test timed out") + .expect("Editor failed"); + + tracing::info!( + ?result, + "Test completed successfully with arrow-prefixed response" + ); + + conductor_handle.abort(); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/empty_conductor_eliza.rs b/src/agent-client-protocol-conductor/tests/empty_conductor_eliza.rs new file mode 100644 index 0000000..a1f500e --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/empty_conductor_eliza.rs @@ -0,0 +1,104 @@ +//! Integration test for conductor with an empty conductor and test agent. +//! +//! This test verifies that: +//! 1. Conductor can orchestrate a chain with an empty conductor as a proxy + test agent +//! 2. Empty conductor (with no components) correctly acts as a passthrough proxy +//! 3. Messages flow correctly through the empty conductor to the agent +//! 4. The full chain works end-to-end + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::{Conductor, ConnectTo, Proxy}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Mock empty conductor component for testing. +/// Creates a nested conductor with no components that acts as a passthrough proxy. +struct MockEmptyConductor; + +impl ConnectTo for MockEmptyConductor { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + // Create an empty conductor with no components - it should act as a passthrough + let empty_components: Vec> = vec![]; + ConnectTo::::connect_to( + ConductorImpl::new_proxy( + "empty-conductor".to_string(), + empty_components, + McpBridgeMode::default(), + ), + client, + ) + .await + } +} + +#[tokio::test] +async fn test_conductor_with_empty_conductor_and_test_agent() +-> Result<(), agent_client_protocol_core::Error> { + // Initialize tracing for debugging + drop( + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("trace")), + ) + .with_test_writer() + .try_init(), + ); + // Create duplex streams for editor <-> conductor communication + let (editor_write, conductor_read) = duplex(8192); + let (conductor_write, editor_read) = duplex(8192); + + // Spawn the conductor + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "outer-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(MockEmptyConductor), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Wait for editor to complete and get the result + let result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + let result = Box::pin(yopo::prompt( + agent_client_protocol_core::ByteStreams::new( + editor_write.compat_write(), + editor_read.compat(), + ), + TestyCommand::Greet.to_prompt(), + )) + .await?; + + tracing::debug!(?result, "Received response from empty conductor chain"); + + // Empty conductor should not modify the response + expect_test::expect![[r#" + "Hello, world!" + "#]] + .assert_debug_eq(&result); + + Ok::(result) + }) + .await + .expect("Test timed out") + .expect("Editor failed"); + + tracing::info!( + ?result, + "Test completed successfully with response from empty conductor chain" + ); + + conductor_handle.abort(); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/initialization_sequence.rs b/src/agent-client-protocol-conductor/tests/initialization_sequence.rs new file mode 100644 index 0000000..d7c1126 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/initialization_sequence.rs @@ -0,0 +1,408 @@ +//! Integration tests for the initialization sequence. +//! +//! These tests verify that: +//! 1. Single-component chains receive `InitializeRequest` (agent mode) +//! 2. Multi-component chains: proxies receive `InitializeProxyRequest` +//! 3. Last component (agent) receives `InitializeRequest` + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::schema::{ + AgentCapabilities, InitializeProxyRequest, InitializeRequest, InitializeResponse, + ProtocolVersion, +}; +use agent_client_protocol_core::{Agent, Client, Conductor, ConnectTo, DynConnectTo, Proxy}; +use agent_client_protocol_test::testy::Testy; +use std::sync::Arc; +use std::sync::Mutex; + +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Test helper to receive a JSON-RPC response +async fn recv( + response: agent_client_protocol_core::SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +/// Tracks what type of initialization request was received +#[derive(Debug, Clone, PartialEq)] +enum InitRequestType { + Initialize, + InitializeProxy, +} + +struct InitConfig { + /// What type of init request was received + received_init_type: Mutex>, +} + +impl InitConfig { + fn new() -> Arc { + Arc::new(Self { + received_init_type: Mutex::new(None), + }) + } + + fn read_init_type(&self) -> Option { + self.received_init_type + .lock() + .expect("not poisoned") + .clone() + } +} + +struct InitComponent { + config: Arc, +} + +impl InitComponent { + fn new(config: &Arc) -> Self { + Self { + config: config.clone(), + } + } +} + +impl ConnectTo for InitComponent { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + let config = self.config; + let config2 = Arc::clone(&config); + + Proxy + .builder() + .name("init-component") + // Handle InitializeProxyRequest (we're a proxy) + .on_receive_request_from( + Client, + async move |request: InitializeProxyRequest, responder, cx| { + *config.received_init_type.lock().expect("unpoisoned") = + Some(InitRequestType::InitializeProxy); + + // Forward InitializeRequest (not InitializeProxyRequest) to successor + cx.send_request_to(agent_client_protocol_core::Agent, request.initialize) + .on_receiving_result(async move |response| { + let response: InitializeResponse = response?; + responder.respond(response) + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + // Handle InitializeRequest (we're the agent) + .on_receive_request_from( + Client, + async move |request: InitializeRequest, responder, _cx| { + *config2.received_init_type.lock().expect("unpoisoned") = + Some(InitRequestType::Initialize); + + // We're the final component, just respond + let response = InitializeResponse::new(request.protocol_version) + .agent_capabilities(AgentCapabilities::new()); + + responder.respond(response) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .connect_to(client) + .await + } +} + +async fn run_test_with_components( + proxies: Vec, + editor_task: impl AsyncFnOnce( + agent_client_protocol_core::ConnectionTo, + ) -> Result<(), agent_client_protocol_core::Error>, +) -> Result<(), agent_client_protocol_core::Error> { + // Set up editor <-> conductor communication + let (editor_out, conductor_in) = duplex(1024); + let (conductor_out, editor_in) = duplex(1024); + + let transport = + agent_client_protocol_core::ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + + agent_client_protocol_core::Client + .builder() + .name("editor-to-connector") + .with_spawned(|_cx| async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxies(proxies), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + )), + ) + .await + }) + .connect_with(transport, editor_task) + .await +} + +#[tokio::test] +async fn test_single_component_gets_initialize_request() +-> Result<(), agent_client_protocol_core::Error> { + // Single component (agent) should receive InitializeRequest - we use ElizaAgent + // which properly handles InitializeRequest + Box::pin(run_test_with_components( + vec![], + async |connection_to_editor| { + let init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await; + + assert!( + init_response.is_ok(), + "Initialize should succeed: {init_response:?}" + ); + + Ok::<(), agent_client_protocol_core::Error>(()) + }, + )) + .await?; + + Ok(()) +} + +#[tokio::test] +async fn test_two_components_proxy_gets_initialize_proxy() +-> Result<(), agent_client_protocol_core::Error> { + // First component (proxy) gets InitializeProxyRequest + // Second component (agent, ElizaAgent) gets InitializeRequest + let component1 = InitConfig::new(); + + Box::pin(run_test_with_components( + vec![InitComponent::new(&component1)], + async |connection_to_editor| { + let init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await; + + assert!( + init_response.is_ok(), + "Initialize should succeed: {init_response:?}" + ); + + Ok::<(), agent_client_protocol_core::Error>(()) + }, + )) + .await?; + + // First component (proxy) should receive InitializeProxyRequest + assert_eq!( + component1.read_init_type(), + Some(InitRequestType::InitializeProxy), + "Proxy component should receive InitializeProxyRequest" + ); + + // Second component (ElizaAgent) receives InitializeRequest implicitly + + Ok(()) +} + +#[tokio::test] +async fn test_three_components_all_proxies_get_initialize_proxy() +-> Result<(), agent_client_protocol_core::Error> { + // First two components (proxies) get InitializeProxyRequest + // Third component (agent, ElizaAgent) gets InitializeRequest + let component1 = InitConfig::new(); + let component2 = InitConfig::new(); + + Box::pin(run_test_with_components( + vec![ + InitComponent::new(&component1), + InitComponent::new(&component2), + ], + async |connection_to_editor| { + let init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await; + + assert!( + init_response.is_ok(), + "Initialize should succeed: {init_response:?}" + ); + + Ok::<(), agent_client_protocol_core::Error>(()) + }, + )) + .await?; + + // First two components (proxies) should receive InitializeProxyRequest + assert_eq!( + component1.read_init_type(), + Some(InitRequestType::InitializeProxy), + "First proxy should receive InitializeProxyRequest" + ); + assert_eq!( + component2.read_init_type(), + Some(InitRequestType::InitializeProxy), + "Second proxy should receive InitializeProxyRequest" + ); + + // Third component (ElizaAgent) receives InitializeRequest implicitly + + Ok(()) +} + +/// A proxy that incorrectly forwards InitializeProxyRequest instead of InitializeRequest. +/// This tests that the conductor rejects such malformed forwarding. +struct BadProxy; + +impl ConnectTo for BadProxy { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + Proxy + .builder() + .name("bad-proxy") + .on_receive_request_from( + Client, + async move |request: InitializeProxyRequest, responder, cx| { + // BUG: forwards InitializeProxyRequest instead of request.initialize + cx.send_request_to(Agent, request) + .on_receiving_result(async move |response| { + let response: InitializeResponse = response?; + responder.respond(response) + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .connect_to(client) + .await + } +} + +/// Run test with explicit proxy and agent DynComponents (for mixing different types) +async fn run_bad_proxy_test( + proxies: Vec>, + agent: DynConnectTo, + editor_task: impl AsyncFnOnce( + agent_client_protocol_core::ConnectionTo, + ) -> Result<(), agent_client_protocol_core::Error>, +) -> Result<(), agent_client_protocol_core::Error> { + let (editor_out, conductor_in) = duplex(1024); + let (conductor_out, editor_in) = duplex(1024); + + let transport = + agent_client_protocol_core::ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + + agent_client_protocol_core::Client + .builder() + .name("editor-to-connector") + .with_spawned(|_cx| async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(agent).proxies(proxies), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + )), + ) + .await + }) + .connect_with(transport, editor_task) + .await +} + +#[tokio::test] +async fn test_conductor_rejects_initialize_proxy_forwarded_to_agent() +-> Result<(), agent_client_protocol_core::Error> { + // BadProxy incorrectly forwards InitializeProxyRequest to the agent. + // The conductor should reject this with an error. + let result = Box::pin(run_bad_proxy_test( + vec![DynConnectTo::new(BadProxy)], + DynConnectTo::new(Testy::new()), + async |connection_to_editor| { + let init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await; + + if let Err(err) = init_response { + assert!( + err.to_string().contains("initialize/proxy"), + "Error should mention initialize/proxy: {err:?}" + ); + } + + Ok::<(), agent_client_protocol_core::Error>(()) + }, + )) + .await; + + match result { + Ok(()) => panic!("Expected error when proxy forwards InitializeProxyRequest to agent"), + Err(err) => { + assert!( + err.to_string().contains("initialize/proxy"), + "Error should mention initialize/proxy: {err:?}" + ); + } + } + + Ok(()) +} + +#[tokio::test] +async fn test_conductor_rejects_initialize_proxy_forwarded_to_proxy() +-> Result<(), agent_client_protocol_core::Error> { + // BadProxy incorrectly forwards InitializeProxyRequest to another proxy. + // The conductor should reject this with an error. + let result = Box::pin(run_bad_proxy_test( + vec![ + DynConnectTo::new(BadProxy), + DynConnectTo::new(InitComponent::new(&InitConfig::new())), // This proxy will receive the bad request + ], + DynConnectTo::new(Testy::new()), // Agent + async |connection_to_editor| { + let init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await; + + // The error may come through recv() or bubble up through the test harness + if let Err(err) = init_response { + assert!( + err.to_string().contains("initialize/proxy"), + "Error should mention initialize/proxy: {err:?}" + ); + } + + Ok::<(), agent_client_protocol_core::Error>(()) + }, + )) + .await; + + // The error might bubble up through run_test_with_components instead + match result { + Ok(()) => panic!("Expected error when proxy forwards InitializeProxyRequest to proxy"), + Err(err) => { + assert!( + err.to_string().contains("initialize/proxy"), + "Error should mention initialize/proxy: {err:?}" + ); + } + } + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/mcp-integration.rs b/src/agent-client-protocol-conductor/tests/mcp-integration.rs new file mode 100644 index 0000000..425ba29 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/mcp-integration.rs @@ -0,0 +1,295 @@ +//! Integration tests for MCP tool routing through proxy components. +//! +//! These tests verify that: +//! 1. Proxy components can provide MCP tools +//! 2. Agent components can discover and invoke those tools +//! 3. Tool invocations route correctly through the proxy + +mod mcp_integration; + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::Agent; +use agent_client_protocol_core::schema::{ + ContentBlock, InitializeRequest, NewSessionRequest, PromptRequest, ProtocolVersion, + SessionNotification, TextContent, +}; +use agent_client_protocol_test::test_binaries; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use futures::{SinkExt, StreamExt, channel::mpsc}; + +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Test helper to receive a JSON-RPC response +async fn recv( + response: agent_client_protocol_core::SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +fn conductor_command() -> Vec { + let binary_path = test_binaries::conductor_binary(); + vec![binary_path.to_string_lossy().to_string()] +} + +async fn run_test_with_mode( + mode: McpBridgeMode, + components: ProxiesAndAgent, + editor_task: impl AsyncFnOnce( + agent_client_protocol_core::ConnectionTo, + ) -> Result<(), agent_client_protocol_core::Error>, +) -> Result<(), agent_client_protocol_core::Error> { + // Initialize tracing for debug output + drop( + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(), + ); + + // Set up editor <-> conductor communication + let (editor_out, conductor_in) = duplex(1024); + let (conductor_out, editor_in) = duplex(1024); + + let transport = + agent_client_protocol_core::ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + + agent_client_protocol_core::Client + .builder() + .name("editor-to-connector") + .with_spawned(|_cx| async move { + Box::pin( + ConductorImpl::new_agent("conductor".to_string(), components, mode).run( + agent_client_protocol_core::ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + ), + ), + ) + .await + }) + .connect_with(transport, editor_task) + .await +} + +/// Test that proxy-provided MCP tools work with stdio bridge mode +#[tokio::test] +async fn test_proxy_provides_mcp_tools_stdio() -> Result<(), agent_client_protocol_core::Error> { + Box::pin(run_test_with_mode( + McpBridgeMode::Stdio { + conductor_command: conductor_command(), + }, + ProxiesAndAgent::new(Testy::new()).proxy(mcp_integration::proxy::ProxyComponent), + async |connection_to_editor| { + // Send initialization request + let init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await; + + assert!( + init_response.is_ok(), + "Initialize should succeed: {init_response:?}" + ); + + // Send session/new request + let session_response = recv( + connection_to_editor + .send_request(NewSessionRequest::new(std::path::PathBuf::from("/"))), + ) + .await; + + assert!( + session_response.is_ok(), + "Session/new should succeed: {session_response:?}" + ); + + let session = session_response.unwrap(); + // ElizACP generates UUID session IDs, just verify it's non-empty + assert!(!session.session_id.0.is_empty()); + + Ok(()) + }, + )) + .await?; + + Ok(()) +} + +/// Test that proxy-provided MCP tools work with HTTP bridge mode +#[tokio::test] +async fn test_proxy_provides_mcp_tools_http() -> Result<(), agent_client_protocol_core::Error> { + Box::pin(run_test_with_mode( + McpBridgeMode::Http, + ProxiesAndAgent::new(Testy::new()).proxy(mcp_integration::proxy::ProxyComponent), + async |connection_to_editor| { + // Send initialization request + let init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await; + + assert!( + init_response.is_ok(), + "Initialize should succeed: {init_response:?}" + ); + + // Send session/new request + let session_response = recv( + connection_to_editor + .send_request(NewSessionRequest::new(std::path::PathBuf::from("/"))), + ) + .await; + + assert!( + session_response.is_ok(), + "Session/new should succeed: {session_response:?}" + ); + + let session = session_response.unwrap(); + // ElizACP generates UUID session IDs, just verify it's non-empty + assert!(!session.session_id.0.is_empty()); + + Ok(()) + }, + )) + .await?; + + Ok(()) +} + +#[tokio::test] +async fn test_agent_handles_prompt() -> Result<(), agent_client_protocol_core::Error> { + // Initialize tracing for debug output + drop( + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_test_writer() + .try_init(), + ); + + // Create channel to collect log events + let (mut log_tx, mut log_rx) = mpsc::unbounded(); + + // Create duplex streams for client <-> conductor communication + let (client_write, conductor_read) = duplex(8192); + let (conductor_write, client_read) = duplex(8192); + + // Spawn the conductor in a background task + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "mcp-integration-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(mcp_integration::proxy::ProxyComponent), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Run the client + let result = Box::pin( + agent_client_protocol_core::Client + .builder() + .name("editor-to-connector") + .on_receive_notification( + { + let mut log_tx = log_tx.clone(); + async move |notification: SessionNotification, + _cx: agent_client_protocol_core::ConnectionTo| { + // Log the notification in debug format + log_tx + .send(format!("{notification:?}")) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + } + }, + agent_client_protocol_core::on_receive_notification!(), + ) + .connect_with( + agent_client_protocol_core::ByteStreams::new( + client_write.compat_write(), + client_read.compat(), + ), + async |connection_to_editor| { + // Initialize + recv( + connection_to_editor + .send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await?; + + // Create session + let session = recv( + connection_to_editor + .send_request(NewSessionRequest::new(std::path::PathBuf::from("/"))), + ) + .await?; + + tracing::debug!(session_id = %session.session_id.0, "Session created"); + + // Send a prompt to call the echo tool + let prompt_response = + recv(connection_to_editor.send_request(PromptRequest::new( + session.session_id.clone(), + vec![ContentBlock::Text(TextContent::new(TestyCommand::CallTool { + server: "test".to_string(), + tool: "echo".to_string(), + params: serde_json::json!({"message": "Hello from the test!"}), + }.to_prompt()))], + ))) + .await?; + + // Log the response + log_tx + .send(format!("{prompt_response:?}")) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error())?; + + Ok(()) + }, + ), + ) + .await; + + conductor_handle.abort(); + result?; + + // Drop the sender and collect all log entries + drop(log_tx); + let mut log_entries = Vec::new(); + while let Some(entry) = log_rx.next().await { + log_entries.push(entry); + } + + // Verify we got a successful tool call response + // The session ID is a UUID generated by ElizACP, so we check for the tool result pattern + assert_eq!(log_entries.len(), 2, "Expected notification + response"); + assert!( + log_entries[0].contains("OK: CallToolResult"), + "Expected successful tool call, got: {}", + log_entries[0] + ); + assert!( + log_entries[0].contains("Echo: Hello from the test!"), + "Expected echo result, got: {}", + log_entries[0] + ); + assert!( + log_entries[1].contains("PromptResponse"), + "Expected prompt response, got: {}", + log_entries[1] + ); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/mcp_integration/mod.rs b/src/agent-client-protocol-conductor/tests/mcp_integration/mod.rs new file mode 100644 index 0000000..44dcc92 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/mcp_integration/mod.rs @@ -0,0 +1 @@ +pub mod proxy; diff --git a/src/agent-client-protocol-conductor/tests/mcp_integration/proxy.rs b/src/agent-client-protocol-conductor/tests/mcp_integration/proxy.rs new file mode 100644 index 0000000..6daaa75 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/mcp_integration/proxy.rs @@ -0,0 +1,50 @@ +//! Proxy component that provides MCP tools + +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::{Conductor, ConnectTo, Proxy}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Parameters for the echo tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoParams { + /// The message to echo back + message: String, +} + +/// Output from the echo tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoOutput { + /// The echoed message + result: String, +} + +pub struct ProxyComponent; + +impl ConnectTo for ProxyComponent { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + let test_server = McpServer::builder("test") + .instructions("A simple test MCP server with an echo tool") + .tool_fn_mut( + "echo", + "Echoes back the input message", + async |params: EchoParams, _context| { + Ok(EchoOutput { + result: format!("Echo: {}", params.message), + }) + }, + agent_client_protocol_core::tool_fn_mut!(), + ) + .build(); + + agent_client_protocol_core::Proxy + .builder() + .name("proxy-component") + .with_mcp_server(test_server) + .connect_to(client) + .await + } +} diff --git a/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain.rs b/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain.rs new file mode 100644 index 0000000..78557bd --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain.rs @@ -0,0 +1,232 @@ +//! Test that MCP server doesn't break the handler chain for NewSessionRequest. +//! +//! This is a regression test for a bug where `McpServer::handle_dispatch` would +//! forward `NewSessionRequest` directly to the agent instead of returning +//! `Handled::No`, which prevented downstream `.on_receive_request_from()` handlers +//! from being invoked. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::schema::{ + AgentCapabilities, InitializeRequest, InitializeResponse, NewSessionRequest, + NewSessionResponse, ProtocolVersion, SessionId, +}; +use agent_client_protocol_core::{Agent, Client, Conductor, ConnectTo, DynConnectTo, Proxy}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; + +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Simple echo tool parameters +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoParams { + message: String, +} + +/// Simple echo tool output +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoOutput { + result: String, +} + +/// Test helper to receive a JSON-RPC response +async fn recv( + response: agent_client_protocol_core::SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +/// Tracks whether the NewSessionRequest handler was invoked +struct HandlerConfig { + new_session_handler_called: AtomicBool, +} + +impl HandlerConfig { + fn new() -> Arc { + Arc::new(Self { + new_session_handler_called: AtomicBool::new(false), + }) + } + + fn was_handler_called(&self) -> bool { + self.new_session_handler_called.load(Ordering::SeqCst) + } +} + +/// A proxy component that has BOTH an MCP server AND a NewSessionRequest handler. +/// The bug was that when both were present, the NewSessionRequest handler was never called. +struct ProxyWithMcpAndHandler { + config: Arc, +} + +impl ConnectTo for ProxyWithMcpAndHandler { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + let config = Arc::clone(&self.config); + + // Create an MCP server with a simple tool + let mcp_server = McpServer::builder("test-server".to_string()) + .instructions("A test MCP server") + .tool_fn_mut( + "echo", + "Echoes back the input", + async |params: EchoParams, _cx| { + Ok(EchoOutput { + result: format!("Echo: {}", params.message), + }) + }, + agent_client_protocol_core::tool_fn_mut!(), + ) + .build(); + + agent_client_protocol_core::Proxy + .builder() + .name("proxy-with-mcp-and-handler") + // Add the MCP server + .with_mcp_server(mcp_server) + // Add a NewSessionRequest handler - this should be invoked! + .on_receive_request_from( + Client, + async move |request: NewSessionRequest, responder, cx| { + // Mark that we were called + config + .new_session_handler_called + .store(true, Ordering::SeqCst); + + // Forward to agent and relay response + cx.send_request_to(Agent, request) + .on_receiving_result(async move |result| { + let response: NewSessionResponse = result?; + responder.respond(response) + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .connect_to(client) + .await + } +} + +/// A simple agent that responds to initialization and session requests +struct SimpleAgent; + +impl ConnectTo for SimpleAgent { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + Agent + .builder() + .name("simple-agent") + .on_receive_request( + async |request: InitializeRequest, responder, _cx| { + responder.respond( + InitializeResponse::new(request.protocol_version) + .agent_capabilities(AgentCapabilities::new()), + ) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_request( + async |_request: NewSessionRequest, responder, _cx| { + responder.respond(NewSessionResponse::new(SessionId::new( + uuid::Uuid::new_v4().to_string(), + ))) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .connect_to(client) + .await + } +} + +async fn run_test( + proxies: Vec>, + agent: DynConnectTo, + editor_task: impl AsyncFnOnce( + agent_client_protocol_core::ConnectionTo, + ) -> Result<(), agent_client_protocol_core::Error>, +) -> Result<(), agent_client_protocol_core::Error> { + let (editor_out, conductor_in) = duplex(1024); + let (conductor_out, editor_in) = duplex(1024); + + let transport = + agent_client_protocol_core::ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + + agent_client_protocol_core::Client + .builder() + .name("editor-to-conductor") + .with_spawned(|_cx| async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(agent).proxies(proxies), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + )), + ) + .await + }) + .connect_with(transport, editor_task) + .await +} + +/// Regression test: NewSessionRequest handler should be invoked even when MCP server is present +#[tokio::test] +async fn test_new_session_handler_invoked_with_mcp_server() +-> Result<(), agent_client_protocol_core::Error> { + let handler_config = HandlerConfig::new(); + let handler_config_clone = Arc::clone(&handler_config); + + let proxy = DynConnectTo::::new(ProxyWithMcpAndHandler { + config: handler_config, + }); + let agent = DynConnectTo::::new(SimpleAgent); + + Box::pin(run_test(vec![proxy], agent, async |connection_to_editor| { + // Initialize first + let _init_response = recv( + connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::LATEST)), + ) + .await?; + + // Create a new session - this should trigger the handler in the proxy + let session_response = + recv(connection_to_editor.send_request(NewSessionRequest::new(PathBuf::from("/tmp")))) + .await?; + + // Verify we got a valid session ID + assert!( + !session_response.session_id.0.is_empty(), + "Should receive a valid session ID" + ); + + Ok::<(), agent_client_protocol_core::Error>(()) + })) + .await?; + + // THE KEY ASSERTION: verify the handler was actually called + assert!( + handler_config_clone.was_handler_called(), + "NewSessionRequest handler should be invoked even when MCP server is in the chain. \ + This is a regression - the MCP server was incorrectly forwarding the request directly \ + to the agent instead of letting it flow through the handler chain." + ); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/nested_arrow_proxy.rs b/src/agent-client-protocol-conductor/tests/nested_arrow_proxy.rs new file mode 100644 index 0000000..8a9395b --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/nested_arrow_proxy.rs @@ -0,0 +1,85 @@ +//! Integration test for conductor with two arrow proxies in sequence. +//! +//! This test verifies that: +//! 1. Multiple arrow proxies work correctly in sequence +//! 2. The '>' prefix is applied multiple times (once per proxy) +//! 3. The full proxy chain works end-to-end +//! +//! Chain structure: +//! test-editor -> conductor -> arrow_proxy1 -> arrow_proxy2 -> test_agent +//! +//! Expected behavior: +//! - arrow_proxy2 adds first '>' to test_agent's response: ">Hello..." +//! - arrow_proxy1 adds second '>' to that: ">>Hello..." +//! +//! Run `just prep-tests` before running this test. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_test::test_binaries::{arrow_proxy_example, testy}; +use agent_client_protocol_test::testy::TestyCommand; +use agent_client_protocol_tokio::AcpAgent; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +#[tokio::test] +async fn test_conductor_with_two_external_arrow_proxies() +-> Result<(), agent_client_protocol_core::Error> { + // Create the component chain: arrow_proxy1 -> arrow_proxy2 -> test_agent + // Uses pre-built binaries to avoid cargo run races during `cargo test --all` + let arrow_proxy1 = AcpAgent::from_args([arrow_proxy_example().to_string_lossy().to_string()])?; + let arrow_proxy2 = AcpAgent::from_args([arrow_proxy_example().to_string_lossy().to_string()])?; + let agent = testy(); + + // Create duplex streams for editor <-> conductor communication + let (editor_write, conductor_read) = duplex(8192); + let (conductor_write, editor_read) = duplex(8192); + + // Spawn the conductor with three components + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(agent) + .proxy(arrow_proxy1) + .proxy(arrow_proxy2), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Wait for editor to complete and get the result + let result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + let result = Box::pin(yopo::prompt( + agent_client_protocol_core::ByteStreams::new( + editor_write.compat_write(), + editor_read.compat(), + ), + TestyCommand::Greet.to_prompt(), + )) + .await?; + + expect_test::expect![[r#" + ">>Hello, world!" + "#]] + .assert_debug_eq(&result); + + Ok::(result) + }) + .await + .expect("Test timed out") + .expect("Editor failed"); + + tracing::info!( + ?result, + "Test completed successfully with double-arrow-prefixed response" + ); + + conductor_handle.abort(); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/nested_conductor.rs b/src/agent-client-protocol-conductor/tests/nested_conductor.rs new file mode 100644 index 0000000..9de033e --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/nested_conductor.rs @@ -0,0 +1,211 @@ +//! Integration test for nested conductors with proxy mode. +//! +//! This test verifies that: +//! 1. Conductors can be nested in proxy chains +//! 2. Inner conductor operates in proxy mode and forwards messages correctly +//! 3. Multiple arrow proxies work correctly through nested conductors +//! 4. The '>' prefix is applied multiple times (once per proxy) +//! +//! Chain structure: +//! test-editor -> outer_conductor -> inner_conductor -> eliza +//! ├─ arrow_proxy1 +//! └─ arrow_proxy2 +//! +//! Expected behavior: +//! - arrow_proxy1 adds first '>' to eliza's response: ">Hello..." +//! - arrow_proxy2 adds second '>' to that: ">>Hello..." +//! - Inner conductor operates in proxy mode, forwarding to eliza +//! - Outer conductor receives the ">>" prefixed response +//! +//! Run `just prep-tests` before running these tests. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::{Conductor, ConnectTo, DynConnectTo}; +use agent_client_protocol_test::arrow_proxy::run_arrow_proxy; +use agent_client_protocol_test::test_binaries::{arrow_proxy_example, conductor_binary, testy}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use agent_client_protocol_tokio::AcpAgent; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Mock arrow proxy component for testing. +/// Runs the arrow proxy logic in-process instead of spawning a subprocess. +struct MockArrowProxy; + +impl ConnectTo for MockArrowProxy { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + run_arrow_proxy(client).await + } +} + +/// Mock inner conductor component for testing. +/// Creates a nested conductor that runs in-process with mock arrow proxies. +struct MockInnerConductor { + num_arrow_proxies: usize, +} + +impl MockInnerConductor { + fn new(num_arrow_proxies: usize) -> Self { + Self { num_arrow_proxies } + } +} + +impl ConnectTo for MockInnerConductor { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + // Create mock arrow proxy components for the inner conductor + // This conductor is ONLY proxies - no actual agent + // Use Serve::serve instead of .run() to get the Serve impl + let mut components: Vec> = Vec::new(); + for _ in 0..self.num_arrow_proxies { + components.push(DynConnectTo::new(MockArrowProxy)); + } + + ConnectTo::::connect_to( + agent_client_protocol_conductor::ConductorImpl::new_proxy( + "inner-conductor".to_string(), + components, + McpBridgeMode::default(), + ), + client, + ) + .await + } +} + +#[tokio::test] +async fn test_nested_conductor_with_arrow_proxies() -> Result<(), agent_client_protocol_core::Error> +{ + // Create the nested component chain using mock components + // Inner conductor will manage: arrow_proxy1 -> arrow_proxy2 -> eliza + // Outer conductor will manage: inner_conductor only + + // Create duplex streams for editor <-> conductor communication + let (editor_write, conductor_read) = duplex(8192); + let (conductor_write, editor_read) = duplex(8192); + + // Spawn the outer conductor with the inner conductor and eliza + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "outer-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(MockInnerConductor::new(2)), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Wait for editor to complete and get the result + let result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + let result = Box::pin(yopo::prompt( + agent_client_protocol_core::ByteStreams::new( + editor_write.compat_write(), + editor_read.compat(), + ), + TestyCommand::Greet.to_prompt(), + )) + .await?; + + tracing::debug!(?result, "Received response from nested conductor chain"); + + expect_test::expect![[r#" + ">>Hello, world!" + "#]] + .assert_debug_eq(&result); + + Ok::(result) + }) + .await + .expect("Test timed out") + .expect("Editor failed"); + + tracing::info!( + ?result, + "Test completed successfully with double-arrow-prefixed response from nested conductor" + ); + + conductor_handle.abort(); + + Ok(()) +} + +#[tokio::test] +async fn test_nested_conductor_with_external_arrow_proxies() +-> Result<(), agent_client_protocol_core::Error> { + // Create the nested component chain using external processes + // Inner conductor spawned as a separate process with two arrow proxies + // Outer conductor manages: inner_conductor -> test agent (both as external processes) + // Uses pre-built binaries to avoid cargo run races during `cargo test --all` + let conductor_path = conductor_binary().to_string_lossy().to_string(); + let arrow_proxy_path = arrow_proxy_example().to_string_lossy().to_string(); + let inner_conductor = AcpAgent::from_args([ + &conductor_path, + "proxy", + &arrow_proxy_path, + &arrow_proxy_path, + ])?; + let agent = testy(); + + // Create duplex streams for editor <-> conductor communication + let (editor_write, conductor_read) = duplex(8192); + let (conductor_write, editor_read) = duplex(8192); + + // Spawn the outer conductor with the inner conductor and eliza as external processes + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "outer-conductor".to_string(), + ProxiesAndAgent::new(agent).proxy(inner_conductor), + McpBridgeMode::default(), + ) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Wait for editor to complete and get the result + let result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + let result = Box::pin(yopo::prompt( + agent_client_protocol_core::ByteStreams::new( + editor_write.compat_write(), + editor_read.compat(), + ), + TestyCommand::Greet.to_prompt(), + )) + .await?; + + tracing::debug!(?result, "Received response from nested conductor chain"); + + expect_test::expect![[r#" + ">>Hello, world!" + "#]] + .assert_debug_eq(&result); + + Ok::(result) + }) + .await + .expect("Test timed out") + .expect("Editor failed"); + + tracing::info!( + ?result, + "Test completed successfully with double-arrow-prefixed response from nested conductor" + ); + + conductor_handle.abort(); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/scoped_mcp_server.rs b/src/agent-client-protocol-conductor/tests/scoped_mcp_server.rs new file mode 100644 index 0000000..6be1552 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/scoped_mcp_server.rs @@ -0,0 +1,148 @@ +//! Test that MCP servers can reference stack-local data. +//! +//! This test demonstrates the new scoped lifetime feature where an MCP tool +//! can capture references to stack-local data (like a Vec) and push to it +//! when the tool is invoked. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::{Agent, Conductor, ConnectTo, Proxy, Role, RunWithConnectionTo}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::sync::Mutex; + +/// Test that an MCP tool can push to a stack-local vector. +/// +/// This validates the scoped lifetime feature - the tool closure captures +/// a reference to `collected_values` which lives on the stack. +#[tokio::test] +async fn test_scoped_mcp_server_through_proxy() -> Result<(), agent_client_protocol_core::Error> { + let conductor = ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(ScopedProxy), + McpBridgeMode::default(), + ); + + let result = Box::pin(yopo::prompt( + conductor, + TestyCommand::CallTool { + server: "test".to_string(), + tool: "push".to_string(), + params: serde_json::json!({"elements": ["Hello", "world"]}), + } + .to_prompt(), + )) + .await?; + + expect_test::expect![[r#" + "OK: CallToolResult { content: [Annotated { raw: Text(RawTextContent { text: \"2\", meta: None }), annotations: None }], structured_content: None, is_error: Some(false), meta: None }" + "#]].assert_debug_eq(&result); + + Ok(()) +} + +/// Test that an MCP tool can push to a stack-local vector through a session. +/// +/// This validates the scoped lifetime feature with session-scoped MCP servers. +/// The MCP server captures a reference to stack-local data that lives for +/// the duration of the session. +#[tokio::test] +async fn test_scoped_mcp_server_through_session() -> Result<(), agent_client_protocol_core::Error> { + // Run the client + Box::pin(agent_client_protocol_core::Client.builder() + .connect_with( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(Testy::new()), + McpBridgeMode::default(), + ), + async |cx| { + // Initialize first + cx.send_request(agent_client_protocol_core::schema::InitializeRequest::new( + agent_client_protocol_core::schema::ProtocolVersion::LATEST, + )) + .block_task() + .await?; + + let collected_values = Mutex::new(Vec::new()); + let result = cx + .build_session(".") + .with_mcp_server(make_mcp_server::(&collected_values))? + .block_task() + .run_until(async |mut active_session| { + active_session + .send_prompt(TestyCommand::CallTool { + server: "test".to_string(), + tool: "push".to_string(), + params: serde_json::json!({"elements": ["Hello", "world"]}), + }.to_prompt())?; + active_session.read_to_string().await + }) + .await?; + + expect_test::expect![[r#" + "OK: CallToolResult { content: [Annotated { raw: Text(RawTextContent { text: \"2\", meta: None }), annotations: None }], structured_content: None, is_error: Some(false), meta: None }" + "#]].assert_debug_eq(&result); + + Ok(()) + }, + )) + .await?; + + Ok(()) +} + +struct ScopedProxy; + +fn make_mcp_server( + values: &Mutex>, +) -> McpServer> { + #[derive(Serialize, Deserialize, JsonSchema)] + struct PushInput { + elements: Vec, + } + + McpServer::builder("test".to_string()) + .instructions("A test MCP server with scoped tool") + .tool_fn_mut( + "push", + "Push a value to the collected values", + async |input: PushInput, _cx| { + let mut values = values.lock().expect("not poisoned"); + values.extend(input.elements); + Ok(values.len()) + }, + agent_client_protocol_core::tool_fn_mut!(), + ) + .tool_fn_mut( + "get", + "Get the collected values", + async |(): (), _cx| { + let values = values.lock().expect("not poisoned"); + Ok(values.clone()) + }, + agent_client_protocol_core::tool_fn_mut!(), + ) + .build() +} + +impl ConnectTo for ScopedProxy { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + // Stack-local data that the MCP tool will push to + let values: Mutex> = Mutex::new(Vec::new()); + + // Build the MCP server that captures a reference to collected_values + let mcp_server = make_mcp_server::(&values); + + Proxy + .builder() + .name("scoped-mcp-server") + .with_mcp_server(mcp_server) + .connect_to(client) + .await + } +} diff --git a/src/agent-client-protocol-conductor/tests/standalone_mcp_server.rs b/src/agent-client-protocol-conductor/tests/standalone_mcp_server.rs new file mode 100644 index 0000000..42c0712 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/standalone_mcp_server.rs @@ -0,0 +1,319 @@ +//! Tests for running McpServer as a standalone MCP server (not part of ACP). +//! +//! These tests verify that `McpServer` can be used directly with MCP clients +//! via the `Component` implementation. + +use agent_client_protocol_core::{ + ByteStreams, ConnectTo, RunWithConnectionTo, mcp_server::McpServer, role::mcp, util::run_until, +}; +use rmcp::{ClientHandler, ServiceExt, model::ClientInfo}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Input for the echo tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoInput { + message: String, +} + +/// Input for the add tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct AddInput { + a: i32, + b: i32, +} + +/// Output for the add tool (structured output) +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct AddOutput { + result: i32, +} + +/// Create a test MCP server with echo and add tools +fn create_test_server() -> McpServer> { + McpServer::builder("test-server") + .instructions("A test MCP server") + .tool_fn( + "echo", + "Echo a message back", + async |input: EchoInput, _cx| Ok(format!("Echo: {}", input.message)), + agent_client_protocol_core::tool_fn!(), + ) + .tool_fn( + "add", + "Add two numbers", + async |input: AddInput, _cx| { + Ok(AddOutput { + result: input.a + input.b, + }) + }, + agent_client_protocol_core::tool_fn!(), + ) + .build() +} + +/// Minimal client handler for rmcp +#[derive(Debug, Clone, Default)] +struct MinimalClientHandler; + +impl ClientHandler for MinimalClientHandler { + fn get_info(&self) -> ClientInfo { + ClientInfo::default() + } +} + +#[tokio::test] +async fn test_standalone_server_list_tools() -> Result<(), agent_client_protocol_core::Error> { + // Create duplex streams for communication + let (server_stream, client_stream) = tokio::io::duplex(8192); + let (server_read, server_write) = tokio::io::split(server_stream); + let (client_read, client_write) = tokio::io::split(client_stream); + + // Create the MCP server + let server = create_test_server(); + + // Wrap client side as ByteStreams (this is what the MCP server will talk to) + let client_as_component = ByteStreams::new(client_write.compat_write(), client_read.compat()); + + Box::pin(run_until( + ConnectTo::::connect_to(server, client_as_component), + async move { + // Create rmcp client on the server side of the duplex (the "other end") + let client = MinimalClientHandler + .serve((server_read, server_write)) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // List tools + let tools_result = client + .list_tools(None) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // Verify we got both tools + assert_eq!(tools_result.tools.len(), 2); + + let tool_names: Vec<&str> = + tools_result.tools.iter().map(|t| t.name.as_ref()).collect(); + assert!(tool_names.contains(&"echo")); + assert!(tool_names.contains(&"add")); + + // Clean up + client + .cancel() + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + Ok(()) + }, + )) + .await +} + +#[tokio::test] +async fn test_standalone_server_call_echo_tool() -> Result<(), agent_client_protocol_core::Error> { + let (server_stream, client_stream) = tokio::io::duplex(8192); + let (server_read, server_write) = tokio::io::split(server_stream); + let (client_read, client_write) = tokio::io::split(client_stream); + + let server = create_test_server(); + let client_as_component = ByteStreams::new(client_write.compat_write(), client_read.compat()); + + Box::pin(run_until( + ConnectTo::::connect_to(server, client_as_component), + async move { + let client = MinimalClientHandler + .serve((server_read, server_write)) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // Call the echo tool + let result = client + .call_tool( + rmcp::model::CallToolRequestParams::new("echo").with_arguments( + serde_json::json!({ "message": "hello world" }) + .as_object() + .unwrap() + .clone(), + ), + ) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // Verify the result + let text = result + .content + .first() + .and_then(|c| c.raw.as_text()) + .map(|t| t.text.as_str()) + .expect("Expected text content"); + + assert_eq!(text, r#""Echo: hello world""#, "Unexpected echo response"); + + client + .cancel() + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + Ok(()) + }, + )) + .await +} + +#[tokio::test] +async fn test_standalone_server_call_add_tool() -> Result<(), agent_client_protocol_core::Error> { + let (server_stream, client_stream) = tokio::io::duplex(8192); + let (server_read, server_write) = tokio::io::split(server_stream); + let (client_read, client_write) = tokio::io::split(client_stream); + + let server = create_test_server(); + let client_as_component = ByteStreams::new(client_write.compat_write(), client_read.compat()); + + Box::pin(run_until( + ConnectTo::::connect_to(server, client_as_component), + async move { + let client = MinimalClientHandler + .serve((server_read, server_write)) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // Call the add tool + let result = client + .call_tool( + rmcp::model::CallToolRequestParams::new("add").with_arguments( + serde_json::json!({ "a": 5, "b": 3 }) + .as_object() + .unwrap() + .clone(), + ), + ) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // The add tool returns structured output (AddOutput) + // Check that we get the expected result + assert!(!result.is_error.unwrap_or(false)); + + // Structured output should have the result + let content = result.content.first().expect("Expected content"); + let text = content.raw.as_text().expect("Expected text content"); + assert!( + text.text.contains('8') || text.text.contains("result"), + "Expected result to contain 8, got: {}", + text.text + ); + + client + .cancel() + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + Ok(()) + }, + )) + .await +} + +#[tokio::test] +async fn test_standalone_server_tool_not_found() -> Result<(), agent_client_protocol_core::Error> { + let (server_stream, client_stream) = tokio::io::duplex(8192); + let (server_read, server_write) = tokio::io::split(server_stream); + let (client_read, client_write) = tokio::io::split(client_stream); + + let server = create_test_server(); + let client_as_component = ByteStreams::new(client_write.compat_write(), client_read.compat()); + + Box::pin(run_until( + ConnectTo::::connect_to(server, client_as_component), + async move { + let client = MinimalClientHandler + .serve((server_read, server_write)) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // Call a non-existent tool + let result = client + .call_tool(rmcp::model::CallToolRequestParams::new("nonexistent")) + .await; + + // Should get an error + assert!(result.is_err(), "Expected error for non-existent tool"); + + client + .cancel() + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + Ok(()) + }, + )) + .await +} + +#[tokio::test] +async fn test_standalone_server_with_disabled_tools() +-> Result<(), agent_client_protocol_core::Error> { + let (server_stream, client_stream) = tokio::io::duplex(8192); + let (server_read, server_write) = tokio::io::split(server_stream); + let (client_read, client_write) = tokio::io::split(client_stream); + + // Create server with echo tool disabled + let server = McpServer::builder("test-server") + .tool_fn( + "echo", + "Echo a message", + async |input: EchoInput, _cx| Ok(format!("Echo: {}", input.message)), + agent_client_protocol_core::tool_fn!(), + ) + .tool_fn( + "add", + "Add two numbers", + async |input: AddInput, _cx| { + Ok(AddOutput { + result: input.a + input.b, + }) + }, + agent_client_protocol_core::tool_fn!(), + ) + .disable_tool("echo")? + .build(); + + let client_as_component = ByteStreams::new(client_write.compat_write(), client_read.compat()); + + Box::pin(run_until( + ConnectTo::::connect_to(server, client_as_component), + async move { + let client = MinimalClientHandler + .serve((server_read, server_write)) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + + // List tools - should only show "add" + let tools_result = client + .list_tools(None) + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + assert_eq!(tools_result.tools.len(), 1); + assert_eq!(tools_result.tools[0].name.as_ref(), "add"); + + // Calling disabled tool should fail + let result = client + .call_tool( + rmcp::model::CallToolRequestParams::new("echo").with_arguments( + serde_json::json!({ "message": "test" }) + .as_object() + .unwrap() + .clone(), + ), + ) + .await; + + assert!(result.is_err(), "Expected error for disabled tool"); + + client + .cancel() + .await + .map_err(agent_client_protocol_core::util::internal_error)?; + Ok(()) + }, + )) + .await +} diff --git a/src/agent-client-protocol-conductor/tests/test_mcp_tool_output_types.rs b/src/agent-client-protocol-conductor/tests/test_mcp_tool_output_types.rs new file mode 100644 index 0000000..3efd32c --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/test_mcp_tool_output_types.rs @@ -0,0 +1,105 @@ +//! Test MCP tools with various output types (string, integer, object) +//! +//! MCP structured output requires JSON objects. This test verifies behavior +//! when tools return non-object types like bare strings or integers. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::{Conductor, ConnectTo, DynConnectTo, Proxy, RunWithConnectionTo}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Empty input for test tools +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EmptyInput {} + +/// Create a proxy with tools that return different types +fn create_test_proxy() -> DynConnectTo { + let mcp_server = McpServer::builder("test_server".to_string()) + .instructions("Test MCP server with various output types") + .tool_fn_mut( + "return_string", + "Returns a bare string", + async |_input: EmptyInput, _context| Ok("hello world".to_string()), + agent_client_protocol_core::tool_fn_mut!(), + ) + .tool_fn_mut( + "return_integer", + "Returns a bare integer", + async |_input: EmptyInput, _context| Ok(42i32), + agent_client_protocol_core::tool_fn_mut!(), + ) + .build(); + + DynConnectTo::new(ProxyWithTestServer { mcp_server }) +} + +struct ProxyWithTestServer> { + mcp_server: McpServer, +} + +impl + 'static + Send> ConnectTo + for ProxyWithTestServer +{ + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + agent_client_protocol_core::Proxy + .builder() + .name("test-proxy") + .with_mcp_server(self.mcp_server) + .connect_to(client) + .await + } +} + +#[tokio::test] +async fn test_tool_returning_string() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_test_proxy()), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "test_server".to_string(), + tool: "return_string".to_string(), + params: serde_json::json!({}), + } + .to_prompt(), + )) + .await?; + + // The result should contain "hello world" somewhere + assert!( + result.contains("hello world"), + "expected 'hello world' in result: {result}" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_tool_returning_integer() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_test_proxy()), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "test_server".to_string(), + tool: "return_integer".to_string(), + params: serde_json::json!({}), + } + .to_prompt(), + )) + .await?; + + // The result should contain "42" somewhere + assert!(result.contains("42"), "expected '42' in result: {result}"); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/test_session_id_in_mcp_tools.rs b/src/agent-client-protocol-conductor/tests/test_session_id_in_mcp_tools.rs new file mode 100644 index 0000000..0a308f4 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/test_session_id_in_mcp_tools.rs @@ -0,0 +1,117 @@ +//! Integration test verifying that MCP tools receive the correct session_id +//! +//! This test verifies the complete flow: +//! 1. Editor creates a session and receives a session_id +//! 2. Proxy provides an MCP server with an echo tool +//! 3. Test agent invokes the tool +//! 4. The tool receives the correct session_id in its context +//! 5. The tool returns the session_id in its response +//! 6. We verify the session_ids match + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::RunWithConnectionTo; +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::{Conductor, ConnectTo, DynConnectTo, Proxy}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Input for the echo tool (null/empty) +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoInput {} + +/// Output from the echo tool containing the session_id +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoOutput { + acp_url: String, +} + +/// Create a proxy that provides an MCP server with a session_id echo tool +fn create_echo_proxy() -> DynConnectTo { + // Create MCP server with an echo tool that returns the session_id + let mcp_server = McpServer::builder("echo_server".to_string()) + .instructions("Test MCP server with session_id echo tool") + .tool_fn_mut( + "echo", + "Returns the current session_id", + async |_input: EchoInput, context| { + Ok(EchoOutput { + acp_url: context.acp_url(), + }) + }, + agent_client_protocol_core::tool_fn_mut!(), + ) + .build(); + + // Create proxy component + DynConnectTo::new(ProxyWithEchoServer { mcp_server }) +} + +struct ProxyWithEchoServer> { + mcp_server: McpServer, +} + +impl + 'static + Send> ConnectTo + for ProxyWithEchoServer +{ + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + agent_client_protocol_core::Proxy + .builder() + .name("echo-proxy") + .with_mcp_server(self.mcp_server) + .connect_to(client) + .await + } +} + +#[tokio::test] +async fn test_list_tools_from_mcp_server() -> Result<(), agent_client_protocol_core::Error> { + use expect_test::expect; + + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_echo_proxy()), + McpBridgeMode::default(), + ), + TestyCommand::ListTools { + server: "echo_server".to_string(), + } + .to_prompt(), + )) + .await?; + + // Check the response using expect_test + expect![[r" + Available tools: + - echo: Returns the current session_id"]] + .assert_eq(&result); + + Ok(()) +} + +#[tokio::test] +async fn test_session_id_delivered_to_mcp_tools() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_echo_proxy()), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "echo_server".to_string(), + tool: "echo".to_string(), + params: serde_json::json!({}), + } + .to_prompt(), + )) + .await?; + + let pattern = regex::Regex::new(r#""acp_url":\s*String\("acp:[0-9a-f-]+"\)"#).unwrap(); + assert!(pattern.is_match(&result), "unexpected result: {result}"); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/test_tool_enable_disable.rs b/src/agent-client-protocol-conductor/tests/test_tool_enable_disable.rs new file mode 100644 index 0000000..48b5bb4 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/test_tool_enable_disable.rs @@ -0,0 +1,271 @@ +//! Integration tests for tool enable/disable functionality +//! +//! These tests verify that `disable_tool`, `enable_tool`, `disable_all_tools`, +//! and `enable_all_tools` correctly filter which tools are visible and callable. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::{Conductor, ConnectTo, DynConnectTo, Proxy, RunWithConnectionTo}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Input for the echo tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EchoInput { + message: String, +} + +/// Input for the greet tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct GreetInput { + name: String, +} + +/// Empty input for simple tools +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct EmptyInput {} + +/// Create a proxy with multiple tools, some disabled via deny-list +fn create_proxy_with_disabled_tool() +-> Result, agent_client_protocol_core::Error> { + let mcp_server = McpServer::builder("test_server".to_string()) + .instructions("Test MCP server with some disabled tools") + .tool_fn( + "echo", + "Echo a message back", + async |input: EchoInput, _context| Ok(format!("Echo: {}", input.message)), + agent_client_protocol_core::tool_fn!(), + ) + .tool_fn( + "greet", + "Greet someone by name", + async |input: GreetInput, _context| Ok(format!("Hello, {}!", input.name)), + agent_client_protocol_core::tool_fn!(), + ) + .tool_fn( + "secret", + "A secret tool that should be disabled", + async |_input: EmptyInput, _context| Ok("This is secret!".to_string()), + agent_client_protocol_core::tool_fn!(), + ) + .disable_tool("secret")? + .build(); + + Ok(DynConnectTo::new(TestProxy { mcp_server })) +} + +/// Create a proxy where all tools are disabled except specific ones (allow-list) +fn create_proxy_with_allowlist() +-> Result, agent_client_protocol_core::Error> { + let mcp_server = McpServer::builder("allowlist_server".to_string()) + .instructions("Test MCP server with allow-list") + .tool_fn( + "echo", + "Echo a message back", + async |input: EchoInput, _context| Ok(format!("Echo: {}", input.message)), + agent_client_protocol_core::tool_fn!(), + ) + .tool_fn( + "greet", + "Greet someone by name", + async |input: GreetInput, _context| Ok(format!("Hello, {}!", input.name)), + agent_client_protocol_core::tool_fn!(), + ) + .tool_fn( + "secret", + "A secret tool", + async |_input: EmptyInput, _context| Ok("This is secret!".to_string()), + agent_client_protocol_core::tool_fn!(), + ) + .disable_all_tools() + .enable_tool("echo")? + .build(); + + Ok(DynConnectTo::new(TestProxy { mcp_server })) +} + +struct TestProxy> { + mcp_server: McpServer, +} + +impl + 'static + Send> ConnectTo for TestProxy { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + agent_client_protocol_core::Proxy + .builder() + .name("test-proxy") + .with_mcp_server(self.mcp_server) + .connect_to(client) + .await + } +} + +// ============================================================================ +// Tests for deny-list (disable specific tools) +// ============================================================================ + +#[tokio::test] +async fn test_list_tools_excludes_disabled() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_proxy_with_disabled_tool()?), + McpBridgeMode::default(), + ), + TestyCommand::ListTools { + server: "test_server".to_string(), + } + .to_prompt(), + )) + .await?; + + // Should contain echo and greet, but NOT secret + assert!(result.contains("echo"), "Expected 'echo' tool in list"); + assert!(result.contains("greet"), "Expected 'greet' tool in list"); + assert!( + !result.contains("secret"), + "Disabled 'secret' tool should not appear in list" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_enabled_tool_can_be_called() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_proxy_with_disabled_tool()?), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "test_server".to_string(), + tool: "echo".to_string(), + params: serde_json::json!({"message": "hello"}), + } + .to_prompt(), + )) + .await?; + + assert!( + result.contains("Echo: hello"), + "Expected echo response, got: {result}" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_disabled_tool_returns_not_found() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_proxy_with_disabled_tool()?), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "test_server".to_string(), + tool: "secret".to_string(), + params: serde_json::json!({}), + } + .to_prompt(), + )) + .await?; + + // Should get an error about tool not found + assert!( + result.contains("not found") || result.contains("error"), + "Expected error for disabled tool, got: {result}" + ); + + Ok(()) +} + +// ============================================================================ +// Tests for allow-list (disable all, enable specific) +// ============================================================================ + +#[tokio::test] +async fn test_allowlist_only_shows_enabled_tools() -> Result<(), agent_client_protocol_core::Error> +{ + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_proxy_with_allowlist()?), + McpBridgeMode::default(), + ), + TestyCommand::ListTools { + server: "allowlist_server".to_string(), + } + .to_prompt(), + )) + .await?; + + // Should only contain echo + assert!(result.contains("echo"), "Expected 'echo' tool in list"); + assert!( + !result.contains("greet"), + "'greet' should not appear (not in allow-list)" + ); + assert!( + !result.contains("secret"), + "'secret' should not appear (not in allow-list)" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_allowlist_enabled_tool_works() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_proxy_with_allowlist()?), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "allowlist_server".to_string(), + tool: "echo".to_string(), + params: serde_json::json!({"message": "allowed"}), + } + .to_prompt(), + )) + .await?; + + assert!( + result.contains("Echo: allowed"), + "Expected echo response, got: {result}" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_allowlist_non_enabled_tool_returns_not_found() +-> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_proxy_with_allowlist()?), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "allowlist_server".to_string(), + tool: "greet".to_string(), + params: serde_json::json!({"name": "World"}), + } + .to_prompt(), + )) + .await?; + + // greet is registered but not enabled, should error + assert!( + result.contains("not found") || result.contains("error"), + "Expected error for non-enabled tool, got: {result}" + ); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/test_tool_fn.rs b/src/agent-client-protocol-conductor/tests/test_tool_fn.rs new file mode 100644 index 0000000..26eff4a --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/test_tool_fn.rs @@ -0,0 +1,78 @@ +//! Integration test for `tool_fn` - stateless concurrent tools +//! +//! This test verifies that `tool_fn` works correctly for stateless tools +//! that don't need mutable state. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::{Conductor, ConnectTo, DynConnectTo, Proxy, RunWithConnectionTo}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +/// Input for the greet tool +#[derive(Debug, Serialize, Deserialize, JsonSchema)] +struct GreetInput { + name: String, +} + +/// Create a proxy that provides an MCP server with a stateless greet tool +fn create_greet_proxy() -> DynConnectTo { + // Create MCP server with a stateless greet tool using tool_fn + let mcp_server = McpServer::builder("greet_server".to_string()) + .instructions("Test MCP server with stateless greet tool") + .tool_fn( + "greet", + "Greet someone by name", + async |input: GreetInput, _context| Ok(format!("Hello, {}!", input.name)), + agent_client_protocol_core::tool_fn!(), + ) + .build(); + + // Create proxy component + DynConnectTo::new(ProxyWithGreetServer { mcp_server }) +} + +struct ProxyWithGreetServer> { + mcp_server: McpServer, +} + +impl + 'static + Send> ConnectTo + for ProxyWithGreetServer +{ + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + Proxy + .builder() + .name("greet-proxy") + .with_mcp_server(self.mcp_server) + .connect_to(client) + .await + } +} + +#[tokio::test] +async fn test_tool_fn_greet() -> Result<(), agent_client_protocol_core::Error> { + let result = Box::pin(yopo::prompt( + ConductorImpl::new_agent( + "test-conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(create_greet_proxy()), + McpBridgeMode::default(), + ), + TestyCommand::CallTool { + server: "greet_server".to_string(), + tool: "greet".to_string(), + params: serde_json::json!({"name": "World"}), + } + .to_prompt(), + )) + .await?; + + expect_test::expect![[r#" + "OK: CallToolResult { content: [Annotated { raw: Text(RawTextContent { text: \"\\\"Hello, World!\\\"\", meta: None }), annotations: None }], structured_content: None, is_error: Some(false), meta: None }" + "#]].assert_debug_eq(&result); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/test_util.rs b/src/agent-client-protocol-conductor/tests/test_util.rs new file mode 100644 index 0000000..cde9f65 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/test_util.rs @@ -0,0 +1,17 @@ +/// Initialize tracing subscriber for tests to output logs to stderr. +/// +/// Call this at the start of any test that needs detailed logging for debugging. +/// Logs will be printed to stderr and only shown when the test fails or when +/// running with `--nocapture`. +pub fn init_test_tracing() { + use tracing_subscriber::{EnvFilter, fmt}; + + drop( + fmt() + .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| { + EnvFilter::new("trace,agent_client_protocol_core::jsonrpc=trace") + })) + .with_test_writer() + .try_init(), + ); +} diff --git a/src/agent-client-protocol-conductor/tests/trace_client_mcp_server.rs b/src/agent-client-protocol-conductor/tests/trace_client_mcp_server.rs new file mode 100644 index 0000000..964e765 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/trace_client_mcp_server.rs @@ -0,0 +1,574 @@ +//! Snapshot test for trace events when a client-hosted MCP server handles tool calls. +//! +//! This test demonstrates the full round-trip flow: +//! 1. Client → Agent: initialize, session/new, session/prompt (left-to-right ACP) +//! 2. Agent → Client: MCP initialize, tools/call (right-to-left MCP, all the way back!) +//! 3. Client → Agent: MCP response +//! 4. Agent → Client: session/update notification, prompt response +//! +//! Unlike trace_mcp_tool_call.rs which tests proxy-hosted MCP servers, this test +//! verifies that MCP requests travel all the way back to the client. + +use agent_client_protocol_conductor::trace::TraceEvent; +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_core::schema::{InitializeRequest, ProtocolVersion}; +use agent_client_protocol_core::{Client, Role, RunWithConnectionTo}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use expect_test::expect; +use futures::StreamExt; +use futures::channel::mpsc; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Mutex; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Normalize events for stable snapshot testing. +/// +/// - Strips timestamps (set to 0.0) +/// - Replaces UUIDs with sequential IDs (id:0, id:1, etc.) +/// - Replaces session IDs with "session:0", etc. +/// - Replaces acp: URLs with "acp:url:0", etc. +/// - Replaces connection_id with "connection:0", etc. +struct EventNormalizer { + id_map: HashMap, + next_id: usize, + session_map: HashMap, + next_session: usize, + acp_url_map: HashMap, + next_acp_url: usize, + connection_map: HashMap, + next_connection: usize, +} + +impl EventNormalizer { + fn new() -> Self { + Self { + id_map: HashMap::new(), + next_id: 0, + session_map: HashMap::new(), + next_session: 0, + acp_url_map: HashMap::new(), + next_acp_url: 0, + connection_map: HashMap::new(), + next_connection: 0, + } + } + + fn normalize_id(&mut self, id: serde_json::Value) -> serde_json::Value { + let id_str = match &id { + serde_json::Value::String(s) => s.clone(), + serde_json::Value::Number(n) => n.to_string(), + _ => return id, + }; + + let normalized = self.id_map.entry(id_str).or_insert_with(|| { + let n = format!("id:{}", self.next_id); + self.next_id += 1; + n + }); + + serde_json::Value::String(normalized.clone()) + } + + fn normalize_session(&mut self, session: Option) -> Option { + session.map(|s| self.normalize_session_id(&s)) + } + + fn normalize_session_id(&mut self, session: &str) -> String { + self.session_map + .entry(session.to_string()) + .or_insert_with(|| { + let n = format!("session:{}", self.next_session); + self.next_session += 1; + n + }) + .clone() + } + + fn normalize_acp_url(&mut self, url: &str) -> String { + self.acp_url_map + .entry(url.to_string()) + .or_insert_with(|| { + let n = format!("acp:url:{}", self.next_acp_url); + self.next_acp_url += 1; + n + }) + .clone() + } + + fn normalize_connection_id(&mut self, id: &str) -> String { + self.connection_map + .entry(id.to_string()) + .or_insert_with(|| { + let n = format!("connection:{}", self.next_connection); + self.next_connection += 1; + n + }) + .clone() + } + + /// Recursively normalize session IDs, acp: URLs, and connection IDs in JSON values. + fn normalize_json(&mut self, value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Object(map) => { + let normalized: serde_json::Map = map + .into_iter() + .map(|(k, v)| { + let v = if k == "sessionId" { + if let serde_json::Value::String(s) = &v { + serde_json::Value::String(self.normalize_session_id(s)) + } else { + self.normalize_json(v) + } + } else if k == "url" || k == "acp_url" { + if let serde_json::Value::String(s) = &v { + if s.starts_with("acp:") || s.starts_with("http://localhost:") { + serde_json::Value::String(self.normalize_acp_url(s)) + } else { + v + } + } else { + self.normalize_json(v) + } + } else if k == "connection_id" { + if let serde_json::Value::String(s) = &v { + serde_json::Value::String(self.normalize_connection_id(s)) + } else { + self.normalize_json(v) + } + } else { + self.normalize_json(v) + }; + (k, v) + }) + .collect(); + serde_json::Value::Object(normalized) + } + serde_json::Value::Array(arr) => { + serde_json::Value::Array(arr.into_iter().map(|v| self.normalize_json(v)).collect()) + } + other => other, + } + } + + fn normalize_events(&mut self, events: Vec) -> Vec { + events + .into_iter() + .map(|event| match event { + TraceEvent::Request(mut r) => { + r.ts = 0.0; + r.id = self.normalize_id(r.id); + r.session = self.normalize_session(r.session); + r.params = self.normalize_json(r.params); + TraceEvent::Request(r) + } + TraceEvent::Response(mut r) => { + r.ts = 0.0; + r.id = self.normalize_id(r.id); + r.payload = self.normalize_json(r.payload); + TraceEvent::Response(r) + } + TraceEvent::Notification(mut n) => { + n.ts = 0.0; + n.session = self.normalize_session(n.session); + n.params = self.normalize_json(n.params); + TraceEvent::Notification(n) + } + _ => panic!("unknown trace event type"), + }) + .collect() + } +} + +/// Create an MCP server with an echo tool for testing. +fn make_echo_mcp_server( + call_count: &Mutex, +) -> McpServer> { + #[derive(Serialize, Deserialize, JsonSchema)] + struct EchoInput { + message: String, + } + + #[derive(Serialize, JsonSchema)] + struct EchoOutput { + echoed: String, + call_number: usize, + } + + McpServer::builder("echo-server".to_string()) + .instructions("A test MCP server hosted by the client") + .tool_fn_mut( + "echo", + "Echoes back the input message", + async |input: EchoInput, _cx| { + let mut count = call_count.lock().expect("not poisoned"); + *count += 1; + Ok(EchoOutput { + echoed: format!("Client echoes: {}", input.message), + call_number: *count, + }) + }, + agent_client_protocol_core::tool_fn_mut!(), + ) + .build() +} + +#[tokio::test] +async fn test_trace_client_mcp_server() -> Result<(), agent_client_protocol_core::Error> { + // Create channel for collecting trace events + let (trace_tx, trace_rx) = mpsc::unbounded(); + + // Create duplex streams for client <-> conductor communication + let (client_write, conductor_read) = duplex(8192); + let (conductor_write, client_read) = duplex(8192); + + // Spawn the conductor with ElizaAgent (no proxies - simple setup) + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(Testy::new()), + McpBridgeMode::default(), + ) + .trace_to(trace_tx) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Run the client with a client-hosted MCP server + let test_result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + Box::pin( + agent_client_protocol_core::Client + .builder() + .name("test-client") + .connect_with( + agent_client_protocol_core::ByteStreams::new( + client_write.compat_write(), + client_read.compat(), + ), + async |cx| { + // Initialize + cx.send_request(InitializeRequest::new(ProtocolVersion::LATEST)) + .block_task() + .await?; + + // Stack-local state that the MCP tool will modify + let call_count = Mutex::new(0usize); + + // Build session with client-hosted MCP server + let result = cx + .build_session(".") + .with_mcp_server(make_echo_mcp_server::(&call_count))? + .block_task() + .run_until(async |mut session| { + // Send prompt that triggers MCP tool call + // The tool call will travel: agent → conductor → client + session.send_prompt(TestyCommand::CallTool { + server: "echo-server".to_string(), + tool: "echo".to_string(), + params: serde_json::json!({"message": "Hello from client test!"}), + }.to_prompt())?; + session.read_to_string().await + }) + .await?; + + // Verify the tool was called + assert_eq!(*call_count.lock().unwrap(), 1); + + // Verify the response contains our echo + assert!(result.contains("Client echoes: Hello from client test!")); + + Ok(()) + }, + ), + ) + .await + }) + .await + .expect("Test timed out"); + + // Abort the conductor to close the trace channel + conductor_handle.abort(); + + // Collect and normalize trace events + let mut normalizer = EventNormalizer::new(); + let events = normalizer.normalize_events(trace_rx.collect().await); + + // Snapshot the trace events + // This should show the full round-trip: + // 1. Client -> Agent: initialize, session/new, session/prompt (left-to-right ACP) + // 2. Agent -> Client: MCP initialize, tools/call (right-to-left MCP - all the way back!) + // 3. Client -> Agent: MCP response + // 4. Agent -> Client: session/update notification, prompt response + expect![[r#" + [ + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Agent", + id: String("id:0"), + method: "initialize", + session: None, + params: Object { + "clientCapabilities": Object { + "auth": Object { + "terminal": Bool(false), + }, + "fs": Object { + "readTextFile": Bool(false), + "writeTextFile": Bool(false), + }, + "terminal": Bool(false), + }, + "protocolVersion": Number(1), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Agent", + to: "Client", + id: String("id:0"), + is_error: false, + payload: Object { + "agentCapabilities": Object { + "auth": Object {}, + "loadSession": Bool(false), + "mcpCapabilities": Object { + "http": Bool(false), + "sse": Bool(false), + }, + "promptCapabilities": Object { + "audio": Bool(false), + "embeddedContext": Bool(false), + "image": Bool(false), + }, + "sessionCapabilities": Object {}, + }, + "authMethods": Array [], + "protocolVersion": Number(1), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Agent", + id: String("id:1"), + method: "session/new", + session: None, + params: Object { + "cwd": String("."), + "mcpServers": Array [ + Object { + "headers": Array [], + "name": String("echo-server"), + "type": String("http"), + "url": String("acp:url:0"), + }, + ], + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Agent", + to: "Client", + id: String("id:2"), + method: "_mcp/connect", + session: None, + params: Object { + "acp_url": String("acp:url:0"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Agent", + to: "Client", + id: String("id:1"), + is_error: false, + payload: Object { + "sessionId": String("session:0"), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Agent", + id: String("id:3"), + method: "session/prompt", + session: None, + params: Object { + "prompt": Array [ + Object { + "text": String("{\"command\":\"call_tool\",\"server\":\"echo-server\",\"tool\":\"echo\",\"params\":{\"message\":\"Hello from client test!\"}}"), + "type": String("text"), + }, + ], + "sessionId": String("session:0"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Client", + to: "Agent", + id: String("id:2"), + is_error: false, + payload: Object { + "connection_id": String("connection:0"), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Mcp, + from: "Agent", + to: "Client", + id: String("id:4"), + method: "initialize", + session: None, + params: Object { + "capabilities": Object {}, + "clientInfo": Object { + "name": String("rmcp"), + "version": String("1.3.0"), + }, + "protocolVersion": String("2025-06-18"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Client", + to: "Agent", + id: String("id:4"), + is_error: false, + payload: Object { + "capabilities": Object { + "tools": Object {}, + }, + "instructions": String("A test MCP server hosted by the client"), + "protocolVersion": String("2025-06-18"), + "serverInfo": Object { + "name": String("rmcp"), + "version": String("1.3.0"), + }, + }, + }, + ), + Notification( + NotificationEvent { + ts: 0.0, + protocol: Mcp, + from: "Agent", + to: "Client", + method: "notifications/initialized", + session: None, + params: Null, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Mcp, + from: "Agent", + to: "Client", + id: String("id:5"), + method: "tools/call", + session: None, + params: Object { + "_meta": Object { + "progressToken": Number(0), + }, + "arguments": Object { + "message": String("Hello from client test!"), + }, + "name": String("echo"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Client", + to: "Agent", + id: String("id:5"), + is_error: false, + payload: Object { + "content": Array [ + Object { + "text": String("{\"call_number\":1,\"echoed\":\"Client echoes: Hello from client test!\"}"), + "type": String("text"), + }, + ], + "isError": Bool(false), + "structuredContent": Object { + "call_number": Number(1), + "echoed": String("Client echoes: Hello from client test!"), + }, + }, + }, + ), + Notification( + NotificationEvent { + ts: 0.0, + protocol: Acp, + from: "Agent", + to: "Client", + method: "session/update", + session: None, + params: Object { + "sessionId": String("session:0"), + "update": Object { + "content": Object { + "text": String("OK: CallToolResult { content: [Annotated { raw: Text(RawTextContent { text: \"{\\\"call_number\\\":1,\\\"echoed\\\":\\\"Client echoes: Hello from client test!\\\"}\", meta: None }), annotations: None }], structured_content: Some(Object {\"call_number\": Number(1), \"echoed\": String(\"Client echoes: Hello from client test!\")}), is_error: Some(false), meta: None }"), + "type": String("text"), + }, + "sessionUpdate": String("agent_message_chunk"), + }, + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Agent", + to: "Client", + id: String("id:3"), + is_error: false, + payload: Object { + "stopReason": String("end_turn"), + }, + }, + ), + ] + "#]] + .assert_debug_eq(&events); + + test_result?; + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/trace_generation.rs b/src/agent-client-protocol-conductor/tests/trace_generation.rs new file mode 100644 index 0000000..9ba3e32 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/trace_generation.rs @@ -0,0 +1,134 @@ +//! Test for trace generation. +//! +//! This test verifies that: +//! 1. Conductor correctly generates trace events when trace_to() is enabled +//! 2. Trace file contains valid JSON lines +//! 3. Events capture the message flow through the conductor +//! +//! Run `just prep-tests` before running this test. + +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_test::test_binaries::{arrow_proxy_example, testy}; +use agent_client_protocol_test::testy::TestyCommand; +use agent_client_protocol_tokio::AcpAgent; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +#[tokio::test] +async fn test_trace_generation() -> Result<(), agent_client_protocol_core::Error> { + // Enable tracing if RUST_LOG is set + drop( + tracing_subscriber::fmt() + .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) + .with_ansi(false) + .try_init(), + ); + // Create a temp file for the trace + let trace_path = std::env::temp_dir().join(format!("trace_test_{}.jsons", std::process::id())); + + // Create the component chain: arrow_proxy -> eliza + // Uses pre-built binaries to avoid cargo run races during `cargo test --all` + let arrow_proxy_agent = + AcpAgent::from_args([arrow_proxy_example().to_string_lossy().to_string()])?; + let eliza_agent = testy(); + + // Create duplex streams for editor <-> conductor communication + let (editor_write, conductor_read) = duplex(8192); + let (conductor_write, editor_read) = duplex(8192); + + let trace_path_clone = trace_path.clone(); + + // Spawn the conductor with tracing enabled + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(eliza_agent).proxy(arrow_proxy_agent), + McpBridgeMode::default(), + ) + .trace_to_path(&trace_path_clone) + .expect("Failed to create trace writer") + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Run a simple prompt through the conductor + let result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + let result = Box::pin(yopo::prompt( + agent_client_protocol_core::ByteStreams::new( + editor_write.compat_write(), + editor_read.compat(), + ), + TestyCommand::Greet.to_prompt(), + )) + .await?; + + Ok::(result) + }) + .await + .expect("Test timed out") + .expect("Editor failed"); + + conductor_handle.abort(); + + // Read and verify the trace file + let trace_content = std::fs::read_to_string(&trace_path).expect("Failed to read trace file"); + + // Parse each line as JSON + let events: Vec = trace_content + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| serde_json::from_str(line).expect("Invalid JSON in trace")) + .collect(); + + println!("Trace file: {}", trace_path.display()); + println!("Generated {} events", events.len()); + for (i, event) in events.iter().enumerate() { + let event_type = event.get("type").and_then(|v| v.as_str()).unwrap_or("?"); + let from = event.get("from").and_then(|v| v.as_str()).unwrap_or("?"); + let to = event.get("to").and_then(|v| v.as_str()).unwrap_or("?"); + let method = event.get("method").and_then(|v| v.as_str()).unwrap_or("-"); + let protocol = event + .get("protocol") + .and_then(|v| v.as_str()) + .unwrap_or("acp"); + println!(" [{i}] {event_type} {from} -> {to} {method} ({protocol})"); + } + + // Verify we got some events + assert!(!events.is_empty(), "Expected trace events, got none"); + + // Verify we have requests and responses + let has_request = events + .iter() + .any(|e| e.get("type").and_then(|v| v.as_str()) == Some("request")); + let has_response = events + .iter() + .any(|e| e.get("type").and_then(|v| v.as_str()) == Some("response")); + + assert!(has_request, "Expected at least one request event"); + assert!(has_response, "Expected at least one response event"); + + // Check that events have required fields + for event in &events { + assert!( + event.get("ts").is_some(), + "Event missing 'ts' field: {event:?}" + ); + assert!( + event.get("type").is_some(), + "Event missing 'type' field: {event:?}" + ); + } + + // Clean up + drop(std::fs::remove_file(&trace_path)); + + println!("Test passed! Response: {result}"); + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/trace_mcp_tool_call.rs b/src/agent-client-protocol-conductor/tests/trace_mcp_tool_call.rs new file mode 100644 index 0000000..83d7e14 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/trace_mcp_tool_call.rs @@ -0,0 +1,558 @@ +//! Snapshot test for trace events when an agent makes an MCP tool call. +//! +//! This test verifies the right-to-left request flow: +//! - Client sends prompt to agent +//! - Agent makes MCP tools/call request back through the conductor +//! - Conductor routes the request to the proxy's MCP server +//! - Response flows back to the agent +//! +//! This captures trace events for the full bidirectional flow. + +mod mcp_integration; + +use agent_client_protocol_conductor::trace::TraceEvent; +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_core::schema::{ + ContentBlock, InitializeRequest, NewSessionRequest, PromptRequest, ProtocolVersion, + SessionNotification, TextContent, +}; +use agent_client_protocol_test::testy::{Testy, TestyCommand}; +use expect_test::expect; +use futures::channel::mpsc; +use futures::{SinkExt, StreamExt}; +use std::collections::HashMap; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Normalize events for stable snapshot testing. +/// +/// - Strips timestamps (set to 0.0) +/// - Replaces UUIDs with sequential IDs (id:0, id:1, etc.) +/// - Replaces session IDs with "session:0", etc. +/// - Replaces acp: URLs with "acp:url:0", etc. +/// - Replaces connection_id with "connection:0", etc. +struct EventNormalizer { + id_map: HashMap, + next_id: usize, + session_map: HashMap, + next_session: usize, + acp_url_map: HashMap, + next_acp_url: usize, + connection_map: HashMap, + next_connection: usize, +} + +impl EventNormalizer { + fn new() -> Self { + Self { + id_map: HashMap::new(), + next_id: 0, + session_map: HashMap::new(), + next_session: 0, + acp_url_map: HashMap::new(), + next_acp_url: 0, + connection_map: HashMap::new(), + next_connection: 0, + } + } + + fn normalize_id(&mut self, id: serde_json::Value) -> serde_json::Value { + let id_str = match &id { + serde_json::Value::String(s) => s.clone(), + serde_json::Value::Number(n) => n.to_string(), + _ => return id, + }; + + let normalized = self.id_map.entry(id_str).or_insert_with(|| { + let n = format!("id:{}", self.next_id); + self.next_id += 1; + n + }); + + serde_json::Value::String(normalized.clone()) + } + + fn normalize_session(&mut self, session: Option) -> Option { + session.map(|s| self.normalize_session_id(&s)) + } + + fn normalize_session_id(&mut self, session: &str) -> String { + self.session_map + .entry(session.to_string()) + .or_insert_with(|| { + let n = format!("session:{}", self.next_session); + self.next_session += 1; + n + }) + .clone() + } + + fn normalize_acp_url(&mut self, url: &str) -> String { + self.acp_url_map + .entry(url.to_string()) + .or_insert_with(|| { + let n = format!("acp:url:{}", self.next_acp_url); + self.next_acp_url += 1; + n + }) + .clone() + } + + fn normalize_connection_id(&mut self, id: &str) -> String { + self.connection_map + .entry(id.to_string()) + .or_insert_with(|| { + let n = format!("connection:{}", self.next_connection); + self.next_connection += 1; + n + }) + .clone() + } + + /// Recursively normalize session IDs, acp: URLs, and connection IDs in JSON values. + fn normalize_json(&mut self, value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Object(map) => { + let normalized: serde_json::Map = map + .into_iter() + .map(|(k, v)| { + let v = if k == "sessionId" { + if let serde_json::Value::String(s) = &v { + serde_json::Value::String(self.normalize_session_id(s)) + } else { + self.normalize_json(v) + } + } else if k == "url" || k == "acp_url" { + if let serde_json::Value::String(s) = &v { + if s.starts_with("acp:") || s.starts_with("http://localhost:") { + serde_json::Value::String(self.normalize_acp_url(s)) + } else { + v + } + } else { + self.normalize_json(v) + } + } else if k == "connection_id" { + if let serde_json::Value::String(s) = &v { + serde_json::Value::String(self.normalize_connection_id(s)) + } else { + self.normalize_json(v) + } + } else { + self.normalize_json(v) + }; + (k, v) + }) + .collect(); + serde_json::Value::Object(normalized) + } + serde_json::Value::Array(arr) => { + serde_json::Value::Array(arr.into_iter().map(|v| self.normalize_json(v)).collect()) + } + other => other, + } + } + + fn normalize_events(&mut self, events: Vec) -> Vec { + events + .into_iter() + .map(|event| match event { + TraceEvent::Request(mut r) => { + r.ts = 0.0; + r.id = self.normalize_id(r.id); + r.session = self.normalize_session(r.session); + r.params = self.normalize_json(r.params); + TraceEvent::Request(r) + } + TraceEvent::Response(mut r) => { + r.ts = 0.0; + r.id = self.normalize_id(r.id); + r.payload = self.normalize_json(r.payload); + TraceEvent::Response(r) + } + TraceEvent::Notification(mut n) => { + n.ts = 0.0; + n.session = self.normalize_session(n.session); + n.params = self.normalize_json(n.params); + TraceEvent::Notification(n) + } + _ => panic!("unknown trace event type"), + }) + .collect() + } +} + +/// Test helper to receive a JSON-RPC response +async fn recv( + response: agent_client_protocol_core::SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +#[tokio::test] +async fn test_trace_mcp_tool_call() -> Result<(), agent_client_protocol_core::Error> { + // Create channel for collecting trace events + let (trace_tx, trace_rx) = mpsc::unbounded(); + + // Create channel to collect notifications (to verify test worked) + let (notif_tx, mut notif_rx) = mpsc::unbounded(); + + // Create duplex streams for client <-> conductor communication + let (client_write, conductor_read) = duplex(8192); + let (conductor_write, client_read) = duplex(8192); + + // Spawn the conductor with: + // - ElizaAgent (deterministic mode) as the agent + // - ProxyComponent that provides the "test" MCP server with echo tool + // - Tracing enabled to capture events + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(Testy::new()).proxy(mcp_integration::proxy::ProxyComponent), + McpBridgeMode::default(), + ) + .trace_to(trace_tx) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Run the client interaction + let test_result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + Box::pin( + agent_client_protocol_core::Client + .builder() + .name("test-client") + .on_receive_notification( + { + let mut notif_tx = notif_tx; + async move |notification: SessionNotification, _cx| { + notif_tx + .send(notification) + .await + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + } + }, + agent_client_protocol_core::on_receive_notification!(), + ) + .connect_with( + agent_client_protocol_core::ByteStreams::new( + client_write.compat_write(), + client_read.compat(), + ), + async |cx| { + // Initialize + recv(cx.send_request(InitializeRequest::new(ProtocolVersion::LATEST))) + .await?; + + // Create session + let session = recv( + cx.send_request(NewSessionRequest::new(std::path::PathBuf::from("/"))), + ) + .await?; + + // Send prompt that triggers MCP tool call + recv(cx.send_request(PromptRequest::new( + session.session_id.clone(), + vec![ContentBlock::Text(TextContent::new(TestyCommand::CallTool { + server: "test".to_string(), + tool: "echo".to_string(), + params: serde_json::json!({"message": "Hello from trace test!"}), + }.to_prompt()))], + ))) + .await?; + + Ok(()) + }, + ), + ) + .await + }) + .await + .expect("Test timed out"); + + // Abort the conductor to close the trace channel + conductor_handle.abort(); + let mut notifications = Vec::new(); + while let Some(notif) = notif_rx.next().await { + notifications.push(notif); + } + assert_eq!(notifications.len(), 1, "Expected one notification"); + + // Collect and normalize trace events + let mut normalizer = EventNormalizer::new(); + let events = normalizer.normalize_events(trace_rx.collect().await); + + // Snapshot the trace events + // This should show: + // 1. Client -> Agent: initialize, session/new, session/prompt (left-to-right) + // 2. Agent -> MCP Server: tools/call (right-to-left, the key part!) + // 3. MCP Server -> Agent: response + // 4. Agent -> Client: notification + response + expect![[r#" + [ + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Proxy(0)", + id: String("id:0"), + method: "_proxy/initialize", + session: None, + params: Object { + "clientCapabilities": Object { + "auth": Object { + "terminal": Bool(false), + }, + "fs": Object { + "readTextFile": Bool(false), + "writeTextFile": Bool(false), + }, + "terminal": Bool(false), + }, + "protocolVersion": Number(1), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Client", + id: String("id:0"), + is_error: false, + payload: Object { + "agentCapabilities": Object { + "auth": Object {}, + "loadSession": Bool(false), + "mcpCapabilities": Object { + "http": Bool(false), + "sse": Bool(false), + }, + "promptCapabilities": Object { + "audio": Bool(false), + "embeddedContext": Bool(false), + "image": Bool(false), + }, + "sessionCapabilities": Object {}, + }, + "authMethods": Array [], + "protocolVersion": Number(1), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Proxy(0)", + id: String("id:1"), + method: "session/new", + session: None, + params: Object { + "cwd": String("/"), + "mcpServers": Array [], + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Proxy(1)", + to: "Proxy(0)", + id: String("id:2"), + method: "_mcp/connect", + session: None, + params: Object { + "acp_url": String("acp:url:0"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Proxy(1)", + id: String("id:2"), + is_error: false, + payload: Object { + "connection_id": String("connection:0"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Client", + id: String("id:1"), + is_error: false, + payload: Object { + "sessionId": String("session:0"), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Proxy(0)", + id: String("id:3"), + method: "session/prompt", + session: None, + params: Object { + "prompt": Array [ + Object { + "text": String("{\"command\":\"call_tool\",\"server\":\"test\",\"tool\":\"echo\",\"params\":{\"message\":\"Hello from trace test!\"}}"), + "type": String("text"), + }, + ], + "sessionId": String("session:0"), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Mcp, + from: "Proxy(1)", + to: "Proxy(0)", + id: String("id:4"), + method: "initialize", + session: None, + params: Object { + "capabilities": Object {}, + "clientInfo": Object { + "name": String("rmcp"), + "version": String("1.3.0"), + }, + "protocolVersion": String("2025-06-18"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Proxy(1)", + id: String("id:4"), + is_error: false, + payload: Object { + "capabilities": Object { + "tools": Object {}, + }, + "instructions": String("A simple test MCP server with an echo tool"), + "protocolVersion": String("2025-06-18"), + "serverInfo": Object { + "name": String("rmcp"), + "version": String("1.3.0"), + }, + }, + }, + ), + Notification( + NotificationEvent { + ts: 0.0, + protocol: Mcp, + from: "Proxy(1)", + to: "Proxy(0)", + method: "notifications/initialized", + session: None, + params: Null, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Mcp, + from: "Proxy(1)", + to: "Proxy(0)", + id: String("id:5"), + method: "tools/call", + session: None, + params: Object { + "_meta": Object { + "progressToken": Number(0), + }, + "arguments": Object { + "message": String("Hello from trace test!"), + }, + "name": String("echo"), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Proxy(1)", + id: String("id:5"), + is_error: false, + payload: Object { + "content": Array [ + Object { + "text": String("{\"result\":\"Echo: Hello from trace test!\"}"), + "type": String("text"), + }, + ], + "isError": Bool(false), + "structuredContent": Object { + "result": String("Echo: Hello from trace test!"), + }, + }, + }, + ), + Notification( + NotificationEvent { + ts: 0.0, + protocol: Acp, + from: "Proxy(1)", + to: "Proxy(0)", + method: "session/update", + session: None, + params: Object { + "sessionId": String("session:0"), + "update": Object { + "content": Object { + "text": String("OK: CallToolResult { content: [Annotated { raw: Text(RawTextContent { text: \"{\\\"result\\\":\\\"Echo: Hello from trace test!\\\"}\", meta: None }), annotations: None }], structured_content: Some(Object {\"result\": String(\"Echo: Hello from trace test!\")}), is_error: Some(false), meta: None }"), + "type": String("text"), + }, + "sessionUpdate": String("agent_message_chunk"), + }, + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Client", + id: String("id:3"), + is_error: false, + payload: Object { + "stopReason": String("end_turn"), + }, + }, + ), + ] + "#]] + .assert_debug_eq(&events); + + test_result?; + + Ok(()) +} diff --git a/src/agent-client-protocol-conductor/tests/trace_snapshot.rs b/src/agent-client-protocol-conductor/tests/trace_snapshot.rs new file mode 100644 index 0000000..b57ea6b --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/trace_snapshot.rs @@ -0,0 +1,323 @@ +//! Snapshot test for trace events from a real yopo interaction. +//! +//! This test runs yopo -> conductor (with arrow_proxy -> test_agent) and +//! captures trace events to a channel for expect_test snapshot verification. +//! +//! Run `just prep-tests` before running this test. + +use agent_client_protocol_conductor::trace::TraceEvent; +use agent_client_protocol_conductor::{ConductorImpl, McpBridgeMode, ProxiesAndAgent}; +use agent_client_protocol_test::test_binaries::{arrow_proxy_example, testy}; +use agent_client_protocol_test::testy::TestyCommand; +use agent_client_protocol_tokio::AcpAgent; +use expect_test::expect; +use futures::StreamExt; +use futures::channel::mpsc; +use std::collections::HashMap; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Normalize events for stable snapshot testing. +/// +/// - Strips timestamps (set to 0.0) +/// - Replaces UUIDs with sequential IDs (id:0, id:1, etc.) +/// - Replaces session IDs with "session:0", etc. +struct EventNormalizer { + id_map: HashMap, + next_id: usize, + session_map: HashMap, + next_session: usize, +} + +impl EventNormalizer { + fn new() -> Self { + Self { + id_map: HashMap::new(), + next_id: 0, + session_map: HashMap::new(), + next_session: 0, + } + } + + fn normalize_id(&mut self, id: serde_json::Value) -> serde_json::Value { + let id_str = match &id { + serde_json::Value::String(s) => s.clone(), + serde_json::Value::Number(n) => n.to_string(), + _ => return id, + }; + + let normalized = self.id_map.entry(id_str).or_insert_with(|| { + let n = format!("id:{}", self.next_id); + self.next_id += 1; + n + }); + + serde_json::Value::String(normalized.clone()) + } + + fn normalize_session(&mut self, session: Option) -> Option { + session.map(|s| self.normalize_session_id(&s)) + } + + fn normalize_session_id(&mut self, session: &str) -> String { + self.session_map + .entry(session.to_string()) + .or_insert_with(|| { + let n = format!("session:{}", self.next_session); + self.next_session += 1; + n + }) + .clone() + } + + /// Recursively normalize session IDs in JSON values. + fn normalize_json(&mut self, value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Object(map) => { + let normalized: serde_json::Map = map + .into_iter() + .map(|(k, v)| { + let v = if k == "sessionId" { + if let serde_json::Value::String(s) = &v { + serde_json::Value::String(self.normalize_session_id(s)) + } else { + self.normalize_json(v) + } + } else { + self.normalize_json(v) + }; + (k, v) + }) + .collect(); + serde_json::Value::Object(normalized) + } + serde_json::Value::Array(arr) => { + serde_json::Value::Array(arr.into_iter().map(|v| self.normalize_json(v)).collect()) + } + other => other, + } + } + + fn normalize_events(&mut self, events: Vec) -> Vec { + events + .into_iter() + .map(|event| match event { + TraceEvent::Request(mut r) => { + r.ts = 0.0; + r.id = self.normalize_id(r.id); + r.session = self.normalize_session(r.session); + r.params = self.normalize_json(r.params); + TraceEvent::Request(r) + } + TraceEvent::Response(mut r) => { + r.ts = 0.0; + r.id = self.normalize_id(r.id); + r.payload = self.normalize_json(r.payload); + TraceEvent::Response(r) + } + TraceEvent::Notification(mut n) => { + n.ts = 0.0; + n.session = self.normalize_session(n.session); + n.params = self.normalize_json(n.params); + TraceEvent::Notification(n) + } + _ => panic!("unknown trace event type"), + }) + .collect() + } +} + +#[tokio::test] +async fn test_trace_snapshot() -> Result<(), agent_client_protocol_core::Error> { + // Create channel for collecting trace events + let (tx, rx) = mpsc::unbounded(); + + // Create the component chain: arrow_proxy -> eliza + // Uses pre-built binaries to avoid cargo run races during `cargo test --all` + let arrow_proxy_agent = + AcpAgent::from_args([arrow_proxy_example().to_string_lossy().to_string()])?; + let eliza_agent = testy(); + + // Create duplex streams for editor <-> conductor communication + let (editor_write, conductor_read) = duplex(8192); + let (conductor_write, editor_read) = duplex(8192); + + // Spawn the conductor with tracing to the channel + let conductor_handle = tokio::spawn(async move { + Box::pin( + ConductorImpl::new_agent( + "conductor".to_string(), + ProxiesAndAgent::new(eliza_agent).proxy(arrow_proxy_agent), + McpBridgeMode::default(), + ) + .trace_to(tx) + .run(agent_client_protocol_core::ByteStreams::new( + conductor_write.compat_write(), + conductor_read.compat(), + )), + ) + .await + }); + + // Run a simple prompt through the conductor + let result = tokio::time::timeout(std::time::Duration::from_secs(30), async move { + Box::pin(yopo::prompt( + agent_client_protocol_core::ByteStreams::new( + editor_write.compat_write(), + editor_read.compat(), + ), + TestyCommand::Greet.to_prompt(), + )) + .await + }) + .await + .expect("Test timed out")?; + + // Abort the conductor to close the trace channel + conductor_handle.abort(); + + // Collect and normalize events + let mut normalizer = EventNormalizer::new(); + let events = normalizer.normalize_events(rx.collect().await); + + // Snapshot the trace events + expect![[r#" + [ + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Proxy(0)", + id: String("id:0"), + method: "_proxy/initialize", + session: None, + params: Object { + "clientCapabilities": Object { + "auth": Object { + "terminal": Bool(false), + }, + "fs": Object { + "readTextFile": Bool(false), + "writeTextFile": Bool(false), + }, + "terminal": Bool(false), + }, + "protocolVersion": Number(1), + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Client", + id: String("id:0"), + is_error: false, + payload: Object { + "agentCapabilities": Object { + "auth": Object {}, + "loadSession": Bool(false), + "mcpCapabilities": Object { + "http": Bool(false), + "sse": Bool(false), + }, + "promptCapabilities": Object { + "audio": Bool(false), + "embeddedContext": Bool(false), + "image": Bool(false), + }, + "sessionCapabilities": Object {}, + }, + "authMethods": Array [], + "protocolVersion": Number(1), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Proxy(0)", + id: String("id:1"), + method: "session/new", + session: None, + params: Object { + "cwd": String("."), + "mcpServers": Array [], + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Client", + id: String("id:1"), + is_error: false, + payload: Object { + "sessionId": String("session:0"), + }, + }, + ), + Request( + RequestEvent { + ts: 0.0, + protocol: Acp, + from: "Client", + to: "Proxy(0)", + id: String("id:2"), + method: "session/prompt", + session: None, + params: Object { + "prompt": Array [ + Object { + "text": String("{\"command\":\"greet\"}"), + "type": String("text"), + }, + ], + "sessionId": String("session:0"), + }, + }, + ), + Notification( + NotificationEvent { + ts: 0.0, + protocol: Acp, + from: "Proxy(1)", + to: "Proxy(0)", + method: "session/update", + session: None, + params: Object { + "sessionId": String("session:0"), + "update": Object { + "content": Object { + "text": String("Hello, world!"), + "type": String("text"), + }, + "sessionUpdate": String("agent_message_chunk"), + }, + }, + }, + ), + Response( + ResponseEvent { + ts: 0.0, + from: "Proxy(0)", + to: "Client", + id: String("id:2"), + is_error: false, + payload: Object { + "stopReason": String("end_turn"), + }, + }, + ), + ] + "#]] + .assert_debug_eq(&events); + + println!("Response: {result}"); + + Ok(()) +} diff --git a/src/agent-client-protocol-cookbook/CHANGELOG.md b/src/agent-client-protocol-cookbook/CHANGELOG.md new file mode 100644 index 0000000..9e0e2bb --- /dev/null +++ b/src/agent-client-protocol-cookbook/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-cookbook/Cargo.toml b/src/agent-client-protocol-cookbook/Cargo.toml new file mode 100644 index 0000000..ef2a3aa --- /dev/null +++ b/src/agent-client-protocol-cookbook/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "agent-client-protocol-cookbook" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "Cookbook of common patterns for building ACP components" +keywords = ["acp", "agent", "proxy", "mcp", "cookbook"] +categories = ["development-tools"] + +[dependencies] +agent-client-protocol-conductor.workspace = true +agent-client-protocol-core.workspace = true +agent-client-protocol-rmcp.workspace = true +agent-client-protocol-tokio.workspace = true +rmcp.workspace = true +schemars.workspace = true +serde.workspace = true +serde_json.workspace = true +tokio.workspace = true +tracing.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-cookbook/src/lib.rs b/src/agent-client-protocol-cookbook/src/lib.rs new file mode 100644 index 0000000..c2f6fa3 --- /dev/null +++ b/src/agent-client-protocol-cookbook/src/lib.rs @@ -0,0 +1,855 @@ +//! Cookbook of common patterns for building ACP components. +//! +//! This crate contains guides and examples for the three main things you can build with ACP: +//! +//! - **Clients** - Connect to an existing agent and send prompts +//! - **Proxies** - Sit between client and agent to add capabilities (like MCP tools) +//! - **Agents** - Respond to prompts with AI-powered responses +//! +//! See the [`agent_client_protocol_core::concepts`] module for detailed explanations of +//! the concepts behind the API. +//! +//! # Building Clients +//! +//! A client connects to an agent, sends requests, and handles responses. Use +//! [`Client.builder()`](agent_client_protocol_core::Client) to build connections. +//! +//! - [`one_shot_prompt`] - Send a single prompt and get a response (simplest pattern) +//! - [`connecting_as_client`] - More details on connection setup and permission handling +//! +//! # Building Proxies +//! +//! A proxy sits between client and agent, intercepting and optionally modifying +//! messages. The most common use case is adding MCP tools. Use [`Proxy.builder()`](agent_client_protocol_core::Proxy) +//! to build proxy connections. +//! +//! **Important:** Proxies don't run standalone—they need the [`agent-client-protocol-conductor`] to +//! orchestrate the connection between client, proxies, and agent. See +//! [`running_proxies_with_conductor`] for how to put the pieces together. +//! +//! - [`global_mcp_server`] - Add tools that work across all sessions +//! - [`per_session_mcp_server`] - Add tools with session-specific state +//! - [`filtering_tools`] - Enable or disable tools dynamically +//! - [`reusable_components`] - Package your proxy as a [`ConnectTo`] for composition +//! - [`running_proxies_with_conductor`] - Run your proxy with an agent +//! +//! [`agent-client-protocol-conductor`]: https://crates.io/crates/agent-client-protocol-conductor +//! +//! # Building Agents +//! +//! An agent receives prompts and generates responses. Use [`Agent.builder()`](agent_client_protocol_core::Agent) +//! to build agent connections. +//! +//! - [`building_an_agent`] - Handle initialization, sessions, and prompts +//! - [`reusable_components`] - Package your agent as a [`ConnectTo`] +//! - [`custom_message_handlers`] - Fine-grained control over message routing +//! +//! [`agent_client_protocol_core::concepts`]: agent_client_protocol_core::concepts +//! [`Client`]: agent_client_protocol_core::Client +//! [`Agent`]: agent_client_protocol_core::Agent +//! [`Proxy`]: agent_client_protocol_core::Proxy +//! [`ConnectTo`]: agent_client_protocol_core::ConnectTo + +pub mod one_shot_prompt { + //! Pattern: You Only Prompt Once. + //! + //! The simplest client pattern: connect to an agent, send one prompt, get the + //! response. This is useful for CLI tools, scripts, or any case where you just + //! need a single interaction with an agent. + //! + //! # Example + //! + //! ``` + //! use agent_client_protocol_core::{Client, Agent, ConnectTo}; + //! use agent_client_protocol_core::schema::{InitializeRequest, ProtocolVersion}; + //! + //! async fn ask_agent( + //! transport: impl ConnectTo + 'static, + //! prompt: &str, + //! ) -> Result { + //! Client.builder() + //! .name("my-client") + //! .connect_with(transport, async |connection| { + //! // Initialize the connection + //! connection.send_request(InitializeRequest::new(ProtocolVersion::LATEST)) + //! .block_task().await?; + //! + //! // Create a session, send prompt, read response + //! let mut session = connection.build_session_cwd()? + //! .block_task() + //! .start_session() + //! .await?; + //! + //! session.send_prompt(prompt)?; + //! session.read_to_string().await + //! }) + //! .await + //! } + //! ``` + //! + //! # How it works + //! + //! 1. **[`connect_with`]** establishes the transport connection and runs your + //! code while handling messages in the background + //! 2. **[`send_request`]** + **[`block_task`]** sends the initialize request + //! and waits for the response + //! 3. **[`build_session_cwd`]** creates a session builder using the current working directory + //! 4. **[`start_session`]** sends the `NewSessionRequest` and returns an + //! [`ActiveSession`] handle + //! 5. **[`send_prompt`]** queues the prompt to send to the agent + //! 6. **[`read_to_string`]** reads all text chunks until the agent finishes + //! + //! # Handling permission requests + //! + //! Most agents will ask for permission before taking actions like running + //! commands or writing files. See [`connecting_as_client`] for how to handle + //! [`RequestPermissionRequest`] messages. + //! + //! [`connect_with`]: agent_client_protocol_core::Builder::connect_with + //! [`send_request`]: agent_client_protocol_core::ConnectionTo::send_request + //! [`block_task`]: agent_client_protocol_core::SentRequest::block_task + //! [`build_session_cwd`]: agent_client_protocol_core::ConnectionTo::build_session_cwd + //! [`start_session`]: agent_client_protocol_core::SessionBuilder::start_session + //! [`ActiveSession`]: agent_client_protocol_core::ActiveSession + //! [`send_prompt`]: agent_client_protocol_core::ActiveSession::send_prompt + //! [`read_to_string`]: agent_client_protocol_core::ActiveSession::read_to_string + //! [`connecting_as_client`]: super::connecting_as_client + //! [`RequestPermissionRequest`]: agent_client_protocol_core::schema::RequestPermissionRequest +} + +pub mod connecting_as_client { + //! Pattern: Connecting as a client. + //! + //! To connect to an ACP agent and send requests, use [`connect_with`]. + //! This runs your code while the connection handles incoming messages + //! in the background. + //! + //! # Basic Example + //! + //! ``` + //! use agent_client_protocol_core::{Client, Agent, ConnectTo}; + //! use agent_client_protocol_core::schema::{InitializeRequest, ProtocolVersion}; + //! + //! async fn connect_to_agent(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + //! Client.builder() + //! .name("my-client") + //! .connect_with(transport, async |connection| { + //! // Initialize the connection + //! connection.send_request(InitializeRequest::new(ProtocolVersion::LATEST)) + //! .block_task().await?; + //! + //! // Create a session and send a prompt + //! connection.build_session_cwd()? + //! .block_task() + //! .run_until(async |mut session| { + //! session.send_prompt("Hello, agent!")?; + //! let response = session.read_to_string().await?; + //! println!("Agent said: {}", response); + //! Ok(()) + //! }) + //! .await + //! }) + //! .await + //! } + //! ``` + //! + //! # Using the Session Builder + //! + //! The [`build_session`] method creates a [`SessionBuilder`] that handles + //! session creation and provides convenient methods for interacting with + //! the session: + //! + //! - [`send_prompt`] - Send a text prompt to the agent + //! - [`read_update`] - Read the next update (text chunk, tool call, etc.) + //! - [`read_to_string`] - Read all text until the turn ends + //! + //! The session builder also supports adding MCP servers with [`with_mcp_server`]. + //! + //! # Handling Permission Requests + //! + //! Agents may send [`RequestPermissionRequest`] to ask for user approval + //! before taking actions. Handle these with [`on_receive_request`]: + //! + //! ```ignore + //! Client.builder() + //! .on_receive_request(async |req: RequestPermissionRequest, responder, _connection| { + //! // Auto-approve by selecting the first option (YOLO mode) + //! let option_id = req.options.first().map(|opt| opt.id.clone()); + //! responder.respond(RequestPermissionResponse { + //! outcome: match option_id { + //! Some(id) => RequestPermissionOutcome::Selected { option_id: id }, + //! None => RequestPermissionOutcome::Cancelled, + //! }, + //! meta: None, + //! }) + //! }, agent_client_protocol_core::on_receive_request!()) + //! .connect_with(transport, async |connection| { /* ... */ }) + //! .await + //! ``` + //! + //! # Note on `block_task` + //! + //! Using [`block_task`] is safe inside `connect_with` because the closure runs + //! as a spawned task, not on the event loop. The event loop continues processing + //! messages (including the response you're waiting for) while your task blocks. + //! + //! [`connect_with`]: agent_client_protocol_core::Builder::connect_with + //! [`block_task`]: agent_client_protocol_core::SentRequest::block_task + //! [`build_session`]: agent_client_protocol_core::ConnectionTo::build_session + //! [`SessionBuilder`]: agent_client_protocol_core::SessionBuilder + //! [`send_prompt`]: agent_client_protocol_core::ActiveSession::send_prompt + //! [`read_update`]: agent_client_protocol_core::ActiveSession::read_update + //! [`read_to_string`]: agent_client_protocol_core::ActiveSession::read_to_string + //! [`with_mcp_server`]: agent_client_protocol_core::SessionBuilder::with_mcp_server + //! [`RequestPermissionRequest`]: agent_client_protocol_core::schema::RequestPermissionRequest + //! [`on_receive_request`]: agent_client_protocol_core::Builder::on_receive_request +} + +pub mod building_an_agent { + //! Pattern: Building an agent. + //! + //! An agent handles prompts and generates responses. At minimum, an agent must: + //! + //! 1. Handle [`InitializeRequest`] to establish the connection + //! 2. Handle [`NewSessionRequest`] to create sessions + //! 3. Handle [`PromptRequest`] to process prompts + //! + //! Use [`Agent.builder()`](agent_client_protocol_core::Agent) to build agent connections. + //! + //! # Minimal Example + //! + //! ``` + //! use agent_client_protocol_core::{Agent, Client, ConnectTo, Dispatch, ConnectionTo}; + //! use agent_client_protocol_core::schema::{ + //! InitializeRequest, InitializeResponse, AgentCapabilities, + //! NewSessionRequest, NewSessionResponse, SessionId, + //! PromptRequest, PromptResponse, StopReason, + //! }; + //! + //! async fn run_agent(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + //! Agent.builder() + //! .name("my-agent") + //! // Handle initialization + //! .on_receive_request(async |req: InitializeRequest, responder, _connection| { + //! responder.respond( + //! InitializeResponse::new(req.protocol_version) + //! .agent_capabilities(AgentCapabilities::new()) + //! ) + //! }, agent_client_protocol_core::on_receive_request!()) + //! // Handle session creation + //! .on_receive_request(async |req: NewSessionRequest, responder, _connection| { + //! responder.respond(NewSessionResponse::new(SessionId::new("session-1"))) + //! }, agent_client_protocol_core::on_receive_request!()) + //! // Handle prompts + //! .on_receive_request(async |req: PromptRequest, responder, connection| { + //! // Send streaming updates via notifications + //! // connection.send_notification(SessionNotification { ... })?; + //! + //! // Return final response + //! responder.respond(PromptResponse::new(StopReason::EndTurn)) + //! }, agent_client_protocol_core::on_receive_request!()) + //! // Reject unknown messages + //! .on_receive_dispatch(async |message: Dispatch, connection: ConnectionTo| { + //! message.respond_with_error(agent_client_protocol_core::Error::method_not_found(), connection) + //! }, agent_client_protocol_core::on_receive_dispatch!()) + //! .connect_to(transport) + //! .await + //! } + //! ``` + //! + //! # Streaming Responses + //! + //! To stream text or other updates to the client, send [`SessionNotification`]s + //! while processing a prompt: + //! + //! ```ignore + //! .on_receive_request(async |req: PromptRequest, responder, connection| { + //! // Stream some text + //! connection.send_notification(SessionNotification { + //! session_id: req.session_id.clone(), + //! update: SessionUpdate::Text(TextUpdate { + //! text: "Hello, ".into(), + //! // ... + //! }), + //! meta: None, + //! })?; + //! + //! connection.send_notification(SessionNotification { + //! session_id: req.session_id.clone(), + //! update: SessionUpdate::Text(TextUpdate { + //! text: "world!".into(), + //! // ... + //! }), + //! meta: None, + //! })?; + //! + //! responder.respond(PromptResponse { + //! stop_reason: StopReason::EndTurn, + //! meta: None, + //! }) + //! }, agent_client_protocol_core::on_receive_request!()) + //! ``` + //! + //! # Requesting Permissions + //! + //! Before taking actions that require user approval (like running commands + //! or writing files), send a [`RequestPermissionRequest`]: + //! + //! ```ignore + //! let response = connection.send_request(RequestPermissionRequest { + //! session_id: session_id.clone(), + //! action: PermissionAction::Bash { command: "rm -rf /".into() }, + //! options: vec![ + //! PermissionOption { id: "allow".into(), label: "Allow".into() }, + //! PermissionOption { id: "deny".into(), label: "Deny".into() }, + //! ], + //! meta: None, + //! }).block_task().await?; + //! + //! match response.outcome { + //! RequestPermissionOutcome::Selected { option_id } if option_id == "allow" => { + //! // User approved, proceed with action + //! } + //! _ => { + //! // User denied or cancelled + //! } + //! } + //! ``` + //! + //! # As a Reusable Component + //! + //! For agents that will be composed with proxies, implement [`ConnectTo`]. + //! See [`reusable_components`] for the pattern. + //! + //! [`InitializeRequest`]: agent_client_protocol_core::schema::InitializeRequest + //! [`NewSessionRequest`]: agent_client_protocol_core::schema::NewSessionRequest + //! [`PromptRequest`]: agent_client_protocol_core::schema::PromptRequest + //! [`SessionNotification`]: agent_client_protocol_core::schema::SessionNotification + //! [`RequestPermissionRequest`]: agent_client_protocol_core::schema::RequestPermissionRequest + //! [`Agent`]: agent_client_protocol_core::Agent + //! [`ConnectTo`]: agent_client_protocol_core::ConnectTo + //! [`reusable_components`]: super::reusable_components +} + +pub mod reusable_components { + //! Pattern: Defining reusable components. + //! + //! When building agents or proxies that will be composed together (for example, + //! with [`agent-client-protocol-conductor`]), define a struct that implements [`ConnectTo`]. + //! This allows your component to be connected to other components in a type-safe way. + //! + //! # Example + //! + //! ``` + //! use agent_client_protocol_core::{ConnectTo, Agent, Client}; + //! use agent_client_protocol_core::schema::{ + //! InitializeRequest, InitializeResponse, AgentCapabilities, + //! }; + //! + //! struct MyAgent { + //! name: String, + //! } + //! + //! impl ConnectTo for MyAgent { + //! async fn connect_to(self, client: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + //! Agent.builder() + //! .name(&self.name) + //! .on_receive_request(async move |req: InitializeRequest, responder, _connection| { + //! responder.respond( + //! InitializeResponse::new(req.protocol_version) + //! .agent_capabilities(AgentCapabilities::new()) + //! ) + //! }, agent_client_protocol_core::on_receive_request!()) + //! .connect_to(client) + //! .await + //! } + //! } + //! + //! let agent = MyAgent { name: "my-agent".into() }; + //! ``` + //! + //! # Important: Don't block the event loop + //! + //! Message handlers run on the event loop. Blocking in a handler prevents the + //! connection from processing new messages. For expensive work: + //! + //! - Use [`ConnectionTo::spawn`] to offload work to a background task + //! - Use [`on_receiving_result`] to schedule work when a response arrives + //! + //! [`ConnectTo`]: agent_client_protocol_core::ConnectTo + //! [`ConnectionTo::spawn`]: agent_client_protocol_core::ConnectionTo::spawn + //! [`on_receiving_result`]: agent_client_protocol_core::SentRequest::on_receiving_result + //! [`agent-client-protocol-conductor`]: https://crates.io/crates/agent-client-protocol-conductor +} + +pub mod custom_message_handlers { + //! Pattern: Custom message handlers. + //! + //! For reusable message handling logic, implement [`HandleDispatchFrom`] and use + //! [`MatchDispatch`] or [`MatchDispatchFrom`] for type-safe dispatching. + //! + //! This is useful when you need to: + //! - Share message handling logic across multiple components + //! - Build complex routing logic that doesn't fit the builder pattern + //! - Integrate with existing handler infrastructure + //! + //! # Example + //! + //! ``` + //! use agent_client_protocol_core::{HandleDispatchFrom, Dispatch, Handled, ConnectionTo, UntypedRole}; + //! use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, AgentCapabilities}; + //! use agent_client_protocol_core::util::MatchDispatch; + //! + //! struct MyHandler; + //! + //! impl HandleDispatchFrom for MyHandler { + //! async fn handle_dispatch_from( + //! &mut self, + //! message: Dispatch, + //! _connection: ConnectionTo, + //! ) -> Result, agent_client_protocol_core::Error> { + //! MatchDispatch::new(message) + //! .if_request(async |req: InitializeRequest, responder| { + //! responder.respond( + //! InitializeResponse::new(req.protocol_version) + //! .agent_capabilities(AgentCapabilities::new()) + //! ) + //! }) + //! .await + //! .done() + //! } + //! + //! fn describe_chain(&self) -> impl std::fmt::Debug { + //! "MyHandler" + //! } + //! } + //! ``` + //! + //! # When to use `MatchDispatch` vs `MatchDispatchFrom` + //! + //! - [`MatchDispatch`] - Use when you don't need peer-aware handling + //! - [`MatchDispatchFrom`] - Use in proxies where messages come from different + //! peers (`Client` vs `Agent`) and may need different handling + //! + //! [`HandleDispatchFrom`]: agent_client_protocol_core::HandleDispatchFrom + //! [`MatchDispatch`]: agent_client_protocol_core::util::MatchDispatch + //! [`MatchDispatchFrom`]: agent_client_protocol_core::util::MatchDispatchFrom +} + +pub mod global_mcp_server { + //! Pattern: Global MCP server in handler chain. + //! + //! Use this pattern when you want a single MCP server that handles tool calls + //! for all sessions. The server is added to the connection's handler chain and + //! automatically injects itself into every `NewSessionRequest` that passes through. + //! + //! # When to use + //! + //! - The MCP server provides stateless tools (no per-session state needed) + //! - You want the simplest setup with minimal boilerplate + //! - Tools don't need access to session-specific context + //! + //! # Using the builder API + //! + //! The simplest way to create an MCP server is with [`McpServer::builder`]: + //! + //! ``` + //! use agent_client_protocol_core::mcp_server::McpServer; + //! use agent_client_protocol_core::{ConnectTo, RunWithConnectionTo, Proxy, Conductor}; + //! use schemars::JsonSchema; + //! use serde::{Deserialize, Serialize}; + //! + //! #[derive(Debug, Deserialize, JsonSchema)] + //! struct EchoParams { message: String } + //! + //! #[derive(Debug, Serialize, JsonSchema)] + //! struct EchoOutput { echoed: String } + //! + //! // Build the MCP server with tools + //! let mcp_server = McpServer::builder("my-tools") + //! .tool_fn("echo", "Echoes the input", + //! async |params: EchoParams, _cx| { + //! Ok(EchoOutput { echoed: params.message }) + //! }, + //! agent_client_protocol_core::tool_fn!()) + //! .build(); + //! + //! // The proxy component is generic over the MCP server's responder type + //! struct MyProxy { + //! mcp_server: McpServer, + //! } + //! + //! impl + Send + 'static> ConnectTo for MyProxy { + //! async fn connect_to(self, conductor: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + //! Proxy.builder() + //! .with_mcp_server(self.mcp_server) + //! .connect_to(conductor) + //! .await + //! } + //! } + //! + //! let proxy = MyProxy { mcp_server }; + //! ``` + //! + //! # Using rmcp + //! + //! If you have an existing [rmcp](https://docs.rs/rmcp) server implementation, + //! use [`McpServer::from_rmcp`] from the `agent-client-protocol-rmcp` crate: + //! + //! ``` + //! use rmcp::{ServerHandler, tool, tool_router, tool_handler}; + //! use rmcp::handler::server::router::tool::ToolRouter; + //! use rmcp::handler::server::wrapper::Parameters; + //! use rmcp::model::*; + //! use agent_client_protocol_core::mcp_server::McpServer; + //! use agent_client_protocol_core::Conductor; + //! use agent_client_protocol_rmcp::McpServerExt; + //! use serde::{Deserialize, Serialize}; + //! + //! #[derive(Debug, Serialize, Deserialize, schemars::JsonSchema)] + //! struct EchoParams { + //! message: String, + //! } + //! + //! #[derive(Clone)] + //! struct MyMcpServer { + //! tool_router: ToolRouter, + //! } + //! + //! impl MyMcpServer { + //! fn new() -> Self { + //! Self { tool_router: Self::tool_router() } + //! } + //! } + //! + //! #[tool_router] + //! impl MyMcpServer { + //! #[tool(description = "Echoes back the input message")] + //! async fn echo(&self, Parameters(params): Parameters) -> Result { + //! Ok(CallToolResult::success(vec![Content::text(format!("Echo: {}", params.message))])) + //! } + //! } + //! + //! #[tool_handler] + //! impl ServerHandler for MyMcpServer { + //! fn get_info(&self) -> ServerInfo { + //! ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + //! .with_protocol_version(ProtocolVersion::V_2024_11_05) + //! .with_server_info(Implementation::from_build_env()) + //! } + //! } + //! + //! // Create an MCP server from the rmcp service + //! let mcp_server = McpServer::::from_rmcp("my-server", MyMcpServer::new); + //! ``` + //! + //! The `from_rmcp` function takes a factory closure that creates a new server + //! instance. This allows each MCP connection to get a fresh server instance. + //! + //! # How it works + //! + //! When you call [`with_mcp_server`], the MCP server is added as a message + //! handler. It: + //! + //! 1. Intercepts `NewSessionRequest` messages and adds its `acp:UUID` URL to the + //! request's `mcp_servers` list + //! 2. Passes the modified request through to the next handler + //! 3. Handles incoming MCP protocol messages (tool calls, etc.) for its URL + //! + //! [`McpServer::builder`]: agent_client_protocol_core::mcp_server::McpServer::builder + //! [`McpServer::from_rmcp`]: agent_client_protocol_rmcp::McpServerExt::from_rmcp + //! [`with_mcp_server`]: agent_client_protocol_core::Builder::with_mcp_server +} + +pub mod per_session_mcp_server { + //! Pattern: Per-session MCP server with workspace context. + //! + //! Use this pattern when each session needs its own MCP server instance + //! with access to session-specific context like the working directory. + //! + //! # When to use + //! + //! - Tools need access to the session's working directory + //! - You want to track active sessions or maintain per-session state + //! - Tools need to customize behavior based on session parameters + //! + //! # Basic pattern with `on_proxy_session_start` + //! + //! The most common pattern intercepts [`NewSessionRequest`], extracts context, + //! creates a per-session MCP server, and uses [`on_proxy_session_start`] to + //! run code after the session is established: + //! + //! ``` + //! use agent_client_protocol_core::mcp_server::McpServer; + //! use agent_client_protocol_core::schema::NewSessionRequest; + //! use agent_client_protocol_core::{Client, Proxy, Conductor, ConnectTo}; + //! + //! async fn run_proxy(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + //! Proxy.builder() + //! .on_receive_request_from(Client, async move |request: NewSessionRequest, responder, connection| { + //! // Extract session context from the request + //! let workspace_path = request.cwd.clone(); + //! + //! // Create tools that capture the workspace path + //! let mcp_server = McpServer::builder("workspace-tools") + //! .tool_fn("get_workspace", "Returns the session's workspace directory", { + //! async move |_params: (), _cx| { + //! Ok(workspace_path.display().to_string()) + //! } + //! }, agent_client_protocol_core::tool_fn!()) + //! .build(); + //! + //! // Build the session and run code after it starts + //! connection.build_session_from(request) + //! .with_mcp_server(mcp_server)? + //! .on_proxy_session_start(responder, async move |session_id| { + //! // This callback runs after the session-id has been sent to the + //! // client but before any further messages from the client or agent + //! // related to this session have been processed. + //! // + //! // You can use this to store the `session_id` before processing + //! // future messages, or to send a first prompt to the agent before + //! // the client has a chance to do so. + //! tracing::info!(%session_id, "Session started"); + //! Ok(()) + //! }) + //! }, agent_client_protocol_core::on_receive_request!()) + //! .connect_to(transport) + //! .await + //! } + //! ``` + //! + //! # How `on_proxy_session_start` works + //! + //! [`on_proxy_session_start`] is the non-blocking way to set up a proxy session: + //! + //! 1. Sends `NewSessionRequest` to the agent + //! 2. When the response arrives, responds to the client automatically + //! 3. Sets up message proxying for the session + //! 4. Runs your callback with the `SessionId` + //! + //! The callback runs after the session is established but doesn't block + //! the message handler. This is ideal for proxies that just need to inject + //! tools and track sessions. + //! + //! # Alternative: blocking with `start_session_proxy` + //! + //! If you need the simpler blocking API (e.g., in a client context where + //! blocking is safe), use [`block_task`] + [`start_session_proxy`]: + //! + //! ``` + //! # use agent_client_protocol_core::mcp_server::McpServer; + //! # use agent_client_protocol_core::schema::NewSessionRequest; + //! # use agent_client_protocol_core::{Client, Proxy, Conductor, ConnectTo}; + //! # async fn run_proxy(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + //! Proxy.builder() + //! .on_receive_request_from(Client, async |request: NewSessionRequest, responder, connection| { + //! let cwd = request.cwd.clone(); + //! let mcp_server = McpServer::builder("tools") + //! .tool_fn("get_cwd", "Returns working directory", { + //! async move |_params: (), _cx| Ok(cwd.display().to_string()) + //! }, agent_client_protocol_core::tool_fn!()) + //! .build(); + //! + //! let session_id = connection.build_session_from(request) + //! .with_mcp_server(mcp_server)? + //! .block_task() + //! .start_session_proxy(responder) + //! .await?; + //! + //! tracing::info!(%session_id, "Session started"); + //! Ok(()) + //! }, agent_client_protocol_core::on_receive_request!()) + //! .connect_to(transport) + //! .await + //! # } + //! ``` + //! + //! For patterns where you need to interact with the session before proxying, + //! use [`start_session`] + [`proxy_remaining_messages`] instead. + //! + //! [`start_session`]: agent_client_protocol_core::SessionBuilder::start_session + //! [`proxy_remaining_messages`]: agent_client_protocol_core::ActiveSession::proxy_remaining_messages + //! + //! [`NewSessionRequest`]: agent_client_protocol_core::schema::NewSessionRequest + //! [`on_proxy_session_start`]: agent_client_protocol_core::SessionBuilder::on_proxy_session_start + //! [`block_task`]: agent_client_protocol_core::SessionBuilder::block_task + //! [`start_session_proxy`]: agent_client_protocol_core::SessionBuilder::start_session_proxy +} + +pub mod filtering_tools { + //! Pattern: Filtering which tools are available. + //! + //! Use [`disable_tool`] and [`enable_tool`] to control which tools are + //! visible to clients. This is useful when: + //! + //! - Some tools should only be available in certain configurations + //! - You want to conditionally expose tools based on runtime settings + //! - You need to restrict access to sensitive tools + //! + //! # Disabling specific tools (deny-list) + //! + //! By default, all registered tools are enabled. Use [`disable_tool`] to + //! hide specific tools: + //! + //! ``` + //! use agent_client_protocol_core::mcp_server::McpServer; + //! use agent_client_protocol_core::{Conductor, RunWithConnectionTo}; + //! use schemars::JsonSchema; + //! use serde::Deserialize; + //! + //! #[derive(Debug, Deserialize, JsonSchema)] + //! struct Params {} + //! + //! fn build_server(enable_admin: bool) -> Result>, agent_client_protocol_core::Error> { + //! let mut builder = McpServer::builder("my-server") + //! .tool_fn("echo", "Echo a message", + //! async |_p: Params, _cx| Ok("echoed"), + //! agent_client_protocol_core::tool_fn!()) + //! .tool_fn("admin", "Admin-only tool", + //! async |_p: Params, _cx| Ok("admin action"), + //! agent_client_protocol_core::tool_fn!()); + //! + //! // Conditionally disable the admin tool + //! if !enable_admin { + //! builder = builder.disable_tool("admin")?; + //! } + //! + //! Ok(builder.build()) + //! } + //! ``` + //! + //! Disabled tools: + //! - Don't appear in `list_tools` responses + //! - Return "tool not found" errors if called directly + //! + //! # Enabling only specific tools (allow-list) + //! + //! Use [`disable_all_tools`] followed by [`enable_tool`] to create an + //! allow-list where only explicitly enabled tools are available: + //! + //! ``` + //! use agent_client_protocol_core::mcp_server::McpServer; + //! use agent_client_protocol_core::{Conductor, RunWithConnectionTo}; + //! use schemars::JsonSchema; + //! use serde::Deserialize; + //! + //! #[derive(Debug, Deserialize, JsonSchema)] + //! struct Params {} + //! + //! fn build_restricted_server() -> Result>, agent_client_protocol_core::Error> { + //! McpServer::builder("restricted-server") + //! .tool_fn("safe", "Safe operation", + //! async |_p: Params, _cx| Ok("safe"), + //! agent_client_protocol_core::tool_fn!()) + //! .tool_fn("dangerous", "Dangerous operation", + //! async |_p: Params, _cx| Ok("danger!"), + //! agent_client_protocol_core::tool_fn!()) + //! .tool_fn("experimental", "Experimental feature", + //! async |_p: Params, _cx| Ok("experimental"), + //! agent_client_protocol_core::tool_fn!()) + //! // Start with all tools disabled + //! .disable_all_tools() + //! // Only enable the safe tool + //! .enable_tool("safe") + //! .map(|b| b.build()) + //! } + //! ``` + //! + //! # Error handling + //! + //! Both [`enable_tool`] and [`disable_tool`] return `Result` and will error + //! if the tool name doesn't match any registered tool. This helps catch typos: + //! + //! ``` + //! use agent_client_protocol_core::mcp_server::McpServer; + //! use agent_client_protocol_core::Conductor; + //! + //! // This will error because "ech" is not a registered tool + //! let result = McpServer::::builder("server") + //! .disable_tool("ech"); // Typo! Should be "echo" + //! + //! assert!(result.is_err()); + //! ``` + //! + //! Calling enable/disable on an already enabled/disabled tool is not an error - + //! the operations are idempotent. + //! + //! [`disable_tool`]: agent_client_protocol_core::mcp_server::McpServerBuilder::disable_tool + //! [`enable_tool`]: agent_client_protocol_core::mcp_server::McpServerBuilder::enable_tool + //! [`disable_all_tools`]: agent_client_protocol_core::mcp_server::McpServerBuilder::disable_all_tools +} + +pub mod running_proxies_with_conductor { + //! Pattern: Running proxies with the conductor. + //! + //! Proxies don't run standalone. To add an MCP server (or other proxy behavior) + //! to an existing agent, you need the **conductor** to orchestrate the connection. + //! + //! The conductor: + //! 1. Accepts connections from clients + //! 2. Chains your proxies together + //! 3. Connects to the final agent + //! 4. Routes messages through the entire chain + //! + //! # Using the `agent-client-protocol-conductor` binary + //! + //! The simplest way to run a proxy is with the [`agent-client-protocol-conductor`] binary. + //! Configure it with a JSON file: + //! + //! ```json + //! { + //! "proxies": [ + //! { "command": ["cargo", "run", "--bin", "my-proxy"] } + //! ], + //! "agent": { "command": ["claude-code", "--agent"] } + //! } + //! ``` + //! + //! Then run: + //! + //! ```bash + //! agent-client-protocol-conductor --config conductor.json + //! ``` + //! + //! # Using the conductor as a library + //! + //! For more control, use [`agent-client-protocol-conductor`] as a library with the `ConductorImpl` type: + //! + //! ```ignore + //! use agent_client_protocol_conductor::{ConductorImpl, ProxiesAndAgent}; + //! + //! // Define your proxy as a ConnectTo + //! let my_proxy = MyProxy::new(); + //! + //! // Spawn the agent process + //! let agent_process = agent_client_protocol_tokio::spawn_process("claude-code", &["--agent"]).await?; + //! + //! // Create the conductor with your proxy chain + //! let conductor = ConductorImpl::new(ProxiesAndAgent { + //! proxies: vec![Box::new(my_proxy)], + //! agent: agent_process, + //! }); + //! + //! // Run the conductor (it will accept client connections on stdin/stdout) + //! conductor.connect_to(client_transport).await?; + //! ``` + //! + //! # Why can't I just connect my proxy directly to an agent? + //! + //! ACP uses a message envelope format for proxy chains. When a proxy sends a + //! message toward the agent, it gets wrapped in a [`SuccessorMessage`] envelope. + //! The conductor handles this wrapping/unwrapping automatically. + //! + //! If you connected directly to an agent, your proxy would send `SuccessorMessage` + //! envelopes that the agent doesn't understand. + //! + //! # Example: Complete proxy with conductor + //! + //! See the [`agent-client-protocol-conductor` tests] for complete working examples of proxies + //! running with the conductor. + //! + //! [`agent-client-protocol-conductor`]: https://crates.io/crates/agent-client-protocol-conductor + //! [`SuccessorMessage`]: agent_client_protocol_core::schema::SuccessorMessage + //! [`agent-client-protocol-conductor` tests]: https://github.com/anthropics/acp-rust-sdk/tree/main/src/agent-client-protocol-conductor/tests +} diff --git a/src/agent-client-protocol-core/CHANGELOG.md b/src/agent-client-protocol-core/CHANGELOG.md new file mode 100644 index 0000000..1d013ff --- /dev/null +++ b/src/agent-client-protocol-core/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-core/Cargo.toml b/src/agent-client-protocol-core/Cargo.toml new file mode 100644 index 0000000..b331dfb --- /dev/null +++ b/src/agent-client-protocol-core/Cargo.toml @@ -0,0 +1,56 @@ +[package] +name = "agent-client-protocol-core" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "Core protocol types and traits for the Agent Client Protocol" +keywords = ["acp", "agent", "protocol", "ai"] +categories = ["development-tools"] + +[features] +default = [] + +# Forward unstable features from agent-client-protocol-schema. +# Enable these to get support for the corresponding unstable ACP methods. +unstable = [ + "unstable_session_model", + "unstable_session_fork", + "unstable_session_resume", + "unstable_session_close", +] +unstable_session_model = ["agent-client-protocol-schema/unstable_session_model"] +unstable_session_fork = ["agent-client-protocol-schema/unstable_session_fork"] +unstable_session_resume = ["agent-client-protocol-schema/unstable_session_resume"] +unstable_session_close = ["agent-client-protocol-schema/unstable_session_close"] + +[dependencies] +agent-client-protocol-schema.workspace = true +agent-client-protocol-derive.workspace = true +anyhow.workspace = true +boxfnonce.workspace = true +futures.workspace = true +futures-concurrency.workspace = true +rustc-hash.workspace = true +jsonrpcmsg.workspace = true +rmcp = { workspace = true, features = ["server"] } +schemars.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true +uuid.workspace = true + +[dev-dependencies] +agent-client-protocol-test.workspace = true +clap.workspace = true +expect-test.workspace = true +shell-words.workspace = true +tokio.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-core/README.md b/src/agent-client-protocol-core/README.md new file mode 100644 index 0000000..1e4c8d0 --- /dev/null +++ b/src/agent-client-protocol-core/README.md @@ -0,0 +1,50 @@ +# agent-client-protocol-core + +Core protocol types and traits for the [Agent Client Protocol (ACP)](https://agentclientprotocol.com/). + +ACP is a protocol for communication between AI agents and their clients (IDEs, CLIs, etc.), +enabling features like tool use, permission requests, and streaming responses. + +## What can you build with this crate? + +- **Clients** that talk to ACP agents (like building your own Claude Code interface) +- **Proxies** that add capabilities to existing agents (like adding custom tools via MCP) +- **Agents** that respond to prompts with AI-powered responses + +## Quick Start: Connecting to an Agent + +The most common use case is connecting to an existing ACP agent as a client: + +```rust +use agent_client_protocol_core::{Client, Agent, ConnectTo}; +use agent_client_protocol_core::schema::{InitializeRequest, ProtocolVersion}; + +Client.builder() + .name("my-client") + .connect_with(transport, async |cx| { + // Initialize the connection + cx.send_request(InitializeRequest::new(ProtocolVersion::LATEST)) + .block_task() + .await?; + + Ok(()) + }) + .await?; +``` + +## Learning More + +See the [crate documentation](https://docs.rs/agent-client-protocol-core) for: + +- **[Cookbook](https://docs.rs/agent-client-protocol-core/latest/agent_client_protocol_core/cookbook/)** — Patterns for building clients, proxies, and agents +- **[Examples](https://github.com/agentclientprotocol/rust-sdk/tree/main/src/agent-client-protocol-core/examples)** — Working code you can run + +## Related Crates + +- **[agent-client-protocol-tokio](../agent-client-protocol-tokio/)** — Tokio utilities for spawning agent processes +- **[agent-client-protocol-derive](../agent-client-protocol-derive/)** — Derive macros for JSON-RPC traits +- **[agent-client-protocol-trace-viewer](../agent-client-protocol-trace-viewer/)** — Interactive trace visualization + +## License + +Apache-2.0 diff --git a/src/agent-client-protocol-core/examples/simple_agent.rs b/src/agent-client-protocol-core/examples/simple_agent.rs new file mode 100644 index 0000000..6dc9e5b --- /dev/null +++ b/src/agent-client-protocol-core/examples/simple_agent.rs @@ -0,0 +1,37 @@ +use agent_client_protocol_core::schema::{ + AgentCapabilities, InitializeRequest, InitializeResponse, +}; +use agent_client_protocol_core::{Agent, Client, ConnectionTo, Dispatch}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +#[tokio::main] +async fn main() -> Result<(), agent_client_protocol_core::Error> { + let agent = Agent + .builder() + .name("my-agent") // for debugging + .on_receive_request( + async move |initialize: InitializeRequest, responder, _connection| { + // Respond to initialize successfully + responder.respond( + InitializeResponse::new(initialize.protocol_version) + .agent_capabilities(AgentCapabilities::new()), + ) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_dispatch( + async move |message: Dispatch, cx: ConnectionTo| { + // Respond to any other message with an error + message.respond_with_error( + agent_client_protocol_core::util::internal_error("TODO"), + cx, + ) + }, + agent_client_protocol_core::on_receive_dispatch!(), + ) + .connect_to(agent_client_protocol_core::ByteStreams::new( + tokio::io::stdout().compat_write(), + tokio::io::stdin().compat(), + )); + Box::pin(agent).await +} diff --git a/src/agent-client-protocol-core/examples/yolo_one_shot_client.rs b/src/agent-client-protocol-core/examples/yolo_one_shot_client.rs new file mode 100644 index 0000000..2acbc5a --- /dev/null +++ b/src/agent-client-protocol-core/examples/yolo_one_shot_client.rs @@ -0,0 +1,162 @@ +//! YOLO one-shot client: A simple ACP client that runs a single prompt against an agent. +//! +//! This is a simplified example showing basic ACP client usage. It only supports +//! simple command strings (not JSON configs or environment variables). +//! +//! For a more full-featured client with JSON config support, see the `yopo` binary crate. +//! +//! # Usage +//! +//! ```bash +//! cargo run --example yolo_one_shot_client -- --command "python my_agent.py" "What is 2+2?" +//! ``` + +use agent_client_protocol_core::schema::{ + ContentBlock, InitializeRequest, NewSessionRequest, PromptRequest, ProtocolVersion, + RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse, + SelectedPermissionOutcome, SessionNotification, TextContent, +}; +use agent_client_protocol_core::{Agent, Client, ConnectionTo}; +use clap::Parser; +use std::path::PathBuf; +use tokio::process::Child; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +#[derive(Parser)] +#[command(name = "yolo-one-shot-client")] +#[command(about = "A simple ACP client for one-shot prompts", long_about = None)] +struct Cli { + /// The command to run the agent (e.g., "python my_agent.py") + #[arg(short, long)] + command: String, + + /// The prompt to send to the agent + prompt: String, +} + +/// Parse a command string into command and args +fn parse_command_string(s: &str) -> Result<(PathBuf, Vec), Box> { + let parts = shell_words::split(s)?; + if parts.is_empty() { + return Err("Command string cannot be empty".into()); + } + let command = PathBuf::from(&parts[0]); + let args = parts[1..].to_vec(); + Ok((command, args)) +} + +/// Spawn a process for the agent and get stdio streams. +fn spawn_agent_process( + command: PathBuf, + args: Vec, +) -> Result< + ( + tokio::process::ChildStdin, + tokio::process::ChildStdout, + Child, + ), + Box, +> { + let mut cmd = tokio::process::Command::new(&command); + cmd.args(&args); + cmd.stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()); + + let mut child = cmd.spawn()?; + let child_stdin = child.stdin.take().ok_or("Failed to open stdin")?; + let child_stdout = child.stdout.take().ok_or("Failed to open stdout")?; + + Ok((child_stdin, child_stdout, child)) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let cli = Cli::parse(); + + // Parse the command string + let (command, args) = parse_command_string(&cli.command)?; + + eprintln!("🚀 Spawning agent: {} {:?}", command.display(), args); + + // Spawn the agent process + let (child_stdin, child_stdout, mut child) = spawn_agent_process(command, args)?; + + // Create transport and connection + let transport = agent_client_protocol_core::ByteStreams::new( + child_stdin.compat_write(), + child_stdout.compat(), + ); + + // Run the client + let client = Client + .builder() + .on_receive_notification( + async move |notification: SessionNotification, _cx| { + // Print session updates to stdout (so 2>/dev/null shows only agent output) + println!("{:?}", notification.update); + Ok(()) + }, + agent_client_protocol_core::on_receive_notification!(), + ) + .on_receive_request( + async move |request: RequestPermissionRequest, responder, _connection| { + // YOLO: Auto-approve all permission requests by selecting the first option + eprintln!("✅ Auto-approving permission request: {request:?}"); + let option_id = request.options.first().map(|opt| opt.option_id.clone()); + if let Some(id) = option_id { + responder.respond(RequestPermissionResponse::new( + RequestPermissionOutcome::Selected(SelectedPermissionOutcome::new(id)), + )) + } else { + eprintln!("⚠️ No options provided in permission request, cancelling"); + responder.respond(RequestPermissionResponse::new( + RequestPermissionOutcome::Cancelled, + )) + } + }, + agent_client_protocol_core::on_receive_request!(), + ) + .connect_with(transport, |connection: ConnectionTo| async move { + // Initialize the agent + eprintln!("🤝 Initializing agent..."); + let init_response = connection + .send_request(InitializeRequest::new(ProtocolVersion::LATEST)) + .block_task() + .await?; + + eprintln!("✓ Agent initialized: {:?}", init_response.agent_info); + + // Create a new session + eprintln!("📝 Creating new session..."); + let new_session_response = connection + .send_request(NewSessionRequest::new( + std::env::current_dir().unwrap_or_else(|_| PathBuf::from("/")), + )) + .block_task() + .await?; + + let session_id = new_session_response.session_id; + eprintln!("✓ Session created"); + + // Send the prompt + eprintln!("💬 Sending prompt: \"{}\"", cli.prompt); + let prompt_response = connection + .send_request(PromptRequest::new( + session_id.clone(), + vec![ContentBlock::Text(TextContent::new(cli.prompt.clone()))], + )) + .block_task() + .await?; + + eprintln!("✅ Agent completed!"); + eprintln!("Stop reason: {:?}", prompt_response.stop_reason); + + Ok(()) + }); + Box::pin(client).await?; + + // Kill the child process when done + drop(child.kill().await); + + Ok(()) +} diff --git a/src/agent-client-protocol-core/src/acp.rs b/src/agent-client-protocol-core/src/acp.rs new file mode 100644 index 0000000..17f8c1d --- /dev/null +++ b/src/agent-client-protocol-core/src/acp.rs @@ -0,0 +1,9 @@ +/// ACP messages sent from agent to client +pub mod agent_to_client; +/// ACP messages sent from client to agent +pub mod client_to_agent; +/// Enum type implementations for ACP message types +mod enum_impls; + +pub use agent_to_client::*; +pub use client_to_agent::*; diff --git a/src/agent-client-protocol-core/src/capabilities.rs b/src/agent-client-protocol-core/src/capabilities.rs new file mode 100644 index 0000000..ee3eff8 --- /dev/null +++ b/src/agent-client-protocol-core/src/capabilities.rs @@ -0,0 +1,188 @@ +//! Capability management for the `_meta.symposium` object in ACP messages. +//! +//! This module provides traits and types for working with capabilities stored in +//! the `_meta.symposium` field of `InitializeRequest` and `InitializeResponse`. +//! +//! # Example +//! +//! ```rust,no_run +//! use agent_client_protocol_core::{MetaCapabilityExt, McpAcpTransport}; +//! # use agent_client_protocol_core::schema::InitializeResponse; +//! # let init_response: InitializeResponse = unimplemented!(); +//! +//! let response = init_response.add_meta_capability(McpAcpTransport); +//! if response.has_meta_capability(McpAcpTransport) { +//! // Agent supports MCP-over-ACP bridging +//! } +//! ``` + +use crate::schema::{InitializeRequest, InitializeResponse}; +use serde_json::json; + +/// Trait for capabilities stored in the `_meta.symposium` object. +/// +/// Capabilities are key-value pairs that signal features or context to components +/// in the proxy chain. Implement this trait to define new capabilities. +pub trait MetaCapability { + /// The key name in the `_meta.symposium` object (e.g., "proxy", "mcp_acp_transport") + fn key(&self) -> &'static str; + + /// The value to set when adding this capability (defaults to `true`) + fn value(&self) -> serde_json::Value { + serde_json::Value::Bool(true) + } +} + +/// The mcp_acp_transport capability - indicates support for MCP-over-ACP bridging. +/// +/// When present in `_meta.symposium.mcp_acp_transport`, signals that the agent +/// supports having MCP servers with `acp:UUID` transport proxied through the conductor. +#[derive(Debug)] +pub struct McpAcpTransport; + +impl MetaCapability for McpAcpTransport { + fn key(&self) -> &'static str { + "mcp_acp_transport" + } +} + +/// Extension trait for checking and modifying capabilities in `InitializeRequest`. +pub trait MetaCapabilityExt { + /// Check if a capability is present in `_meta.symposium` + fn has_meta_capability(&self, capability: impl MetaCapability) -> bool; + + /// Add a capability to `_meta.symposium`, creating the structure if needed + #[must_use] + fn add_meta_capability(self, capability: impl MetaCapability) -> Self; + + /// Remove a capability from `_meta.symposium` if present + #[must_use] + fn remove_meta_capability(self, capability: impl MetaCapability) -> Self; +} + +impl MetaCapabilityExt for InitializeRequest { + fn has_meta_capability(&self, capability: impl MetaCapability) -> bool { + self.client_capabilities + .meta + .as_ref() + .and_then(|meta| meta.get("symposium")) + .and_then(|symposium| symposium.get(capability.key())) + .is_some() + } + + fn add_meta_capability(mut self, capability: impl MetaCapability) -> Self { + let meta = self + .client_capabilities + .meta + .get_or_insert_with(Default::default); + + let symposium = meta.entry("symposium").or_insert_with(|| json!({})); + + if let Some(symposium_obj) = symposium.as_object_mut() { + symposium_obj.insert("version".to_string(), json!("1.0")); + symposium_obj.insert(capability.key().to_string(), capability.value()); + } + + self + } + + fn remove_meta_capability(mut self, capability: impl MetaCapability) -> Self { + if let Some(ref mut meta) = self.client_capabilities.meta + && let Some(symposium) = meta.get_mut("symposium") + && let Some(symposium_obj) = symposium.as_object_mut() + { + symposium_obj.remove(capability.key()); + } + self + } +} + +impl MetaCapabilityExt for InitializeResponse { + fn has_meta_capability(&self, capability: impl MetaCapability) -> bool { + self.agent_capabilities + .meta + .as_ref() + .and_then(|meta| meta.get("symposium")) + .and_then(|symposium| symposium.get(capability.key())) + .is_some() + } + + fn add_meta_capability(mut self, capability: impl MetaCapability) -> Self { + let meta = self + .agent_capabilities + .meta + .get_or_insert_with(Default::default); + + let symposium = meta.entry("symposium").or_insert_with(|| json!({})); + + if let Some(symposium_obj) = symposium.as_object_mut() { + symposium_obj.insert("version".to_string(), json!("1.0")); + symposium_obj.insert(capability.key().to_string(), capability.value()); + } + + self + } + + fn remove_meta_capability(mut self, capability: impl MetaCapability) -> Self { + if let Some(ref mut meta) = self.agent_capabilities.meta + && let Some(symposium) = meta.get_mut("symposium") + && let Some(symposium_obj) = symposium.as_object_mut() + { + symposium_obj.remove(capability.key()); + } + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::schema::{ClientCapabilities, ProtocolVersion}; + use serde_json::json; + + #[test] + fn test_add_capability_to_request() { + let request = InitializeRequest::new(ProtocolVersion::LATEST); + + let request = request.add_meta_capability(McpAcpTransport); + + assert!(request.has_meta_capability(McpAcpTransport)); + assert_eq!( + request.client_capabilities.meta.as_ref().unwrap()["symposium"]["mcp_acp_transport"], + json!(true) + ); + } + + #[test] + fn test_remove_capability_from_request() { + let mut meta = serde_json::Map::new(); + meta.insert( + "symposium".to_string(), + json!({ + "version": "1.0", + "mcp_acp_transport": true + }), + ); + let client_capabilities = ClientCapabilities::new().meta(meta); + + let request = InitializeRequest::new(ProtocolVersion::LATEST) + .client_capabilities(client_capabilities); + + let request = request.remove_meta_capability(McpAcpTransport); + + assert!(!request.has_meta_capability(McpAcpTransport)); + } + + #[test] + fn test_add_capability_to_response() { + let response = InitializeResponse::new(ProtocolVersion::LATEST); + + let response = response.add_meta_capability(McpAcpTransport); + + assert!(response.has_meta_capability(McpAcpTransport)); + assert_eq!( + response.agent_capabilities.meta.as_ref().unwrap()["symposium"]["mcp_acp_transport"], + json!(true) + ); + } +} diff --git a/src/agent-client-protocol-core/src/component.rs b/src/agent-client-protocol-core/src/component.rs new file mode 100644 index 0000000..9ee8a11 --- /dev/null +++ b/src/agent-client-protocol-core/src/component.rs @@ -0,0 +1,257 @@ +//! ConnectTo abstraction for agents and proxies. +//! +//! This module provides the [`ConnectTo`] trait that defines the interface for things +//! that can be run as part of a conductor's chain - agents, proxies, or any ACP-speaking component. +//! +//! ## Usage +//! +//! Components connect to other components, creating a chain of message processors. +//! The type parameter `R` is the role that this component connects to (its counterpart). +//! +//! To implement a component, implement the `connect_to` method: +//! +//! ```rust,ignore +//! use agent_client_protocol_core::ConnectTo; +//! use agent_client_protocol_core::Client; +//! +//! struct MyAgent { +//! // configuration fields +//! } +//! +//! // An agent connects to clients +//! impl ConnectTo for MyAgent { +//! async fn connect_to(self, client: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! agent_client_protocol_core::Agent.builder() +//! .name("my-agent") +//! // configure handlers here +//! .connect_to(client) +//! .await +//! } +//! } +//! ``` + +use futures::future::BoxFuture; +use std::{fmt::Debug, future::Future, marker::PhantomData}; + +use crate::{Channel, role::Role}; + +/// A component that can exchange JSON-RPC messages to an endpoint playing the role `R` +/// (e.g., an ACP [`Agent`](`crate::role::acp::Agent`) or an MCP [`Server`](`crate::role::mcp::Server`)). +/// +/// This trait represents anything that can communicate via JSON-RPC messages over channels - +/// agents, proxies, in-process connections, or any ACP-speaking component. +/// +/// The type parameter `R` is the role that this component serves (its counterpart). +/// For example: +/// - An agent implements `Serve` - it serves clients +/// - A proxy implements `Serve` - it serves conductors +/// - Transports like `Channel` implement `Serve` for all `R` since they're role-agnostic +/// +/// # Component Types +/// +/// The trait is implemented by several built-in types representing different communication patterns: +/// +/// - **[`ByteStreams`]**: A component communicating over byte streams (stdin/stdout, sockets, etc.) +/// - **[`Channel`]**: A component communicating via in-process message channels (for testing or direct connections) +/// - **[`AcpAgent`]**: An external agent running in a separate process with stdio communication +/// - **Custom components**: Proxies, transformers, or any ACP-aware service +/// +/// # Two Ways to Serve +/// +/// Components can be used in two ways: +/// +/// 1. **`serve(client)`** - Serve by forwarding to another component (most components implement this) +/// 2. **`into_server()`** - Convert into a channel endpoint and server future (base cases implement this) +/// +/// Most components only need to implement `serve(client)` - the `into_server()` method has a default +/// implementation that creates an intermediate channel and calls `serve`. +/// +/// # Implementation Example +/// +/// ```rust,ignore +/// use agent_client_protocol_core::{Serve, role::Client}; +/// +/// struct MyAgent { +/// config: AgentConfig, +/// } +/// +/// impl Serve for MyAgent { +/// async fn serve(self, client: impl Serve) -> Result<(), agent_client_protocol_core::Error> { +/// // Set up connection that forwards to client +/// agent_client_protocol_core::Agent.builder() +/// .name("my-agent") +/// .on_receive_request(async |req: MyRequest, cx| { +/// // Handle request +/// cx.respond(MyResponse { status: "ok".into() }) +/// }) +/// .serve(client) +/// .await +/// } +/// } +/// ``` +/// +/// # Heterogeneous Collections +/// +/// For storing different component types in the same collection, use [`DynConnectTo`]: +/// +/// ```rust,ignore +/// use agent_client_protocol_core::Client; +/// +/// let components: Vec> = vec![ +/// DynConnectTo::new(proxy1), +/// DynConnectTo::new(proxy2), +/// DynConnectTo::new(agent), +/// ]; +/// ``` +/// +/// [`ByteStreams`]: crate::ByteStreams +/// [`AcpAgent`]: https://docs.rs/agent-client-protocol-tokio/latest/agent_client_protocol_tokio/struct.AcpAgent.html +/// [`Builder`]: crate::Builder +pub trait ConnectTo: Send + 'static { + /// Serve this component by forwarding to a client component. + /// + /// Most components implement this method to set up their connection and + /// forward messages to the provided client. + /// + /// # Arguments + /// + /// * `client` - The component to forward messages to (implements `Serve`) + /// + /// # Returns + /// + /// A future that resolves when the component stops serving, either successfully + /// or with an error. The future must be `Send`. + fn connect_to( + self, + client: impl ConnectTo, + ) -> impl Future> + Send; + + /// Convert this component into a channel endpoint and server future. + /// + /// This method returns: + /// - A `Channel` that can be used to communicate with this component + /// - A `BoxFuture` that runs the component's server logic + /// + /// The default implementation creates an intermediate channel pair and calls `serve` + /// on one endpoint while returning the other endpoint for the caller to use. + /// + /// Base cases like `Channel` and `ByteStreams` override this to avoid unnecessary copying. + /// + /// # Returns + /// + /// A tuple of `(Channel, BoxFuture)` where the channel is for the caller to use + /// and the future must be spawned to run the server. + fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) + where + Self: Sized, + { + let (channel_a, channel_b) = Channel::duplex(); + let future = Box::pin(self.connect_to(channel_b)); + (channel_a, future) + } +} + +/// Type-erased connect trait for object-safe dynamic dispatch. +/// +/// This trait is internal and used by [`DynConnectTo`]. Users should implement +/// [`ConnectTo`] instead, which is automatically converted to `ErasedConnectTo` +/// via a blanket implementation. +trait ErasedConnectTo: Send { + fn type_name(&self) -> String; + + fn connect_to_erased( + self: Box, + client: Box>, + ) -> BoxFuture<'static, Result<(), crate::Error>>; + + fn into_channel_and_future_erased( + self: Box, + ) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>); +} + +/// Blanket implementation: any `Serve` can be type-erased. +impl, R: Role> ErasedConnectTo for C { + fn type_name(&self) -> String { + std::any::type_name::().to_string() + } + + fn connect_to_erased( + self: Box, + client: Box>, + ) -> BoxFuture<'static, Result<(), crate::Error>> { + Box::pin(async move { + (*self) + .connect_to(DynConnectTo { + inner: client, + _marker: PhantomData, + }) + .await + }) + } + + fn into_channel_and_future_erased( + self: Box, + ) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { + (*self).into_channel_and_future() + } +} + +/// A dynamically-typed component for heterogeneous collections. +/// +/// This type wraps any [`ConnectTo`] implementation and provides dynamic dispatch, +/// allowing you to store different component types in the same collection. +/// +/// The type parameter `R` is the role that all components in the +/// collection serve (their counterpart). +/// +/// # Examples +/// +/// ```rust,ignore +/// use agent_client_protocol_core::{DynConnectTo, Client}; +/// +/// let components: Vec> = vec![ +/// DynConnectTo::new(Proxy1), +/// DynConnectTo::new(Proxy2), +/// DynConnectTo::new(Agent), +/// ]; +/// ``` +pub struct DynConnectTo { + inner: Box>, + _marker: PhantomData, +} + +impl DynConnectTo { + /// Create a new `DynConnectTo` from any type implementing [`ConnectTo`]. + pub fn new>(component: C) -> Self { + Self { + inner: Box::new(component), + _marker: PhantomData, + } + } + + /// Returns the type name of the wrapped component. + #[must_use] + pub fn type_name(&self) -> String { + self.inner.type_name() + } +} + +impl ConnectTo for DynConnectTo { + async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { + self.inner + .connect_to_erased(Box::new(client) as Box>) + .await + } + + fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { + self.inner.into_channel_and_future_erased() + } +} + +impl Debug for DynConnectTo { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DynServe") + .field("type_name", &self.type_name()) + .finish() + } +} diff --git a/src/agent-client-protocol-core/src/concepts/acp_basics.rs b/src/agent-client-protocol-core/src/concepts/acp_basics.rs new file mode 100644 index 0000000..d0f2a6b --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/acp_basics.rs @@ -0,0 +1,75 @@ +//! The roles in ACP: clients, agents, proxies, and conductors. +//! +//! The Agent-Client Protocol defines how AI agents communicate with the +//! applications that use them. Understanding the four roles helps you choose +//! the right approach for what you're building. +//! +//! # Clients +//! +//! A **client** is an application that wants to use an AI agent. Examples include: +//! +//! - IDEs like VS Code or JetBrains +//! - Command-line tools +//! - Web applications +//! - Automation scripts +//! +//! Clients send prompts to agents and receive responses. They can also handle +//! requests from the agent (like permission requests or tool approvals). +//! +//! # Agents +//! +//! An **agent** is an AI-powered service that responds to prompts. Examples include: +//! +//! - Claude Code +//! - Custom agents built with language models +//! +//! Agents receive prompts, process them (typically by calling an LLM), and stream +//! back responses. They may also request permissions, invoke tools, or ask for +//! user input during processing. +//! +//! # Proxies +//! +//! A **proxy** sits between a client and an agent, intercepting and potentially +//! modifying messages in both directions. Proxies are useful for: +//! +//! - Adding tools via MCP (Model Context Protocol) servers +//! - Injecting system prompts or context +//! - Logging and debugging +//! - Filtering or transforming messages +//! +//! Proxies can be chained - you can have multiple proxies between a client and +//! an agent, each adding its own capabilities. +//! +//! # Conductors +//! +//! A **conductor** orchestrates a chain of proxies with a final agent. It: +//! +//! - Spawns and manages proxy processes +//! - Routes messages through the chain +//! - Handles initialization and shutdown +//! +//! The [`agent-client-protocol-conductor`] crate provides a conductor implementation. Most users +//! don't need to implement conductors themselves - they just configure which +//! proxies to use. +//! +//! # Message Flow +//! +//! Messages flow through the system like this: +//! +//! ```text +//! Client <-> Proxy 1 <-> Proxy 2 <-> ... <-> Agent +//! ^ ^ +//! | | +//! +------ Conductor manages -------+ +//! ``` +//! +//! Each arrow represents a bidirectional connection. Requests flow toward the +//! agent, responses flow back toward the client, and notifications can flow +//! in either direction. +//! +//! # Next Steps +//! +//! Now that you understand the roles, see [Connections and Links](super::connections) +//! to learn how to establish connections in code. +//! +//! [`agent-client-protocol-conductor`]: https://crates.io/crates/agent-client-protocol-conductor diff --git a/src/agent-client-protocol-core/src/concepts/callbacks.rs b/src/agent-client-protocol-core/src/concepts/callbacks.rs new file mode 100644 index 0000000..231c117 --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/callbacks.rs @@ -0,0 +1,127 @@ +//! Handling incoming messages with `on_receive_*` callbacks. +//! +//! So far we've seen how to *send* messages. But ACP is bidirectional - the +//! remote peer can also send messages to you. Use callbacks to handle them. +//! +//! # Handling Requests +//! +//! Use `on_receive_request` to handle incoming requests that expect a response: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::{ValidateRequest, ValidateResponse}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Client.builder() +//! .on_receive_request(async |req: ValidateRequest, responder, cx| { +//! // Process the request +//! let is_valid = req.data.len() > 0; +//! +//! // Send the response +//! responder.respond(ValidateResponse { is_valid, error: None }) +//! }, agent_client_protocol_core::on_receive_request!()) +//! .connect_with(transport, async |cx| { Ok(()) }) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Your callback receives three arguments: +//! - The request payload (e.g., `PermissionRequest`) +//! - A [`Responder`] for sending the response +//! - A [`ConnectionTo`] for sending other messages +//! +//! # Handling Notifications +//! +//! Use `on_receive_notification` for fire-and-forget messages that don't need +//! a response: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::StatusUpdate; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Client.builder() +//! .on_receive_notification(async |notif: StatusUpdate, cx| { +//! println!("Status: {}", notif.message); +//! Ok(()) +//! }, agent_client_protocol_core::on_receive_notification!()) +//! # .connect_with(transport, async |_| Ok(())).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # The Request Context +//! +//! The [`Responder`] lets you send a response to the request: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::{MyRequest, MyResponse}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder() +//! # .on_receive_request(async |req: MyRequest, responder, cx| { +//! // Send a successful response +//! responder.respond(MyResponse { status: "ok".into() })?; +//! # Ok(()) +//! # }, agent_client_protocol_core::on_receive_request!()) +//! # .connect_with(transport, async |_| Ok(())).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Or send an error: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::{MyRequest, MyResponse}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder() +//! # .on_receive_request(async |req: MyRequest, responder, cx| { +//! responder.respond_with_error(agent_client_protocol_core::Error::invalid_params())?; +//! # Ok(()) +//! # }, agent_client_protocol_core::on_receive_request!()) +//! # .connect_with(transport, async |_| Ok(())).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! You must send exactly one response per request. If your callback returns +//! without responding, an error response is sent automatically. +//! +//! # Multiple Handlers +//! +//! You can register multiple handlers. They're tried in order until one +//! handles the message: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::{ValidateRequest, ValidateResponse, ExecuteRequest, ExecuteResponse}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Client.builder() +//! .on_receive_request(async |req: ValidateRequest, responder, cx| { +//! // Handle validation requests +//! responder.respond(ValidateResponse { is_valid: true, error: None }) +//! }, agent_client_protocol_core::on_receive_request!()) +//! .on_receive_request(async |req: ExecuteRequest, responder, cx| { +//! // Handle execution requests +//! responder.respond(ExecuteResponse { result: "done".into() }) +//! }, agent_client_protocol_core::on_receive_request!()) +//! # .connect_with(transport, async |_| Ok(())).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Ordering Guarantees +//! +//! Callbacks run inside the dispatch loop and block further message processing +//! until they complete. This gives you ordering guarantees but also means you +//! need to be careful about deadlocks. +//! +//! See [Ordering](super::ordering) for the full details. +//! +//! # Next Steps +//! +//! - [Explicit Peers](super::peers) - Use `_from` variants to specify the source peer +//! - [Ordering](super::ordering) - Understand dispatch loop semantics +//! +//! [`Responder`]: crate::Responder +//! [`ConnectionTo`]: crate::ConnectionTo diff --git a/src/agent-client-protocol-core/src/concepts/connections.rs b/src/agent-client-protocol-core/src/concepts/connections.rs new file mode 100644 index 0000000..d93893c --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/connections.rs @@ -0,0 +1,120 @@ +//! Establishing connections using role types and connection builders. +//! +//! To communicate over ACP, you need to establish a connection. This involves +//! choosing a **role type** that matches your role and using a **connection builder** +//! to configure and run the connection. +//! +//! # Choosing a Role Type +//! +//! Your role type determines what messages you can send and who you can send them to. +//! Choose based on what you're building: +//! +//! | You are building... | Use this role type | +//! |---------------------|-------------------| +//! | A client that talks to an agent | [`Client`] | +//! | An agent that responds to clients | [`Agent`] | +//! | A proxy in a conductor chain | [`Proxy`] | +//! +//! # The Connection Builder Pattern +//! +//! Every role type has a `builder()` method that returns a connection builder. +//! The builder lets you configure handlers, then connect to a transport: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Client.builder() +//! .name("my-client") +//! .connect_with(transport, async |cx| { +//! // Use `cx` to send requests and handle responses +//! Ok(()) +//! }) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # The Connection Context +//! +//! Inside `connect_with`, you receive a [`ConnectionTo`] (connection context) that +//! lets you interact with the remote peer: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_core::schema::{InitializeRequest, ProtocolVersion}; +//! # use agent_client_protocol_test::StatusUpdate; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! // Send a request and wait for the response +//! let response = cx.send_request(InitializeRequest::new(ProtocolVersion::LATEST)) +//! .block_task() +//! .await?; +//! +//! // Send a notification (fire-and-forget) +//! cx.send_notification(StatusUpdate { message: "hello".into() })?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Sending Requests +//! +//! When you call `send_request()`, you get back a [`SentRequest`] that represents +//! the pending response. You have two main ways to handle it: +//! +//! ## Option 1: Block and wait +//! +//! Use `block_task()` when you need the response before continuing: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::MyRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! let response = cx.send_request(MyRequest {}) +//! .block_task() +//! .await?; +//! // Use response here +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! ## Option 2: Schedule a callback +//! +//! Use `on_receiving_result()` when you want to handle the response asynchronously: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::MyRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.send_request(MyRequest {}) +//! .on_receiving_result(async |result| { +//! match result { +//! Ok(response) => { /* handle success */ } +//! Err(error) => { /* handle error */ } +//! } +//! Ok(()) +//! })?; +//! // Continues immediately, callback runs when response arrives +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! See [Ordering](super::ordering) for important details about how these differ. +//! +//! # Next Steps +//! +//! - [Sessions](super::sessions) - Create multi-turn conversations +//! - [Callbacks](super::callbacks) - Handle incoming requests from the remote peer +//! +//! [`Client`]: crate::Client +//! [`Agent`]: crate::Agent +//! [`Proxy`]: crate::Proxy +//! [`ConnectionTo`]: crate::ConnectionTo +//! [`SentRequest`]: crate::SentRequest diff --git a/src/agent-client-protocol-core/src/concepts/error_handling.rs b/src/agent-client-protocol-core/src/concepts/error_handling.rs new file mode 100644 index 0000000..f1cc0b9 --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/error_handling.rs @@ -0,0 +1,116 @@ +//! Error handling patterns in agent-client-protocol-core. +//! +//! This chapter explains how errors work in agent-client-protocol-core callbacks and the difference +//! between *protocol errors* (sent to the peer) and *connection errors* (which +//! shut down the connection). +//! +//! # Callback Return Types +//! +//! Almost all agent-client-protocol-core callbacks return `Result<_, crate::Error>`. What happens +//! when you return an `Err` depends on the context: +//! +//! **Returning `Err` from a callback shuts down the connection.** +//! +//! This is appropriate for truly unrecoverable situations—internal bugs, +//! resource exhaustion, or when you want to terminate the connection. +//! But most of the time, you want to send an error *to the peer* while +//! keeping the connection alive. +//! +//! # Sending Protocol Errors +//! +//! To send an error response to a request (without closing the connection), +//! use the request context's `respond` method: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::{ValidateRequest, ValidateResponse}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Client.builder() +//! .on_receive_request(async |request: ValidateRequest, responder, _cx| { +//! if request.data.is_empty() { +//! // Send error to peer, keep connection alive +//! responder.respond_with_error(agent_client_protocol_core::Error::invalid_params())?; +//! return Ok(()); +//! } +//! +//! // Process valid request... +//! responder.respond(ValidateResponse { is_valid: true, error: None })?; +//! Ok(()) +//! }, agent_client_protocol_core::on_receive_request!()) +//! # .connect_with(transport, async |_| Ok(())).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! For sending error notifications (one-way error messages), use +//! [`send_error_notification`][crate::ConnectionTo::send_error_notification]: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.send_error_notification(agent_client_protocol_core::Error::internal_error() +//! .data("Something went wrong"))?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # The `into_internal_error` Helper +//! +//! When working with external libraries that return their own error types, +//! you need to convert them to `agent_client_protocol_core::Error`. The +//! [`Error::into_internal_error`][crate::Error::into_internal_error] method +//! provides a convenient way to do this: +//! +//! ``` +//! use agent_client_protocol_core::Error; +//! +//! # fn example() -> Result<(), agent_client_protocol_core::Error> { +//! # let data = "hello"; +//! # let path = "/tmp/test.txt"; +//! // Convert any error type to agent_client_protocol_core::Error +//! let value = serde_json::to_value(&data) +//! .map_err(Error::into_internal_error)?; +//! +//! // Or with a file operation +//! let contents = std::fs::read_to_string(path) +//! .map_err(Error::into_internal_error); +//! # Ok(()) +//! # } +//! ``` +//! +//! This wraps the original error's message in an internal error, which is +//! appropriate for unexpected failures. For expected error conditions that +//! should be communicated to the peer, create specific error types instead. +//! +//! # Error Types +//! +//! The [`Error`][crate::Error] type provides factory methods for common +//! JSON-RPC error codes: +//! +//! - [`Error::parse_error()`][crate::Error::parse_error] - Invalid JSON +//! - [`Error::invalid_request()`][crate::Error::invalid_request] - Malformed request +//! - [`Error::method_not_found()`][crate::Error::method_not_found] - Unknown method +//! - [`Error::invalid_params()`][crate::Error::invalid_params] - Bad parameters +//! - [`Error::internal_error()`][crate::Error::internal_error] - Server error +//! +//! You can add context with `.data()`: +//! +//! ``` +//! let error = agent_client_protocol_core::Error::invalid_params() +//! .data(serde_json::json!({ +//! "field": "timeout", +//! "reason": "must be positive" +//! })); +//! ``` +//! +//! # Summary +//! +//! | Situation | What to do | +//! |-----------|------------| +//! | Send error response to request | `responder.respond(Err(error))` then `Ok(())` | +//! | Send error notification | `cx.send_error_notification(error)` then `Ok(())` | +//! | Shut down connection | Return `Err(error)` from callback | +//! | Convert external error | `.map_err(Error::into_internal_error)?` | diff --git a/src/agent-client-protocol-core/src/concepts/mod.rs b/src/agent-client-protocol-core/src/concepts/mod.rs new file mode 100644 index 0000000..5c32b48 --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/mod.rs @@ -0,0 +1,40 @@ +//! Core concepts for understanding and using agent-client-protocol-core. +//! +//! This module provides detailed explanations of the key concepts you need +//! to work effectively with the agent-client-protocol-core SDK. Read these in order for a progressive +//! introduction, or jump to specific topics as needed. +//! +//! # Table of Contents +//! +//! 1. [ACP Basics][`crate::concepts::acp_basics`] - The roles in the protocol: clients, +//! agents, proxies, and conductors. +//! +//! 2. [Connections][`crate::concepts::connections`] - How to establish connections +//! using link types and connection builders. +//! +//! 3. [Sessions][`crate::concepts::sessions`] - Creating and managing sessions for +//! multi-turn conversations with agents. +//! +//! 4. [Callbacks][`crate::concepts::callbacks`] - Handling incoming messages with +//! `on_receive_*` methods. +//! +//! 5. [Explicit Peers][`crate::concepts::peers`] - Using `_to` and `_from` variants +//! when you need to specify which peer you're communicating with. +//! +//! 6. [Ordering][`crate::concepts::ordering`] - How the dispatch loop processes +//! messages and what ordering guarantees you get. +//! +//! 7. [Proxies and Conductors][`crate::concepts::proxies`] - Building proxies that +//! intercept and modify messages between clients and agents. +//! +//! 8. [Error Handling][`crate::concepts::error_handling`] - Protocol errors vs +//! connection errors, and how to handle them. + +pub mod acp_basics; +pub mod callbacks; +pub mod connections; +pub mod error_handling; +pub mod ordering; +pub mod peers; +pub mod proxies; +pub mod sessions; diff --git a/src/agent-client-protocol-core/src/concepts/ordering.rs b/src/agent-client-protocol-core/src/concepts/ordering.rs new file mode 100644 index 0000000..602b844 --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/ordering.rs @@ -0,0 +1,174 @@ +//! Message ordering, concurrency, and the dispatch loop. +//! +//! Understanding how agent-client-protocol-core processes messages is key to writing correct code. +//! This chapter explains the dispatch loop and the ordering guarantees you +//! can rely on. +//! +//! # The Dispatch Loop +//! +//! Each connection has a central **dispatch loop** that processes incoming +//! messages one at a time. When a message arrives, it is passed to your +//! handlers in order until one claims it. +//! +//! The key property: **the dispatch loop waits for each handler to complete +//! before processing the next message.** This gives you sequential ordering +//! guarantees within a single connection. +//! +//! # `on_*` Methods Block the Loop +//! +//! Methods whose names begin with `on_` register callbacks that run inside +//! the dispatch loop. When your callback is invoked, the loop is blocked +//! until your callback completes. +//! +//! This includes: +//! - [`on_receive_request`] and [`on_receive_notification`] +//! - [`on_receiving_result`] and [`on_receiving_ok_result`] +//! - [`on_session_start`] and [`on_proxy_session_start`] +//! +//! This means: +//! - No other messages are processed while your callback runs +//! - You can safely do setup before "releasing" control back to the loop +//! - Messages are processed in the order they arrive +//! +//! # Deadlock Risk +//! +//! Because `on_*` callbacks block the dispatch loop, it's easy to create +//! deadlocks. The most common pattern: +//! +//! ```ignore +//! // DEADLOCK: This blocks the loop waiting for a response, +//! // but the response can't arrive because the loop is blocked! +//! builder.on_receive_request(async |request: MyRequest, responder, cx| { +//! let response = cx.send_request(SomeRequest { ... }) +//! .block_task() // <-- Waits for response +//! .await?; // <-- But response can never arrive! +//! responder.respond(response) +//! }, on_receive_request!()); +//! ``` +//! +//! The response can never arrive because the dispatch loop is blocked waiting +//! for your callback to complete. +//! +//! # `block_task` vs `on_receiving_result` +//! +//! When you send a request, you get a [`SentRequest`] with two ways to handle it: +//! +//! ## `block_task()` - Acks immediately, you process later +//! +//! Use this in spawned tasks where you need to wait for the response: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::MyRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.spawn({ +//! let cx = cx.clone(); +//! async move { +//! // Safe: we're in a spawned task, not blocking the dispatch loop +//! let response = cx.send_request(MyRequest {}) +//! .block_task() +//! .await?; +//! // Process response... +//! Ok(()) +//! } +//! })?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! The dispatch loop continues immediately after delivering the response. +//! Your code receives the response and can take as long as it wants. +//! +//! ## `on_receiving_result()` - Your callback blocks the loop +//! +//! Use this when you need ordering guarantees: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::MyRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.send_request(MyRequest {}) +//! .on_receiving_result(async |result| { +//! // Dispatch loop is blocked until this completes +//! let response = result?; +//! // Do something with response... +//! Ok(()) +//! })?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! The dispatch loop waits for your callback to complete before processing +//! the next message. Use this when you need to ensure no other messages +//! are processed until you've handled the response. +//! +//! # Escaping the Loop: `spawn` +//! +//! Use [`spawn`] to run work outside the dispatch loop: +//! +//! ```ignore +//! builder.on_receive_request(async |request: MyRequest, responder, cx| { +//! cx.spawn(async move { +//! // This runs outside the loop - other messages may be processed +//! let response = cx.send_request(SomeRequest { ... }) +//! .block_task() +//! .await?; +//! // ... +//! Ok(()) +//! })?; +//! responder.respond(MyResponse { ... }) // Return immediately +//! }, on_receive_request!()); +//! ``` +//! +//! # `run_until` Methods +//! +//! Methods named `run_until` (like on session builders) run in a spawned task, +//! so awaiting them won't cause deadlocks: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.build_session_cwd()? +//! .block_task() +//! .run_until(async |mut session| { +//! // Safe to await here - we're in a spawned task +//! session.send_prompt("Hello")?; +//! let response = session.read_to_string().await?; +//! Ok(()) +//! }) +//! .await?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Summary +//! +//! | Pattern | Blocks Loop? | Use When | +//! |---------|--------------|----------| +//! | `on_*` callback | Yes | Quick decisions, need ordering | +//! | `on_receiving_result` | Yes | Need to process response before next message | +//! | `block_task()` | No | In spawned tasks, need response value | +//! | `spawn(...)` | No | Long-running work, don't need ordering | +//! | `block_task().run_until(...)` | No | Session-scoped work | +//! +//! # Next Steps +//! +//! - [Proxies and Conductors](super::proxies) - Building message interceptors +//! +//! [`on_receive_request`]: crate::Builder::on_receive_request +//! [`on_receive_notification`]: crate::Builder::on_receive_notification +//! [`on_receiving_result`]: crate::SentRequest::on_receiving_result +//! [`on_receiving_ok_result`]: crate::SentRequest::on_receiving_ok_result +//! [`on_session_start`]: crate::SessionBuilder::on_session_start +//! [`on_proxy_session_start`]: crate::SessionBuilder::on_proxy_session_start +//! [`SentRequest`]: crate::SentRequest +//! [`spawn`]: crate::ConnectionTo::spawn diff --git a/src/agent-client-protocol-core/src/concepts/peers.rs b/src/agent-client-protocol-core/src/concepts/peers.rs new file mode 100644 index 0000000..5bdfa56 --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/peers.rs @@ -0,0 +1,106 @@ +//! Explicit peers: using `_to` and `_from` variants. +//! +//! So far, we've used methods like `send_request` and `on_receive_request` +//! without specifying *who* we're sending to or receiving from. That's because +//! each role type has a **default peer**. +//! +//! # Default Peers +//! +//! For simple role types, there's only one peer to talk to: +//! +//! | Role Type | Default Peer | +//! |-----------|--------------| +//! | [`Client`] | The agent | +//! | [`Agent`] | The client | +//! +//! So when you write: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_core::schema::{InitializeRequest, ProtocolVersion}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! // As a client +//! cx.send_request(InitializeRequest::new(ProtocolVersion::LATEST)); +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! The request automatically goes to the agent, because that's the only peer +//! a client can talk to. +//! +//! # Explicit Peer Methods +//! +//! Every method has an explicit variant that takes a peer argument: +//! +//! | Default method | Explicit variant | +//! |----------------|------------------| +//! | `send_request` | `send_request_to(peer, request)` | +//! | `send_notification` | `send_notification_to(peer, request)` | +//! | `on_receive_request` | `on_receive_request_from(peer, callback)` | +//! | `on_receive_notification` | `on_receive_notification_from(peer, callback)` | +//! +//! For simple role types, the explicit form is equivalent: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_test::MyRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! # let req = MyRequest {}; +//! // These are equivalent for Client: +//! cx.send_request(req.clone()); +//! cx.send_request_to(Agent, req); +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Why Explicit Peers Matter +//! +//! Explicit peers become essential when working with proxies. A proxy sits +//! between a client and an agent, so it has *two* peers: +//! +//! - [`Client`] - the client (or previous proxy in the chain) +//! - [`Agent`] - the agent (or next proxy in the chain) +//! +//! When writing proxy code, you need to specify which direction: +//! +//! ``` +//! # use agent_client_protocol_core::{Proxy, Client, Agent, Conductor, ConnectTo}; +//! # use agent_client_protocol_test::MyRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Proxy.builder() +//! // Receive a request from the client +//! .on_receive_request_from(Client, async |req: MyRequest, responder, cx| { +//! // Forward it to the agent +//! cx.send_request_to(Agent, req) +//! .forward_response_to(responder) +//! }, agent_client_protocol_core::on_receive_request!()) +//! .connect_to(transport) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! See [Proxies and Conductors](super::proxies) for more on building proxies. +//! +//! # Available Peer Types +//! +//! | Peer Type | Represents | +//! |-----------|------------| +//! | [`Client`] | The client direction | +//! | [`Agent`] | The agent direction | +//! | [`Conductor`] | The conductor (for proxies) | +//! +//! # Next Steps +//! +//! - [Ordering](super::ordering) - Understand dispatch loop semantics +//! - [Proxies and Conductors](super::proxies) - Build proxies that use explicit peers +//! +//! [`Client`]: crate::Client +//! [`Agent`]: crate::Agent +//! [`Conductor`]: crate::Conductor diff --git a/src/agent-client-protocol-core/src/concepts/proxies.rs b/src/agent-client-protocol-core/src/concepts/proxies.rs new file mode 100644 index 0000000..c8857df --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/proxies.rs @@ -0,0 +1,144 @@ +//! Building proxies that intercept and modify messages. +//! +//! A **proxy** sits between a client and an agent, intercepting messages +//! in both directions. This is how you add capabilities like MCP tools, +//! logging, or message transformation. +//! +//! # The Proxy Role Type +//! +//! Proxies use the [`Proxy`] role type, which has two peers: +//! +//! - [`Client`] - messages from/to the client direction +//! - [`Agent`] - messages from/to the agent direction +//! +//! Unlike simpler links, there's no default peer - you must always specify +//! which direction you're communicating with. +//! +//! # Default Forwarding +//! +//! By default, [`Proxy`] forwards all messages it doesn't handle. +//! This means a minimal proxy that does nothing is just: +//! +//! ``` +//! # use agent_client_protocol_core::{Proxy, Conductor, ConnectTo}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Proxy.builder() +//! .connect_to(transport) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! All messages pass through unchanged. +//! +//! # Intercepting Messages +//! +//! To intercept specific messages, use `on_receive_*_from` with explicit peers: +//! +//! ``` +//! # use agent_client_protocol_core::{Proxy, Client, Agent, Conductor, ConnectTo}; +//! # use agent_client_protocol_test::ProcessRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Proxy.builder() +//! // Intercept requests from the client +//! .on_receive_request_from(Client, async |req: ProcessRequest, responder, cx| { +//! // Modify the request +//! let modified = ProcessRequest { +//! data: format!("prefix: {}", req.data), +//! }; +//! +//! // Forward to agent and relay the response back +//! cx.send_request_to(Agent, modified) +//! .forward_response_to(responder) +//! }, agent_client_protocol_core::on_receive_request!()) +//! .connect_to(transport) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Messages you don't handle are forwarded automatically. +//! +//! # Adding MCP Servers +//! +//! A common use case is adding tools via MCP. You can add them globally +//! (available in all sessions) or per-session. +//! +//! ## Global MCP Server +//! +//! ``` +//! # use agent_client_protocol_core::{Proxy, Conductor, ConnectTo}; +//! # use agent_client_protocol_core::mcp_server::McpServer; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # let my_mcp_server = McpServer::::builder("tools").build(); +//! Proxy.builder() +//! .with_mcp_server(my_mcp_server) +//! .connect_to(transport) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! ## Per-Session MCP Server +//! +//! ``` +//! # use agent_client_protocol_core::{Proxy, Client, Conductor, ConnectTo}; +//! # use agent_client_protocol_core::schema::NewSessionRequest; +//! # use agent_client_protocol_core::mcp_server::McpServer; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Proxy.builder() +//! .on_receive_request_from(Client, async |req: NewSessionRequest, responder, cx| { +//! let my_mcp_server = McpServer::::builder("tools").build(); +//! cx.build_session_from(req) +//! .with_mcp_server(my_mcp_server)? +//! .on_proxy_session_start(responder, async |session_id| { +//! // Session started with MCP server attached +//! Ok(()) +//! }) +//! }, agent_client_protocol_core::on_receive_request!()) +//! .connect_to(transport) +//! .await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # The Conductor +//! +//! Proxies don't run standalone - they're orchestrated by a **conductor**. +//! The conductor: +//! +//! - Spawns proxy processes +//! - Chains them together +//! - Connects the final proxy to the agent +//! +//! The [`agent-client-protocol-conductor`] crate provides a conductor binary. You configure +//! it with a list of proxies to run. +//! +//! # Proxy Chains +//! +//! Multiple proxies can be chained: +//! +//! ```text +//! Client <-> Proxy A <-> Proxy B <-> Agent +//! ``` +//! +//! Each proxy sees messages from its perspective: +//! - `Client` is "toward the client" (Proxy A, or conductor if first) +//! - `Agent` is "toward the agent" (Proxy B, or agent if last) +//! +//! Messages flow through each proxy in order. Each can inspect, modify, +//! or handle messages before they continue. +//! +//! # Summary +//! +//! | Task | Approach | +//! |------|----------| +//! | Forward everything | Just `connect_to(transport)` | +//! | Intercept specific messages | `on_receive_*_from` with explicit peers | +//! | Add global tools | `with_mcp_server` on builder | +//! | Add per-session tools | `with_mcp_server` on session builder | +//! +//! [`Proxy`]: crate::Proxy +//! [`Client`]: crate::Client +//! [`Agent`]: crate::Agent +//! [`agent-client-protocol-conductor`]: https://crates.io/crates/agent-client-protocol-conductor diff --git a/src/agent-client-protocol-core/src/concepts/sessions.rs b/src/agent-client-protocol-core/src/concepts/sessions.rs new file mode 100644 index 0000000..dd0f8bd --- /dev/null +++ b/src/agent-client-protocol-core/src/concepts/sessions.rs @@ -0,0 +1,130 @@ +//! Creating and managing sessions for multi-turn conversations. +//! +//! A **session** represents a multi-turn conversation with an agent. Within a +//! session, you can send prompts, receive responses, and the agent maintains +//! context across turns. +//! +//! # Creating a Session +//! +//! Use the session builder to create a new session: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.build_session_cwd()? // Use current working directory +//! .block_task() // Mark as blocking +//! .run_until(async |session| { +//! // Use the session here +//! Ok(()) +//! }) +//! .await?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! Or specify a custom working directory: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.build_session("/path/to/project") +//! .block_task() +//! .run_until(async |session| { Ok(()) }) +//! .await?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Sending Prompts +//! +//! Inside `run_until`, you get an [`ActiveSession`] that lets you interact +//! with the agent: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # Client.builder().connect_with(transport, async |cx| { +//! # cx.build_session_cwd()?.block_task() +//! .run_until(async |mut session| { +//! // Send a prompt +//! session.send_prompt("What is 2 + 2?")?; +//! +//! // Read the complete response as a string +//! let response = session.read_to_string().await?; +//! println!("{}", response); +//! +//! // Send another prompt in the same session +//! session.send_prompt("And what is 3 + 3?")?; +//! let response = session.read_to_string().await?; +//! +//! Ok(()) +//! }) +//! # .await?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! # Adding MCP Servers +//! +//! You can attach MCP (Model Context Protocol) servers to a session to provide +//! tools to the agent: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_core::mcp_server::McpServer; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! # let my_mcp_server = McpServer::::builder("tools").build(); +//! # Client.builder().connect_with(transport, async |cx| { +//! cx.build_session_cwd()? +//! .with_mcp_server(my_mcp_server)? +//! .block_task() +//! .run_until(async |session| { Ok(()) }) +//! .await?; +//! # Ok(()) +//! # }).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! See the cookbook for detailed MCP server examples. +//! +//! # Non-Blocking Session Start +//! +//! If you're inside an `on_receive_*` callback and need to start a session, +//! use `on_session_start` instead of `block_task().run_until()`: +//! +//! ``` +//! # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +//! # use agent_client_protocol_core::schema::NewSessionRequest; +//! # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Client.builder() +//! .on_receive_request(async |req: NewSessionRequest, responder, cx| { +//! cx.build_session_from(req) +//! .on_session_start(async |session| { +//! // Handle the session +//! Ok(()) +//! })?; +//! Ok(()) +//! }, agent_client_protocol_core::on_receive_request!()) +//! # .connect_with(transport, async |_| Ok(())).await?; +//! # Ok(()) +//! # } +//! ``` +//! +//! This follows the same ordering guarantees as other `on_*` methods - see +//! [Ordering](super::ordering) for details. +//! +//! # Next Steps +//! +//! - [Callbacks](super::callbacks) - Handle incoming requests +//! - [Ordering](super::ordering) - Understand when to use `block_task` vs `on_*` +//! +//! [`ActiveSession`]: crate::ActiveSession diff --git a/src/agent-client-protocol-core/src/cookbook.rs b/src/agent-client-protocol-core/src/cookbook.rs new file mode 100644 index 0000000..689cb83 --- /dev/null +++ b/src/agent-client-protocol-core/src/cookbook.rs @@ -0,0 +1,12 @@ +//! Cookbook of common patterns for building ACP components. +//! +//! This module has moved to the [`agent_client_protocol_cookbook`] crate, which provides +//! more comprehensive examples with full access to the entire agent-client-protocol-core ecosystem. +//! +//! See [`agent_client_protocol_cookbook`] for: +//! +//! - **Clients** - Connect to an existing agent and send prompts +//! - **Proxies** - Sit between client and agent to add capabilities (like MCP tools) +//! - **Agents** - Respond to prompts with AI-powered responses +//! +//! [`agent_client_protocol_cookbook`]: https://docs.rs/agent-client-protocol-cookbook diff --git a/src/agent-client-protocol-core/src/handler.rs b/src/agent-client-protocol-core/src/handler.rs new file mode 100644 index 0000000..e232ef9 --- /dev/null +++ b/src/agent-client-protocol-core/src/handler.rs @@ -0,0 +1,7 @@ +//! Handler types for building custom JSON-RPC message handlers. +//! +//! This module contains the handler types used by [`Builder`](crate::Builder) +//! to process incoming messages. Most users won't need to use these types directly, +//! as the builder methods on `Builder` handle the construction automatically. + +pub use crate::jsonrpc::{HandleDispatchFrom, handlers::NullHandler}; diff --git a/src/agent-client-protocol-core/src/jsonrpc.rs b/src/agent-client-protocol-core/src/jsonrpc.rs new file mode 100644 index 0000000..a6ea537 --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc.rs @@ -0,0 +1,3428 @@ +//! Core JSON-RPC server support. + +use agent_client_protocol_schema::SessionId; +// Re-export jsonrpcmsg for use in public API +pub use jsonrpcmsg; + +// Types re-exported from crate root +use serde::{Deserialize, Serialize}; +use std::fmt::Debug; +use std::panic::Location; +use std::pin::pin; +use uuid::Uuid; + +use boxfnonce::SendBoxFnOnce; +use futures::channel::{mpsc, oneshot}; +use futures::future::{self, BoxFuture, Either}; +use futures::{AsyncRead, AsyncWrite, StreamExt}; + +mod dynamic_handler; +pub(crate) mod handlers; +mod incoming_actor; +mod outgoing_actor; +pub(crate) mod run; +mod task_actor; +mod transport_actor; + +use crate::jsonrpc::dynamic_handler::DynamicHandlerMessage; +pub use crate::jsonrpc::handlers::NullHandler; +use crate::jsonrpc::handlers::{ChainedHandler, NamedHandler}; +use crate::jsonrpc::handlers::{MessageHandler, NotificationHandler, RequestHandler}; +use crate::jsonrpc::outgoing_actor::{OutgoingMessageTx, send_raw_message}; +use crate::jsonrpc::run::SpawnedRun; +use crate::jsonrpc::run::{ChainRun, NullRun, RunWithConnectionTo}; +use crate::jsonrpc::task_actor::{Task, TaskTx}; +use crate::mcp_server::McpServer; +use crate::role::HasPeer; +use crate::role::Role; +use crate::util::json_cast; +use crate::{Agent, Client, ConnectTo, RoleId}; + +/// Handlers process incoming JSON-RPC messages on a connection. +/// +/// When messages arrive, they flow through a chain of handlers. Each handler can +/// either **claim** the message (handle it) or **decline** it (pass to the next handler). +/// +/// # Message Flow +/// +/// Messages flow through three layers of handlers in order: +/// +/// ```text +/// ┌─────────────────────────────────────────────────────────────────┐ +/// │ Incoming Message │ +/// └─────────────────────────────────────────────────────────────────┘ +/// │ +/// ▼ +/// ┌─────────────────────────────────────────────────────────────────┐ +/// │ 1. User Handlers (registered via on_receive_request, etc.) │ +/// │ - Tried in registration order │ +/// │ - First handler to return Handled::Yes claims the message │ +/// └─────────────────────────────────────────────────────────────────┘ +/// │ Handled::No +/// ▼ +/// ┌─────────────────────────────────────────────────────────────────┐ +/// │ 2. Dynamic Handlers (added at runtime) │ +/// │ - Used for session-specific message handling │ +/// │ - Added via ConnectionTo::add_dynamic_handler │ +/// └─────────────────────────────────────────────────────────────────┘ +/// │ Handled::No +/// ▼ +/// ┌─────────────────────────────────────────────────────────────────┐ +/// │ 3. Role Default Handler │ +/// │ - Fallback based on the connection's Role │ +/// │ - Handles protocol-level messages (e.g., proxy forwarding) │ +/// └─────────────────────────────────────────────────────────────────┘ +/// │ Handled::No +/// ▼ +/// ┌─────────────────────────────────────────────────────────────────┐ +/// │ Unhandled: Error response sent (or queued if retry=true) │ +/// └─────────────────────────────────────────────────────────────────┘ +/// ``` +/// +/// # The `Handled` Return Value +/// +/// Each handler returns [`Handled`] to indicate whether it processed the message: +/// +/// - **`Handled::Yes`** - Message was handled. No further handlers are invoked. +/// - **`Handled::No { message, retry }`** - Message was not handled. The message +/// (possibly modified) is passed to the next handler in the chain. +/// +/// For convenience, handlers can return `()` which is equivalent to `Handled::Yes`. +/// +/// # The Retry Mechanism +/// +/// The `retry` flag in `Handled::No` controls what happens when no handler claims a message: +/// +/// - **`retry: false`** (default) - Send a "method not found" error response immediately. +/// - **`retry: true`** - Queue the message and retry it when new dynamic handlers are added. +/// +/// This mechanism exists because of a timing issue with sessions: when a `session/new` +/// response is being processed, the dynamic handler for that session hasn't been registered +/// yet, but `session/update` notifications for that session may already be arriving. +/// By setting `retry: true`, these early notifications are queued until the session's +/// dynamic handler is added. +/// +/// # Handler Registration +/// +/// Most users register handlers using the builder methods on [`Builder`]: +/// +/// ``` +/// # use agent_client_protocol_core::{Agent, Client, ConnectTo}; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, AgentCapabilities}; +/// # use agent_client_protocol_test::StatusUpdate; +/// # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +/// Agent.builder() +/// .on_receive_request(async |req: InitializeRequest, responder, cx| { +/// responder.respond( +/// InitializeResponse::new(req.protocol_version) +/// .agent_capabilities(AgentCapabilities::new()), +/// ) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .on_receive_notification(async |notif: StatusUpdate, cx| { +/// // Process notification +/// Ok(()) +/// }, agent_client_protocol_core::on_receive_notification!()) +/// .connect_to(transport) +/// .await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// The type parameter on the closure determines which messages are dispatched to it. +/// Messages that don't match the type are automatically passed to the next handler. +/// +/// # Implementing Custom Handlers +/// +/// For advanced use cases, you can implement `HandleMessageAs` directly: +/// +/// ```ignore +/// struct MyHandler; +/// +/// impl HandleMessageAs for MyHandler { +/// +/// async fn handle_dispatch( +/// &mut self, +/// message: Dispatch, +/// cx: ConnectionTo, +/// ) -> Result, Error> { +/// if message.method() == "my/custom/method" { +/// // Handle it +/// Ok(Handled::Yes) +/// } else { +/// // Pass to next handler +/// Ok(Handled::No { message, retry: false }) +/// } +/// } +/// +/// fn describe_chain(&self) -> impl std::fmt::Debug { +/// "MyHandler" +/// } +/// } +/// ``` +/// +/// # Important: Handlers Must Not Block +/// +/// The connection processes messages on a single async task. While a handler is running, +/// no other messages can be processed. For expensive operations, use [`ConnectionTo::spawn`] +/// to run work concurrently: +/// +/// ``` +/// # use agent_client_protocol_core::{Client, Agent, ConnectTo}; +/// # use agent_client_protocol_test::{expensive_operation, ProcessComplete}; +/// # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +/// # Client.builder().connect_with(transport, async |cx| { +/// cx.spawn({ +/// let connection = cx.clone(); +/// async move { +/// let result = expensive_operation("data").await?; +/// connection.send_notification(ProcessComplete { result })?; +/// Ok(()) +/// } +/// })?; +/// # Ok(()) +/// # }).await?; +/// # Ok(()) +/// # } +/// ``` +#[allow(async_fn_in_trait)] +/// A handler for incoming JSON-RPC messages. +/// +/// This trait is implemented by types that can process incoming messages on a connection. +/// Handlers are registered with a [`Builder`] and are called in order until +/// one claims the message. +/// +/// The type parameter `R` is the role this handler plays - who I am. +/// For an agent handler, `R = Agent` (I handle messages as an agent). +/// For a client handler, `R = Client` (I handle messages as a client). +pub trait HandleDispatchFrom: Send { + /// Attempt to claim an incoming message (request or notification). + /// + /// # Important: do not block + /// + /// The server will not process new messages until this handler returns. + /// You should avoid blocking in this callback unless you wish to block the server (e.g., for rate limiting). + /// The recommended approach to manage expensive operations is to the [`ConnectionTo::spawn`] method available on the message context. + /// + /// # Parameters + /// + /// * `message` - The incoming message to handle. + /// * `connection` - The connection, used to send messages and access connection state. + /// + /// # Returns + /// + /// * `Ok(Handled::Yes)` if the message was claimed. It will not be propagated further. + /// * `Ok(Handled::No(message))` if not; the (possibly changed) message will be passed to the remaining handlers. + /// * `Err` if an internal error occurs (this will bring down the server). + fn handle_dispatch_from( + &mut self, + message: Dispatch, + connection: ConnectionTo, + ) -> impl Future, crate::Error>> + Send; + + /// Returns a debug description of the registered handlers for diagnostics. + fn describe_chain(&self) -> impl std::fmt::Debug; +} + +impl HandleDispatchFrom for &mut H +where + H: HandleDispatchFrom, +{ + fn handle_dispatch_from( + &mut self, + message: Dispatch, + cx: ConnectionTo, + ) -> impl Future, crate::Error>> + Send { + H::handle_dispatch_from(self, message, cx) + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + H::describe_chain(self) + } +} + +/// A JSON-RPC connection that can act as either a server, client, or both. +/// +/// [`Builder`] provides a builder-style API for creating JSON-RPC servers and clients. +/// You start by calling `Role.builder()` (e.g., `Client.builder()`), then add message +/// handlers, and finally drive the connection with either [`connect_to`](Builder::connect_to) +/// or [`connect_with`](Builder::connect_with), providing a component implementation +/// (e.g., [`ByteStreams`] for byte streams). +/// +/// # JSON-RPC Primer +/// +/// JSON-RPC 2.0 has two fundamental message types: +/// +/// * **Requests** - Messages that expect a response. They have an `id` field that gets +/// echoed back in the response so the sender can correlate them. +/// * **Notifications** - Fire-and-forget messages with no `id` field. The sender doesn't +/// expect or receive a response. +/// +/// # Type-Driven Message Dispatch +/// +/// The handler registration methods use Rust's type system to determine which messages +/// to handle. The type parameter you provide controls what gets dispatched to your handler: +/// +/// ## Single Message Types +/// +/// The simplest case - handle one specific message type: +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, SessionNotification}; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// connection +/// .on_receive_request(async |req: InitializeRequest, responder, cx| { +/// // Handle only InitializeRequest messages +/// responder.respond(InitializeResponse::make()) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .on_receive_notification(async |notif: SessionNotification, cx| { +/// // Handle only SessionUpdate notifications +/// Ok(()) +/// }, agent_client_protocol_core::on_receive_notification!()) +/// # .connect_to(agent_client_protocol_test::MockTransport).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// ## Enum Message Types +/// +/// You can also handle multiple related messages with a single handler by defining an enum +/// that implements the appropriate trait ([`JsonRpcRequest`] or [`JsonRpcNotification`]): +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # use agent_client_protocol_core::{JsonRpcRequest, JsonRpcMessage, UntypedMessage}; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, PromptRequest, PromptResponse}; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// // Define an enum for multiple request types +/// #[derive(Debug, Clone)] +/// enum MyRequests { +/// Initialize(InitializeRequest), +/// Prompt(PromptRequest), +/// } +/// +/// // Implement JsonRpcRequest for your enum +/// # impl JsonRpcMessage for MyRequests { +/// # fn matches_method(_method: &str) -> bool { false } +/// # fn method(&self) -> &str { "myRequests" } +/// # fn to_untyped_message(&self) -> Result { todo!() } +/// # fn parse_message(_method: &str, _params: &impl serde::Serialize) -> Result { Err(agent_client_protocol_core::Error::method_not_found()) } +/// # } +/// impl JsonRpcRequest for MyRequests { type Response = serde_json::Value; } +/// +/// // Handle all variants in one place +/// connection.on_receive_request(async |req: MyRequests, responder, cx| { +/// match req { +/// MyRequests::Initialize(init) => { responder.respond(serde_json::json!({})) } +/// MyRequests::Prompt(prompt) => { responder.respond(serde_json::json!({})) } +/// } +/// }, agent_client_protocol_core::on_receive_request!()) +/// # .connect_to(agent_client_protocol_test::MockTransport).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// ## Mixed Message Types +/// +/// For enums containing both requests AND notifications, use [`on_receive_dispatch`](Self::on_receive_dispatch): +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # use agent_client_protocol_core::Dispatch; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, SessionNotification}; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// // on_receive_dispatch receives Dispatch which can be either a request or notification +/// connection.on_receive_dispatch(async |msg: Dispatch, _cx| { +/// match msg { +/// Dispatch::Request(req, responder) => { +/// responder.respond(InitializeResponse::make()) +/// } +/// Dispatch::Notification(notif) => { +/// Ok(()) +/// } +/// Dispatch::Response(result, router) => { +/// // Forward response to its destination +/// router.respond_with_result(result) +/// } +/// } +/// }, agent_client_protocol_core::on_receive_dispatch!()) +/// # .connect_to(agent_client_protocol_test::MockTransport).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// # Handler Registration +/// +/// Register handlers using these methods (listed from most common to most flexible): +/// +/// * [`on_receive_request`](Self::on_receive_request) - Handle JSON-RPC requests (messages expecting responses) +/// * [`on_receive_notification`](Self::on_receive_notification) - Handle JSON-RPC notifications (fire-and-forget) +/// * [`on_receive_dispatch`](Self::on_receive_dispatch) - Handle enums containing both requests and notifications +/// * [`with_handler`](Self::with_handler) - Low-level primitive for maximum flexibility +/// +/// ## Handler Ordering +/// +/// Handlers are tried in the order you register them. The first handler that claims a message +/// (by matching its type) will process it. Subsequent handlers won't see that message: +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # use agent_client_protocol_core::{Dispatch, UntypedMessage}; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, PromptRequest, PromptResponse}; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// connection +/// .on_receive_request(async |req: InitializeRequest, responder, cx| { +/// // This runs first for InitializeRequest +/// responder.respond(InitializeResponse::make()) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .on_receive_request(async |req: PromptRequest, responder, cx| { +/// // This runs first for PromptRequest +/// responder.respond(PromptResponse::make()) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .on_receive_dispatch(async |msg: Dispatch, cx| { +/// // This runs for any message not handled above +/// msg.respond_with_error(agent_client_protocol_core::util::internal_error("unknown method"), cx) +/// }, agent_client_protocol_core::on_receive_dispatch!()) +/// # .connect_to(agent_client_protocol_test::MockTransport).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// # Event Loop and Concurrency +/// +/// Understanding the event loop is critical for writing correct handlers. +/// +/// ## The Event Loop +/// +/// [`Builder`] runs all handler callbacks on a single async task - the event loop. +/// While a handler is running, **the server cannot receive new messages**. This means +/// any blocking or expensive work in your handlers will stall the entire connection. +/// +/// To avoid blocking the event loop, use [`ConnectionTo::spawn`] to offload serious +/// work to concurrent tasks: +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// connection.on_receive_request(async |req: AnalyzeRequest, responder, cx| { +/// // Clone cx for the spawned task +/// cx.spawn({ +/// let connection = cx.clone(); +/// async move { +/// let result = expensive_analysis(&req.data).await?; +/// connection.send_notification(AnalysisComplete { result })?; +/// Ok(()) +/// } +/// })?; +/// +/// // Respond immediately without blocking +/// responder.respond(AnalysisStarted { job_id: 42 }) +/// }, agent_client_protocol_core::on_receive_request!()) +/// # .connect_to(agent_client_protocol_test::MockTransport).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// Note that the entire connection runs within one async task, so parallelism must be +/// managed explicitly using [`spawn`](ConnectionTo::spawn). +/// +/// ## The Connection Context +/// +/// Handler callbacks receive a context object (`cx`) for interacting with the connection: +/// +/// * **For request handlers** - [`Responder`] provides [`respond`](Responder::respond) +/// to send the response, plus methods to send other messages +/// * **For notification handlers** - [`ConnectionTo`] provides methods to send messages +/// and spawn tasks +/// +/// Both context types support: +/// * [`send_request`](ConnectionTo::send_request) - Send requests to the other side +/// * [`send_notification`](ConnectionTo::send_notification) - Send notifications +/// * [`spawn`](ConnectionTo::spawn) - Run tasks concurrently without blocking the event loop +/// +/// The [`SentRequest`] returned by `send_request` provides methods like +/// [`on_receiving_result`](SentRequest::on_receiving_result) that help you +/// avoid accidentally blocking the event loop while waiting for responses. +/// +/// # Driving the Connection +/// +/// After adding handlers, you must drive the connection using one of two modes: +/// +/// ## Server Mode: `connect_to()` +/// +/// Use [`connect_to`](Self::connect_to) when you only need to respond to incoming messages: +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// connection +/// .on_receive_request(async |req: MyRequest, responder, cx| { +/// responder.respond(MyResponse { status: "ok".into() }) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .connect_to(MockTransport) // Runs until connection closes or error occurs +/// .await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// The connection will process incoming messages and invoke your handlers until the +/// connection is closed or an error occurs. +/// +/// ## Client Mode: `connect_with()` +/// +/// Use [`connect_with`](Self::connect_with) when you need to both handle incoming messages +/// AND send your own requests/notifications: +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # use agent_client_protocol_core::schema::InitializeRequest; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// connection +/// .on_receive_request(async |req: MyRequest, responder, cx| { +/// responder.respond(MyResponse { status: "ok".into() }) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .connect_with(MockTransport, async |cx| { +/// // You can send requests to the other side +/// let response = cx.send_request(InitializeRequest::make()) +/// .block_task() +/// .await?; +/// +/// // And send notifications +/// cx.send_notification(StatusUpdate { message: "ready".into() })?; +/// +/// Ok(()) +/// }) +/// .await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// The connection will serve incoming messages in the background while your client closure +/// runs. When the closure returns, the connection shuts down. +/// +/// # Example: Complete Agent +/// +/// ```no_run +/// # use agent_client_protocol_core::UntypedRole; +/// # use agent_client_protocol_core::{Builder}; +/// # use agent_client_protocol_core::ByteStreams; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, PromptRequest, PromptResponse, SessionNotification}; +/// # use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// let transport = ByteStreams::new( +/// tokio::io::stdout().compat_write(), +/// tokio::io::stdin().compat(), +/// ); +/// +/// UntypedRole.builder() +/// .name("my-agent") // Optional: for debugging logs +/// .on_receive_request(async |init: InitializeRequest, responder, cx| { +/// let response: InitializeResponse = todo!(); +/// responder.respond(response) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .on_receive_request(async |prompt: PromptRequest, responder, cx| { +/// // You can send notifications while processing a request +/// let notif: SessionNotification = todo!(); +/// cx.send_notification(notif)?; +/// +/// // Then respond to the request +/// let response: PromptResponse = todo!(); +/// responder.respond(response) +/// }, agent_client_protocol_core::on_receive_request!()) +/// .connect_to(transport) +/// .await?; +/// # Ok(()) +/// # } +/// ``` +#[must_use] +#[derive(Debug)] +pub struct Builder +where + Handler: HandleDispatchFrom, + Runner: RunWithConnectionTo, +{ + /// My role. + host: Host, + + /// Name of the connection, used in tracing logs. + name: Option, + + /// Handler for incoming messages. + handler: Handler, + + /// Responder for background tasks. + responder: Runner, +} + +impl Builder { + /// Create a new connection builder for the given role. + /// This type follows a builder pattern; use other methods to configure and then invoke + /// [`Self::connect_to`] (to use as a server) or [`Self::connect_with`] to use as a client. + pub fn new(role: Host) -> Self { + Self { + host: role, + name: None, + handler: NullHandler, + responder: NullRun, + } + } +} + +impl Builder +where + Handler: HandleDispatchFrom, +{ + /// Create a new connection builder with the given handler. + pub fn new_with(role: Host, handler: Handler) -> Self { + Self { + host: role, + name: None, + handler, + responder: NullRun, + } + } +} + +impl< + Host: Role, + Handler: HandleDispatchFrom, + Runner: RunWithConnectionTo, +> Builder +{ + /// Set the "name" of this connection -- used only for debugging logs. + pub fn name(mut self, name: impl ToString) -> Self { + self.name = Some(name.to_string()); + self + } + + /// Merge another [`Builder`] into this one. + /// + /// Prefer [`Self::on_receive_request`] or [`Self::on_receive_notification`]. + /// This is a low-level method that is not intended for general use. + pub fn with_connection_builder( + self, + other: Builder< + Host, + impl HandleDispatchFrom, + impl RunWithConnectionTo, + >, + ) -> Builder< + Host, + impl HandleDispatchFrom, + impl RunWithConnectionTo, + > { + Builder { + host: self.host, + name: self.name, + handler: ChainedHandler::new( + self.handler, + NamedHandler::new(other.name, other.handler), + ), + responder: ChainRun::new(self.responder, other.responder), + } + } + + /// Add a new [`HandleDispatchFrom`] to the chain. + /// + /// Prefer [`Self::on_receive_request`] or [`Self::on_receive_notification`]. + /// This is a low-level method that is not intended for general use. + pub fn with_handler( + self, + handler: impl HandleDispatchFrom, + ) -> Builder, Runner> { + Builder { + host: self.host, + name: self.name, + handler: ChainedHandler::new(self.handler, handler), + responder: self.responder, + } + } + + /// Add a new [`RunWithConnectionTo`] to the chain. + pub fn with_responder( + self, + responder: Run1, + ) -> Builder> + where + Run1: RunWithConnectionTo, + { + Builder { + host: self.host, + name: self.name, + handler: self.handler, + responder: ChainRun::new(self.responder, responder), + } + } + + /// Enqueue a task to run once the connection is actively serving traffic. + #[track_caller] + pub fn with_spawned( + self, + task: F, + ) -> Builder> + where + F: FnOnce(ConnectionTo) -> Fut + Send, + Fut: Future> + Send, + { + let location = Location::caller(); + self.with_responder(SpawnedRun::new(location, task)) + } + + /// Register a handler for messages that can be either requests OR notifications. + /// + /// Use this when you want to handle an enum type that contains both request and + /// notification variants. Your handler receives a [`Dispatch`] which + /// is an enum with two variants: + /// + /// - `Dispatch::Request(request, responder)` - A request with its response context + /// - `Dispatch::Notification(notification)` - A notification + /// - `Dispatch::Response(result, router)` - A response to a request we sent + /// + /// # Example + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # use agent_client_protocol_core::Dispatch; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// # let connection = mock_connection(); + /// connection.on_receive_dispatch(async |message: Dispatch, _cx| { + /// match message { + /// Dispatch::Request(req, responder) => { + /// // Handle request and send response + /// responder.respond(MyResponse { status: "ok".into() }) + /// } + /// Dispatch::Notification(notif) => { + /// // Handle notification (no response needed) + /// Ok(()) + /// } + /// Dispatch::Response(result, router) => { + /// // Forward response to its destination + /// router.respond_with_result(result) + /// } + /// } + /// }, agent_client_protocol_core::on_receive_dispatch!()) + /// # .connect_to(agent_client_protocol_test::MockTransport).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// For most use cases, prefer [`on_receive_request`](Self::on_receive_request) or + /// [`on_receive_notification`](Self::on_receive_notification) which provide cleaner APIs + /// for handling requests or notifications separately. + /// + /// # Ordering + /// + /// This callback runs inside the dispatch loop and blocks further message processing + /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on + /// ordering guarantees and how to avoid deadlocks. + pub fn on_receive_dispatch( + self, + op: F, + to_future_hack: ToFut, + ) -> Builder, Runner> + where + Host::Counterpart: HasPeer, + Req: JsonRpcRequest, + Notif: JsonRpcNotification, + F: AsyncFnMut( + Dispatch, + ConnectionTo, + ) -> Result + + Send, + T: IntoHandled>, + ToFut: Fn( + &mut F, + Dispatch, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, + { + let handler = MessageHandler::new( + self.host.counterpart(), + self.host.counterpart(), + op, + to_future_hack, + ); + self.with_handler(handler) + } + + /// Register a handler for JSON-RPC requests of type `Req`. + /// + /// Your handler receives two arguments: + /// 1. The request (type `Req`) + /// 2. A [`Responder`] for sending the response + /// + /// The request context allows you to: + /// - Send the response with [`Responder::respond`] + /// - Send notifications to the client with [`ConnectionTo::send_notification`] + /// - Send requests to the client with [`ConnectionTo::send_request`] + /// + /// # Example + /// + /// ```ignore + /// # use agent_client_protocol_core::UntypedRole; + /// # use agent_client_protocol_core::{Builder}; + /// # use agent_client_protocol_core::schema::{PromptRequest, PromptResponse, SessionNotification}; + /// # fn example(connection: Builder>) { + /// connection.on_receive_request(async |request: PromptRequest, responder, cx| { + /// // Send a notification while processing + /// let notif: SessionNotification = todo!(); + /// cx.send_notification(notif)?; + /// + /// // Do some work... + /// let result = todo!("process the prompt"); + /// + /// // Send the response + /// let response: PromptResponse = todo!(); + /// responder.respond(response) + /// }, agent_client_protocol_core::on_receive_request!()); + /// # } + /// ``` + /// + /// # Type Parameter + /// + /// `Req` can be either a single request type or an enum of multiple request types. + /// See the [type-driven dispatch](Self#type-driven-message-dispatch) section for details. + /// + /// # Ordering + /// + /// This callback runs inside the dispatch loop and blocks further message processing + /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on + /// ordering guarantees and how to avoid deadlocks. + pub fn on_receive_request( + self, + op: F, + to_future_hack: ToFut, + ) -> Builder, Runner> + where + Host::Counterpart: HasPeer, + F: AsyncFnMut( + Req, + Responder, + ConnectionTo, + ) -> Result + + Send, + T: IntoHandled<(Req, Responder)>, + ToFut: Fn( + &mut F, + Req, + Responder, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, + { + let handler = RequestHandler::new( + self.host.counterpart(), + self.host.counterpart(), + op, + to_future_hack, + ); + self.with_handler(handler) + } + + /// Register a handler for JSON-RPC notifications of type `Notif`. + /// + /// Notifications are fire-and-forget messages that don't expect a response. + /// Your handler receives: + /// 1. The notification (type `Notif`) + /// 2. A [`ConnectionTo`] for sending messages to the other side + /// + /// Unlike request handlers, you cannot send a response (notifications don't have IDs), + /// but you can still send your own requests and notifications using the context. + /// + /// # Example + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// # let connection = mock_connection(); + /// connection.on_receive_notification(async |notif: SessionUpdate, cx| { + /// // Process the notification + /// update_session_state(¬if)?; + /// + /// // Optionally send a notification back + /// cx.send_notification(StatusUpdate { + /// message: "Acknowledged".into(), + /// })?; + /// + /// Ok(()) + /// }, agent_client_protocol_core::on_receive_notification!()) + /// # .connect_to(agent_client_protocol_test::MockTransport).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Type Parameter + /// + /// `Notif` can be either a single notification type or an enum of multiple notification types. + /// See the [type-driven dispatch](Self#type-driven-message-dispatch) section for details. + /// + /// # Ordering + /// + /// This callback runs inside the dispatch loop and blocks further message processing + /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on + /// ordering guarantees and how to avoid deadlocks. + pub fn on_receive_notification( + self, + op: F, + to_future_hack: ToFut, + ) -> Builder, Runner> + where + Host::Counterpart: HasPeer, + Notif: JsonRpcNotification, + F: AsyncFnMut(Notif, ConnectionTo) -> Result + Send, + T: IntoHandled<(Notif, ConnectionTo)>, + ToFut: Fn( + &mut F, + Notif, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, + { + let handler = NotificationHandler::new( + self.host.counterpart(), + self.host.counterpart(), + op, + to_future_hack, + ); + self.with_handler(handler) + } + + /// Register a handler for messages from a specific peer. + /// + /// This is similar to [`on_receive_dispatch`](Self::on_receive_dispatch), but allows + /// specifying the source peer explicitly. This is useful when receiving messages + /// from a peer that requires message transformation (e.g., unwrapping `SuccessorMessage` + /// envelopes when receiving from an agent via a proxy). + /// + /// For the common case of receiving from the default counterpart, use + /// [`on_receive_dispatch`](Self::on_receive_dispatch) instead. + /// + /// # Ordering + /// + /// This callback runs inside the dispatch loop and blocks further message processing + /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on + /// ordering guarantees and how to avoid deadlocks. + pub fn on_receive_dispatch_from< + Req: JsonRpcRequest, + Notif: JsonRpcNotification, + Peer: Role, + F, + T, + ToFut, + >( + self, + peer: Peer, + op: F, + to_future_hack: ToFut, + ) -> Builder, Runner> + where + Host::Counterpart: HasPeer, + F: AsyncFnMut( + Dispatch, + ConnectionTo, + ) -> Result + + Send, + T: IntoHandled>, + ToFut: Fn( + &mut F, + Dispatch, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, + { + let handler = MessageHandler::new(self.host.counterpart(), peer, op, to_future_hack); + self.with_handler(handler) + } + + /// Register a handler for JSON-RPC requests from a specific peer. + /// + /// This is similar to [`on_receive_request`](Self::on_receive_request), but allows + /// specifying the source peer explicitly. This is useful when receiving messages + /// from a peer that requires message transformation (e.g., unwrapping `SuccessorRequest` + /// envelopes when receiving from an agent via a proxy). + /// + /// For the common case of receiving from the default counterpart, use + /// [`on_receive_request`](Self::on_receive_request) instead. + /// + /// # Example + /// + /// ```ignore + /// use agent_client_protocol_core::Agent; + /// use agent_client_protocol_core::schema::InitializeRequest; + /// + /// // Conductor receiving from agent direction - messages will be unwrapped from SuccessorMessage + /// connection.on_receive_request_from(Agent, async |req: InitializeRequest, responder, cx| { + /// // Handle the request + /// responder.respond(InitializeResponse::make()) + /// }) + /// ``` + /// + /// # Ordering + /// + /// This callback runs inside the dispatch loop and blocks further message processing + /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on + /// ordering guarantees and how to avoid deadlocks. + pub fn on_receive_request_from( + self, + peer: Peer, + op: F, + to_future_hack: ToFut, + ) -> Builder, Runner> + where + Host::Counterpart: HasPeer, + F: AsyncFnMut( + Req, + Responder, + ConnectionTo, + ) -> Result + + Send, + T: IntoHandled<(Req, Responder)>, + ToFut: Fn( + &mut F, + Req, + Responder, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, + { + let handler = RequestHandler::new(self.host.counterpart(), peer, op, to_future_hack); + self.with_handler(handler) + } + + /// Register a handler for JSON-RPC notifications from a specific peer. + /// + /// This is similar to [`on_receive_notification`](Self::on_receive_notification), but allows + /// specifying the source peer explicitly. This is useful when receiving messages + /// from a peer that requires message transformation (e.g., unwrapping `SuccessorNotification` + /// envelopes when receiving from an agent via a proxy). + /// + /// For the common case of receiving from the default counterpart, use + /// [`on_receive_notification`](Self::on_receive_notification) instead. + /// + /// # Ordering + /// + /// This callback runs inside the dispatch loop and blocks further message processing + /// until it completes. See the [`ordering`](crate::concepts::ordering) module for details on + /// ordering guarantees and how to avoid deadlocks. + pub fn on_receive_notification_from( + self, + peer: Peer, + op: F, + to_future_hack: ToFut, + ) -> Builder, Runner> + where + Host::Counterpart: HasPeer, + F: AsyncFnMut(Notif, ConnectionTo) -> Result + Send, + T: IntoHandled<(Notif, ConnectionTo)>, + ToFut: Fn( + &mut F, + Notif, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, + { + let handler = NotificationHandler::new(self.host.counterpart(), peer, op, to_future_hack); + self.with_handler(handler) + } + + /// Add an MCP server that will be added to all new sessions that are proxied through this connection. + /// + /// Only applicable to proxies. + pub fn with_mcp_server( + self, + mcp_server: McpServer>, + ) -> Builder< + Host, + impl HandleDispatchFrom, + impl RunWithConnectionTo, + > + where + Host::Counterpart: HasPeer + HasPeer, + { + let (handler, responder) = mcp_server.into_handler_and_responder(); + self.with_handler(handler).with_responder(responder) + } + + /// Run in server mode with the provided transport. + /// + /// This drives the connection by continuously processing messages from the transport + /// and dispatching them to your registered handlers. The connection will run until: + /// - The transport closes (e.g., EOF on byte streams) + /// - An error occurs + /// - One of your handlers returns an error + /// + /// The transport is responsible for serializing and deserializing `jsonrpcmsg::Message` + /// values to/from the underlying I/O mechanism (byte streams, channels, etc.). + /// + /// Use this mode when you only need to respond to incoming messages and don't need + /// to initiate your own requests. If you need to send requests to the other side, + /// use [`connect_with`](Self::connect_with) instead. + /// + /// # Example: Byte Stream Transport + /// + /// ```no_run + /// # use agent_client_protocol_core::UntypedRole; + /// # use agent_client_protocol_core::{Builder}; + /// # use agent_client_protocol_core::ByteStreams; + /// # use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// let transport = ByteStreams::new( + /// tokio::io::stdout().compat_write(), + /// tokio::io::stdin().compat(), + /// ); + /// + /// UntypedRole.builder() + /// .on_receive_request(async |req: MyRequest, responder, cx| { + /// responder.respond(MyResponse { status: "ok".into() }) + /// }, agent_client_protocol_core::on_receive_request!()) + /// .connect_to(transport) + /// .await?; + /// # Ok(()) + /// # } + /// ``` + pub async fn connect_to( + self, + transport: impl ConnectTo + 'static, + ) -> Result<(), crate::Error> { + self.connect_with(transport, async move |_cx| future::pending().await) + .await + } + + /// Run the connection until the provided closure completes. + /// + /// This drives the connection by: + /// 1. Running your registered handlers in the background to process incoming messages + /// 2. Executing your `main_fn` closure with a [`ConnectionTo`] for sending requests/notifications + /// + /// The connection stays active until your `main_fn` returns, then shuts down gracefully. + /// If the connection closes unexpectedly before `main_fn` completes, this returns an error. + /// + /// Use this mode when you need to initiate communication (send requests/notifications) + /// in addition to responding to incoming messages. For server-only mode where you just + /// respond to messages, use [`connect_to`](Self::connect_to) instead. + /// + /// # Example + /// + /// ```no_run + /// # use agent_client_protocol_core::UntypedRole; + /// # use agent_client_protocol_core::{Builder}; + /// # use agent_client_protocol_core::ByteStreams; + /// # use agent_client_protocol_core::schema::InitializeRequest; + /// # use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// let transport = ByteStreams::new( + /// tokio::io::stdout().compat_write(), + /// tokio::io::stdin().compat(), + /// ); + /// + /// UntypedRole.builder() + /// .on_receive_request(async |req: MyRequest, responder, cx| { + /// // Handle incoming requests in the background + /// responder.respond(MyResponse { status: "ok".into() }) + /// }, agent_client_protocol_core::on_receive_request!()) + /// .connect_with(transport, async |cx| { + /// // Initialize the protocol + /// let init_response = cx.send_request(InitializeRequest::make()) + /// .block_task() + /// .await?; + /// + /// // Send more requests... + /// let result = cx.send_request(MyRequest {}) + /// .block_task() + /// .await?; + /// + /// // When this closure returns, the connection shuts down + /// Ok(()) + /// }) + /// .await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Parameters + /// + /// - `main_fn`: Your client logic. Receives a [`ConnectionTo`] for sending messages. + /// + /// # Errors + /// + /// Returns an error if the connection closes before `main_fn` completes. + pub async fn connect_with( + self, + transport: impl ConnectTo + 'static, + main_fn: impl AsyncFnOnce(ConnectionTo) -> Result, + ) -> Result { + let (_, future) = self.into_connection_and_future(transport, main_fn); + future.await + } + + /// Helper that returns a [`ConnectionTo`] and a future that runs this connection until `main_fn` returns. + fn into_connection_and_future( + self, + transport: impl ConnectTo + 'static, + main_fn: impl AsyncFnOnce(ConnectionTo) -> Result, + ) -> ( + ConnectionTo, + impl Future>, + ) { + let Self { + name, + handler, + responder, + host: me, + } = self; + + let (outgoing_tx, outgoing_rx) = mpsc::unbounded(); + let (new_task_tx, new_task_rx) = mpsc::unbounded(); + let (dynamic_handler_tx, dynamic_handler_rx) = mpsc::unbounded(); + let connection = ConnectionTo::new( + me.counterpart(), + outgoing_tx, + new_task_tx, + dynamic_handler_tx, + ); + + // Convert transport into server - this returns a channel for us to use + // and a future that runs the transport + let transport_component = crate::DynConnectTo::new(transport); + let (transport_channel, transport_future) = transport_component.into_channel_and_future(); + let spawn_result = connection.spawn(transport_future); + + // Destructure the channel endpoints + let Channel { + rx: transport_incoming_rx, + tx: transport_outgoing_tx, + } = transport_channel; + + let (reply_tx, reply_rx) = mpsc::unbounded(); + + let future = crate::util::instrument_with_connection_name(name, { + let connection = connection.clone(); + async move { + let () = spawn_result?; + + let background = async { + futures::try_join!( + // Protocol layer: OutgoingMessage → jsonrpcmsg::Message + outgoing_actor::outgoing_protocol_actor( + outgoing_rx, + reply_tx.clone(), + transport_outgoing_tx, + ), + // Protocol layer: jsonrpcmsg::Message → handler/reply routing + incoming_actor::incoming_protocol_actor( + me.counterpart(), + &connection, + transport_incoming_rx, + dynamic_handler_rx, + reply_rx, + handler, + ), + task_actor::task_actor(new_task_rx, &connection), + responder.run_with_connection_to(connection.clone()), + )?; + Ok(()) + }; + + crate::util::run_until(background, main_fn(connection.clone())).await + } + }); + + (connection, future) + } +} + +impl ConnectTo for Builder +where + R: Role, + H: HandleDispatchFrom + 'static, + Run: RunWithConnectionTo + 'static, +{ + async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { + Builder::connect_to(self, client).await + } +} + +/// The payload sent through the response oneshot channel. +/// +/// Includes the response value and an optional ack channel for dispatch loop +/// synchronization. +pub(crate) struct ResponsePayload { + /// The response result - either the JSON value or an error. + pub(crate) result: Result, + + /// Optional acknowledgment channel for dispatch loop synchronization. + /// + /// When present, the receiver must send on this channel to signal that + /// response processing is complete, allowing the dispatch loop to continue + /// to the next message. + /// + /// This is `None` for error paths where the response is sent directly + /// (e.g., when the outgoing channel is broken) rather than through the + /// normal dispatch loop flow. + pub(crate) ack_tx: Option>, +} + +impl std::fmt::Debug for ResponsePayload { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResponsePayload") + .field("result", &self.result) + .field("ack_tx", &self.ack_tx.as_ref().map(|_| "...")) + .finish() + } +} + +/// Message sent to the incoming actor for reply subscription management. +enum ReplyMessage { + /// Subscribe to receive a response for the given request id. + /// When a response with this id arrives, it will be sent through the oneshot + /// along with an ack channel that must be signaled when processing is complete. + /// The method name is stored to allow routing responses through typed handlers. + Subscribe { + id: jsonrpcmsg::Id, + + /// id of the peer this request was sent to + role_id: RoleId, + + /// (original) method of the request -- the actual request may have been transformed + /// to a successor method, but this will reflect the method of the wrapped request + method: String, + + sender: oneshot::Sender, + }, +} + +impl std::fmt::Debug for ReplyMessage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ReplyMessage::Subscribe { id, method, .. } => f + .debug_struct("Subscribe") + .field("id", id) + .field("method", method) + .finish(), + } + } +} + +/// Messages send to be serialized over the transport. +#[derive(Debug)] +enum OutgoingMessage { + /// Send a request to the server. + Request { + /// id assigned to this request (generated by sender) + id: jsonrpcmsg::Id, + + /// the original method + method: String, + + /// the peer we sent this to + role_id: RoleId, + + /// the message to send; this may have a distinct method + /// depending on the peer + untyped: UntypedMessage, + + /// where to send the response when it arrives (includes ack channel) + response_tx: oneshot::Sender, + }, + + /// Send a notification to the server. + Notification { + /// the message to send; this may have a distinct method + /// depending on the peer + untyped: UntypedMessage, + }, + + /// Send a response to a message from the server + Response { + id: jsonrpcmsg::Id, + + response: Result, + }, + + /// Send a generalized error message + Error { error: crate::Error }, +} + +/// Return type from JrHandler; indicates whether the request was handled or not. +#[must_use] +#[derive(Debug)] +pub enum Handled { + /// The message was handled + Yes, + + /// The message was not handled; returns the original value. + /// + /// If `retry` is true, + No { + /// The message to be passed to subsequent handlers + /// (typically the original message, but it may have been + /// mutated.) + message: T, + + /// If true, request the message to be queued and retried with + /// dynamic handlers as they are added. + /// + /// This is used for managing session updates since the dynamic + /// handler for a session cannot be added until the response to the + /// new session request has been processed and there may be updates + /// that get processed at the same time. + retry: bool, + }, +} + +/// Trait for converting handler return values into [`Handled`]. +/// +/// This trait allows handlers to return either `()` (which becomes `Handled::Yes`) +/// or an explicit `Handled` value for more control over handler propagation. +pub trait IntoHandled { + /// Convert this value into a `Handled`. + fn into_handled(self) -> Handled; +} + +impl IntoHandled for () { + fn into_handled(self) -> Handled { + Handled::Yes + } +} + +impl IntoHandled for Handled { + fn into_handled(self) -> Handled { + self + } +} + +/// Connection context for sending messages and spawning tasks. +/// +/// This is the primary handle for interacting with the JSON-RPC connection from +/// within handler callbacks. You can use it to: +/// +/// * Send requests and notifications to the other side +/// * Spawn concurrent tasks that run alongside the connection +/// * Respond to requests (via [`Responder`] which wraps this) +/// +/// # Cloning +/// +/// `ConnectionTo` is cheaply cloneable - all clones refer to the same underlying connection. +/// This makes it easy to share across async tasks. +/// +/// # Event Loop and Concurrency +/// +/// Handler callbacks run on the event loop, which means the connection cannot process new +/// messages while your handler is running. Use [`spawn`](Self::spawn) to offload any +/// expensive or blocking work to concurrent tasks. +/// +/// See the [Event Loop and Concurrency](Builder#event-loop-and-concurrency) section +/// for more details. +#[derive(Clone, Debug)] +pub struct ConnectionTo { + counterpart: Counterpart, + message_tx: OutgoingMessageTx, + task_tx: TaskTx, + dynamic_handler_tx: mpsc::UnboundedSender>, +} + +impl ConnectionTo { + fn new( + counterpart: Counterpart, + message_tx: mpsc::UnboundedSender, + task_tx: mpsc::UnboundedSender, + dynamic_handler_tx: mpsc::UnboundedSender>, + ) -> Self { + Self { + counterpart, + message_tx, + task_tx, + dynamic_handler_tx, + } + } + + /// Return the counterpart role this connection is talking to. + pub fn counterpart(&self) -> Counterpart { + self.counterpart.clone() + } + + /// Spawns a task that will run so long as the JSON-RPC connection is being served. + /// + /// This is the primary mechanism for offloading expensive work from handler callbacks + /// to avoid blocking the event loop. Spawned tasks run concurrently with the connection, + /// allowing the server to continue processing messages. + /// + /// # Event Loop + /// + /// Handler callbacks run on the event loop, which cannot process new messages while + /// your handler is running. Use `spawn` for any expensive operations: + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// # let connection = mock_connection(); + /// connection.on_receive_request(async |req: ProcessRequest, responder, cx| { + /// // Clone cx for the spawned task + /// cx.spawn({ + /// let connection = cx.clone(); + /// async move { + /// let result = expensive_operation(&req.data).await?; + /// connection.send_notification(ProcessComplete { result })?; + /// Ok(()) + /// } + /// })?; + /// + /// // Respond immediately + /// responder.respond(ProcessResponse { result: "started".into() }) + /// }, agent_client_protocol_core::on_receive_request!()) + /// # .connect_to(agent_client_protocol_test::MockTransport).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// If the spawned task returns an error, the entire server will shut down. + #[track_caller] + pub fn spawn( + &self, + task: impl IntoFuture, IntoFuture: Send + 'static>, + ) -> Result<(), crate::Error> { + let location = std::panic::Location::caller(); + let task = task.into_future(); + Task::new(location, task).spawn(&self.task_tx) + } + + /// Spawn a JSON-RPC connection in the background and return a [`ConnectionTo`] for sending messages to it. + /// + /// This is useful for creating multiple connections that communicate with each other, + /// such as implementing proxy patterns or connecting to multiple backend services. + /// + /// # Arguments + /// + /// - `builder`: The connection builder with handlers configured + /// - `transport`: The transport component to connect to + /// + /// # Returns + /// + /// A `ConnectionTo` that you can use to send requests and notifications to the spawned connection. + /// + /// # Example: Proxying to a backend connection + /// + /// ``` + /// # use agent_client_protocol_core::UntypedRole; + /// # use agent_client_protocol_core::{Builder, ConnectionTo}; + /// # use agent_client_protocol_test::*; + /// # async fn example(cx: ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { + /// // Set up a backend connection builder + /// let backend = UntypedRole.builder() + /// .on_receive_request(async |req: MyRequest, responder, _cx| { + /// responder.respond(MyResponse { status: "ok".into() }) + /// }, agent_client_protocol_core::on_receive_request!()); + /// + /// // Spawn it and get a context to send requests to it + /// let backend_connection = cx.spawn_connection(backend, MockTransport)?; + /// + /// // Now you can forward requests to the backend + /// let response = backend_connection.send_request(MyRequest {}).block_task().await?; + /// # Ok(()) + /// # } + /// ``` + #[track_caller] + pub fn spawn_connection( + &self, + builder: Builder< + R, + impl HandleDispatchFrom + 'static, + impl RunWithConnectionTo + 'static, + >, + transport: impl ConnectTo + 'static, + ) -> Result, crate::Error> { + let (connection, future) = + builder.into_connection_and_future(transport, |_| std::future::pending()); + Task::new(std::panic::Location::caller(), future).spawn(&self.task_tx)?; + Ok(connection) + } + + /// Send a request/notification and forward the response appropriately. + /// + /// The request context's response type matches the request's response type, + /// enabling type-safe message forwarding. + pub fn send_proxied_message, Notif: JsonRpcNotification>( + &self, + message: Dispatch, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.send_proxied_message_to(self.counterpart(), message) + } + + /// Send a request/notification and forward the response appropriately. + /// + /// The request context's response type matches the request's response type, + /// enabling type-safe message forwarding. + pub fn send_proxied_message_to< + Peer: Role, + Req: JsonRpcRequest, + Notif: JsonRpcNotification, + >( + &self, + peer: Peer, + message: Dispatch, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + match message { + Dispatch::Request(request, responder) => self + .send_request_to(peer, request) + .forward_response_to(responder), + Dispatch::Notification(notification) => self.send_notification_to(peer, notification), + Dispatch::Response(result, router) => { + // Responses are forwarded directly to their destination + router.respond_with_result(result) + } + } + } + + /// Send an outgoing request and return a [`SentRequest`] for handling the reply. + /// + /// The returned [`SentRequest`] provides methods for receiving the response without + /// blocking the event loop: + /// + /// * [`on_receiving_result`](SentRequest::on_receiving_result) - Schedule + /// a callback to run when the response arrives (doesn't block the event loop) + /// * [`block_task`](SentRequest::block_task) - Block the current task until the response + /// arrives (only safe in spawned tasks, not in handlers) + /// + /// # Anti-Footgun Design + /// + /// The API intentionally makes it difficult to block on the result directly to prevent + /// the common mistake of blocking the event loop while waiting for a response: + /// + /// ```compile_fail + /// # use agent_client_protocol_test::*; + /// # async fn example(cx: agent_client_protocol_core::ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { + /// // ❌ This doesn't compile - prevents blocking the event loop + /// let response = cx.send_request(MyRequest {}).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example(cx: agent_client_protocol_core::ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { + /// // ✅ Option 1: Schedule callback (safe in handlers) + /// cx.send_request(MyRequest {}) + /// .on_receiving_result(async |result| { + /// // Handle the response + /// Ok(()) + /// })?; + /// + /// // ✅ Option 2: Block in spawned task (safe because task is concurrent) + /// cx.spawn({ + /// let cx = cx.clone(); + /// async move { + /// let response = cx.send_request(MyRequest {}) + /// .block_task() + /// .await?; + /// // Process response... + /// Ok(()) + /// } + /// })?; + /// # Ok(()) + /// # } + /// ``` + /// Send an outgoing request to the default counterpart peer. + /// + /// This is a convenience method that sends to the counterpart role `R`. + /// For explicit control over the target peer, use [`send_request_to`](Self::send_request_to). + pub fn send_request(&self, request: Req) -> SentRequest + where + Counterpart: HasPeer, + { + self.send_request_to(self.counterpart.clone(), request) + } + + /// Send an outgoing request to a specific peer. + /// + /// The message will be transformed according to the [`HasPeer`](crate::role::HasPeer) + /// implementation before being sent. + pub fn send_request_to( + &self, + peer: Peer, + request: Req, + ) -> SentRequest + where + Counterpart: HasPeer, + { + let method = request.method().to_string(); + let id = jsonrpcmsg::Id::String(uuid::Uuid::new_v4().to_string()); + let (response_tx, response_rx) = oneshot::channel(); + let role_id = peer.role_id(); + let remote_style = self.counterpart.remote_style(peer); + match remote_style.transform_outgoing_message(request) { + Ok(untyped) => { + // Transform the message for the target role + let message = OutgoingMessage::Request { + id: id.clone(), + method: method.clone(), + role_id, + untyped, + response_tx, + }; + + match self.message_tx.unbounded_send(message) { + Ok(()) => (), + Err(error) => { + let OutgoingMessage::Request { + method, + response_tx, + .. + } = error.into_inner() + else { + unreachable!(); + }; + + response_tx + .send(ResponsePayload { + result: Err(crate::util::internal_error(format!( + "failed to send outgoing request `{method}" + ))), + ack_tx: None, + }) + .unwrap(); + } + } + } + + Err(err) => { + response_tx + .send(ResponsePayload { + result: Err(crate::util::internal_error(format!( + "failed to create untyped request for `{method}`: {err}" + ))), + ack_tx: None, + }) + .unwrap(); + } + } + + SentRequest::new(id, method.clone(), self.task_tx.clone(), response_rx) + .map(move |json| ::from_value(&method, json)) + } + + /// Send an outgoing notification to the default counterpart peer (no reply expected). + /// + /// Notifications are fire-and-forget messages that don't have IDs and don't expect responses. + /// This method sends the notification immediately and returns. + /// + /// This is a convenience method that sends to the counterpart role `R`. + /// For explicit control over the target peer, use [`send_notification_to`](Self::send_notification_to). + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example(cx: agent_client_protocol_core::ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { + /// cx.send_notification(StatusUpdate { + /// message: "Processing...".into(), + /// })?; + /// # Ok(()) + /// # } + /// ``` + pub fn send_notification( + &self, + notification: N, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.send_notification_to(self.counterpart.clone(), notification) + } + + /// Send an outgoing notification to a specific peer (no reply expected). + /// + /// The message will be transformed according to the [`HasPeer`](crate::role::HasPeer) + /// implementation before being sent. + pub fn send_notification_to( + &self, + peer: Peer, + notification: N, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + let remote_style = self.counterpart.remote_style(peer); + tracing::debug!( + role = std::any::type_name::(), + peer = std::any::type_name::(), + notification_type = std::any::type_name::(), + ?remote_style, + original_method = notification.method(), + "send_notification_to" + ); + let transformed = remote_style.transform_outgoing_message(notification)?; + tracing::debug!( + transformed_method = %transformed.method, + "send_notification_to transformed" + ); + send_raw_message( + &self.message_tx, + OutgoingMessage::Notification { + untyped: transformed, + }, + ) + } + + /// Send an error notification (no reply expected). + pub fn send_error_notification(&self, error: crate::Error) -> Result<(), crate::Error> { + send_raw_message(&self.message_tx, OutgoingMessage::Error { error }) + } + + /// Register a dynamic message handler, used to intercept messages specific to a particular session + /// or some similar modal thing. + /// + /// Dynamic message handlers are called first for every incoming message. + /// + /// If they decline to handle the message, then the message is passed to the regular registered handlers. + /// + /// The handler will stay registered until the returned registration guard is dropped. + pub fn add_dynamic_handler( + &self, + handler: impl HandleDispatchFrom + 'static, + ) -> Result, crate::Error> { + let uuid = Uuid::new_v4(); + self.dynamic_handler_tx + .unbounded_send(DynamicHandlerMessage::AddDynamicHandler( + uuid, + Box::new(handler), + )) + .map_err(crate::util::internal_error)?; + + Ok(DynamicHandlerRegistration::new(uuid, self.clone())) + } + + fn remove_dynamic_handler(&self, uuid: Uuid) { + // Ignore errors + drop( + self.dynamic_handler_tx + .unbounded_send(DynamicHandlerMessage::RemoveDynamicHandler(uuid)), + ); + } +} + +#[derive(Clone, Debug)] +pub struct DynamicHandlerRegistration { + uuid: Uuid, + cx: ConnectionTo, +} + +impl DynamicHandlerRegistration { + fn new(uuid: Uuid, cx: ConnectionTo) -> Self { + Self { uuid, cx } + } + + /// Prevents the dynamic handler from being removed when dropped. + pub fn run_indefinitely(self) { + std::mem::forget(self); + } +} + +impl Drop for DynamicHandlerRegistration { + fn drop(&mut self) { + self.cx.remove_dynamic_handler(self.uuid); + } +} + +/// The context to respond to an incoming request. +/// +/// This context is provided to request handlers and serves a dual role: +/// +/// 1. **Respond to the request** - Use [`respond`](Self::respond) or +/// [`respond_with_result`](Self::respond_with_result) to send the response +/// 2. **Send other messages** - Use the [`ConnectionTo`] parameter passed to your +/// handler, which provides [`send_request`](`ConnectionTo::send_request`), +/// [`send_notification`](`ConnectionTo::send_notification`), and +/// [`spawn`](`ConnectionTo::spawn`) +/// +/// # Example +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// connection.on_receive_request(async |req: ProcessRequest, responder, cx| { +/// // Send a notification while processing +/// cx.send_notification(StatusUpdate { +/// message: "processing".into(), +/// })?; +/// +/// // Do some work... +/// let result = process(&req.data)?; +/// +/// // Respond to the request +/// responder.respond(ProcessResponse { result }) +/// }, agent_client_protocol_core::on_receive_request!()) +/// # .connect_to(agent_client_protocol_test::MockTransport).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// # Event Loop Considerations +/// +/// Like all handlers, request handlers run on the event loop. Use +/// [`spawn`](ConnectionTo::spawn) for expensive operations to avoid blocking +/// the connection. +/// +/// See the [Event Loop and Concurrency](Builder#event-loop-and-concurrency) +/// section for more details. +#[must_use] +pub struct Responder { + /// The method of the request. + method: String, + + /// The `id` of the message we are replying to. + id: jsonrpcmsg::Id, + + /// Function to send the response to its destination. + /// + /// For incoming requests: serializes to JSON and sends over the wire. + /// For incoming responses: sends to the waiting oneshot channel. + send_fn: SendBoxFnOnce<'static, (Result,), Result<(), crate::Error>>, +} + +impl std::fmt::Debug for Responder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Responder") + .field("method", &self.method) + .field("id", &self.id) + .field("response_type", &std::any::type_name::()) + .finish_non_exhaustive() + } +} + +impl Responder { + /// Create a new request context for an incoming request. + /// + /// The response will be serialized to JSON and sent over the wire. + fn new(message_tx: OutgoingMessageTx, method: String, id: jsonrpcmsg::Id) -> Self { + let id_clone = id.clone(); + Self { + method, + id, + send_fn: SendBoxFnOnce::new( + move |response: Result| { + send_raw_message( + &message_tx, + OutgoingMessage::Response { + id: id_clone, + response, + }, + ) + }, + ), + } + } + + /// Cast this request context to a different response type. + /// + /// The provided type `T` will be serialized to JSON before sending. + pub fn cast(self) -> Responder { + self.wrap_params(move |method, value| match value { + Ok(value) => T::into_json(value, method), + Err(e) => Err(e), + }) + } +} + +impl Responder { + /// Method of the incoming request + #[must_use] + pub fn method(&self) -> &str { + &self.method + } + + /// ID of the incoming request/response as a JSON value + #[must_use] + pub fn id(&self) -> serde_json::Value { + crate::util::id_to_json(&self.id) + } + + /// Convert to a `Responder` that expects a JSON value + /// and which checks (dynamically) that the JSON value it receives + /// can be converted to `T`. + pub fn erase_to_json(self) -> Responder { + self.wrap_params(|method, value| T::from_value(method, value?)) + } + + /// Return a new Responder with a different method name. + pub fn wrap_method(self, method: String) -> Responder { + Responder { + method, + id: self.id, + send_fn: self.send_fn, + } + } + + /// Return a new Responder that expects a response of type U. + /// + /// `wrap_fn` will be invoked with the method name and the result to transform + /// type `U` into type `T` before sending. + pub fn wrap_params( + self, + wrap_fn: impl FnOnce(&str, Result) -> Result + Send + 'static, + ) -> Responder { + let method = self.method.clone(); + Responder { + method: self.method, + id: self.id, + send_fn: SendBoxFnOnce::new(move |input: Result| { + let t_value = wrap_fn(&method, input); + self.send_fn.call(t_value) + }), + } + } + + /// Respond to the JSON-RPC request with either a value (`Ok`) or an error (`Err`). + pub fn respond_with_result( + self, + response: Result, + ) -> Result<(), crate::Error> { + tracing::debug!(id = ?self.id, "respond called"); + self.send_fn.call(response) + } + + /// Respond to the JSON-RPC request with a value. + pub fn respond(self, response: T) -> Result<(), crate::Error> { + self.respond_with_result(Ok(response)) + } + + /// Respond to the JSON-RPC request with an internal error containing a message. + pub fn respond_with_internal_error(self, message: impl ToString) -> Result<(), crate::Error> { + self.respond_with_error(crate::util::internal_error(message)) + } + + /// Respond to the JSON-RPC request with an error. + pub fn respond_with_error(self, error: crate::Error) -> Result<(), crate::Error> { + tracing::debug!(id = ?self.id, ?error, "respond_with_error called"); + self.respond_with_result(Err(error)) + } +} + +/// Context for handling an incoming JSON-RPC response. +/// +/// This is the response-side counterpart to [`Responder`]. While `Responder` handles +/// incoming requests (where you send a response over the wire), `ResponseRouter` handles +/// incoming responses (where you route the response to a local task waiting for it). +/// +/// Both are fundamentally "sinks" that push the message through a `send_fn`, but they +/// represent different points in the message lifecycle and carry different metadata. +#[must_use] +pub struct ResponseRouter { + /// The method of the original request. + method: String, + + /// The `id` of the original request. + id: jsonrpcmsg::Id, + + /// The RoleId to which the original request was sent + /// (and hence from which the reply is expected). + role_id: RoleId, + + /// Function to send the response to the waiting task. + send_fn: SendBoxFnOnce<'static, (Result,), Result<(), crate::Error>>, +} + +impl std::fmt::Debug for ResponseRouter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResponseRouter") + .field("method", &self.method) + .field("id", &self.id) + .field("response_type", &std::any::type_name::()) + .finish_non_exhaustive() + } +} + +impl ResponseRouter { + /// Create a new response context for routing a response to a local awaiter. + /// + /// When `respond_with_result` is called, the response is sent through the oneshot + /// channel to the code that originally sent the request. + pub(crate) fn new( + method: String, + id: jsonrpcmsg::Id, + role_id: RoleId, + sender: oneshot::Sender, + ) -> Self { + Self { + method, + id, + role_id, + send_fn: SendBoxFnOnce::new( + move |response: Result| { + sender + .send(ResponsePayload { + result: response, + ack_tx: None, + }) + .map_err(|_| { + crate::util::internal_error("failed to send response, receiver dropped") + }) + }, + ), + } + } + + /// Cast this response context to a different response type. + /// + /// The provided type `T` will be serialized to JSON before sending. + pub fn cast(self) -> ResponseRouter { + self.wrap_params(move |method, value| match value { + Ok(value) => T::into_json(value, method), + Err(e) => Err(e), + }) + } +} + +impl ResponseRouter { + /// Method of the original request + #[must_use] + pub fn method(&self) -> &str { + &self.method + } + + /// ID of the original request as a JSON value + #[must_use] + pub fn id(&self) -> serde_json::Value { + crate::util::id_to_json(&self.id) + } + + /// The peer to which the original request was sent. + /// + /// This is the peer from which we expect to receive the response. + #[must_use] + pub fn role_id(&self) -> RoleId { + self.role_id.clone() + } + + /// Convert to a `ResponseRouter` that expects a JSON value + /// and which checks (dynamically) that the JSON value it receives + /// can be converted to `T`. + pub fn erase_to_json(self) -> ResponseRouter { + self.wrap_params(|method, value| T::from_value(method, value?)) + } + + /// Return a new ResponseRouter that expects a response of type U. + /// + /// `wrap_fn` will be invoked with the method name and the result to transform + /// type `U` into type `T` before sending. + fn wrap_params( + self, + wrap_fn: impl FnOnce(&str, Result) -> Result + Send + 'static, + ) -> ResponseRouter { + let method = self.method.clone(); + ResponseRouter { + method: self.method, + id: self.id, + role_id: self.role_id, + send_fn: SendBoxFnOnce::new(move |input: Result| { + let t_value = wrap_fn(&method, input); + self.send_fn.call(t_value) + }), + } + } + + /// Complete the response by sending the result to the waiting task. + pub fn respond_with_result( + self, + response: Result, + ) -> Result<(), crate::Error> { + tracing::debug!(id = ?self.id, "response routed to awaiter"); + self.send_fn.call(response) + } + + /// Complete the response by sending a value to the waiting task. + pub fn respond(self, response: T) -> Result<(), crate::Error> { + self.respond_with_result(Ok(response)) + } + + /// Complete the response by sending an internal error to the waiting task. + pub fn respond_with_internal_error(self, message: impl ToString) -> Result<(), crate::Error> { + self.respond_with_error(crate::util::internal_error(message)) + } + + /// Complete the response by sending an error to the waiting task. + pub fn respond_with_error(self, error: crate::Error) -> Result<(), crate::Error> { + tracing::debug!(id = ?self.id, ?error, "error routed to awaiter"); + self.respond_with_result(Err(error)) + } +} + +/// Common bounds for any JSON-RPC message. +/// +/// # Derive Macro +/// +/// For simple message types, you can use the `JsonRpcRequest` or `JsonRpcNotification` derive macros +/// which will implement both `JsonRpcMessage` and the respective trait. See [`JsonRpcRequest`] and +/// [`JsonRpcNotification`] for examples. +pub trait JsonRpcMessage: 'static + Debug + Sized + Send + Clone { + /// Check if this message type matches the given method name. + fn matches_method(method: &str) -> bool; + + /// The method name for the message. + fn method(&self) -> &str; + + /// Convert this message into an untyped message. + fn to_untyped_message(&self) -> Result; + + /// Parse this type from a method name and parameters. + /// + /// Returns an error if the method doesn't match or deserialization fails. + /// Callers should use `matches_method` first to check if this type handles the method. + fn parse_message(method: &str, params: &impl Serialize) -> Result; +} + +/// Defines the "payload" of a successful response to a JSON-RPC request. +/// +/// # Derive Macro +/// +/// Use `#[derive(JsonRpcResponse)]` to automatically implement this trait: +/// +/// ```ignore +/// use agent_client_protocol_core::JsonRpcResponse; +/// use serde::{Serialize, Deserialize}; +/// +/// #[derive(Debug, Serialize, Deserialize, JsonRpcResponse)] +/// #[response(method = "_hello")] +/// struct HelloResponse { +/// greeting: String, +/// } +/// ``` +pub trait JsonRpcResponse: 'static + Debug + Sized + Send + Clone { + /// Convert this message into a JSON value. + fn into_json(self, method: &str) -> Result; + + /// Parse a JSON value into the response type. + fn from_value(method: &str, value: serde_json::Value) -> Result; +} + +impl JsonRpcResponse for serde_json::Value { + fn from_value(_method: &str, value: serde_json::Value) -> Result { + Ok(value) + } + + fn into_json(self, _method: &str) -> Result { + Ok(self) + } +} + +/// A struct that represents a notification (JSON-RPC message that does not expect a response). +/// +/// # Derive Macro +/// +/// Use `#[derive(JsonRpcNotification)]` to automatically implement both `JsonRpcMessage` and `JsonRpcNotification`: +/// +/// ```ignore +/// use agent_client_protocol_core::JsonRpcNotification; +/// use serde::{Serialize, Deserialize}; +/// +/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)] +/// #[notification(method = "_ping")] +/// struct PingNotification { +/// timestamp: u64, +/// } +/// ``` +pub trait JsonRpcNotification: JsonRpcMessage {} + +/// A struct that represents a request (JSON-RPC message expecting a response). +/// +/// # Derive Macro +/// +/// Use `#[derive(JsonRpcRequest)]` to automatically implement both `JsonRpcMessage` and `JsonRpcRequest`: +/// +/// ```ignore +/// use agent_client_protocol_core::{JsonRpcRequest, JsonRpcResponse}; +/// use serde::{Serialize, Deserialize}; +/// +/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcRequest)] +/// #[request(method = "_hello", response = HelloResponse)] +/// struct HelloRequest { +/// name: String, +/// } +/// +/// #[derive(Debug, Serialize, Deserialize, JsonRpcResponse)] +/// struct HelloResponse { +/// greeting: String, +/// } +/// ``` +pub trait JsonRpcRequest: JsonRpcMessage { + /// The type of data expected in response. + type Response: JsonRpcResponse; +} + +/// An enum capturing an in-flight request or notification. +/// In the case of a request, also includes the context used to respond to the request. +/// +/// Type parameters allow specifying the concrete request and notification types. +/// By default, both are `UntypedMessage` for dynamic dispatch. +/// The request context's response type matches the request's response type. +#[derive(Debug)] +pub enum Dispatch { + /// Incoming request and the context where the response should be sent. + Request(Req, Responder), + + /// Incoming notification. + Notification(Notif), + + /// Incoming response to a request we sent. + /// + /// The first field is the response result (success or error from the remote). + /// The second field is the context for forwarding the response to its destination + /// (typically a waiting oneshot channel). + Response( + Result, + ResponseRouter, + ), +} + +impl Dispatch { + /// Map the request and notification types to new types. + /// + /// Note: Response variants are passed through unchanged since they don't + /// contain a parseable message payload. + pub fn map( + self, + map_request: impl FnOnce(Req, Responder) -> (Req1, Responder), + map_notification: impl FnOnce(Notif) -> Notif1, + ) -> Dispatch + where + Req1: JsonRpcRequest, + Notif1: JsonRpcMessage, + { + match self { + Dispatch::Request(request, responder) => { + let (new_request, new_responder) = map_request(request, responder); + Dispatch::Request(new_request, new_responder) + } + Dispatch::Notification(notification) => { + let new_notification = map_notification(notification); + Dispatch::Notification(new_notification) + } + Dispatch::Response(result, router) => Dispatch::Response(result, router), + } + } + + /// Respond to the message with an error. + /// + /// If this message is a request, this error becomes the reply to the request. + /// + /// If this message is a notification, the error is sent as a notification. + /// + /// If this message is a response, the error is forwarded to the waiting handler. + pub fn respond_with_error( + self, + error: crate::Error, + cx: ConnectionTo, + ) -> Result<(), crate::Error> { + match self { + Dispatch::Request(_, responder) => responder.respond_with_error(error), + Dispatch::Notification(_) => cx.send_error_notification(error), + Dispatch::Response(_, responder) => responder.respond_with_error(error), + } + } + + /// Convert to a `Responder` that expects a JSON value + /// and which checks (dynamically) that the JSON value it receives + /// can be converted to `T`. + /// + /// Note: Response variants cannot be erased since their payload is already + /// parsed. This returns an error for Response variants. + pub fn erase_to_json(self) -> Result { + match self { + Dispatch::Request(response, responder) => Ok(Dispatch::Request( + response.to_untyped_message()?, + responder.erase_to_json(), + )), + Dispatch::Notification(notification) => { + Ok(Dispatch::Notification(notification.to_untyped_message()?)) + } + Dispatch::Response(_, _) => Err(crate::util::internal_error( + "cannot erase Response variant to JSON", + )), + } + } + + /// Convert the message in self to an untyped message. + /// + /// Note: Response variants don't have an untyped message representation. + /// This returns an error for Response variants. + pub fn to_untyped_message(&self) -> Result { + match self { + Dispatch::Request(request, _) => request.to_untyped_message(), + Dispatch::Notification(notification) => notification.to_untyped_message(), + Dispatch::Response(_, _) => Err(crate::util::internal_error( + "Response variant has no untyped message representation", + )), + } + } + + /// Convert self to an untyped message context. + /// + /// Note: Response variants cannot be converted. This returns an error for Response variants. + pub fn into_untyped_dispatch(self) -> Result { + match self { + Dispatch::Request(request, responder) => Ok(Dispatch::Request( + request.to_untyped_message()?, + responder.erase_to_json(), + )), + Dispatch::Notification(notification) => { + Ok(Dispatch::Notification(notification.to_untyped_message()?)) + } + Dispatch::Response(_, _) => Err(crate::util::internal_error( + "cannot convert Response variant to untyped message context", + )), + } + } + + /// Returns the request ID if this is a request or response, None if notification. + pub fn id(&self) -> Option { + match self { + Dispatch::Request(_, cx) => Some(cx.id()), + Dispatch::Notification(_) => None, + Dispatch::Response(_, cx) => Some(cx.id()), + } + } + + /// Returns the method of the message. + /// + /// For requests and notifications, this is the method from the message payload. + /// For responses, this is the method of the original request. + pub fn method(&self) -> &str { + match self { + Dispatch::Request(msg, _) => msg.method(), + Dispatch::Notification(msg) => msg.method(), + Dispatch::Response(_, cx) => cx.method(), + } + } +} + +impl Dispatch { + /// Attempts to parse `self` into a typed message context. + /// + /// # Returns + /// + /// * `Ok(Ok(typed))` if this is a request/notification of the given types + /// * `Ok(Err(self))` if not + /// * `Err` if has the correct method for the given types but parsing fails + #[tracing::instrument(skip(self), fields(Request = ?std::any::type_name::(), Notif = ?std::any::type_name::()), level = "trace", ret)] + pub(crate) fn into_typed_dispatch( + self, + ) -> Result, Dispatch>, crate::Error> { + tracing::debug!( + message = ?self, + "into_typed_dispatch" + ); + match self { + Dispatch::Request(message, responder) => { + if Req::matches_method(&message.method) { + match Req::parse_message(&message.method, &message.params) { + Ok(req) => { + tracing::trace!(?req, "parsed ok"); + Ok(Ok(Dispatch::Request(req, responder.cast()))) + } + Err(err) => { + tracing::trace!(?err, "parse error"); + Err(err) + } + } + } else { + tracing::trace!("method doesn't match"); + Ok(Err(Dispatch::Request(message, responder))) + } + } + + Dispatch::Notification(message) => { + if Notif::matches_method(&message.method) { + match Notif::parse_message(&message.method, &message.params) { + Ok(notif) => { + tracing::trace!(?notif, "parse ok"); + Ok(Ok(Dispatch::Notification(notif))) + } + Err(err) => { + tracing::trace!(?err, "parse error"); + Err(err) + } + } + } else { + tracing::trace!("method doesn't match"); + Ok(Err(Dispatch::Notification(message))) + } + } + + Dispatch::Response(result, cx) => { + let method = cx.method(); + if Req::matches_method(method) { + // Parse the response result + let typed_result = match result { + Ok(value) => { + match ::from_value(method, value) { + Ok(parsed) => { + tracing::trace!(?parsed, "parse ok"); + Ok(parsed) + } + Err(err) => { + tracing::trace!(?err, "parse error"); + return Err(err); + } + } + } + Err(err) => { + tracing::trace!("error, passthrough"); + Err(err) + } + }; + Ok(Ok(Dispatch::Response(typed_result, cx.cast()))) + } else { + tracing::trace!("method doesn't match"); + Ok(Err(Dispatch::Response(result, cx))) + } + } + } + } + + /// True if this message has a field with the given name. + /// + /// Returns `false` for Response variants. + #[must_use] + pub fn has_field(&self, field_name: &str) -> bool { + self.message() + .and_then(|m| m.params().get(field_name)) + .is_some() + } + + /// Returns true if this message has a session-id field. + /// + /// Returns `false` for Response variants. + pub(crate) fn has_session_id(&self) -> bool { + self.has_field("sessionId") + } + + /// Extract the ACP session-id from this message (if any). + /// + /// Returns `Ok(None)` for Response variants. + pub(crate) fn get_session_id(&self) -> Result, crate::Error> { + let Some(message) = self.message() else { + return Ok(None); + }; + let Some(value) = message.params().get("sessionId") else { + return Ok(None); + }; + let session_id = serde_json::from_value(value.clone())?; + Ok(Some(session_id)) + } + + /// Try to parse this as a notification of the given type. + /// + /// # Returns + /// + /// * `Ok(Ok(typed))` if this is a request/notification of the given types + /// * `Ok(Err(self))` if not + /// * `Err` if has the correct method for the given types but parsing fails + pub fn into_notification( + self, + ) -> Result, crate::Error> { + match self { + Dispatch::Notification(msg) => { + if !N::matches_method(&msg.method) { + return Ok(Err(Dispatch::Notification(msg))); + } + match N::parse_message(&msg.method, &msg.params) { + Ok(n) => Ok(Ok(n)), + Err(err) => Err(err), + } + } + Dispatch::Request(..) | Dispatch::Response(..) => Ok(Err(self)), + } + } + + /// Try to parse this as a request of the given type. + /// + /// # Returns + /// + /// * `Ok(Ok(typed))` if this is a request/notification of the given types + /// * `Ok(Err(self))` if not + /// * `Err` if has the correct method for the given types but parsing fails + pub fn into_request( + self, + ) -> Result), Dispatch>, crate::Error> { + match self { + Dispatch::Request(msg, responder) => { + if !Req::matches_method(&msg.method) { + return Ok(Err(Dispatch::Request(msg, responder))); + } + match Req::parse_message(&msg.method, &msg.params) { + Ok(req) => Ok(Ok((req, responder.cast()))), + Err(err) => Err(err), + } + } + Dispatch::Notification(..) | Dispatch::Response(..) => Ok(Err(self)), + } + } +} + +impl Dispatch { + /// Returns the message payload for requests and notifications. + /// + /// Returns `None` for Response variants since they don't contain a message payload. + pub fn message(&self) -> Option<&M> { + match self { + Dispatch::Request(msg, _) | Dispatch::Notification(msg) => Some(msg), + Dispatch::Response(_, _) => None, + } + } + + /// Map the request/notification message. + /// + /// Response variants pass through unchanged. + pub(crate) fn try_map_message( + self, + map_message: impl FnOnce(M) -> Result, + ) -> Result, crate::Error> { + match self { + Dispatch::Request(request, cx) => Ok(Dispatch::Request(map_message(request)?, cx)), + Dispatch::Notification(notification) => { + Ok(Dispatch::::Notification(map_message(notification)?)) + } + Dispatch::Response(result, cx) => Ok(Dispatch::Response(result, cx)), + } + } +} + +/// An incoming JSON message without any typing. Can be a request or a notification. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct UntypedMessage { + /// The JSON-RPC method name + pub method: String, + /// The JSON-RPC parameters as a raw JSON value + pub params: serde_json::Value, +} + +impl UntypedMessage { + /// Returns an untyped message with the given method and parameters. + pub fn new(method: &str, params: impl Serialize) -> Result { + let params = serde_json::to_value(params)?; + Ok(Self { + method: method.to_string(), + params, + }) + } + + /// Returns the method name + #[must_use] + pub fn method(&self) -> &str { + &self.method + } + + /// Returns the parameters as a JSON value + #[must_use] + pub fn params(&self) -> &serde_json::Value { + &self.params + } + + /// Consumes this message and returns the method and params + #[must_use] + pub fn into_parts(self) -> (String, serde_json::Value) { + (self.method, self.params) + } + + /// Convert `self` to a JSON-RPC message. + pub(crate) fn into_jsonrpc_msg( + self, + id: Option, + ) -> Result { + let Self { method, params } = self; + Ok(jsonrpcmsg::Request::new_v2(method, json_cast(params)?, id)) + } +} + +impl JsonRpcMessage for UntypedMessage { + fn matches_method(_method: &str) -> bool { + // UntypedMessage matches any method - it's the untyped fallback + true + } + + fn method(&self) -> &str { + &self.method + } + + fn to_untyped_message(&self) -> Result { + Ok(self.clone()) + } + + fn parse_message(method: &str, params: &impl Serialize) -> Result { + UntypedMessage::new(method, params) + } +} + +impl JsonRpcRequest for UntypedMessage { + type Response = serde_json::Value; +} + +impl JsonRpcNotification for UntypedMessage {} + +/// Represents a pending response of type `R` from an outgoing request. +/// +/// Returned by [`ConnectionTo::send_request`], this type provides methods for handling +/// the response without blocking the event loop. The API is intentionally designed to make +/// it difficult to accidentally block. +/// +/// # Anti-Footgun Design +/// +/// You cannot directly `.await` a `SentRequest`. Instead, you must choose how to handle +/// the response: +/// +/// ## Option 1: Schedule a Callback (Safe in Handlers) +/// +/// Use [`on_receiving_result`](Self::on_receiving_result) to schedule a task +/// that runs when the response arrives. This doesn't block the event loop: +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # async fn example(cx: agent_client_protocol_core::ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { +/// cx.send_request(MyRequest {}) +/// .on_receiving_result(async |result| { +/// match result { +/// Ok(response) => { +/// // Handle successful response +/// Ok(()) +/// } +/// Err(error) => { +/// // Handle error +/// Err(error) +/// } +/// } +/// })?; +/// # Ok(()) +/// # } +/// ``` +/// +/// ## Option 2: Block in a Spawned Task (Safe Only in `spawn`) +/// +/// Use [`block_task`](Self::block_task) to block until the response arrives, but **only** +/// in a spawned task (never in a handler): +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # async fn example(cx: agent_client_protocol_core::ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { +/// // ✅ Safe: Spawned task runs concurrently +/// cx.spawn({ +/// let cx = cx.clone(); +/// async move { +/// let response = cx.send_request(MyRequest {}) +/// .block_task() +/// .await?; +/// // Process response... +/// Ok(()) +/// } +/// })?; +/// # Ok(()) +/// # } +/// ``` +/// +/// ```no_run +/// # use agent_client_protocol_test::*; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// # let connection = mock_connection(); +/// // ❌ NEVER do this in a handler - blocks the event loop! +/// connection.on_receive_request(async |req: MyRequest, responder, cx| { +/// let response = cx.send_request(MyRequest {}) +/// .block_task() // This will deadlock! +/// .await?; +/// responder.respond(response) +/// }, agent_client_protocol_core::on_receive_request!()) +/// # .connect_to(agent_client_protocol_test::MockTransport).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// # Why This Design? +/// +/// If you block the event loop while waiting for a response, the connection cannot process +/// the incoming response message, creating a deadlock. This API design prevents that footgun +/// by making blocking explicit and encouraging non-blocking patterns. +pub struct SentRequest { + id: jsonrpcmsg::Id, + method: String, + task_tx: TaskTx, + response_rx: oneshot::Receiver, + to_result: Box Result + Send>, +} + +impl Debug for SentRequest { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SentRequest") + .field("id", &self.id) + .field("method", &self.method) + .field("task_tx", &self.task_tx) + .field("response_rx", &self.response_rx) + .finish_non_exhaustive() + } +} + +impl SentRequest { + fn new( + id: jsonrpcmsg::Id, + method: String, + task_tx: mpsc::UnboundedSender, + response_rx: oneshot::Receiver, + ) -> Self { + Self { + id, + method, + response_rx, + task_tx, + to_result: Box::new(Ok), + } + } +} + +impl SentRequest { + /// The id of the outgoing request. + #[must_use] + pub fn id(&self) -> serde_json::Value { + crate::util::id_to_json(&self.id) + } + + /// The method of the request this is in response to. + #[must_use] + pub fn method(&self) -> &str { + &self.method + } + + /// Create a new response that maps the result of the response to a new type. + pub fn map( + self, + map_fn: impl Fn(T) -> Result + 'static + Send, + ) -> SentRequest { + SentRequest { + id: self.id, + method: self.method, + response_rx: self.response_rx, + task_tx: self.task_tx, + to_result: Box::new(move |value| map_fn((self.to_result)(value)?)), + } + } + + /// Forward the response (success or error) to a request context when it arrives. + /// + /// This is a convenience method for proxying messages between connections. When the + /// response arrives, it will be automatically sent to the provided request context, + /// whether it's a successful response or an error. + /// + /// # Example: Proxying requests + /// + /// ``` + /// # use agent_client_protocol_core::UntypedRole; + /// # use agent_client_protocol_core::{Builder, ConnectionTo}; + /// # use agent_client_protocol_test::*; + /// # async fn example(cx: ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { + /// // Set up backend connection builder + /// let backend = UntypedRole.builder() + /// .on_receive_request(async |req: MyRequest, responder, cx| { + /// responder.respond(MyResponse { status: "ok".into() }) + /// }, agent_client_protocol_core::on_receive_request!()); + /// + /// // Spawn backend and get a context to send to it + /// let backend_connection = cx.spawn_connection(backend, MockTransport)?; + /// + /// // Set up proxy that forwards requests to backend + /// UntypedRole.builder() + /// .on_receive_request({ + /// let backend_connection = backend_connection.clone(); + /// async move |req: MyRequest, responder, cx| { + /// // Forward the request to backend and proxy the response back + /// backend_connection.send_request(req) + /// .forward_response_to(responder)?; + /// Ok(()) + /// } + /// }, agent_client_protocol_core::on_receive_request!()); + /// # Ok(()) + /// # } + /// ``` + /// + /// # Type Safety + /// + /// The request context's response type must match the request's response type, + /// ensuring type-safe message forwarding. + /// + /// # When to Use + /// + /// Use this when: + /// - You're implementing a proxy or gateway pattern + /// - You want to forward responses without processing them + /// - The response types match between the outgoing request and incoming request + /// + /// This is equivalent to calling `on_receiving_result` and manually forwarding + /// the result, but more concise. + pub fn forward_response_to(self, responder: Responder) -> Result<(), crate::Error> + where + T: Send, + { + self.on_receiving_result(async move |result| responder.respond_with_result(result)) + } + + /// Block the current task until the response is received. + /// + /// **Warning:** This method blocks the current async task. It is **only safe** to use + /// in spawned tasks created with [`ConnectionTo::spawn`]. Using it directly in a + /// handler callback will deadlock the connection. + /// + /// # Safe Usage (in spawned tasks) + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// # let connection = mock_connection(); + /// connection.on_receive_request(async |req: MyRequest, responder, cx| { + /// // Spawn a task to handle the request + /// cx.spawn({ + /// let connection = cx.clone(); + /// async move { + /// // Safe: We're in a spawned task, not blocking the event loop + /// let response = connection.send_request(OtherRequest {}) + /// .block_task() + /// .await?; + /// + /// // Process the response... + /// Ok(()) + /// } + /// })?; + /// + /// // Respond immediately + /// responder.respond(MyResponse { status: "ok".into() }) + /// }, agent_client_protocol_core::on_receive_request!()) + /// # .connect_to(agent_client_protocol_test::MockTransport).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Unsafe Usage (in handlers - will deadlock!) + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// # let connection = mock_connection(); + /// connection.on_receive_request(async |req: MyRequest, responder, cx| { + /// // ❌ DEADLOCK: Handler blocks event loop, which can't process the response + /// let response = cx.send_request(OtherRequest {}) + /// .block_task() + /// .await?; + /// + /// responder.respond(MyResponse { status: response.value }) + /// }, agent_client_protocol_core::on_receive_request!()) + /// # .connect_to(agent_client_protocol_test::MockTransport).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # When to Use + /// + /// Use this method when: + /// - You're in a spawned task (via [`ConnectionTo::spawn`]) + /// - You need the response value to proceed with your logic + /// - Linear control flow is more natural than callbacks + /// + /// For handler callbacks, use [`on_receiving_result`](Self::on_receiving_result) instead. + pub async fn block_task(self) -> Result + where + T: Send, + { + match self.response_rx.await { + Ok(ResponsePayload { + result: Ok(json_value), + ack_tx, + }) => { + // Ack immediately - we're in a spawned task, so the dispatch loop + // can continue while we process the value. + if let Some(tx) = ack_tx { + let _ = tx.send(()); + } + match (self.to_result)(json_value) { + Ok(value) => Ok(value), + Err(err) => Err(err), + } + } + Ok(ResponsePayload { + result: Err(err), + ack_tx, + }) => { + if let Some(tx) = ack_tx { + let _ = tx.send(()); + } + Err(err) + } + Err(err) => Err(crate::util::internal_error(format!( + "response to `{}` never received: {}", + self.method, err + ))), + } + } + + /// Schedule an async task to run when a successful response is received. + /// + /// This is a convenience wrapper around [`on_receiving_result`](Self::on_receiving_result) + /// for the common pattern of forwarding errors to a request context while only processing + /// successful responses. + /// + /// # Behavior + /// + /// - If the response is `Ok(value)`, your task receives the value and the request context + /// - If the response is `Err(error)`, the error is automatically sent to `responder` + /// and your task is not called + /// + /// # Example: Chaining requests + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// # let connection = mock_connection(); + /// connection.on_receive_request(async |req: ValidateRequest, responder, cx| { + /// // Send initial request + /// cx.send_request(ValidateRequest { data: req.data.clone() }) + /// .on_receiving_ok_result(responder, async |validation, responder| { + /// // Only runs if validation succeeded + /// if validation.is_valid { + /// // Respond to original request + /// responder.respond(ValidateResponse { is_valid: true, error: None }) + /// } else { + /// responder.respond_with_error(agent_client_protocol_core::util::internal_error("validation failed")) + /// } + /// })?; + /// + /// Ok(()) + /// }, agent_client_protocol_core::on_receive_request!()) + /// # .connect_to(agent_client_protocol_test::MockTransport).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Ordering + /// + /// Like [`on_receiving_result`](Self::on_receiving_result), the callback blocks the + /// dispatch loop until it completes. See the [`ordering`](crate::concepts::ordering) module + /// for details. + /// + /// # When to Use + /// + /// Use this when: + /// - You need to respond to a request based on another request's result + /// - You want errors to automatically propagate to the request context + /// - You only care about the success case + /// + /// For more control over error handling, use [`on_receiving_result`](Self::on_receiving_result). + #[track_caller] + pub fn on_receiving_ok_result( + self, + responder: Responder, + task: impl FnOnce(T, Responder) -> F + 'static + Send, + ) -> Result<(), crate::Error> + where + F: Future> + 'static + Send, + T: Send, + { + self.on_receiving_result(async move |result| match result { + Ok(value) => task(value, responder).await, + Err(err) => responder.respond_with_error(err), + }) + } + + /// Schedule an async task to run when the response is received. + /// + /// This is the recommended way to handle responses in handler callbacks, as it doesn't + /// block the event loop. The task will be spawned automatically when the response arrives. + /// + /// # Example: Handle response in callback + /// + /// ```no_run + /// # use agent_client_protocol_test::*; + /// # async fn example() -> Result<(), agent_client_protocol_core::Error> { + /// # let connection = mock_connection(); + /// connection.on_receive_request(async |req: MyRequest, responder, cx| { + /// // Send a request and schedule a callback for the response + /// cx.send_request(QueryRequest { id: 22 }) + /// .on_receiving_result({ + /// let connection = cx.clone(); + /// async move |result| { + /// match result { + /// Ok(response) => { + /// println!("Got response: {:?}", response); + /// // Can send more messages here + /// connection.send_notification(QueryComplete {})?; + /// Ok(()) + /// } + /// Err(error) => { + /// eprintln!("Request failed: {}", error); + /// Err(error) + /// } + /// } + /// } + /// })?; + /// + /// // Handler continues immediately without waiting + /// responder.respond(MyResponse { status: "processing".into() }) + /// }, agent_client_protocol_core::on_receive_request!()) + /// # .connect_to(agent_client_protocol_test::MockTransport).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Ordering + /// + /// The callback runs as a spawned task, but the dispatch loop waits for it to complete + /// before processing the next message. This gives you ordering guarantees: no other + /// messages will be processed while your callback runs. + /// + /// This differs from [`block_task`](Self::block_task), which signals completion immediately + /// upon receiving the response (before your code processes it). + /// + /// See the [`ordering`](crate::concepts::ordering) module for details on ordering guarantees + /// and how to avoid deadlocks. + /// + /// # Error Handling + /// + /// If the scheduled task returns `Err`, the entire server will shut down. Make sure to handle + /// errors appropriately within your task. + /// + /// # When to Use + /// + /// Use this method when: + /// - You're in a handler callback (not a spawned task) + /// - You want ordering guarantees (no other messages processed during your callback) + /// - You need to do async work before "releasing" control back to the dispatch loop + /// + /// For spawned tasks where you don't need ordering guarantees, consider [`block_task`](Self::block_task). + #[track_caller] + pub fn on_receiving_result( + self, + task: impl FnOnce(Result) -> F + 'static + Send, + ) -> Result<(), crate::Error> + where + F: Future> + 'static + Send, + T: Send, + { + let task_tx = self.task_tx.clone(); + let method = self.method; + let response_rx = self.response_rx; + let to_result = self.to_result; + let location = Location::caller(); + + Task::new(location, async move { + match response_rx.await { + Ok(ResponsePayload { result, ack_tx }) => { + // Convert the result using to_result for Ok values + let typed_result = match result { + Ok(json_value) => to_result(json_value), + Err(err) => Err(err), + }; + + // Run the user's callback + let outcome = task(typed_result).await; + + // Ack AFTER the callback completes - this is the key difference + // from block_task. The dispatch loop waits for this ack. + if let Some(tx) = ack_tx { + let _ = tx.send(()); + } + + outcome + } + Err(err) => Err(crate::util::internal_error(format!( + "response to `{method}` never received: {err}" + ))), + } + }) + .spawn(&task_tx) + } +} + +// ============================================================================ +// IntoJrConnectionTransport Implementations +// ============================================================================ + +/// A component that communicates over line streams. +/// +/// `Lines` implements the [`ConnectTo`] trait for any pair of line-based streams +/// (a `Stream>` for incoming and a `Sink` for outgoing), +/// handling serialization of JSON-RPC messages to/from newline-delimited JSON. +/// +/// This is a lower-level primitive than [`ByteStreams`] that enables interception and +/// transformation of individual lines before they are parsed or after they are serialized. +/// This is particularly useful for debugging, logging, or implementing custom line-based +/// protocols. +/// +/// # Use Cases +/// +/// - **Line-by-line logging**: Intercept and log each line before parsing +/// - **Custom protocols**: Transform lines before/after JSON-RPC processing +/// - **Debugging**: Inspect raw message strings +/// - **Line filtering**: Skip or modify specific messages +/// +/// Most users should use [`ByteStreams`] instead, which provides a simpler interface +/// for byte-based I/O. +/// +/// [`ConnectTo`]: crate::ConnectTo +#[derive(Debug)] +pub struct Lines { + /// Outgoing line sink (where we write serialized JSON-RPC messages) + pub outgoing: OutgoingSink, + /// Incoming line stream (where we read and parse JSON-RPC messages) + pub incoming: IncomingStream, +} + +impl Lines +where + OutgoingSink: futures::Sink + Send + 'static, + IncomingStream: futures::Stream> + Send + 'static, +{ + /// Create a new line stream transport. + pub fn new(outgoing: OutgoingSink, incoming: IncomingStream) -> Self { + Self { outgoing, incoming } + } +} + +impl ConnectTo for Lines +where + OutgoingSink: futures::Sink + Send + 'static, + IncomingStream: futures::Stream> + Send + 'static, +{ + async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { + let (channel, serve_self) = ConnectTo::::into_channel_and_future(self); + match futures::future::select(Box::pin(client.connect_to(channel)), serve_self).await { + Either::Left((result, _)) | Either::Right((result, _)) => result, + } + } + + fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { + let Self { outgoing, incoming } = self; + + // Create a channel pair for the client to use + let (channel_for_caller, channel_for_lines) = Channel::duplex(); + + // Create the server future that runs the line stream actors + let server_future = Box::pin(async move { + let Channel { rx, tx } = channel_for_lines; + + // Run both actors concurrently + let outgoing_future = transport_actor::transport_outgoing_lines_actor(rx, outgoing); + let incoming_future = transport_actor::transport_incoming_lines_actor(incoming, tx); + + // Wait for both to complete + futures::try_join!(outgoing_future, incoming_future)?; + + Ok(()) + }); + + (channel_for_caller, server_future) + } +} + +/// A component that communicates over byte streams (stdin/stdout, sockets, pipes, etc.). +/// +/// `ByteStreams` implements the [`ConnectTo`] trait for any pair of `AsyncRead` and `AsyncWrite` +/// streams, handling serialization of JSON-RPC messages to/from newline-delimited JSON. +/// This is the standard way to communicate with external processes or network connections. +/// +/// # Use Cases +/// +/// - **Stdio communication**: Connect to agents or proxies via stdin/stdout +/// - **Network sockets**: TCP, Unix domain sockets, or other stream-based protocols +/// - **Named pipes**: Cross-process communication on the same machine +/// - **File I/O**: Reading from and writing to file descriptors +/// +/// # Example +/// +/// Connecting to an agent via stdio: +/// +/// ```no_run +/// use agent_client_protocol_core::UntypedRole; +/// # use agent_client_protocol_core::{ByteStreams}; +/// use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; +/// +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// let component = ByteStreams::new( +/// tokio::io::stdout().compat_write(), +/// tokio::io::stdin().compat(), +/// ); +/// +/// // Use as a component in a connection +/// agent_client_protocol_core::UntypedRole.builder() +/// .name("my-client") +/// .connect_to(component) +/// .await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// [`ConnectTo`]: crate::ConnectTo +#[derive(Debug)] +pub struct ByteStreams { + /// Outgoing byte stream (where we write serialized messages) + pub outgoing: OB, + /// Incoming byte stream (where we read and parse messages) + pub incoming: IB, +} + +impl ByteStreams +where + OB: AsyncWrite + Send + 'static, + IB: AsyncRead + Send + 'static, +{ + /// Create a new byte stream transport. + pub fn new(outgoing: OB, incoming: IB) -> Self { + Self { outgoing, incoming } + } +} + +impl ConnectTo for ByteStreams +where + OB: AsyncWrite + Send + 'static, + IB: AsyncRead + Send + 'static, +{ + async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { + let (channel, serve_self) = ConnectTo::::into_channel_and_future(self); + match futures::future::select(pin!(client.connect_to(channel)), serve_self).await { + Either::Left((result, _)) | Either::Right((result, _)) => result, + } + } + + fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { + use futures::AsyncBufReadExt; + use futures::AsyncWriteExt; + use futures::io::BufReader; + let Self { outgoing, incoming } = self; + + // Convert byte streams to line streams + // Box both streams to satisfy Unpin requirements + let incoming_lines = Box::pin(BufReader::new(incoming).lines()); + + // Create a sink that writes lines (with newlines) to the outgoing byte stream + // We need to Box the writer since it may not be Unpin + let outgoing_sink = + futures::sink::unfold(Box::pin(outgoing), async move |mut writer, line: String| { + let mut bytes = line.into_bytes(); + bytes.push(b'\n'); + writer.write_all(&bytes).await?; + Ok::<_, std::io::Error>(writer) + }); + + // Delegate to Lines component + ConnectTo::::into_channel_and_future(Lines::new(outgoing_sink, incoming_lines)) + } +} + +/// A channel endpoint representing one side of a bidirectional message channel. +/// +/// `Channel` represents a single endpoint's view of a bidirectional communication channel. +/// Each endpoint has: +/// - `rx`: A receiver for incoming messages (or errors) from the counterpart +/// - `tx`: A sender for outgoing messages (or errors) to the counterpart +/// +/// # Example +/// +/// ```no_run +/// # use agent_client_protocol_core::UntypedRole; +/// # use agent_client_protocol_core::{Channel, Builder}; +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// // Create a pair of connected channels +/// let (channel_a, channel_b) = Channel::duplex(); +/// +/// // Each channel can be used by a different component +/// UntypedRole.builder() +/// .name("connection-a") +/// .connect_to(channel_a) +/// .await?; +/// # Ok(()) +/// # } +/// ``` +#[derive(Debug)] +pub struct Channel { + /// Receives messages (or errors) from the counterpart. + pub rx: mpsc::UnboundedReceiver>, + /// Sends messages (or errors) to the counterpart. + pub tx: mpsc::UnboundedSender>, +} + +impl Channel { + /// Create a pair of connected channel endpoints. + /// + /// Returns two `Channel` instances that are connected to each other: + /// - Messages sent via `channel_a.tx` are received on `channel_b.rx` + /// - Messages sent via `channel_b.tx` are received on `channel_a.rx` + /// + /// # Returns + /// + /// A tuple `(channel_a, channel_b)` of connected channel endpoints. + #[must_use] + pub fn duplex() -> (Self, Self) { + // Create channels: A sends Result which B receives as Message + let (a_tx, b_rx) = mpsc::unbounded(); + let (b_tx, a_rx) = mpsc::unbounded(); + + let channel_a = Self { rx: a_rx, tx: a_tx }; + let channel_b = Self { rx: b_rx, tx: b_tx }; + + (channel_a, channel_b) + } + + /// Copy messages from `rx` to `tx`. + /// + /// # Returns + /// + /// A `Result` indicating success or failure. + pub async fn copy(mut self) -> Result<(), crate::Error> { + while let Some(msg) = self.rx.next().await { + self.tx + .unbounded_send(msg) + .map_err(crate::util::internal_error)?; + } + Ok(()) + } +} + +impl ConnectTo for Channel { + async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { + let (client_channel, client_serve) = client.into_channel_and_future(); + + match futures::try_join!( + Channel { + rx: client_channel.rx, + tx: self.tx + } + .copy(), + Channel { + rx: self.rx, + tx: client_channel.tx + } + .copy(), + client_serve + ) { + Ok(((), (), ())) => Ok(()), + Err(err) => Err(err), + } + } + + fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { + (self, Box::pin(future::ready(Ok(())))) + } +} diff --git a/src/agent-client-protocol-core/src/jsonrpc/connection.rs b/src/agent-client-protocol-core/src/jsonrpc/connection.rs new file mode 100644 index 0000000..e69de29 diff --git a/src/agent-client-protocol-core/src/jsonrpc/dynamic_handler.rs b/src/agent-client-protocol-core/src/jsonrpc/dynamic_handler.rs new file mode 100644 index 0000000..177348d --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc/dynamic_handler.rs @@ -0,0 +1,55 @@ +use futures::future::BoxFuture; +use uuid::Uuid; + +use crate::role::Role; +use crate::{ConnectionTo, Dispatch, HandleDispatchFrom, Handled}; + +/// Internal dyn-safe wrapper around `HandleMessageAs` +/// +/// The type parameter `R` is the role's counterpart (who we connect to). +pub(crate) trait DynHandleDispatchFrom: Send { + fn dyn_handle_dispatch_from( + &mut self, + message: Dispatch, + cx: ConnectionTo, + ) -> BoxFuture<'_, Result, crate::Error>>; + + fn dyn_describe_chain(&self) -> String; +} + +impl> DynHandleDispatchFrom + for H +{ + fn dyn_handle_dispatch_from( + &mut self, + message: Dispatch, + cx: ConnectionTo, + ) -> BoxFuture<'_, Result, crate::Error>> { + Box::pin(HandleDispatchFrom::handle_dispatch_from(self, message, cx)) + } + + fn dyn_describe_chain(&self) -> String { + format!("{:?}", H::describe_chain(self)) + } +} + +/// Messages used to add/remove dynamic handlers +pub(crate) enum DynamicHandlerMessage { + AddDynamicHandler(Uuid, Box>), + RemoveDynamicHandler(Uuid), +} + +impl std::fmt::Debug for DynamicHandlerMessage { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::AddDynamicHandler(arg0, arg1) => f + .debug_tuple("AddDynamicHandler") + .field(arg0) + .field(&arg1.dyn_describe_chain()) + .finish(), + Self::RemoveDynamicHandler(arg0) => { + f.debug_tuple("RemoveDynamicHandler").field(arg0).finish() + } + } + } +} diff --git a/src/agent-client-protocol-core/src/jsonrpc/handlers.rs b/src/agent-client-protocol-core/src/jsonrpc/handlers.rs new file mode 100644 index 0000000..3575ef0 --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc/handlers.rs @@ -0,0 +1,519 @@ +use crate::jsonrpc::{HandleDispatchFrom, Handled, IntoHandled, JsonRpcResponse}; + +use crate::role::{HasPeer, Role, handle_incoming_dispatch}; +use crate::{ConnectionTo, Dispatch, JsonRpcNotification, JsonRpcRequest, UntypedMessage}; +// Types re-exported from crate root +use super::Responder; +use std::marker::PhantomData; +use std::ops::AsyncFnMut; + +/// Null handler that accepts no messages. +#[derive(Debug)] +pub struct NullHandler; + +impl NullHandler { + /// Creates a new null handler. + #[must_use] + pub fn new() -> Self { + Self + } +} + +impl Default for NullHandler { + fn default() -> Self { + Self::new() + } +} + +impl HandleDispatchFrom for NullHandler { + fn describe_chain(&self) -> impl std::fmt::Debug { + "(null)" + } + + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + _cx: ConnectionTo, + ) -> Result, crate::Error> { + Ok(Handled::No { + message, + retry: false, + }) + } +} + +/// Handler for typed request messages +pub struct RequestHandler< + Counterpart: Role, + Peer: Role, + Req: JsonRpcRequest = UntypedMessage, + F = (), + ToFut = (), +> { + counterpart: Counterpart, + peer: Peer, + handler: F, + to_future_hack: ToFut, + phantom: PhantomData, +} + +impl + RequestHandler +{ + /// Creates a new request handler + pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self { + Self { + counterpart, + peer, + handler, + to_future_hack, + phantom: PhantomData, + } + } +} + +impl HandleDispatchFrom + for RequestHandler +where + Counterpart: HasPeer, + Req: JsonRpcRequest, + F: AsyncFnMut( + Req, + Responder, + ConnectionTo, + ) -> Result + + Send, + T: crate::IntoHandled<(Req, Responder)>, + ToFut: Fn( + &mut F, + Req, + Responder, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, +{ + fn describe_chain(&self) -> impl std::fmt::Debug { + std::any::type_name::() + } + + async fn handle_dispatch_from( + &mut self, + dispatch: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + handle_incoming_dispatch( + self.counterpart.clone(), + self.peer.clone(), + dispatch, + connection, + async |dispatch, connection| { + match dispatch { + Dispatch::Request(message, responder) => { + tracing::debug!( + request_type = std::any::type_name::(), + message = ?message, + "RequestHandler::handle_request" + ); + if Req::matches_method(&message.method) { + match Req::parse_message(&message.method, &message.params) { + Ok(req) => { + tracing::trace!( + ?req, + "RequestHandler::handle_request: parse completed" + ); + let typed_responder = responder.cast(); + let result = (self.to_future_hack)( + &mut self.handler, + req, + typed_responder, + connection, + ) + .await?; + match result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: (request, responder), + retry, + } => { + // Handler returned the request back, convert to untyped + let untyped = request.to_untyped_message()?; + Ok(Handled::No { + message: Dispatch::Request( + untyped, + responder.erase_to_json(), + ), + retry, + }) + } + } + } + Err(err) => { + tracing::trace!( + ?err, + "RequestHandler::handle_request: parse errored" + ); + Err(err) + } + } + } else { + tracing::trace!("RequestHandler::handle_request: method doesn't match"); + Ok(Handled::No { + message: Dispatch::Request(message, responder), + retry: false, + }) + } + } + + Dispatch::Notification(..) | Dispatch::Response(..) => Ok(Handled::No { + message: dispatch, + retry: false, + }), + } + }, + ) + .await + } +} + +/// Handler for typed notification messages +pub struct NotificationHandler< + Counterpart: Role, + Peer: Role, + Notif: JsonRpcNotification = UntypedMessage, + F = (), + ToFut = (), +> { + counterpart: Counterpart, + peer: Peer, + handler: F, + to_future_hack: ToFut, + phantom: PhantomData, +} + +impl + NotificationHandler +{ + /// Creates a new notification handler + pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self { + Self { + counterpart, + peer, + handler, + to_future_hack, + phantom: PhantomData, + } + } +} + +impl HandleDispatchFrom + for NotificationHandler +where + Counterpart: HasPeer, + Notif: JsonRpcNotification, + F: AsyncFnMut(Notif, ConnectionTo) -> Result + Send, + T: crate::IntoHandled<(Notif, ConnectionTo)>, + ToFut: Fn( + &mut F, + Notif, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, +{ + fn describe_chain(&self) -> impl std::fmt::Debug { + std::any::type_name::() + } + + async fn handle_dispatch_from( + &mut self, + dispatch: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + handle_incoming_dispatch( + self.counterpart.clone(), + self.peer.clone(), + dispatch, + connection, + async |dispatch, connection| { + match dispatch { + Dispatch::Notification(message) => { + tracing::debug!( + request_type = std::any::type_name::(), + message = ?message, + "NotificationHandler::handle_dispatch" + ); + if Notif::matches_method(&message.method) { + match Notif::parse_message(&message.method, &message.params) { + Ok(notif) => { + tracing::trace!( + ?notif, + "NotificationHandler::handle_notification: parse completed" + ); + let result = + (self.to_future_hack)(&mut self.handler, notif, connection) + .await?; + match result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: (notification, _cx), + retry, + } => { + // Handler returned the notification back, convert to untyped + let untyped = notification.to_untyped_message()?; + Ok(Handled::No { + message: Dispatch::Notification(untyped), + retry, + }) + } + } + } + Err(err) => { + tracing::trace!( + ?err, + "NotificationHandler::handle_notification: parse errored" + ); + Err(err) + } + } + } else { + tracing::trace!( + "NotificationHandler::handle_notification: method doesn't match" + ); + Ok(Handled::No { + message: Dispatch::Notification(message), + retry: false, + }) + } + } + + Dispatch::Request(..) | Dispatch::Response(..) => Ok(Handled::No { + message: dispatch, + retry: false, + }), + } + }, + ) + .await + } +} + +/// Handler that handles both requests and notifications of specific types. +pub struct MessageHandler< + Counterpart: Role, + Peer: Role, + Req: JsonRpcRequest = UntypedMessage, + Notif: JsonRpcNotification = UntypedMessage, + F = (), + ToFut = (), +> { + counterpart: Counterpart, + peer: Peer, + handler: F, + to_future_hack: ToFut, + phantom: PhantomData)>, +} + +impl + MessageHandler +{ + /// Creates a new message handler + pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self { + Self { + counterpart, + peer, + handler, + to_future_hack, + phantom: PhantomData, + } + } +} + +impl + HandleDispatchFrom for MessageHandler +where + Counterpart: HasPeer, + F: AsyncFnMut(Dispatch, ConnectionTo) -> Result + + Send, + T: IntoHandled>, + ToFut: Fn( + &mut F, + Dispatch, + ConnectionTo, + ) -> crate::BoxFuture<'_, Result> + + Send + + Sync, +{ + fn describe_chain(&self) -> impl std::fmt::Debug { + format!( + "({}, {})", + std::any::type_name::(), + std::any::type_name::() + ) + } + + async fn handle_dispatch_from( + &mut self, + dispatch: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + handle_incoming_dispatch( + self.counterpart.clone(), + self.peer.clone(), + dispatch, + connection, + async |dispatch, connection| match dispatch.into_typed_dispatch::()? { + Ok(typed_dispatch) => { + let result = + (self.to_future_hack)(&mut self.handler, typed_dispatch, connection) + .await?; + match result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: Dispatch::Request(request, responder), + retry, + } => { + let untyped = request.to_untyped_message()?; + Ok(Handled::No { + message: Dispatch::Request(untyped, responder.erase_to_json()), + retry, + }) + } + Handled::No { + message: Dispatch::Notification(notification), + retry, + } => { + let untyped = notification.to_untyped_message()?; + Ok(Handled::No { + message: Dispatch::Notification(untyped), + retry, + }) + } + Handled::No { + message: Dispatch::Response(result, responder), + retry, + } => { + let method = responder.method(); + let untyped_result = match result { + Ok(response) => response.into_json(method).map(Ok), + Err(err) => Ok(Err(err)), + }?; + Ok(Handled::No { + message: Dispatch::Response( + untyped_result, + responder.erase_to_json(), + ), + retry, + }) + } + } + } + + Err(dispatch) => Ok(Handled::No { + message: dispatch, + retry: false, + }), + }, + ) + .await + } +} + +/// Wraps a handler with an optional name for tracing/debugging. +pub struct NamedHandler { + name: Option, + handler: H, +} + +impl NamedHandler { + /// Creates a new named handler + pub fn new(name: Option, handler: H) -> Self { + Self { name, handler } + } +} + +impl> HandleDispatchFrom + for NamedHandler +{ + fn describe_chain(&self) -> impl std::fmt::Debug { + format!( + "NamedHandler({:?}, {:?})", + self.name, + self.handler.describe_chain() + ) + } + + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + if let Some(name) = &self.name { + crate::util::instrumented_with_connection_name( + name.clone(), + self.handler.handle_dispatch_from(message, connection), + ) + .await + } else { + self.handler.handle_dispatch_from(message, connection).await + } + } +} + +/// Chains two handlers together, trying the first handler and falling back to the second +pub struct ChainedHandler { + handler1: H1, + handler2: H2, +} + +impl ChainedHandler { + /// Creates a new chain handler + pub fn new(handler1: H1, handler2: H2) -> Self { + Self { handler1, handler2 } + } +} + +impl HandleDispatchFrom for ChainedHandler +where + H1: HandleDispatchFrom, + H2: HandleDispatchFrom, +{ + fn describe_chain(&self) -> impl std::fmt::Debug { + format!( + "{:?}, {:?}", + self.handler1.describe_chain(), + self.handler2.describe_chain() + ) + } + + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + match self + .handler1 + .handle_dispatch_from(message, connection.clone()) + .await? + { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message, + retry: retry1, + } => match self + .handler2 + .handle_dispatch_from(message, connection) + .await? + { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message, + retry: retry2, + } => Ok(Handled::No { + message, + retry: retry1 | retry2, + }), + }, + } + } +} diff --git a/src/agent-client-protocol-core/src/jsonrpc/incoming_actor.rs b/src/agent-client-protocol-core/src/jsonrpc/incoming_actor.rs new file mode 100644 index 0000000..6d4375b --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc/incoming_actor.rs @@ -0,0 +1,332 @@ +// Types re-exported from crate root +use std::collections::HashMap; + +use futures::StreamExt as _; +use futures::channel::mpsc; +use futures::channel::oneshot; +use futures_concurrency::stream::StreamExt as _; +use rustc_hash::FxHashMap; +use uuid::Uuid; + +use crate::Dispatch; +use crate::RoleId; +use crate::UntypedMessage; +use crate::jsonrpc::ConnectionTo; +use crate::jsonrpc::HandleDispatchFrom; +use crate::jsonrpc::ReplyMessage; +use crate::jsonrpc::Responder; +use crate::jsonrpc::ResponseRouter; +use crate::jsonrpc::dynamic_handler::DynHandleDispatchFrom; +use crate::jsonrpc::dynamic_handler::DynamicHandlerMessage; + +use crate::role::Role; + +use super::Handled; + +struct PendingReply { + method: String, + role_id: RoleId, + sender: oneshot::Sender, +} + +/// Incoming protocol actor: The central dispatch loop for a connection. +/// +/// This actor handles JSON-RPC protocol semantics: +/// - Routes responses to pending request awaiters +/// - Routes requests/notifications to registered handlers +/// - Converts jsonrpcmsg::Request to UntypedMessage for handlers +/// - Manages reply subscriptions from outgoing requests +/// +/// This is the protocol layer - it has no knowledge of how messages arrived. +/// +/// The type parameter `MyRole` is the role of this endpoint (e.g., `Agent`). +/// Messages are received from `MyRole::Counterpart` (e.g., `Client`). +pub(super) async fn incoming_protocol_actor( + counterpart: Counterpart, + connection: &ConnectionTo, + transport_rx: mpsc::UnboundedReceiver>, + dynamic_handler_rx: mpsc::UnboundedReceiver>, + reply_rx: mpsc::UnboundedReceiver, + mut handler: impl HandleDispatchFrom, +) -> Result<(), crate::Error> { + let mut my_rx = transport_rx + .map(IncomingProtocolMsg::Transport) + .merge(dynamic_handler_rx.map(IncomingProtocolMsg::DynamicHandler)) + .merge(reply_rx.map(IncomingProtocolMsg::Reply)); + + let mut dynamic_handlers: FxHashMap>> = + FxHashMap::default(); + let mut pending_messages: Vec = vec![]; + + // Map from request ID to (method, sender) for response dispatch. + // Keys are JSON values because jsonrpcmsg::Id doesn't implement Eq. + // The method is stored to allow routing responses through typed handlers. + let mut pending_replies: HashMap = HashMap::new(); + + while let Some(message_result) = my_rx.next().await { + tracing::trace!(message = ?message_result, actor = "incoming_protocol_actor"); + match message_result { + IncomingProtocolMsg::Reply(message) => match message { + ReplyMessage::Subscribe { + id, + role_id, + method, + sender, + } => { + tracing::trace!(?id, %method, "incoming_actor: subscribing to response"); + let id = serde_json::to_value(&id).unwrap(); + pending_replies.insert( + id, + PendingReply { + method, + role_id, + sender, + }, + ); + } + }, + + IncomingProtocolMsg::DynamicHandler(message) => match message { + DynamicHandlerMessage::AddDynamicHandler(uuid, mut handler) => { + // Before adding the new handler, give it a chance to process + // any pending messages. + let mut new_pending_messages = vec![]; + for pending_message in pending_messages { + tracing::trace!(method = pending_message.method(), handler = ?handler.dyn_describe_chain(), "Retrying message"); + match handler + .dyn_handle_dispatch_from(pending_message, connection.clone()) + .await? + { + Handled::Yes => { + tracing::trace!("Message handled"); + } + Handled::No { + message: m, + retry: _, + } => { + tracing::trace!(method = m.method(), handler = ?handler.dyn_describe_chain(), "Message not handled"); + new_pending_messages.push(m); + } + } + } + pending_messages = new_pending_messages; + + // Add handler so it will be used for future incoming messages. + dynamic_handlers.insert(uuid, handler); + } + DynamicHandlerMessage::RemoveDynamicHandler(uuid) => { + dynamic_handlers.remove(&uuid); + } + }, + + IncomingProtocolMsg::Transport(message) => match message { + Ok(message) => match message { + jsonrpcmsg::Message::Request(request) => { + tracing::trace!(method = %request.method, id = ?request.id, "Handling request"); + let dispatch = dispatch_from_request(connection, request); + dispatch_dispatch( + counterpart.clone(), + connection, + dispatch, + &mut dynamic_handlers, + &mut handler, + &mut pending_messages, + ) + .await?; + } + jsonrpcmsg::Message::Response(response) => { + tracing::trace!(id = ?response.id, has_result = response.result.is_some(), has_error = response.error.is_some(), "Handling response"); + if let Some(id) = response.id { + let result = if let Some(value) = response.result { + Ok(value) + } else if let Some(error) = response.error { + // Convert jsonrpcmsg::Error to crate::Error + Err(crate::Error::new(error.code, error.message).data(error.data)) + } else { + // Response with neither result nor error - treat as null result + Ok(serde_json::Value::Null) + }; + + let id_json = serde_json::to_value(&id).unwrap(); + if let Some(pending_reply) = pending_replies.remove(&id_json) { + // Route the response through the handler chain + let dispatch = dispatch_from_response(id, pending_reply, result); + dispatch_dispatch( + counterpart.clone(), + connection, + dispatch, + &mut dynamic_handlers, + &mut handler, + &mut pending_messages, + ) + .await?; + } else { + tracing::warn!( + ?id, + "incoming_actor: received response for unknown id, no subscriber found" + ); + } + } + } + }, + Err(error) => { + // Parse error from transport - send error notification back to remote + tracing::warn!(?error, "Transport parse error, sending error notification"); + connection.send_error_notification(error)?; + } + }, + } + } + Ok(()) +} + +#[derive(Debug)] +enum IncomingProtocolMsg { + Transport(Result), + DynamicHandler(DynamicHandlerMessage), + Reply(ReplyMessage), +} + +/// Dispatches a JSON-RPC request to the handler. +/// Report an error back to the server if it does not get handled. +fn dispatch_from_request( + connection: &ConnectionTo, + request: jsonrpcmsg::Request, +) -> Dispatch { + let message = UntypedMessage::new(&request.method, &request.params).expect("well-formed JSON"); + + match &request.id { + Some(id) => Dispatch::Request( + message, + Responder::new( + connection.message_tx.clone(), + request.method.clone(), + id.clone(), + ), + ), + None => Dispatch::Notification(message), + } +} + +/// Dispatches a JSON-RPC response through the handler chain. +/// +/// This allows handlers to intercept and process responses before they reach +/// the awaiting code. The default behavior is to forward the response to the +/// local awaiter via the oneshot channel. +fn dispatch_from_response( + id: jsonrpcmsg::Id, + pending_reply: PendingReply, + result: Result, +) -> Dispatch { + let PendingReply { + method, + role_id, + sender, + } = pending_reply; + + // Create a Dispatch::Response with a ResponseRouter that routes to the oneshot + let router = ResponseRouter::new(method.clone(), id.clone(), role_id, sender); + Dispatch::Response(result, router) +} + +#[tracing::instrument( + skip(connection, dispatch, dynamic_handlers, handler, pending_messages), + fields(method = dispatch.method()), + level = "trace", +)] +async fn dispatch_dispatch( + counterpart: Counterpart, + connection: &ConnectionTo, + mut dispatch: Dispatch, + dynamic_handlers: &mut FxHashMap>>, + handler: &mut impl HandleDispatchFrom, + pending_messages: &mut Vec, +) -> Result<(), crate::Error> { + tracing::trace!(?dispatch, "dispatch_dispatch"); + + let mut retry_any = false; + + let id = dispatch.id(); + let method = dispatch.method().to_string(); + + // First, apply the handlers given by the user. + tracing::trace!(handler = ?handler.describe_chain(), "Attempting handler chain"); + match handler + .handle_dispatch_from(dispatch, connection.clone()) + .await? + { + Handled::Yes => { + tracing::trace!(?method, ?id, handler = ?handler.describe_chain(), "Handler accepted message"); + return Ok(()); + } + + Handled::No { message: m, retry } => { + tracing::trace!(?method, ?id, handler = ?handler.describe_chain(), "Handler declined message"); + dispatch = m; + retry_any |= retry; + } + } + + // Next, apply any dynamic handlers. + for dynamic_handler in dynamic_handlers.values_mut() { + tracing::trace!(handler = ?dynamic_handler.dyn_describe_chain(), "Attempting dynamic handler"); + match dynamic_handler + .dyn_handle_dispatch_from(dispatch, connection.clone()) + .await? + { + Handled::Yes => { + tracing::trace!(?method, ?id, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler accepted message"); + return Ok(()); + } + + Handled::No { message: m, retry } => { + tracing::trace!(?method, ?id, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler declined message"); + retry_any |= retry; + dispatch = m; + } + } + } + + // Finally, apply the default handler for the role. + tracing::trace!(role = ?counterpart, "Attempting default handler"); + match counterpart + .default_handle_dispatch_from(dispatch, connection.clone()) + .await? + { + Handled::Yes => { + tracing::trace!(?method, handler = "default", "Role accepted message"); + return Ok(()); + } + Handled::No { message: m, retry } => { + tracing::trace!(?method, handler = "default", "Role declined message"); + dispatch = m; + retry_any |= retry; + } + } + + // If the message was never handled, check whether the retry flag was set. + // If so, enqueue it for later processing. Else, reject it. + if retry_any { + tracing::debug!( + ?method, + "Retrying message as new dynamic handlers are added" + ); + pending_messages.push(dispatch); + Ok(()) + } else { + match dispatch { + Dispatch::Request(..) | Dispatch::Notification(_) => { + tracing::info!(?method, "Rejecting message with error, no handler"); + let method = dispatch.method().to_string(); + dispatch.respond_with_error( + crate::Error::method_not_found().data(method), + connection.clone(), + ) + } + Dispatch::Response(result, router) => { + tracing::trace!(?method, "Forwarding response"); + router.respond_with_result(result) + } + } + } +} diff --git a/src/agent-client-protocol-core/src/jsonrpc/outgoing_actor.rs b/src/agent-client-protocol-core/src/jsonrpc/outgoing_actor.rs new file mode 100644 index 0000000..7eee713 --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc/outgoing_actor.rs @@ -0,0 +1,101 @@ +// Types re-exported from crate root +use futures::StreamExt as _; +use futures::channel::mpsc; + +use crate::jsonrpc::OutgoingMessage; +use crate::jsonrpc::ReplyMessage; + +pub type OutgoingMessageTx = mpsc::UnboundedSender; + +pub(crate) fn send_raw_message( + tx: &OutgoingMessageTx, + message: OutgoingMessage, +) -> Result<(), crate::Error> { + tracing::debug!(?message, ?tx, "send_raw_message"); + tx.unbounded_send(message) + .map_err(crate::util::internal_error) +} + +/// Outgoing protocol actor: Converts application-level OutgoingMessage to protocol-level jsonrpcmsg::Message. +/// +/// This actor handles JSON-RPC protocol semantics: +/// - Subscribes to reply_actor for response correlation +/// - Converts OutgoingMessage variants to jsonrpcmsg::Message +/// +/// This is the protocol layer - it has no knowledge of how messages are transported. +pub(super) async fn outgoing_protocol_actor( + mut outgoing_rx: mpsc::UnboundedReceiver, + reply_tx: mpsc::UnboundedSender, + transport_tx: mpsc::UnboundedSender>, +) -> Result<(), crate::Error> { + while let Some(message) = outgoing_rx.next().await { + tracing::debug!(?message, "outgoing_protocol_actor"); + + // Create the message to be sent over the transport + let json_rpc_message = match message { + OutgoingMessage::Request { + id, + role_id, + method, + untyped, + response_tx, + } => { + // Record where the reply should be sent once it arrives. + reply_tx + .unbounded_send(ReplyMessage::Subscribe { + id: id.clone(), + role_id, + method, + sender: response_tx, + }) + .map_err(crate::Error::into_internal_error)?; + + jsonrpcmsg::Message::Request(untyped.into_jsonrpc_msg(Some(id))?) + } + OutgoingMessage::Notification { untyped } => { + let msg = untyped.into_jsonrpc_msg(None)?; + jsonrpcmsg::Message::Request(msg) + } + OutgoingMessage::Response { + id, + response: Ok(value), + } => { + tracing::debug!(?id, "Sending success response"); + jsonrpcmsg::Message::Response(jsonrpcmsg::Response::success_v2(value, Some(id))) + } + OutgoingMessage::Response { + id, + response: Err(error), + } => { + tracing::warn!(?id, ?error, "Sending error response"); + // Convert crate::Error to jsonrpcmsg::Error + let jsonrpc_error = jsonrpcmsg::Error { + code: error.code.into(), + message: error.message, + data: error.data, + }; + jsonrpcmsg::Message::Response(jsonrpcmsg::Response::error_v2( + jsonrpc_error, + Some(id), + )) + } + OutgoingMessage::Error { error } => { + // Convert crate::Error to jsonrpcmsg::Error + let jsonrpc_error = jsonrpcmsg::Error { + code: error.code.into(), + message: error.message, + data: error.data, + }; + // Response with id: None means this is an error notification that couldn't be + // correlated to a specific request (e.g., parse error before we could read the id) + jsonrpcmsg::Message::Response(jsonrpcmsg::Response::error_v2(jsonrpc_error, None)) + } + }; + + // Send to transport layer (wrapped in Ok since transport expects Result) + transport_tx + .unbounded_send(Ok(json_rpc_message)) + .map_err(crate::Error::into_internal_error)?; + } + Ok(()) +} diff --git a/src/agent-client-protocol-core/src/jsonrpc/run.rs b/src/agent-client-protocol-core/src/jsonrpc/run.rs new file mode 100644 index 0000000..ff07f62 --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc/run.rs @@ -0,0 +1,104 @@ +//! Run trait for background tasks that run alongside a connection. +//! +//! Run implementations are composable background tasks that run while a connection is active. +//! They're used for things like MCP tool handlers that need to receive calls through +//! channels and invoke user-provided closures. + +use std::future::Future; + +use crate::{ConnectionTo, role::Role}; + +/// A background task that runs alongside a connection. +/// +/// `RunIn` means "run in the context of being role R". The task receives +/// a `ConnectionTo` for communicating with the other side. +/// +/// Implementations are composed using [`ChainRun`] and run in parallel +/// when the connection is active. +pub trait RunWithConnectionTo: Send { + /// Run this task to completion. + fn run_with_connection_to( + self, + cx: ConnectionTo, + ) -> impl Future> + Send; +} + +/// A no-op RunIn that completes immediately. +#[derive(Debug, Default)] +pub struct NullRun; + +impl RunWithConnectionTo for NullRun { + async fn run_with_connection_to( + self, + _cx: ConnectionTo, + ) -> Result<(), crate::Error> { + Ok(()) + } +} + +/// Chains two RunIn implementations to run in parallel. +#[derive(Debug)] +pub struct ChainRun { + a: A, + b: B, +} + +impl ChainRun { + /// Create a new chained RunIn from two RunIn implementations. + pub fn new(a: A, b: B) -> Self { + Self { a, b } + } +} + +impl RunWithConnectionTo for ChainRun +where + A: RunWithConnectionTo, + B: RunWithConnectionTo, +{ + async fn run_with_connection_to( + self, + cx: ConnectionTo, + ) -> Result<(), crate::Error> { + // Box the futures to avoid stack overflow with deeply nested RunIn chains + let a_fut = Box::pin(self.a.run_with_connection_to(cx.clone())); + let b_fut = Box::pin(self.b.run_with_connection_to(cx.clone())); + let ((), ()) = futures::future::try_join(a_fut, b_fut).await?; + Ok(()) + } +} + +/// A RunIn created from a closure via [`with_spawned`](crate::Builder::with_spawned). +pub struct SpawnedRun { + task_fn: F, + location: &'static std::panic::Location<'static>, +} + +impl SpawnedRun { + /// Create a new spawned RunIn from a closure. + pub fn new(location: &'static std::panic::Location<'static>, task_fn: F) -> Self { + Self { task_fn, location } + } +} + +impl RunWithConnectionTo for SpawnedRun +where + Counterpart: Role, + F: FnOnce(ConnectionTo) -> Fut + Send, + Fut: Future> + Send, +{ + async fn run_with_connection_to( + self, + connection: ConnectionTo, + ) -> Result<(), crate::Error> { + let location = self.location; + (self.task_fn)(connection).await.map_err(|err| { + let data = err.data.clone(); + err.data(serde_json::json! { + { + "spawned_at": format!("{}:{}:{}", location.file(), location.line(), location.column()), + "data": data, + } + }) + }) + } +} diff --git a/src/agent-client-protocol-core/src/jsonrpc/task_actor.rs b/src/agent-client-protocol-core/src/jsonrpc/task_actor.rs new file mode 100644 index 0000000..a885f15 --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc/task_actor.rs @@ -0,0 +1,61 @@ +use std::panic::Location; + +use futures::{FutureExt, channel::mpsc, future::BoxFuture}; + +use crate::ConnectionTo; +use crate::role::Role; +use crate::util::process_stream_concurrently; + +pub type TaskTx = mpsc::UnboundedSender; + +#[must_use] +pub(crate) struct Task { + future: BoxFuture<'static, Result<(), crate::Error>>, +} + +impl Task { + pub fn new( + location: &'static Location<'static>, + task_future: impl IntoFuture, IntoFuture: Send + 'static>, + ) -> Self { + let task_future = task_future.into_future(); + Task { + future: futures::FutureExt::map( + task_future, + |result| match result { + Ok(()) => Ok(()), + Err(err) => { + let data = err.data.clone(); + Err(err.data(serde_json::json! { + { + "spawned_at": format!("{}:{}:{}", location.file(), location.line(), location.column()), + "data": data, + } + })) + } + }, + ) + .boxed() + } + } + + pub fn spawn(self, task_tx: &TaskTx) -> Result<(), crate::Error> { + task_tx + .unbounded_send(self) + .map_err(crate::util::internal_error)?; + Ok(()) + } +} + +/// The "task actor" manages dynamically spawned tasks. +pub(super) async fn task_actor( + task_rx: mpsc::UnboundedReceiver, + _cx: &ConnectionTo, +) -> Result<(), crate::Error> { + process_stream_concurrently( + task_rx, + async |task| task.future.await, + |a, b| Box::pin(a(b)), + ) + .await +} diff --git a/src/agent-client-protocol-core/src/jsonrpc/transport_actor.rs b/src/agent-client-protocol-core/src/jsonrpc/transport_actor.rs new file mode 100644 index 0000000..c252f1f --- /dev/null +++ b/src/agent-client-protocol-core/src/jsonrpc/transport_actor.rs @@ -0,0 +1,119 @@ +use std::pin::pin; + +// Types re-exported from crate root +use futures::StreamExt as _; +use futures::channel::mpsc; + +/// Transport outgoing actor for line streams: Serializes jsonrpcmsg::Message and yields lines. +/// +/// This is a line-based variant of `transport_outgoing_actor` that works with a Sink +/// instead of an AsyncWrite byte stream. This enables interception of lines before they are +/// written to the underlying transport. +/// +/// This actor handles transport mechanics: +/// - Unwraps Result from the channel +/// - Serializes jsonrpcmsg::Message to JSON strings +/// - Yields newline-terminated strings +/// - Handles serialization errors +/// +/// This is the transport layer - it has no knowledge of protocol semantics (IDs, correlation, etc.). +pub(super) async fn transport_outgoing_lines_actor( + mut transport_rx: mpsc::UnboundedReceiver>, + outgoing_lines: impl futures::Sink, +) -> Result<(), crate::Error> { + use futures::SinkExt; + let mut outgoing_lines = pin!(outgoing_lines); + + while let Some(message_result) = transport_rx.next().await { + // Unwrap the Result - errors here would be from the channel itself + let json_rpc_message = message_result?; + match serde_json::to_string(&json_rpc_message) { + Ok(line) => { + tracing::trace!(message = %line, "Sending JSON-RPC message"); + outgoing_lines + .send(line) + .await + .map_err(crate::Error::into_internal_error)?; + } + + Err(serialization_error) => { + match json_rpc_message { + jsonrpcmsg::Message::Request(_request) => { + // If we failed to serialize a request, + // just ignore it. + // + // Q: (Maybe it'd be nice to "reply" with an error?) + tracing::error!( + ?serialization_error, + "Failed to serialize request, ignoring" + ); + } + jsonrpcmsg::Message::Response(response) => { + // If we failed to serialize a *response*, + // send an error in response. + tracing::error!(?serialization_error, id = ?response.id, "Failed to serialize response, sending internal_error instead"); + // Convert crate::Error to jsonrpcmsg::Error + let acp_error = crate::Error::internal_error(); + let jsonrpc_error = jsonrpcmsg::Error { + code: acp_error.code.into(), + message: acp_error.message, + data: acp_error.data, + }; + let error_line = serde_json::to_string(&jsonrpcmsg::Response::error( + jsonrpc_error, + response.id, + )) + .unwrap(); + outgoing_lines + .send(error_line) + .await + .map_err(crate::Error::into_internal_error)?; + } + } + } + } + } + Ok(()) +} + +/// Transport incoming actor for line streams: Parses lines into jsonrpcmsg::Message. +/// +/// This is a line-based variant of `transport_incoming_actor` that works with a +/// Stream> instead of an AsyncRead byte stream. This enables +/// interception of lines before they are parsed. +/// +/// This actor handles transport mechanics: +/// - Reads lines from the stream +/// - Parses to jsonrpcmsg::Message +/// - Handles parse errors +/// +/// This is the transport layer - it has no knowledge of protocol semantics. +pub(super) async fn transport_incoming_lines_actor( + incoming_lines: impl futures::Stream>, + transport_tx: mpsc::UnboundedSender>, +) -> Result<(), crate::Error> { + let mut incoming_lines = pin!(incoming_lines); + while let Some(line_result) = incoming_lines.next().await { + let line = line_result.map_err(crate::Error::into_internal_error)?; + tracing::trace!(message = %line, "Received JSON-RPC message"); + + let message: Result = serde_json::from_str(&line); + match message { + Ok(msg) => { + transport_tx + .unbounded_send(Ok(msg)) + .map_err(crate::Error::into_internal_error)?; + } + Err(_) => { + transport_tx + .unbounded_send(Err(crate::Error::parse_error().data(serde_json::json!( + { + "line": &line + } + )))) + .map_err(crate::Error::into_internal_error)?; + } + } + } + Ok(()) +} diff --git a/src/agent-client-protocol-core/src/lib.rs b/src/agent-client-protocol-core/src/lib.rs new file mode 100644 index 0000000..4193d1e --- /dev/null +++ b/src/agent-client-protocol-core/src/lib.rs @@ -0,0 +1,210 @@ +#![deny(missing_docs)] + +//! # agent-client-protocol-core -- the Symposium Agent Client Protocol (ACP) SDK +//! +//! **agent-client-protocol-core** is a Rust SDK for building [Agent-Client Protocol (ACP)][acp] applications. +//! ACP is a protocol for communication between AI agents and their clients (IDEs, CLIs, etc.), +//! enabling features like tool use, permission requests, and streaming responses. +//! +//! [acp]: https://agentclientprotocol.com/ +//! +//! ## What can you build with agent-client-protocol-core? +//! +//! - **Clients** that talk to ACP agents (like building your own Claude Code interface) +//! - **Proxies** that add capabilities to existing agents (like adding custom tools via MCP) +//! - **Agents** that respond to prompts with AI-powered responses +//! +//! ## Quick Start: Connecting to an Agent +//! +//! The most common use case is connecting to an existing ACP agent as a client. +//! Here's a minimal example that initializes a connection, creates a session, +//! and sends a prompt: +//! +//! ```no_run +//! use agent_client_protocol_core::Client; +//! use agent_client_protocol_core::schema::{InitializeRequest, ProtocolVersion}; +//! +//! # async fn run(transport: impl agent_client_protocol_core::ConnectTo) -> Result<(), agent_client_protocol_core::Error> { +//! Client.builder() +//! .name("my-client") +//! .connect_with(transport, async |cx| { +//! // Step 1: Initialize the connection +//! cx.send_request(InitializeRequest::new(ProtocolVersion::LATEST)) +//! .block_task().await?; +//! +//! // Step 2: Create a session and send a prompt +//! cx.build_session_cwd()? +//! .block_task() +//! .run_until(async |mut session| { +//! session.send_prompt("What is 2 + 2?")?; +//! let response = session.read_to_string().await?; +//! println!("{}", response); +//! Ok(()) +//! }) +//! .await +//! }) +//! .await +//! # } +//! ``` +//! +//! For a complete working example, see [`yolo_one_shot_client.rs`][yolo]. +//! +//! [yolo]: https://github.com/agentclientprotocol/rust-sdk/blob/main/src/agent-client-protocol-core/examples/yolo_one_shot_client.rs +//! +//! ## Cookbook +//! +//! The [`agent_client_protocol_cookbook`] crate contains practical guides and examples: +//! +//! - Connecting as a client +//! - Global MCP server +//! - Per-session MCP server with workspace context +//! - Building agents and reusable components +//! - Running proxies with the conductor +//! +//! [`agent_client_protocol_cookbook`]: https://docs.rs/agent-client-protocol-cookbook +//! +//! ## Core Concepts +//! +//! The [`concepts`] module provides detailed explanations of how agent-client-protocol-core works, +//! including connections, sessions, callbacks, ordering guarantees, and more. +//! +//! ## Related Crates +//! +//! - [`agent-client-protocol-tokio`] - Tokio utilities for spawning agent processes +//! - [`agent-client-protocol-conductor`] - Binary for running proxy chains +//! +//! [`agent-client-protocol-tokio`]: https://crates.io/crates/agent-client-protocol-tokio +//! [`agent-client-protocol-conductor`]: https://crates.io/crates/agent-client-protocol-conductor + +/// Capability management for the `_meta.symposium` object +mod capabilities; +/// Component abstraction for agents and proxies +pub mod component; +/// Core concepts for understanding and using agent-client-protocol-core +pub mod concepts; +/// Cookbook of common patterns for building ACP components +pub mod cookbook; +/// JSON-RPC handler types for building custom message handlers +pub mod handler; +/// JSON-RPC connection and handler infrastructure +mod jsonrpc; +/// MCP server support for providing MCP tools over ACP +pub mod mcp_server; +/// Role types for ACP connections +pub mod role; +/// ACP protocol schema types - all message types, requests, responses, and supporting types +pub mod schema; +/// Utility functions and types +pub mod util; + +pub use capabilities::*; + +/// JSON-RPC message types. +/// +/// This module re-exports types from the `jsonrpcmsg` crate that are transitively +/// reachable through the public API (e.g., via [`Channel`]). +/// +/// Users of the `agent-client-protocol-core` crate can use these types without adding a direct dependency +/// on `jsonrpcmsg`. +pub mod jsonrpcmsg { + pub use jsonrpcmsg::{Error, Id, Message, Params, Request, Response}; +} + +pub use jsonrpc::{ + Builder, ByteStreams, Channel, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, + IntoHandled, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Lines, + NullHandler, Responder, ResponseRouter, SentRequest, UntypedMessage, + run::{ChainRun, NullRun, RunWithConnectionTo}, +}; + +pub use role::{ + Role, RoleId, UntypedRole, + acp::{Agent, Client, Conductor, Proxy}, +}; + +pub use component::{ConnectTo, DynConnectTo}; + +// Re-export BoxFuture for implementing Component traits +pub use futures::future::BoxFuture; + +// Re-export the six primary message enum types at the root +pub use schema::{ + AgentNotification, AgentRequest, AgentResponse, ClientNotification, ClientRequest, + ClientResponse, +}; + +// Re-export commonly used infrastructure types for convenience +pub use schema::{Error, ErrorCode}; + +// Re-export derive macros for custom JSON-RPC types +pub use agent_client_protocol_derive::{JsonRpcNotification, JsonRpcRequest, JsonRpcResponse}; + +mod session; +pub use session::*; + +/// This is a hack that must be given as the final argument of +/// [`McpServerBuilder::tool_fn`](`crate::mcp_server::McpServerBuilder::tool_fn`) when defining tools. +/// Look away, lest ye be blinded by its vileness! +/// +/// Fine, if you MUST know, it's a horrific workaround for not having +/// [return-type notation](https://github.com/rust-lang/rust/issues/109417) +/// and for [this !@$#!%! bug](https://github.com/rust-lang/rust/issues/110338). +/// Trust me, the need for it hurts me more than it hurts you. --nikomatsakis +#[macro_export] +macro_rules! tool_fn_mut { + () => { + |func, params, context| Box::pin(func(params, context)) + }; +} + +/// This is a hack that must be given as the final argument of +/// [`McpServerBuilder::tool_fn`](`crate::mcp_server::McpServerBuilder::tool_fn`) when defining stateless concurrent tools. +/// See [`tool_fn_mut!`] for the gory details. +#[macro_export] +macro_rules! tool_fn { + () => { + |func, params, context| Box::pin(func(params, context)) + }; +} + +/// This macro is used for the value of the `to_future_hack` parameter of +/// [`Builder::on_receive_request`] and [`Builder::on_receive_request_from`]. +/// +/// It expands to `|f, req, responder, cx| Box::pin(f(req, responder, cx))`. +/// +/// This is needed until [return-type notation](https://github.com/rust-lang/rust/issues/109417) +/// is stabilized. +#[macro_export] +macro_rules! on_receive_request { + () => { + |f: &mut _, req, responder, cx| Box::pin(f(req, responder, cx)) + }; +} + +/// This macro is used for the value of the `to_future_hack` parameter of +/// [`Builder::on_receive_notification`] and [`Builder::on_receive_notification_from`]. +/// +/// It expands to `|f, notif, cx| Box::pin(f(notif, cx))`. +/// +/// This is needed until [return-type notation](https://github.com/rust-lang/rust/issues/109417) +/// is stabilized. +#[macro_export] +macro_rules! on_receive_notification { + () => { + |f: &mut _, notif, cx| Box::pin(f(notif, cx)) + }; +} + +/// This macro is used for the value of the `to_future_hack` parameter of +/// [`Builder::on_receive_dispatch`] and [`Builder::on_receive_dispatch_from`]. +/// +/// It expands to `|f, dispatch, cx| Box::pin(f(dispatch, cx))`. +/// +/// This is needed until [return-type notation](https://github.com/rust-lang/rust/issues/109417) +/// is stabilized. +#[macro_export] +macro_rules! on_receive_dispatch { + () => { + |f: &mut _, dispatch, cx| Box::pin(f(dispatch, cx)) + }; +} diff --git a/src/agent-client-protocol-core/src/mcp_server/active_session.rs b/src/agent-client-protocol-core/src/mcp_server/active_session.rs new file mode 100644 index 0000000..7bb1c5c --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/active_session.rs @@ -0,0 +1,247 @@ +use futures::channel::mpsc; +use futures::{SinkExt, StreamExt}; +use rustc_hash::FxHashMap; + +use crate::mcp_server::{McpConnectionTo, McpServerConnect}; +use crate::role; +use crate::role::HasPeer; +use crate::schema::{ + McpConnectRequest, McpConnectResponse, McpDisconnectNotification, McpOverAcpMessage, +}; +use crate::util::MatchDispatchFrom; +use crate::{ + Agent, Channel, ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Responder, + Role, UntypedMessage, +}; +use std::sync::Arc; + +/// The message handler for an MCP server offered to a particular session. +/// This is added as a 'dynamic' handler to the connection context +/// (see [`ConnectionTo::add_dynamic_handler`]) and handles MCP-over-ACP messages +/// with the appropriate ACP url. +pub(super) struct McpActiveSession { + /// The ACP URL created for this session + acp_url: String, + + /// The MCP server we are managing + mcp_connect: Arc>, + + /// Active connections to MCP server tasks + connections: FxHashMap>, +} + +impl McpActiveSession +where + Counterpart: HasPeer, +{ + pub fn new(acp_url: String, mcp_connect: Arc>) -> Self { + Self { + acp_url, + mcp_connect, + connections: FxHashMap::default(), + } + } + + /// Handle connection requests for our MCP server by creating a new connection. + /// A *connection* is an actual running instance of this MCP server. + fn handle_connect_request( + &mut self, + request: McpConnectRequest, + responder: Responder, + acp_connection: &ConnectionTo, + ) -> Result)>, crate::Error> { + // Check that this is for our MCP server + if request.acp_url != self.acp_url { + return Ok(Handled::No { + message: (request, responder), + retry: false, + }); + } + + // Create a unique connection ID and a channel for future communication + let connection_id = format!("mcp-over-acp-connection:{}", uuid::Uuid::new_v4()); + let (mcp_server_tx, mut mcp_server_rx) = mpsc::channel(128); + self.connections + .insert(connection_id.clone(), mcp_server_tx); + + // Create connected channel pair for client-server communication + let (client_channel, server_channel) = Channel::duplex(); + + // Create client-side handler that wraps messages and forwards to successor + let client_component = { + let connection_id = connection_id.clone(); + let acp_connection = acp_connection.clone(); + + role::mcp::Client + .builder() + .on_receive_dispatch( + async move |message: Dispatch, _mcp_connection| { + // Wrap the message in McpOverAcp{Request,Notification} and forward to successor + let wrapped = message.map( + |request, responder| { + ( + McpOverAcpMessage { + connection_id: connection_id.clone(), + message: request, + meta: None, + }, + responder, + ) + }, + |notification| McpOverAcpMessage { + connection_id: connection_id.clone(), + message: notification, + meta: None, + }, + ); + acp_connection.send_proxied_message_to(Agent, wrapped) + }, + crate::on_receive_dispatch!(), + ) + .with_spawned(move |mcp_connection| async move { + // Messages we pull off this channel were sent from the agent. + // Forward them back to the MCP server. + while let Some(msg) = mcp_server_rx.next().await { + mcp_connection.send_proxied_message_to(role::mcp::Server, msg)?; + } + Ok(()) + }) + }; + + // Get the MCP server component + let spawned_server = self.mcp_connect.connect(McpConnectionTo { + acp_url: request.acp_url.clone(), + connection: acp_connection.clone(), + }); + + // Spawn both sides of the connection + let spawn_results = acp_connection + .spawn(async move { client_component.connect_to(client_channel).await }) + .and_then(|()| { + // Spawn the MCP server serving the server channel + acp_connection.spawn(async move { spawned_server.connect_to(server_channel).await }) + }); + + match spawn_results { + Ok(()) => { + responder.respond(McpConnectResponse { + connection_id, + meta: None, + })?; + Ok(Handled::Yes) + } + + Err(err) => { + responder.respond_with_error(err)?; + Ok(Handled::Yes) + } + } + } + + /// Forward MCP-over-ACP requests to the connection. + async fn handle_mcp_over_acp_request( + &mut self, + request: McpOverAcpMessage, + responder: Responder, + ) -> Result< + Handled<( + McpOverAcpMessage, + Responder, + )>, + crate::Error, + > { + // Check if we have a registered server with the given URL. If not, don't try to handle the request. + let Some(mcp_server_tx) = self.connections.get_mut(&request.connection_id) else { + return Ok(Handled::No { + message: (request, responder), + retry: false, + }); + }; + + mcp_server_tx + .send(Dispatch::Request(request.message, responder)) + .await + .map_err(crate::Error::into_internal_error)?; + + Ok(Handled::Yes) + } + + /// Forward MCP-over-ACP notifications to the connection. + async fn handle_mcp_over_acp_notification( + &mut self, + notification: McpOverAcpMessage, + ) -> Result>, crate::Error> { + // Check if we have a registered server with the given URL. If not, don't try to handle the request. + let Some(mcp_server_tx) = self.connections.get_mut(¬ification.connection_id) else { + return Ok(Handled::No { + message: notification, + retry: false, + }); + }; + + mcp_server_tx + .send(Dispatch::Notification(notification.message)) + .await + .map_err(crate::Error::into_internal_error)?; + + Ok(Handled::Yes) + } + + /// Disconnect a connection. + fn handle_mcp_disconnect_notification( + &mut self, + successor_notification: McpDisconnectNotification, + ) -> Handled { + // Remove connection if we have it. Otherwise, do not handle the notification. + if self + .connections + .remove(&successor_notification.connection_id) + .is_some() + { + Handled::Yes + } else { + Handled::No { + message: successor_notification, + retry: false, + } + } + } +} + +impl HandleDispatchFrom for McpActiveSession +where + Counterpart: HasPeer, +{ + fn describe_chain(&self) -> impl std::fmt::Debug { + "McpServerSession" + } + + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + MatchDispatchFrom::new(message, &connection) + // MCP connect requests come from the Agent direction (wrapped in SuccessorMessage) + .if_request_from(Agent, async |request, responder| { + self.handle_connect_request(request, responder, &connection) + }) + .await + // MCP over ACP requests come from the Agent direction + .if_request_from(Agent, async |request, responder| { + self.handle_mcp_over_acp_request(request, responder).await + }) + .await + // MCP over ACP notifications come from the Agent direction + .if_notification_from(Agent, async |notification| { + self.handle_mcp_over_acp_notification(notification).await + }) + .await + // MCP disconnect notifications come from the Agent direction + .if_notification_from(Agent, async |notification| { + Ok(self.handle_mcp_disconnect_notification(notification)) + }) + .await + .done() + } +} diff --git a/src/agent-client-protocol-core/src/mcp_server/builder.rs b/src/agent-client-protocol-core/src/mcp_server/builder.rs new file mode 100644 index 0000000..6bd9891 --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/builder.rs @@ -0,0 +1,629 @@ +//! MCP server builder for creating MCP servers. + +use std::{collections::HashSet, marker::PhantomData, pin::pin, sync::Arc}; + +use futures::{ + SinkExt, + channel::{mpsc, oneshot}, + future::{BoxFuture, Either}, +}; +use futures_concurrency::future::TryJoin; +use rustc_hash::FxHashMap; + +/// Tracks which tools are enabled. +/// +/// - `DenyList`: All tools enabled except those in the set (default) +/// - `AllowList`: Only tools in the set are enabled +#[derive(Clone, Debug)] +pub enum EnabledTools { + /// All tools enabled except those in the deny set. + DenyList(HashSet), + /// Only tools in the allow set are enabled. + AllowList(HashSet), +} + +impl Default for EnabledTools { + fn default() -> Self { + EnabledTools::DenyList(HashSet::new()) + } +} + +impl EnabledTools { + /// Check if a tool is enabled. + #[must_use] + pub fn is_enabled(&self, name: &str) -> bool { + match self { + EnabledTools::DenyList(deny) => !deny.contains(name), + EnabledTools::AllowList(allow) => allow.contains(name), + } + } +} +use rmcp::{ + ErrorData, ServerHandler, + handler::server::tool::{schema_for_output, schema_for_type}, + model::{CallToolResult, ListToolsResult, Tool}, +}; +use schemars::JsonSchema; +use serde::{Serialize, de::DeserializeOwned}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +use super::{McpConnectionTo, McpTool}; +use crate::{ + ByteStreams, ConnectTo, DynConnectTo, + jsonrpc::run::{ChainRun, NullRun, RunWithConnectionTo}, + mcp_server::{ + McpServer, McpServerConnect, + responder::{ToolCall, ToolFnMutResponder, ToolFnResponder}, + }, + role::{self, Role}, +}; + +/// Builder for creating MCP servers with tools. +/// +/// Use [`McpServer::builder`] to create a new builder, then chain methods to +/// configure the server and call [`build`](Self::build) to create the server. +/// +/// # Example +/// +/// ```rust,ignore +/// let server = McpServer::builder("my-server".to_string()) +/// .instructions("A helpful assistant") +/// .tool(EchoTool) +/// .tool_fn( +/// "greet", +/// "Greet someone by name", +/// async |input: GreetInput, _cx| Ok(format!("Hello, {}!", input.name)), +/// agent_client_protocol_core::tool_fn!(), +/// ) +/// .build(); +/// ``` +#[derive(Debug)] +pub struct McpServerBuilder +where + Responder: RunWithConnectionTo, +{ + phantom: PhantomData, + name: String, + data: McpServerData, + responder: Responder, +} + +#[derive(Debug)] +struct McpServerData { + instructions: Option, + tool_models: Vec, + tools: FxHashMap>, + enabled_tools: EnabledTools, +} + +/// A registered tool with its metadata. +struct RegisteredTool { + tool: Arc>, + /// Whether this tool returns structured output (i.e., has an output_schema). + has_structured_output: bool, +} + +impl std::fmt::Debug for RegisteredTool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RegisteredTool") + .field("has_structured_output", &self.has_structured_output) + .finish_non_exhaustive() + } +} + +impl Default for McpServerData { + fn default() -> Self { + Self { + instructions: None, + tool_models: Vec::new(), + tools: FxHashMap::default(), + enabled_tools: EnabledTools::default(), + } + } +} + +impl McpServerBuilder { + pub(super) fn new(name: String) -> Self { + Self { + name, + phantom: PhantomData, + data: McpServerData::default(), + responder: NullRun, + } + } +} + +impl McpServerBuilder +where + Responder: RunWithConnectionTo, +{ + /// Set the server instructions that are provided to the client. + #[must_use] + pub fn instructions(mut self, instructions: impl ToString) -> Self { + self.data.instructions = Some(instructions.to_string()); + self + } + + /// Add a tool to the server. + #[must_use] + pub fn tool(mut self, tool: impl McpTool + 'static) -> Self { + let tool_model = make_tool_model(&tool); + let has_structured_output = tool_model.output_schema.is_some(); + self.data.tool_models.push(tool_model); + self.data.tools.insert( + tool.name(), + RegisteredTool { + tool: make_erased_mcp_tool(tool), + has_structured_output, + }, + ); + self + } + + /// Disable all tools. After calling this, only tools explicitly enabled + /// with [`enable_tool`](Self::enable_tool) will be available. + #[must_use] + pub fn disable_all_tools(mut self) -> Self { + self.data.enabled_tools = EnabledTools::AllowList(HashSet::new()); + self + } + + /// Enable all tools. After calling this, all tools will be available + /// except those explicitly disabled with [`disable_tool`](Self::disable_tool). + #[must_use] + pub fn enable_all_tools(mut self) -> Self { + self.data.enabled_tools = EnabledTools::DenyList(HashSet::new()); + self + } + + /// Disable a specific tool by name. + /// + /// Returns an error if the tool is not registered. + pub fn disable_tool(mut self, name: &str) -> Result { + if !self.data.tools.contains_key(name) { + return Err(crate::Error::invalid_request().data(format!("unknown tool: {name}"))); + } + match &mut self.data.enabled_tools { + EnabledTools::DenyList(deny) => { + deny.insert(name.to_string()); + } + EnabledTools::AllowList(allow) => { + allow.remove(name); + } + } + Ok(self) + } + + /// Enable a specific tool by name. + /// + /// Returns an error if the tool is not registered. + pub fn enable_tool(mut self, name: &str) -> Result { + if !self.data.tools.contains_key(name) { + return Err(crate::Error::invalid_request().data(format!("unknown tool: {name}"))); + } + match &mut self.data.enabled_tools { + EnabledTools::DenyList(deny) => { + deny.remove(name); + } + EnabledTools::AllowList(allow) => { + allow.insert(name.to_string()); + } + } + Ok(self) + } + + /// Private fn: adds the tool but also adds a responder that will be + /// run while the MCP server is active. + fn tool_with_responder( + self, + tool: impl McpTool + 'static, + tool_responder: impl RunWithConnectionTo, + ) -> McpServerBuilder> { + let this = self.tool(tool); + McpServerBuilder { + phantom: PhantomData, + name: this.name, + data: this.data, + responder: ChainRun::new(this.responder, tool_responder), + } + } + + /// Convenience wrapper for defining a "single-threaded" tool without having to create a struct. + /// By "single-threaded", we mean that only one invocation of the tool can be running at a time. + /// Typically agents invoke a tool once per session and then block waiting for the result, + /// so this is fine, but they could attempt to run multiple invocations concurrently, in which + /// case those invocations would be serialized. + /// + /// # Parameters + /// + /// * `name`: The name of the tool. + /// * `description`: The description of the tool. + /// * `func`: The function that implements the tool. Use an async closure like `async |args, cx| { .. }`. + /// + /// # Examples + /// + /// ```rust,ignore + /// McpServer::builder("my-server") + /// .tool_fn_mut( + /// "greet", + /// "Greet someone by name", + /// async |input: GreetInput, _cx| Ok(format!("Hello, {}!", input.name)), + /// ) + /// ``` + pub fn tool_fn_mut( + self, + name: impl ToString, + description: impl ToString, + func: F, + tool_future_hack: impl for<'a> Fn( + &'a mut F, + P, + McpConnectionTo, + ) -> BoxFuture<'a, Result> + + Send + + 'static, + ) -> McpServerBuilder> + where + P: JsonSchema + DeserializeOwned + 'static + Send, + Ret: JsonSchema + Serialize + 'static + Send, + F: AsyncFnMut(P, McpConnectionTo) -> Result + Send, + { + let (call_tx, call_rx) = mpsc::channel(128); + self.tool_with_responder( + ToolFnTool { + name: name.to_string(), + description: description.to_string(), + call_tx, + }, + ToolFnMutResponder { + func, + call_rx, + tool_future_fn: Box::new(tool_future_hack), + }, + ) + } + + /// Convenience wrapper for defining a stateless tool that can run concurrently. + /// Unlike [`tool_fn_mut`](Self::tool_fn_mut), multiple invocations of this tool can run + /// at the same time since the function is `Fn` rather than `FnMut`. + /// + /// # Parameters + /// + /// * `name`: The name of the tool. + /// * `description`: The description of the tool. + /// * `func`: The function that implements the tool. Use an async closure like `async |args, cx| { .. }`. + /// + /// # Examples + /// + /// ```rust,ignore + /// McpServer::builder("my-server") + /// .tool_fn( + /// "greet", + /// "Greet someone by name", + /// async |input: GreetInput, _cx| Ok(format!("Hello, {}!", input.name)), + /// ) + /// ``` + pub fn tool_fn( + self, + name: impl ToString, + description: impl ToString, + func: F, + tool_future_hack: impl for<'a> Fn( + &'a F, + P, + McpConnectionTo, + ) -> BoxFuture<'a, Result> + + Send + + Sync + + 'static, + ) -> McpServerBuilder> + where + P: JsonSchema + DeserializeOwned + 'static + Send, + Ret: JsonSchema + Serialize + 'static + Send, + F: AsyncFn(P, McpConnectionTo) -> Result + + Send + + Sync + + 'static, + { + let (call_tx, call_rx) = mpsc::channel(128); + self.tool_with_responder( + ToolFnTool { + name: name.to_string(), + description: description.to_string(), + call_tx, + }, + ToolFnResponder { + func, + call_rx, + tool_future_fn: Box::new(tool_future_hack), + }, + ) + } + + /// Create an MCP server from this builder. + /// + /// This builder can be attached to new sessions (see [`SessionBuilder::with_mcp_server`](`crate::SessionBuilder::with_mcp_server`)) + /// or served up as part of a proxy (see [`Builder::with_mcp_server`](`crate::Builder::with_mcp_server`)). + pub fn build(self) -> McpServer { + McpServer::new( + McpServerBuilt { + name: self.name, + data: Arc::new(self.data), + }, + self.responder, + ) + } +} + +struct McpServerBuilt { + name: String, + data: Arc>, +} + +impl McpServerConnect for McpServerBuilt { + fn name(&self) -> String { + self.name.clone() + } + + fn connect( + &self, + mcp_connection: McpConnectionTo, + ) -> DynConnectTo { + DynConnectTo::new(McpServerConnection { + data: self.data.clone(), + mcp_connection, + }) + } +} + +/// An MCP server instance connected to the ACP framework. +pub(crate) struct McpServerConnection { + data: Arc>, + mcp_connection: McpConnectionTo, +} + +impl ConnectTo for McpServerConnection { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), crate::Error> { + // Create tokio byte streams that rmcp expects + let (mcp_server_stream, mcp_client_stream) = tokio::io::duplex(8192); + let (mcp_server_read, mcp_server_write) = tokio::io::split(mcp_server_stream); + let (mcp_client_read, mcp_client_write) = tokio::io::split(mcp_client_stream); + + let run_client = async { + // Connect byte_streams to the provided client + let byte_streams = + ByteStreams::new(mcp_client_write.compat_write(), mcp_client_read.compat()); + drop( + as ConnectTo>::connect_to( + byte_streams, + client, + ) + .await, + ); + Ok(()) + }; + + let run_server = async { + // Run the rmcp server with the server side of the duplex stream + let running_server = rmcp::ServiceExt::serve(self, (mcp_server_read, mcp_server_write)) + .await + .map_err(crate::Error::into_internal_error)?; + + // Wait for the server to finish + running_server + .waiting() + .await + .map(|_quit_reason| ()) + .map_err(crate::Error::into_internal_error) + }; + + (run_client, run_server).try_join().await?; + Ok(()) + } +} + +impl ServerHandler for McpServerConnection { + async fn call_tool( + &self, + request: rmcp::model::CallToolRequestParams, + context: rmcp::service::RequestContext, + ) -> Result { + // Lookup the tool definition, erroring if not found or disabled + let Some(registered) = self.data.tools.get(&request.name[..]) else { + return Err(rmcp::model::ErrorData::invalid_params( + format!("tool `{}` not found", request.name), + None, + )); + }; + + // Treat disabled tools as not found + if !self.data.enabled_tools.is_enabled(&request.name) { + return Err(rmcp::model::ErrorData::invalid_params( + format!("tool `{}` not found", request.name), + None, + )); + } + + // Convert input into JSON + let serde_value = serde_json::to_value(request.arguments).expect("valid json"); + + // Execute the user's tool, unless cancellation occurs + let has_structured_output = registered.has_structured_output; + match futures::future::select( + registered + .tool + .call_tool(serde_value, self.mcp_connection.clone()), + pin!(context.ct.cancelled()), + ) + .await + { + // If completed successfully + Either::Left((m, _)) => match m { + Ok(result) => { + // Use structured output only if the tool declared an output_schema + if has_structured_output { + Ok(CallToolResult::structured(result)) + } else { + Ok(CallToolResult::success(vec![rmcp::model::Content::text( + result.to_string(), + )])) + } + } + Err(error) => Err(to_rmcp_error(error)), + }, + + // If cancelled + Either::Right(((), _)) => { + Err(rmcp::ErrorData::internal_error("operation cancelled", None)) + } + } + } + + async fn list_tools( + &self, + _request: Option, + _context: rmcp::service::RequestContext, + ) -> Result { + // Return only enabled tools + let tools: Vec<_> = self + .data + .tool_models + .iter() + .filter(|t| self.data.enabled_tools.is_enabled(&t.name)) + .cloned() + .collect(); + Ok(ListToolsResult::with_all_items(tools)) + } + + fn get_info(&self) -> rmcp::model::ServerInfo { + // Basic server info + let base = rmcp::model::ServerInfo::new( + rmcp::model::ServerCapabilities::builder() + .enable_tools() + .build(), + ) + .with_server_info(rmcp::model::Implementation::default()) + .with_protocol_version(rmcp::model::ProtocolVersion::default()); + + if let Some(instr) = self.data.instructions.clone() { + base.with_instructions(instr) + } else { + base + } + } +} + +/// Erased version of the MCP tool trait that is dyn-compatible. +trait ErasedMcpTool: Send + Sync { + fn call_tool( + &self, + input: serde_json::Value, + connection: McpConnectionTo, + ) -> BoxFuture<'_, Result>; +} + +/// Create an `rmcp` tool model from our [`McpTool`] trait. +fn make_tool_model>(tool: &M) -> Tool { + let mut tool = rmcp::model::Tool::new( + tool.name(), + tool.description(), + schema_for_type::(), + ) + .with_execution(rmcp::model::ToolExecution::new()); + + if let Ok(schema) = schema_for_output::() { + // schema_for_output returns Err for non-object types (strings, integers, etc.) + // since MCP structured output requires JSON objects. We set + // output_schema to None for these tools, signaling unstructured output. + tool = tool.with_raw_output_schema(schema); + } + + tool +} + +/// Create a [`ErasedMcpTool`] from a [`McpTool`], erasing the type details. +fn make_erased_mcp_tool<'s, R: Role, M: McpTool + 's>( + tool: M, +) -> Arc + 's> { + struct ErasedMcpToolImpl { + tool: M, + } + + impl ErasedMcpTool for ErasedMcpToolImpl + where + R: Role, + M: McpTool, + { + fn call_tool( + &self, + input: serde_json::Value, + context: McpConnectionTo, + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let input = serde_json::from_value(input).map_err(crate::util::internal_error)?; + serde_json::to_value(self.tool.call_tool(input, context).await?) + .map_err(crate::util::internal_error) + }) + } + } + + Arc::new(ErasedMcpToolImpl { tool }) +} + +/// Convert a [`crate::Error`] into an [`rmcp::ErrorData`]. +fn to_rmcp_error(error: crate::Error) -> rmcp::ErrorData { + rmcp::ErrorData { + code: rmcp::model::ErrorCode(error.code.into()), + message: error.message.into(), + data: error.data, + } +} + +/// MCP tool used for `tool_fn` and `tooL_fn_mut`. +/// Each time it is invoked, it sends a `ToolCall` message to `call_tx`. +struct ToolFnTool { + name: String, + description: String, + call_tx: mpsc::Sender>, +} + +impl McpTool for ToolFnTool +where + R: Role, + P: JsonSchema + DeserializeOwned + 'static + Send, + Ret: JsonSchema + Serialize + 'static + Send, +{ + type Input = P; + type Output = Ret; + + fn name(&self) -> String { + self.name.clone() + } + + fn description(&self) -> String { + self.description.clone() + } + + async fn call_tool( + &self, + params: P, + mcp_connection: McpConnectionTo, + ) -> Result { + let (result_tx, result_rx) = oneshot::channel(); + + self.call_tx + .clone() + .send(ToolCall { + params, + mcp_connection, + result_tx, + }) + .await + .map_err(crate::util::internal_error)?; + + result_rx.await.map_err(crate::util::internal_error)? + } +} diff --git a/src/agent-client-protocol-core/src/mcp_server/connect.rs b/src/agent-client-protocol-core/src/mcp_server/connect.rs new file mode 100644 index 0000000..7fa614b --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/connect.rs @@ -0,0 +1,72 @@ +use std::sync::Arc; + +use crate::{ + DynConnectTo, + mcp_server::McpConnectionTo, + role::{self, Role}, +}; + +/// Trait for types that can create MCP server connections. +/// +/// Implement this trait to create custom MCP servers. Each call to [`connect`](Self::connect) +/// should return a new [`ConnectTo`](crate::ConnectTo) that serves MCP requests for a single +/// connection. +/// +/// # Example +/// +/// ```rust,ignore +/// use agent_client_protocol_core::mcp_server::{McpServerConnect, McpConnectionTo}; +/// use agent_client_protocol_core::{DynConnectTo, role::Role}; +/// +/// struct MyMcpServer { +/// name: String, +/// } +/// +/// impl McpServerConnect for MyMcpServer { +/// fn name(&self) -> String { +/// self.name.clone() +/// } +/// +/// fn connect(&self, cx: McpConnectionTo) -> DynConnectTo { +/// // Create and return a component that handles MCP requests +/// DynConnectTo::new(MyMcpComponent::new(cx)) +/// } +/// } +/// ``` +pub trait McpServerConnect: Send + Sync + 'static { + /// The name of the MCP server, used to identify it in session responses. + fn name(&self) -> String; + + /// Create a component to service a new MCP connection. + /// + /// This is called each time an agent connects to this MCP server. The returned + /// component will handle MCP protocol messages for that connection. + /// + /// The [`McpConnectionTo`] provides access to the ACP connection context and the + /// server's ACP URL. + fn connect(&self, cx: McpConnectionTo) -> DynConnectTo; +} + +impl> McpServerConnect + for Box +{ + fn name(&self) -> String { + S::name(self) + } + + fn connect(&self, cx: McpConnectionTo) -> DynConnectTo { + S::connect(self, cx) + } +} + +impl> McpServerConnect + for Arc +{ + fn name(&self) -> String { + S::name(self) + } + + fn connect(&self, cx: McpConnectionTo) -> DynConnectTo { + S::connect(self, cx) + } +} diff --git a/src/agent-client-protocol-core/src/mcp_server/context.rs b/src/agent-client-protocol-core/src/mcp_server/context.rs new file mode 100644 index 0000000..7905f20 --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/context.rs @@ -0,0 +1,22 @@ +use crate::{ConnectionTo, role::Role}; + +/// Context about the ACP and MCP connection available to an MCP server. +#[derive(Clone, Debug)] +pub struct McpConnectionTo { + pub(super) acp_url: String, + pub(super) connection: ConnectionTo, +} + +impl McpConnectionTo { + /// The `acp:UUID` that was given. + pub fn acp_url(&self) -> String { + self.acp_url.clone() + } + + /// The host connection context. + /// + /// If this MCP server is hosted inside of an ACP context, this will be the ACP connection context. + pub fn connection_to(&self) -> ConnectionTo { + self.connection.clone() + } +} diff --git a/src/agent-client-protocol-core/src/mcp_server/mod.rs b/src/agent-client-protocol-core/src/mcp_server/mod.rs new file mode 100644 index 0000000..35c14f8 --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/mod.rs @@ -0,0 +1,60 @@ +//! MCP server support for providing MCP tools over ACP. +//! +//! This module provides the infrastructure for building MCP servers that +//! integrate with ACP connections. +//! +//! ## Quick Start +//! +//! ```rust,ignore +//! use agent_client_protocol_core::mcp_server::{McpServer, McpTool}; +//! +//! // Create an MCP server with tools +//! let server = McpServer::builder("my-server".to_string()) +//! .instructions("A helpful assistant") +//! .tool(MyTool) +//! .build(); +//! +//! // Use the server as a handler on your connection +//! Proxy.builder() +//! .with_handler(server) +//! .serve(client) +//! .await?; +//! ``` +//! +//! ## Custom MCP Server Implementations +//! +//! You can implement [`McpServerConnect`](`crate::mcp_server::McpServerConnect`) to create custom MCP servers: +//! +//! ```rust,ignore +//! use agent_client_protocol_core::mcp_server::{McpServer, McpServerConnect, McpContext}; +//! use agent_client_protocol_core::{DynComponent, JrLink}; +//! +//! struct MyCustomServer; +//! +//! impl McpServerConnect for MyCustomServer { +//! fn name(&self) -> String { +//! "my-custom-server".to_string() +//! } +//! +//! fn connect(&self, cx: McpContext) -> DynComponent { +//! // Return a component that serves MCP requests +//! DynComponent::new(my_mcp_component) +//! } +//! } +//! +//! let server = McpServer::new(MyCustomServer); +//! ``` + +mod active_session; +mod builder; +mod connect; +mod context; +mod responder; +mod server; +mod tool; + +pub use builder::{EnabledTools, McpServerBuilder}; +pub use connect::McpServerConnect; +pub use context::McpConnectionTo; +pub use server::McpServer; +pub use tool::McpTool; diff --git a/src/agent-client-protocol-core/src/mcp_server/responder.rs b/src/agent-client-protocol-core/src/mcp_server/responder.rs new file mode 100644 index 0000000..aed19d8 --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/responder.rs @@ -0,0 +1,146 @@ +//! MCP-specific responder types. + +use futures::{ + StreamExt, + channel::{mpsc, oneshot}, + future::BoxFuture, +}; + +use crate::{ + ConnectionTo, jsonrpc::run::RunWithConnectionTo, mcp_server::McpConnectionTo, role::Role, +}; + +/// A tool call request sent through the channel. +pub(super) struct ToolCall { + pub(crate) params: P, + pub(crate) mcp_connection: McpConnectionTo, + pub(crate) result_tx: futures::channel::oneshot::Sender>, +} + +/// Responder for a `tool_fn` closure that receives tool calls through a channel +/// and invokes the user's async function. +pub(super) struct ToolFnMutResponder { + pub(crate) func: F, + pub(crate) call_rx: mpsc::Receiver>, + pub(crate) tool_future_fn: Box< + dyn for<'a> Fn( + &'a mut F, + P, + McpConnectionTo, + ) -> BoxFuture<'a, Result> + + Send, + >, +} + +impl RunWithConnectionTo + for ToolFnMutResponder +where + Counterpart: Role, + Counterpart1: Role, + P: Send, + R: Send, + F: Send, +{ + async fn run_with_connection_to( + self, + _connection: ConnectionTo, + ) -> Result<(), crate::Error> { + let ToolFnMutResponder { + mut func, + mut call_rx, + tool_future_fn, + } = self; + while let Some(ToolCall { + params, + mcp_connection, + result_tx, + }) = call_rx.next().await + { + let result = tool_future_fn(&mut func, params, mcp_connection).await; + result_tx + .send(result) + .map_err(|_| crate::util::internal_error("failed to send MCP result"))?; + } + Ok(()) + } +} + +/// Responder for a `tool_fn` closure that receives tool calls through a channel +/// and invokes the user's async function concurrently. +pub(super) struct ToolFnResponder { + pub(crate) func: F, + pub(crate) call_rx: mpsc::Receiver>, + pub(crate) tool_future_fn: Box< + dyn for<'a> Fn( + &'a F, + P, + McpConnectionTo, + ) -> BoxFuture<'a, Result> + + Send + + Sync, + >, +} + +impl RunWithConnectionTo + for ToolFnResponder +where + Counterpart: Role, + Counterpart1: Role, + P: Send, + R: Send, + F: Send + Sync, +{ + async fn run_with_connection_to( + self, + _connection: ConnectionTo, + ) -> Result<(), crate::Error> { + let ToolFnResponder { + func, + call_rx, + tool_future_fn, + } = self; + crate::util::process_stream_concurrently( + call_rx, + async |tool_call| { + fn hack<'a, F, P, R, MyRole>( + func: &'a F, + params: P, + mcp_connection: McpConnectionTo, + tool_future_fn: &'a ( + dyn Fn( + &'a F, + P, + McpConnectionTo, + ) -> BoxFuture<'a, Result> + + Send + + Sync + ), + result_tx: oneshot::Sender>, + ) -> BoxFuture<'a, ()> + where + MyRole: Role, + P: Send, + R: Send, + F: Send + Sync, + { + Box::pin(async move { + let result = tool_future_fn(func, params, mcp_connection).await; + // Ignore send errors - the receiver may have been dropped + drop(result_tx.send(result)); + }) + } + + let ToolCall { + params, + mcp_connection, + result_tx, + } = tool_call; + + hack(&func, params, mcp_connection, &*tool_future_fn, result_tx).await; + Ok(()) + }, + |a, b| Box::pin(a(b)), + ) + .await + } +} diff --git a/src/agent-client-protocol-core/src/mcp_server/server.rs b/src/agent-client-protocol-core/src/mcp_server/server.rs new file mode 100644 index 0000000..bdac9f8 --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/server.rs @@ -0,0 +1,256 @@ +//! MCP server builder for creating MCP servers. + +use std::{marker::PhantomData, sync::Arc}; + +use agent_client_protocol_schema::NewSessionRequest; +use futures::{StreamExt, channel::mpsc}; +use uuid::Uuid; + +use crate::{ + Agent, Client, ConnectTo, ConnectionTo, Dispatch, DynConnectTo, HandleDispatchFrom, Handled, + Role, + jsonrpc::{ + DynamicHandlerRegistration, + run::{NullRun, RunWithConnectionTo}, + }, + mcp_server::{ + McpConnectionTo, McpServerConnect, active_session::McpActiveSession, + builder::McpServerBuilder, + }, + role::{self, HasPeer}, + util::MatchDispatchFrom, +}; + +/// An MCP server that can be attached to ACP connections. +/// +/// `McpServer` wraps an [`McpServerConnect`](`super::McpServerConnect`) implementation and can be used either: +/// - As a message handler via [`Builder::with_handler`](`crate::Builder::with_handler`), automatically +/// attaching to new sessions +/// - Manually for more control +/// +/// # Creating an MCP Server +/// +/// Use [`McpServer::builder`] to create a server with tools: +/// +/// ```rust,ignore +/// let server = McpServer::builder("my-server".to_string()) +/// .instructions("A helpful assistant") +/// .tool(MyTool) +/// .build(); +/// ``` +/// +/// Or implement [`McpServerConnect`](`super::McpServerConnect`) for custom server behavior: +/// +/// ```rust,ignore +/// let server = McpServer::new(MyCustomServerConnect); +/// ``` +pub struct McpServer { + /// The host role that is serving up this MCP server + phantom: PhantomData, + + /// The ACP URL we assigned for this mcp server; always unique + acp_url: String, + + /// The "connect" instance + connect: Arc>, + + /// The "responder" is a task that should be run alongside the message handler. + /// Some futures direct messages back through channels to this future which actually + /// handles responding to the client. + /// + /// This is how we bridge the gap between the rmcp implementation, + /// which requires `'static`, and our APIs, which do not. + responder: Run, +} + +impl std::fmt::Debug + for McpServer +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("McpServer") + .field("phantom", &self.phantom) + .field("acp_url", &self.acp_url) + .field("responder", &self.responder) + .finish_non_exhaustive() + } +} + +impl McpServer { + /// Create an empty server with no content. + pub fn builder(name: impl ToString) -> McpServerBuilder { + McpServerBuilder::new(name.to_string()) + } +} + +impl McpServer +where + Run: RunWithConnectionTo, +{ + /// Create an MCP server from something that implements the [`McpServerConnect`](`super::McpServerConnect`) trait. + /// + /// # See also + /// + /// See [`Self::builder`] to construct MCP servers from Rust code. + pub fn new(c: impl McpServerConnect, responder: Run) -> Self { + McpServer { + phantom: PhantomData, + acp_url: format!("acp:{}", Uuid::new_v4()), + connect: Arc::new(c), + responder, + } + } + + /// Split this MCP server into the message handler and a future that must be run while the handler is active. + pub(crate) fn into_handler_and_responder(self) -> (McpNewSessionHandler, Run) + where + Counterpart: HasPeer, + { + let Self { + phantom: _, + acp_url, + connect, + responder, + } = self; + (McpNewSessionHandler::new(acp_url, connect), responder) + } +} + +/// Message handler created from a [`McpServer`]. +pub(crate) struct McpNewSessionHandler +where + Counterpart: HasPeer, +{ + acp_url: String, + connect: Arc>, + active_session: McpActiveSession, +} + +impl McpNewSessionHandler +where + Counterpart: HasPeer, +{ + pub fn new(acp_url: String, connect: Arc>) -> Self { + Self { + active_session: McpActiveSession::new(acp_url.clone(), connect.clone()), + acp_url, + connect, + } + } + + /// Modify the new session request to include this MCP server. + fn modify_new_session_request(&self, request: &mut NewSessionRequest) { + request.mcp_servers.push(crate::schema::McpServer::Http( + crate::schema::McpServerHttp::new(self.connect.name(), self.acp_url.clone()), + )); + } +} + +impl McpNewSessionHandler +where + Counterpart: HasPeer, +{ + /// Attach this server to the new session, spawning off a dynamic handler that will + /// manage requests coming from this session. + /// + /// # Return value + /// + /// Returns a [`DynamicHandlerRegistration`] for the handler that intercepts messages + /// related to this MCP server. Once the value is dropped, the MCP server messages + /// will no longer be received, so you need to keep this value alive as long as the session + /// is in use. You can also invoke [`DynamicHandlerRegistration::run_indefinitely`] + /// if you want to keep the handler running indefinitely. + pub fn into_dynamic_handler( + self, + request: &mut NewSessionRequest, + cx: &ConnectionTo, + ) -> Result, crate::Error> + where + Counterpart: HasPeer, + { + self.modify_new_session_request(request); + cx.add_dynamic_handler(self.active_session) + } +} + +impl HandleDispatchFrom for McpNewSessionHandler +where + Counterpart: HasPeer + HasPeer, +{ + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + cx: ConnectionTo, + ) -> Result, crate::Error> { + MatchDispatchFrom::new(message, &cx) + .if_request_from(Client, async |mut request: NewSessionRequest, responder| { + self.modify_new_session_request(&mut request); + Ok(Handled::No { + message: (request, responder), + retry: false, + }) + }) + .await + .otherwise_delegate(&mut self.active_session) + .await + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + format!("McpServer({})", self.connect.name()) + } +} + +impl ConnectTo for McpServer +where + Run: RunWithConnectionTo + 'static, +{ + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), crate::Error> { + let Self { + acp_url, + connect, + responder, + phantom: _, + } = self; + + let (tx, mut rx) = mpsc::unbounded(); + + role::mcp::Server + .builder() + .with_responder(responder) + .on_receive_dispatch( + async |message_from_client: Dispatch, _cx| { + tx.unbounded_send(message_from_client) + .map_err(|_| crate::util::internal_error("nobody listening to mcp server")) + }, + crate::on_receive_dispatch!(), + ) + .with_spawned(async move |connection_to_client| { + let spawned_server: DynConnectTo = + connect.connect(McpConnectionTo { + acp_url, + connection: connection_to_client.clone(), + }); + + let client = role::mcp::Client + .builder() + .on_receive_dispatch( + async |message_from_server: Dispatch, _| { + // when we receive a message from the server, fwd to the client + connection_to_client.send_proxied_message(message_from_server) + }, + crate::on_receive_dispatch!(), + ) + .connect_with(spawned_server, async |connection_to_server| { + while let Some(message_from_client) = rx.next().await { + connection_to_server.send_proxied_message(message_from_client)?; + } + Ok(()) + }); + Box::pin(client).await + }) + .connect_to(client) + .await + } +} diff --git a/src/agent-client-protocol-core/src/mcp_server/tool.rs b/src/agent-client-protocol-core/src/mcp_server/tool.rs new file mode 100644 index 0000000..5c26f04 --- /dev/null +++ b/src/agent-client-protocol-core/src/mcp_server/tool.rs @@ -0,0 +1,82 @@ +//! MCP tool trait for defining tools. + +use schemars::JsonSchema; +use serde::{Serialize, de::DeserializeOwned}; + +use crate::role::Role; + +use super::McpConnectionTo; + +/// Trait for defining MCP tools. +/// +/// Implement this trait to create a tool that can be registered with an MCP server. +/// The tool's input and output types must implement JSON Schema for automatic +/// documentation. +/// +/// # Example +/// +/// ```rust,ignore +/// use agent_client_protocol_core::mcp_server::{McpTool, McpContext}; +/// use schemars::JsonSchema; +/// use serde::{Deserialize, Serialize}; +/// +/// #[derive(JsonSchema, Deserialize)] +/// struct EchoInput { +/// message: String, +/// } +/// +/// #[derive(JsonSchema, Serialize)] +/// struct EchoOutput { +/// echoed: String, +/// } +/// +/// struct EchoTool; +/// +/// impl McpTool for EchoTool { +/// type Input = EchoInput; +/// type Output = EchoOutput; +/// +/// fn name(&self) -> String { +/// "echo".to_string() +/// } +/// +/// fn description(&self) -> String { +/// "Echoes back the input message".to_string() +/// } +/// +/// async fn call_tool( +/// &self, +/// input: EchoInput, +/// _context: McpContext, +/// ) -> Result { +/// Ok(EchoOutput { +/// echoed: format!("Echo: {}", input.message), +/// }) +/// } +/// } +/// ``` +pub trait McpTool: Send + Sync { + /// The type of input the tool accepts. + type Input: JsonSchema + DeserializeOwned + Send + 'static; + + /// The type of output the tool produces. + type Output: JsonSchema + Serialize + Send + 'static; + + /// The name of the tool + fn name(&self) -> String; + + /// A description of what the tool does + fn description(&self) -> String; + + /// A human-readable title for the tool + fn title(&self) -> Option { + None + } + + /// Define the tool's behavior. You can implement this with an `async fn`. + fn call_tool( + &self, + input: Self::Input, + context: McpConnectionTo, + ) -> impl Future> + Send; +} diff --git a/src/agent-client-protocol-core/src/role.rs b/src/agent-client-protocol-core/src/role.rs new file mode 100644 index 0000000..e1c28aa --- /dev/null +++ b/src/agent-client-protocol-core/src/role.rs @@ -0,0 +1,307 @@ +//! Role types for ACP connections. +//! +//! Roles represent the logical identity of an endpoint in an ACP connection. +//! Each role has a counterpart (who it connects to) and may have multiple peers +//! (who it can exchange messages with). + +use std::{any::TypeId, fmt::Debug, future::Future, hash::Hash}; + +use serde::{Deserialize, Serialize}; + +use crate::schema::{METHOD_SUCCESSOR_MESSAGE, SuccessorMessage}; +use crate::util::json_cast; +use crate::{Builder, ConnectionTo, Dispatch, Handled, JsonRpcMessage, UntypedMessage}; + +/// Roles for the ACP protocol. +pub mod acp; + +/// Roles for the MCP protocol. +pub mod mcp; + +/// The role that an endpoint plays in an ACP connection. +/// +/// Roles are the fundamental building blocks of ACP's type system: +/// - [`acp::Client`] connects to [`acp::Agent`] +/// - [`acp::Agent`] connects to [`acp::Client`] +/// - [`acp::Proxy`] connects to [`acp::Conductor`] +/// - [`acp::Conductor`] connects to [`acp::Proxy`] +/// +/// Each role determines: +/// - Who the counterpart is (via [`Role::Counterpart`]) +/// - How unhandled messages are processed (via `Role::default_message_handler`) +pub trait Role: Debug + Clone + Send + Sync + 'static + Eq + Ord + Hash { + /// The role that this endpoint connects to. + /// + /// For example: + /// - `Client::Counterpart = Agent` + /// - `Agent::Counterpart = Client` + /// - `Proxy::Counterpart = Conductor` + /// - `Conductor::Counterpart = Proxy` + type Counterpart: Role; + + /// Creates a new builder playing this role. + fn builder(self) -> Builder + where + Self: Sized, + { + Builder::new(self) + } + + /// Returns a unique identifier for this role. + fn role_id(&self) -> RoleId; + + /// Method invoked when there is no defined message handler. + fn default_handle_dispatch_from( + &self, + message: Dispatch, + connection: ConnectionTo, + ) -> impl Future, crate::Error>> + Send; + + /// Returns the counterpart role. + fn counterpart(&self) -> Self::Counterpart; +} + +/// Declares that a role can send messages to a specific peer. +/// +/// Most roles only communicate with their counterpart, but some (like [`acp::Proxy`]) +/// can communicate with multiple peers: +/// - `Proxy: HasPeer` - proxy can send/receive from clients +/// - `Proxy: HasPeer` - proxy can send/receive from agents +/// - `Proxy: HasPeer` - proxy can send/receive from its conductor +/// +/// The [`RemoteStyle`] determines how messages are transformed: +/// - [`RemoteStyle::Counterpart`] - pass through unchanged +/// - [`RemoteStyle::Predecessor`] - pass through, but reject wrapped messages +/// - [`RemoteStyle::Successor`] - wrap in a [`SuccessorMessage`] envelope +/// +/// [`SuccessorMessage`]: crate::schema::SuccessorMessage +pub trait HasPeer: Role { + /// Returns the remote style for sending to this peer. + fn remote_style(&self, peer: Peer) -> RemoteStyle; +} + +/// Describes how messages are transformed when sent to a remote peer. +#[derive(Clone, Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum RemoteStyle { + /// Pass each message through exactly as it is. + Counterpart, + + /// Only messages not wrapped in successor. + Predecessor, + + /// Wrap messages in a [`SuccessorMessage`] envelope. + Successor, +} + +impl RemoteStyle { + pub(crate) fn transform_outgoing_message( + &self, + msg: M, + ) -> Result { + match self { + RemoteStyle::Counterpart | RemoteStyle::Predecessor => msg.to_untyped_message(), + RemoteStyle::Successor => SuccessorMessage { + message: msg, + meta: None, + } + .to_untyped_message(), + } + } +} + +/// Unique identifier for a role instance. +/// +/// Used to identify the source/destination of messages when multiple +/// peers are possible on a single connection. +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)] +#[non_exhaustive] +pub enum RoleId { + /// Singleton role identified by type name and type ID. + Singleton(&'static str, TypeId), +} + +impl RoleId { + /// Create the role ID for a singleton role type. + pub fn from_singleton(_role: &R) -> RoleId + where + R: Role + Default, + { + RoleId::Singleton(std::any::type_name::(), TypeId::of::()) + } +} + +// ============================================================================ +// Role implementations +// ============================================================================ + +pub(crate) async fn handle_incoming_dispatch( + counterpart: Counterpart, + peer: Peer, + dispatch: Dispatch, + connection: ConnectionTo, + handle_dispatch: impl AsyncFnOnce( + Dispatch, + ConnectionTo, + ) -> Result, crate::Error>, +) -> Result, crate::Error> +where + Counterpart: Role + HasPeer, + Peer: Role, +{ + tracing::trace!( + method = %dispatch.method(), + ?counterpart, + ?peer, + ?dispatch, + "handle_incoming_dispatch: enter" + ); + + // Responses are different from other messages. + // + // For normal incoming messages, messages from non-default + // peers are tagged with special method names and carry + // special payload that have be "unwrapped". + // + // For responses, the payload is untouched. The response + // carries an `id` and we use this `id` to look up information + // on the request that was sent to determine which peer it was + // directed at (and therefore which peer sent us the response). + if let Dispatch::Response(_, router) = &dispatch { + tracing::trace!( + response_role_id = ?router.role_id(), + peer_role_id = ?peer.role_id(), + "handle_incoming_dispatch: response" + ); + + if router.role_id() == peer.role_id() { + return handle_dispatch(dispatch, connection).await; + } + return Ok(Handled::No { + message: dispatch, + retry: false, + }); + } + + // Handle other messages by looking at the 'remote style' + let method = dispatch.method(); + match counterpart.remote_style(peer) { + RemoteStyle::Counterpart => { + // "Counterpart" is the default peer, no special checks required. + tracing::trace!("handle_incoming_dispatch: Counterpart style, passing through"); + handle_dispatch(dispatch, connection).await + } + RemoteStyle::Predecessor => { + // "Predecessor" is the default peer, no special checks required. + tracing::trace!("handle_incoming_dispatch: Predecessor style, passing through"); + if method == METHOD_SUCCESSOR_MESSAGE { + // Methods coming from the successor are not coming from + // our counterpart. + Ok(Handled::No { + message: dispatch, + retry: false, + }) + } else { + handle_dispatch(dispatch, connection).await + } + } + RemoteStyle::Successor => { + // Successor style means we have to look for a special method name. + if method != METHOD_SUCCESSOR_MESSAGE { + tracing::trace!( + method, + expected = METHOD_SUCCESSOR_MESSAGE, + "handle_incoming_dispatch: Successor style but method doesn't match, returning Handled::No" + ); + return Ok(Handled::No { + message: dispatch, + retry: false, + }); + } + + tracing::trace!( + "handle_incoming_dispatch: Successor style, unwrapping SuccessorMessage" + ); + + // The outer message has method="_proxy/successor" and params containing the inner message. + // We need to deserialize the params (not the whole message) to extract the inner UntypedMessage. + let untyped_message = dispatch.message().ok_or_else(|| { + crate::util::internal_error( + "Response variant cannot be unwrapped as SuccessorMessage", + ) + })?; + let SuccessorMessage { message, meta } = json_cast(untyped_message.params())?; + let successor_dispatch = dispatch.try_map_message(|_| Ok(message))?; + tracing::trace!( + unwrapped_method = %successor_dispatch.method(), + "handle_incoming_dispatch: unwrapped to inner message" + ); + match handle_dispatch(successor_dispatch, connection).await? { + Handled::Yes => { + tracing::trace!( + "handle_incoming_dispatch: inner handler returned Handled::Yes" + ); + Ok(Handled::Yes) + } + + Handled::No { + message: successor_dispatch, + retry, + } => { + tracing::trace!( + "handle_incoming_dispatch: inner handler returned Handled::No, re-wrapping" + ); + Ok(Handled::No { + message: successor_dispatch.try_map_message(|message| { + SuccessorMessage { message, meta }.to_untyped_message() + })?, + retry, + }) + } + } + } + } +} + +/// A dummy role you can use to exchange JSON-RPC messages without any knowledge of the underlying protocol. +/// Don't sue this. +#[derive( + Copy, Clone, Default, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, +)] +pub struct UntypedRole; + +impl UntypedRole { + /// Creates a new builder for a connection from this role. + pub fn builder(self) -> Builder { + Builder::new(self) + } +} + +impl Role for UntypedRole { + type Counterpart = UntypedRole; + + fn role_id(&self) -> RoleId { + RoleId::from_singleton(self) + } + + async fn default_handle_dispatch_from( + &self, + message: Dispatch, + _connection: ConnectionTo, + ) -> Result, crate::Error> { + Ok(Handled::No { + message, + retry: false, + }) + } + + fn counterpart(&self) -> Self::Counterpart { + *self + } +} + +impl HasPeer for UntypedRole { + fn remote_style(&self, _peer: UntypedRole) -> RemoteStyle { + RemoteStyle::Counterpart + } +} diff --git a/src/agent-client-protocol-core/src/role/acp.rs b/src/agent-client-protocol-core/src/role/acp.rs new file mode 100644 index 0000000..3ffb334 --- /dev/null +++ b/src/agent-client-protocol-core/src/role/acp.rs @@ -0,0 +1,302 @@ +use std::{fmt::Debug, hash::Hash}; + +use agent_client_protocol_schema::{NewSessionRequest, NewSessionResponse, SessionId}; + +use crate::jsonrpc::{Builder, handlers::NullHandler, run::NullRun}; +use crate::role::{HasPeer, RemoteStyle}; +use crate::schema::{InitializeProxyRequest, InitializeRequest, METHOD_INITIALIZE_PROXY}; +use crate::util::MatchDispatchFrom; +use crate::{ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Role, RoleId}; + +/// The client role - typically an IDE or CLI that controls an agent. +/// +/// Clients send prompts and receive responses from agents. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Client; + +impl Role for Client { + type Counterpart = Agent; + + async fn default_handle_dispatch_from( + &self, + message: Dispatch, + _connection: ConnectionTo, + ) -> Result, crate::Error> { + Ok(Handled::No { + message, + retry: false, + }) + } + + fn role_id(&self) -> RoleId { + RoleId::from_singleton(self) + } + + fn counterpart(&self) -> Self::Counterpart { + Agent + } +} + +impl Client { + /// Create a connection builder for a client. + pub fn builder(self) -> Builder { + Builder::new(self) + } + + /// Connect to `agent` and run `main_fn` with the [`ConnectionTo`]. + /// Returns the result of `main_fn` (or an error if something goes wrong). + /// + /// Equivalent to `self.builder().connect_with(agent, main_fn)`. + pub async fn connect_with( + self, + agent: impl ConnectTo, + main_fn: impl AsyncFnOnce(ConnectionTo) -> Result, + ) -> Result { + self.builder().connect_with(agent, main_fn).await + } +} + +impl HasPeer for Client { + fn remote_style(&self, _peer: Client) -> RemoteStyle { + RemoteStyle::Counterpart + } +} + +/// The agent role - typically an LLM that responds to prompts. +/// +/// Agents receive prompts from clients and respond with answers, +/// potentially invoking tools along the way. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Agent; + +impl Role for Agent { + type Counterpart = Client; + + fn role_id(&self) -> RoleId { + RoleId::from_singleton(self) + } + + fn counterpart(&self) -> Self::Counterpart { + Client + } + + async fn default_handle_dispatch_from( + &self, + message: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + MatchDispatchFrom::new(message, &connection) + .if_message_from(Agent, async |message: Dispatch| { + // Subtle: messages that have a session-id field + // should be captured by a dynamic message handler + // for that session -- but there is a race condition + // between the dynamic handler being added and + // possible updates. Therefore, we "retry" all such + // messages, so that they will be resent as new handlers + // are added. + let retry = message.has_session_id(); + Ok(Handled::No { message, retry }) + }) + .await + .done() + } +} + +impl Agent { + /// Create a connection builder for an agent. + pub fn builder(self) -> Builder { + Builder::new(self) + } +} + +impl HasPeer for Agent { + fn remote_style(&self, _peer: Agent) -> RemoteStyle { + RemoteStyle::Counterpart + } +} + +/// The proxy role - an intermediary that can intercept and modify messages. +/// +/// Proxies sit between a client and an agent (or another proxy), and can: +/// - Add tools via MCP servers +/// - Filter or transform messages +/// - Inject additional context +/// +/// Proxies connect to a [`Conductor`] which orchestrates the proxy chain. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Proxy; + +impl Role for Proxy { + type Counterpart = Conductor; + + async fn default_handle_dispatch_from( + &self, + message: crate::Dispatch, + _connection: crate::ConnectionTo, + ) -> Result, crate::Error> { + Ok(Handled::No { + message, + retry: false, + }) + } + + fn role_id(&self) -> RoleId { + RoleId::from_singleton(self) + } + + fn counterpart(&self) -> Self::Counterpart { + Conductor + } +} + +impl Proxy { + /// Create a connection builder for a proxy. + pub fn builder(self) -> Builder { + Builder::new(self) + } +} + +impl HasPeer for Proxy { + fn remote_style(&self, _peer: Proxy) -> RemoteStyle { + RemoteStyle::Counterpart + } +} + +/// The conductor role - orchestrates proxy chains. +/// +/// Conductors manage connections between clients, proxies, and agents, +/// routing messages through the appropriate proxy chain. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Conductor; + +impl Role for Conductor { + type Counterpart = Proxy; + + fn role_id(&self) -> RoleId { + RoleId::from_singleton(self) + } + + fn counterpart(&self) -> Self::Counterpart { + Proxy + } + + async fn default_handle_dispatch_from( + &self, + message: Dispatch, + cx: ConnectionTo, + ) -> Result, crate::Error> { + // Handle various special messages: + MatchDispatchFrom::new(message, &cx) + .if_request_from(Client, async |_req: InitializeRequest, responder| { + responder.respond_with_error(crate::Error::invalid_request().data(format!( + "proxies must be initialized with `{METHOD_INITIALIZE_PROXY}`" + ))) + }) + .await + // Initialize Proxy coming from the client -- forward to the agent but + // convert into a regular initialize. + .if_request_from( + Client, + async |request: InitializeProxyRequest, responder| { + let InitializeProxyRequest { initialize } = request; + cx.send_request_to(Agent, initialize) + .forward_response_to(responder) + }, + ) + .await + // New session coming from the client -- proxy to the agent + // and add a dynamic handler for that session-id. + .if_request_from(Client, async |request: NewSessionRequest, responder| { + cx.send_request_to(Agent, request).on_receiving_result({ + let cx = cx.clone(); + async move |result| { + if let Ok(NewSessionResponse { session_id, .. }) = &result { + cx.add_dynamic_handler(ProxySessionMessages::new(session_id.clone()))? + .run_indefinitely(); + } + responder.respond_with_result(result) + } + }) + }) + .await + // Incoming message from the client -- forward to the agent + .if_message_from(Client, async |message: Dispatch| { + cx.send_proxied_message_to(Agent, message) + }) + .await + // Incoming message from the agent -- forward to the client + .if_message_from(Agent, async |message: Dispatch| { + cx.send_proxied_message_to(Client, message) + }) + .await + .done() + } +} + +impl Conductor { + /// Create a connection builder for a conductor. + pub fn builder(self) -> Builder { + Builder::new(self) + } +} + +impl HasPeer for Conductor { + fn remote_style(&self, _peer: Client) -> RemoteStyle { + RemoteStyle::Predecessor + } +} + +impl HasPeer for Conductor { + fn remote_style(&self, _peer: Agent) -> RemoteStyle { + RemoteStyle::Successor + } +} + +/// Dynamic handler that proxies session messages from Agent to Client. +/// +/// This is used internally to handle session message routing after a +/// `session.new` request has been forwarded. +pub(crate) struct ProxySessionMessages { + session_id: SessionId, +} + +impl ProxySessionMessages { + /// Create a new proxy handler for the given session. + pub fn new(session_id: SessionId) -> Self { + Self { session_id } + } +} + +impl HandleDispatchFrom for ProxySessionMessages +where + Counterpart: HasPeer + HasPeer, +{ + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + connection: ConnectionTo, + ) -> Result, crate::Error> { + MatchDispatchFrom::new(message, &connection) + .if_message_from(Agent, async |message| { + // If this is for our session-id, proxy it to the client. + if let Some(session_id) = message.get_session_id()? + && session_id == self.session_id + { + connection.send_proxied_message_to(Client, message)?; + return Ok(Handled::Yes); + } + + // Otherwise, leave it alone. + Ok(Handled::No { + message, + retry: false, + }) + }) + .await + .done() + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + format!("ProxySessionMessages({})", self.session_id) + } +} diff --git a/src/agent-client-protocol-core/src/role/mcp.rs b/src/agent-client-protocol-core/src/role/mcp.rs new file mode 100644 index 0000000..508c0f4 --- /dev/null +++ b/src/agent-client-protocol-core/src/role/mcp.rs @@ -0,0 +1,90 @@ +//! MCP (Model Context Protocol) role types. +//! +//! These roles are used for MCP connections, which are separate from ACP but +//! use the same underlying connection infrastructure. + +use crate::{ + Handled, RoleId, + jsonrpc::{Builder, handlers::NullHandler, run::NullRun}, + role::{HasPeer, RemoteStyle, Role}, +}; + +/// The MCP client role - connects to MCP servers to access tools and resources. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Client; + +impl Role for Client { + type Counterpart = Server; + + fn role_id(&self) -> RoleId { + RoleId::from_singleton(self) + } + + fn counterpart(&self) -> Self::Counterpart { + Server + } + + async fn default_handle_dispatch_from( + &self, + message: crate::Dispatch, + _connection: crate::ConnectionTo, + ) -> Result, crate::Error> { + Ok(Handled::No { + message, + retry: false, + }) + } +} + +impl Client { + /// Create a connection builder for an MCP client. + pub fn builder(self) -> Builder { + Builder::new(self) + } +} + +impl HasPeer for Client { + fn remote_style(&self, _peer: Client) -> RemoteStyle { + RemoteStyle::Counterpart + } +} + +/// The MCP server role - provides tools and resources to MCP clients. +#[derive(Debug, Default, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Server; + +impl Role for Server { + type Counterpart = Client; + + fn role_id(&self) -> RoleId { + RoleId::from_singleton(self) + } + + fn counterpart(&self) -> Self::Counterpart { + Client + } + + async fn default_handle_dispatch_from( + &self, + message: crate::Dispatch, + _connection: crate::ConnectionTo, + ) -> Result, crate::Error> { + Ok(Handled::No { + message, + retry: false, + }) + } +} + +impl Server { + /// Create a connection builder for an MCP server. + pub fn builder(self) -> Builder { + Builder::new(self) + } +} + +impl HasPeer for Server { + fn remote_style(&self, _peer: Server) -> RemoteStyle { + RemoteStyle::Counterpart + } +} diff --git a/src/agent-client-protocol-core/src/schema/agent_to_client/mod.rs b/src/agent-client-protocol-core/src/schema/agent_to_client/mod.rs new file mode 100644 index 0000000..2af4bcf --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/agent_to_client/mod.rs @@ -0,0 +1,3 @@ +// Agent → Client message implementations +mod notifications; +mod requests; diff --git a/src/agent-client-protocol-core/src/schema/agent_to_client/notifications.rs b/src/agent-client-protocol-core/src/schema/agent_to_client/notifications.rs new file mode 100644 index 0000000..a321593 --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/agent_to_client/notifications.rs @@ -0,0 +1,3 @@ +use crate::schema::SessionNotification; + +impl_jsonrpc_notification!(SessionNotification, "session/update"); diff --git a/src/agent-client-protocol-core/src/schema/agent_to_client/requests.rs b/src/agent-client-protocol-core/src/schema/agent_to_client/requests.rs new file mode 100644 index 0000000..b258aed --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/agent_to_client/requests.rs @@ -0,0 +1,44 @@ +use crate::schema::{ + CreateTerminalRequest, CreateTerminalResponse, KillTerminalRequest, KillTerminalResponse, + ReadTextFileRequest, ReadTextFileResponse, ReleaseTerminalRequest, ReleaseTerminalResponse, + RequestPermissionRequest, RequestPermissionResponse, TerminalOutputRequest, + TerminalOutputResponse, WaitForTerminalExitRequest, WaitForTerminalExitResponse, + WriteTextFileRequest, WriteTextFileResponse, +}; + +impl_jsonrpc_request!( + RequestPermissionRequest, + RequestPermissionResponse, + "session/request_permission" +); +impl_jsonrpc_request!( + WriteTextFileRequest, + WriteTextFileResponse, + "fs/write_text_file" +); +impl_jsonrpc_request!( + ReadTextFileRequest, + ReadTextFileResponse, + "fs/read_text_file" +); +impl_jsonrpc_request!( + CreateTerminalRequest, + CreateTerminalResponse, + "terminal/create" +); +impl_jsonrpc_request!( + TerminalOutputRequest, + TerminalOutputResponse, + "terminal/output" +); +impl_jsonrpc_request!( + ReleaseTerminalRequest, + ReleaseTerminalResponse, + "terminal/release" +); +impl_jsonrpc_request!( + WaitForTerminalExitRequest, + WaitForTerminalExitResponse, + "terminal/wait_for_exit" +); +impl_jsonrpc_request!(KillTerminalRequest, KillTerminalResponse, "terminal/kill"); diff --git a/src/agent-client-protocol-core/src/schema/client_to_agent/mod.rs b/src/agent-client-protocol-core/src/schema/client_to_agent/mod.rs new file mode 100644 index 0000000..b150749 --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/client_to_agent/mod.rs @@ -0,0 +1,3 @@ +// Client → Agent message implementations +mod notifications; +mod requests; diff --git a/src/agent-client-protocol-core/src/schema/client_to_agent/notifications.rs b/src/agent-client-protocol-core/src/schema/client_to_agent/notifications.rs new file mode 100644 index 0000000..8162c23 --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/client_to_agent/notifications.rs @@ -0,0 +1,3 @@ +use crate::schema::CancelNotification; + +impl_jsonrpc_notification!(CancelNotification, "session/cancel"); diff --git a/src/agent-client-protocol-core/src/schema/client_to_agent/requests.rs b/src/agent-client-protocol-core/src/schema/client_to_agent/requests.rs new file mode 100644 index 0000000..5215391 --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/client_to_agent/requests.rs @@ -0,0 +1,49 @@ +use crate::schema::{ + AuthenticateRequest, AuthenticateResponse, InitializeRequest, InitializeResponse, + ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, + NewSessionRequest, NewSessionResponse, PromptRequest, PromptResponse, + SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, SetSessionModeRequest, + SetSessionModeResponse, +}; +#[cfg(feature = "unstable_session_close")] +use crate::schema::{CloseSessionRequest, CloseSessionResponse}; +#[cfg(feature = "unstable_session_fork")] +use crate::schema::{ForkSessionRequest, ForkSessionResponse}; +#[cfg(feature = "unstable_session_resume")] +use crate::schema::{ResumeSessionRequest, ResumeSessionResponse}; +#[cfg(feature = "unstable_session_model")] +use crate::schema::{SetSessionModelRequest, SetSessionModelResponse}; + +impl_jsonrpc_request!(InitializeRequest, InitializeResponse, "initialize"); +impl_jsonrpc_request!(AuthenticateRequest, AuthenticateResponse, "authenticate"); +impl_jsonrpc_request!(LoadSessionRequest, LoadSessionResponse, "session/load"); +impl_jsonrpc_request!(ListSessionsRequest, ListSessionsResponse, "session/list"); +impl_jsonrpc_request!(NewSessionRequest, NewSessionResponse, "session/new"); +impl_jsonrpc_request!(PromptRequest, PromptResponse, "session/prompt"); +impl_jsonrpc_request!( + SetSessionModeRequest, + SetSessionModeResponse, + "session/set_mode" +); +impl_jsonrpc_request!( + SetSessionConfigOptionRequest, + SetSessionConfigOptionResponse, + "session/set_config_option" +); + +#[cfg(feature = "unstable_session_model")] +impl_jsonrpc_request!( + SetSessionModelRequest, + SetSessionModelResponse, + "session/set_model" +); +#[cfg(feature = "unstable_session_fork")] +impl_jsonrpc_request!(ForkSessionRequest, ForkSessionResponse, "session/fork"); +#[cfg(feature = "unstable_session_resume")] +impl_jsonrpc_request!( + ResumeSessionRequest, + ResumeSessionResponse, + "session/resume" +); +#[cfg(feature = "unstable_session_close")] +impl_jsonrpc_request!(CloseSessionRequest, CloseSessionResponse, "session/close"); diff --git a/src/agent-client-protocol-core/src/schema/enum_impls.rs b/src/agent-client-protocol-core/src/schema/enum_impls.rs new file mode 100644 index 0000000..c7fb06c --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/enum_impls.rs @@ -0,0 +1,54 @@ +//! JsonRpcMessage and JsonRpcNotification/JsonRpcRequest implementations for +//! the ACP enum types from agent-client-protocol-schema. + +use crate::schema::{AgentNotification, AgentRequest, ClientNotification, ClientRequest}; + +// ============================================================================ +// Agent side (messages that agents receive) +// ============================================================================ + +impl_jsonrpc_request_enum!(ClientRequest { + InitializeRequest => "initialize", + AuthenticateRequest => "authenticate", + NewSessionRequest => "session/new", + LoadSessionRequest => "session/load", + ListSessionsRequest => "session/list", + #[cfg(feature = "unstable_session_fork")] + ForkSessionRequest => "session/fork", + #[cfg(feature = "unstable_session_resume")] + ResumeSessionRequest => "session/resume", + #[cfg(feature = "unstable_session_close")] + CloseSessionRequest => "session/close", + SetSessionModeRequest => "session/set_mode", + SetSessionConfigOptionRequest => "session/set_config_option", + PromptRequest => "session/prompt", + #[cfg(feature = "unstable_session_model")] + SetSessionModelRequest => "session/set_model", + [ext] ExtMethodRequest, +}); + +impl_jsonrpc_notification_enum!(ClientNotification { + CancelNotification => "session/cancel", + [ext] ExtNotification, +}); + +// ============================================================================ +// Client side (messages that clients/editors receive) +// ============================================================================ + +impl_jsonrpc_request_enum!(AgentRequest { + WriteTextFileRequest => "fs/write_text_file", + ReadTextFileRequest => "fs/read_text_file", + RequestPermissionRequest => "session/request_permission", + CreateTerminalRequest => "terminal/create", + TerminalOutputRequest => "terminal/output", + ReleaseTerminalRequest => "terminal/release", + WaitForTerminalExitRequest => "terminal/wait_for_exit", + KillTerminalRequest => "terminal/kill", + [ext] ExtMethodRequest, +}); + +impl_jsonrpc_notification_enum!(AgentNotification { + SessionNotification => "session/update", + [ext] ExtNotification, +}); diff --git a/src/agent-client-protocol-core/src/schema/mod.rs b/src/agent-client-protocol-core/src/schema/mod.rs new file mode 100644 index 0000000..30f6d93 --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/mod.rs @@ -0,0 +1,232 @@ +//! ACP protocol schema types and message implementations. +//! +//! This module contains all the types from the Agent-Client Protocol schema, +//! including requests, responses, notifications, and supporting types. +//! All types are re-exported flatly from this module. + +// --------------------------------------------------------------------------- +// Macros for implementing JsonRpc traits on schema types +// --------------------------------------------------------------------------- + +/// Implement `JsonRpcMessage`, `JsonRpcRequest`, and `JsonRpcResponse` for a +/// request/response pair from the schema crate. +/// +/// ```ignore +/// impl_jsonrpc_request!(PromptRequest, PromptResponse, "session/prompt"); +/// ``` +macro_rules! impl_jsonrpc_request { + ($req:ty, $resp:ty, $method:literal) => { + impl $crate::JsonRpcMessage for $req { + fn matches_method(method: &str) -> bool { + method == $method + } + + fn method(&self) -> &str { + $method + } + + fn to_untyped_message(&self) -> Result<$crate::UntypedMessage, $crate::Error> { + $crate::UntypedMessage::new($method, self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if method != $method { + return Err($crate::Error::method_not_found()); + } + $crate::util::json_cast(params) + } + } + + impl $crate::JsonRpcRequest for $req { + type Response = $resp; + } + + impl $crate::JsonRpcResponse for $resp { + fn into_json(self, _method: &str) -> Result { + serde_json::to_value(self).map_err($crate::Error::into_internal_error) + } + + fn from_value(_method: &str, value: serde_json::Value) -> Result { + $crate::util::json_cast(&value) + } + } + }; +} + +/// Implement `JsonRpcMessage` and `JsonRpcNotification` for a notification type +/// from the schema crate. +/// +/// ```ignore +/// impl_jsonrpc_notification!(CancelNotification, "session/cancel"); +/// ``` +macro_rules! impl_jsonrpc_notification { + ($notif:ty, $method:literal) => { + impl $crate::JsonRpcMessage for $notif { + fn matches_method(method: &str) -> bool { + method == $method + } + + fn method(&self) -> &str { + $method + } + + fn to_untyped_message(&self) -> Result<$crate::UntypedMessage, $crate::Error> { + $crate::UntypedMessage::new($method, self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if method != $method { + return Err($crate::Error::method_not_found()); + } + $crate::util::json_cast(params) + } + } + + impl $crate::JsonRpcNotification for $notif {} + }; +} + +/// Implement `JsonRpcMessage` and `JsonRpcRequest` for an enum that dispatches +/// across multiple request types, with an extension method fallback. +/// +/// Variants can optionally have `#[cfg(...)]` attributes for conditional compilation. +/// +/// ```ignore +/// impl_jsonrpc_request_enum!(ClientRequest { +/// InitializeRequest => "initialize", +/// PromptRequest => "session/prompt", +/// #[cfg(feature = "unstable_session_model")] +/// SetSessionModelRequest => "session/set_model", +/// [ext] ExtMethodRequest, +/// }); +/// ``` +macro_rules! impl_jsonrpc_request_enum { + ($enum:ty { + $( $(#[$meta:meta])* $variant:ident => $method:literal, )* + [ext] $ext_variant:ident, + }) => { + impl $crate::JsonRpcMessage for $enum { + fn matches_method(_method: &str) -> bool { + true + } + + fn method(&self) -> &str { + match self { + $( $(#[$meta])* Self::$variant(_) => $method, )* + Self::$ext_variant(ext) => &ext.method, + _ => "_unknown", + } + } + + fn to_untyped_message(&self) -> Result<$crate::UntypedMessage, $crate::Error> { + $crate::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + match method { + $( $(#[$meta])* $method => $crate::util::json_cast(params).map(Self::$variant), )* + _ => { + if let Some(custom_method) = method.strip_prefix('_') { + $crate::util::json_cast(params).map( + |ext_req: $crate::schema::ExtRequest| { + Self::$ext_variant($crate::schema::ExtRequest::new( + custom_method.to_string(), + ext_req.params, + )) + }, + ) + } else { + Err($crate::Error::method_not_found()) + } + } + } + } + } + + impl $crate::JsonRpcRequest for $enum { + type Response = serde_json::Value; + } + }; +} + +/// Implement `JsonRpcMessage` and `JsonRpcNotification` for an enum that +/// dispatches across multiple notification types, with an extension fallback. +/// +/// Variants can optionally have `#[cfg(...)]` attributes for conditional compilation. +/// +/// ```ignore +/// impl_jsonrpc_notification_enum!(AgentNotification { +/// SessionNotification => "session/update", +/// [ext] ExtNotification, +/// }); +/// ``` +macro_rules! impl_jsonrpc_notification_enum { + ($enum:ty { + $( $(#[$meta:meta])* $variant:ident => $method:literal, )* + [ext] $ext_variant:ident, + }) => { + impl $crate::JsonRpcMessage for $enum { + fn matches_method(_method: &str) -> bool { + true + } + + fn method(&self) -> &str { + match self { + $( $(#[$meta])* Self::$variant(_) => $method, )* + Self::$ext_variant(ext) => &ext.method, + _ => "_unknown", + } + } + + fn to_untyped_message(&self) -> Result<$crate::UntypedMessage, $crate::Error> { + $crate::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + match method { + $( $(#[$meta])* $method => $crate::util::json_cast(params).map(Self::$variant), )* + _ => { + if let Some(custom_method) = method.strip_prefix('_') { + $crate::util::json_cast(params).map( + |ext_notif: $crate::schema::ExtNotification| { + Self::$ext_variant($crate::schema::ExtNotification::new( + custom_method.to_string(), + ext_notif.params, + )) + }, + ) + } else { + Err($crate::Error::method_not_found()) + } + } + } + } + } + + impl $crate::JsonRpcNotification for $enum {} + }; +} + +// Internal organization +mod agent_to_client; +mod client_to_agent; +mod enum_impls; +mod proxy_protocol; + +// Re-export everything from agent_client_protocol_schema +pub use agent_client_protocol_schema::*; + +// Re-export proxy/MCP protocol types +pub use proxy_protocol::*; diff --git a/src/agent-client-protocol-core/src/schema/proxy_protocol.rs b/src/agent-client-protocol-core/src/schema/proxy_protocol.rs new file mode 100644 index 0000000..38b7109 --- /dev/null +++ b/src/agent-client-protocol-core/src/schema/proxy_protocol.rs @@ -0,0 +1,207 @@ +//! Protocol types for proxy and MCP-over-ACP communication. +//! +//! These types are intended to become part of the ACP protocol specification. + +use crate::{JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, UntypedMessage}; +use agent_client_protocol_schema::InitializeResponse; +use serde::{Deserialize, Serialize}; + +// ============================================================================= +// Successor forwarding protocol +// ============================================================================= + +/// JSON-RPC method name for successor forwarding. +pub const METHOD_SUCCESSOR_MESSAGE: &str = "_proxy/successor"; + +/// A message being sent to the successor component. +/// +/// Used in `_proxy/successor` when the proxy wants to forward a message downstream. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SuccessorMessage { + /// The message to be sent to the successor component. + #[serde(flatten)] + pub message: M, + + /// Optional metadata + #[serde(skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +impl JsonRpcMessage for SuccessorMessage { + fn matches_method(method: &str) -> bool { + method == METHOD_SUCCESSOR_MESSAGE + } + + fn method(&self) -> &str { + METHOD_SUCCESSOR_MESSAGE + } + + fn to_untyped_message(&self) -> Result { + UntypedMessage::new( + METHOD_SUCCESSOR_MESSAGE, + SuccessorMessage { + message: self.message.to_untyped_message()?, + meta: self.meta.clone(), + }, + ) + } + + fn parse_message(method: &str, params: &impl Serialize) -> Result { + if method != METHOD_SUCCESSOR_MESSAGE { + return Err(crate::Error::method_not_found()); + } + let outer = crate::util::json_cast::<_, SuccessorMessage>(params)?; + if !M::matches_method(&outer.message.method) { + return Err(crate::Error::method_not_found()); + } + let inner = M::parse_message(&outer.message.method, &outer.message.params)?; + Ok(SuccessorMessage { + message: inner, + meta: outer.meta, + }) + } +} + +impl JsonRpcRequest for SuccessorMessage { + type Response = Req::Response; +} + +impl JsonRpcNotification for SuccessorMessage {} + +// ============================================================================= +// MCP-over-ACP protocol +// ============================================================================= + +/// JSON-RPC method name for MCP connect requests +pub const METHOD_MCP_CONNECT_REQUEST: &str = "_mcp/connect"; + +/// Creates a new MCP connection. This is equivalent to "running the command". +#[derive(Debug, Clone, Serialize, Deserialize, crate::JsonRpcRequest)] +#[request(method = "_mcp/connect", response = McpConnectResponse, crate = crate)] +pub struct McpConnectRequest { + /// The ACP URL to connect to (e.g., "acp:uuid") + pub acp_url: String, + + /// Optional metadata + #[serde(skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +/// Response to an MCP connect request +#[derive(Debug, Clone, Serialize, Deserialize, crate::JsonRpcResponse)] +#[response(crate = crate)] +pub struct McpConnectResponse { + /// Unique identifier for the established MCP connection + pub connection_id: String, + + /// Optional metadata + #[serde(skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +/// JSON-RPC method name for MCP disconnect notifications +pub const METHOD_MCP_DISCONNECT_NOTIFICATION: &str = "_mcp/disconnect"; + +/// Disconnects the MCP connection. +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize, crate::JsonRpcNotification)] +#[notification(method = "_mcp/disconnect", crate = crate)] +pub struct McpDisconnectNotification { + /// The id of the connection to disconnect. + pub connection_id: String, + + /// Optional metadata + #[serde(skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +/// JSON-RPC method name for MCP requests over ACP +pub const METHOD_MCP_MESSAGE: &str = "_mcp/message"; + +/// An MCP request sent via ACP. This could be an MCP-server-to-MCP-client request +/// (in which case it goes from the ACP client to the ACP agent, +/// note the reversal of roles) or an MCP-client-to-MCP-server request +/// (in which case it goes from the ACP agent to the ACP client). +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct McpOverAcpMessage { + /// id given in response to `_mcp/connect` request. + pub connection_id: String, + + /// Request to be sent to the MCP server or client. + #[serde(flatten)] + pub message: M, + + /// Optional metadata + #[serde(skip_serializing_if = "Option::is_none")] + pub meta: Option, +} + +impl JsonRpcMessage for McpOverAcpMessage { + fn matches_method(method: &str) -> bool { + method == METHOD_MCP_MESSAGE + } + + fn method(&self) -> &str { + METHOD_MCP_MESSAGE + } + + fn to_untyped_message(&self) -> Result { + let message = self.message.to_untyped_message()?; + UntypedMessage::new( + METHOD_MCP_MESSAGE, + McpOverAcpMessage { + connection_id: self.connection_id.clone(), + message, + meta: self.meta.clone(), + }, + ) + } + + fn parse_message(method: &str, params: &impl Serialize) -> Result { + if method != METHOD_MCP_MESSAGE { + return Err(crate::Error::method_not_found()); + } + let outer = crate::util::json_cast::<_, McpOverAcpMessage>(params)?; + if !M::matches_method(&outer.message.method) { + return Err(crate::Error::method_not_found()); + } + let inner = M::parse_message(&outer.message.method, &outer.message.params)?; + Ok(McpOverAcpMessage { + connection_id: outer.connection_id, + message: inner, + meta: outer.meta, + }) + } +} + +impl JsonRpcRequest for McpOverAcpMessage { + type Response = R::Response; +} + +impl JsonRpcNotification for McpOverAcpMessage {} + +// ============================================================================= +// Proxy initialization protocol +// ============================================================================= + +/// JSON-RPC method name for proxy initialization. +pub const METHOD_INITIALIZE_PROXY: &str = "_proxy/initialize"; + +/// Initialize request for proxy components. +/// +/// This is sent to components that have a successor in the chain. +/// Components that receive this (instead of `InitializeRequest`) know they +/// are operating as a proxy and should forward messages to their successor. +#[derive(Debug, Clone, Serialize, Deserialize, crate::JsonRpcRequest)] +#[request(method = "_proxy/initialize", response = InitializeResponse, crate = crate)] +pub struct InitializeProxyRequest { + /// The underlying initialize request data. + #[serde(flatten)] + pub initialize: agent_client_protocol_schema::InitializeRequest, +} + +impl From for InitializeProxyRequest { + fn from(initialize: agent_client_protocol_schema::InitializeRequest) -> Self { + Self { initialize } + } +} diff --git a/src/agent-client-protocol-core/src/session.rs b/src/agent-client-protocol-core/src/session.rs new file mode 100644 index 0000000..9908c00 --- /dev/null +++ b/src/agent-client-protocol-core/src/session.rs @@ -0,0 +1,771 @@ +use std::{future::Future, marker::PhantomData, path::Path}; + +use agent_client_protocol_schema::{ + ContentBlock, ContentChunk, NewSessionRequest, NewSessionResponse, PromptRequest, + PromptResponse, SessionModeState, SessionNotification, SessionUpdate, StopReason, +}; +use futures::channel::mpsc; +use tokio::sync::oneshot; + +use crate::{ + Agent, Client, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Responder, Role, + jsonrpc::{ + DynamicHandlerRegistration, + run::{ChainRun, NullRun, RunWithConnectionTo}, + }, + mcp_server::McpServer, + role::{HasPeer, acp::ProxySessionMessages}, + schema::SessionId, + util::{MatchDispatch, MatchDispatchFrom, run_until}, +}; + +/// Marker type indicating the session builder will block the current task. +#[derive(Debug)] +pub struct Blocking; +impl SessionBlockState for Blocking {} + +/// Marker type indicating the session builder will not block the current task. +#[derive(Debug)] +pub struct NonBlocking; +impl SessionBlockState for NonBlocking {} + +/// Trait for marker types that indicate blocking vs blocking API. +/// See [`SessionBuilder::block_task`]. +pub trait SessionBlockState: Send + 'static + Sync + std::fmt::Debug {} + +impl ConnectionTo +where + Counterpart: HasPeer, +{ + /// Session builder for a new session request. + pub fn build_session(&self, cwd: impl AsRef) -> SessionBuilder { + SessionBuilder::new(self, NewSessionRequest::new(cwd.as_ref())) + } + + /// Session builder using the current working directory. + /// + /// This is a convenience wrapper around [`build_session`](Self::build_session) + /// that uses [`std::env::current_dir`] to get the working directory. + /// + /// Returns an error if the current directory cannot be determined. + pub fn build_session_cwd(&self) -> Result, crate::Error> { + let cwd = std::env::current_dir().map_err(|e| { + crate::Error::internal_error().data(format!("cannot get current directory: {e}")) + })?; + Ok(self.build_session(cwd)) + } + + /// Session builder starting from an existing request. + /// + /// Use this when you've intercepted a `session.new` request and want to + /// modify it (e.g., inject MCP servers) before forwarding. + pub fn build_session_from( + &self, + request: NewSessionRequest, + ) -> SessionBuilder { + SessionBuilder::new(self, request) + } + + /// Given a session response received from the agent, + /// attach a handler to process messages related to this session + /// and let you access them. + /// + /// Normally you would not use this method directly but would + /// instead use [`Self::build_session`] and then [`SessionBuilder::start_session`]. + /// + /// The vector `dynamic_handler_registrations` contains any dynamic + /// handle registrations associated with this session (e.g., from MCP servers). + /// You can simply pass `Default::default()` if not applicable. + pub fn attach_session<'responder>( + &self, + response: NewSessionResponse, + mcp_handler_registrations: Vec>, + ) -> Result, crate::Error> { + let NewSessionResponse { + session_id, + modes, + meta, + .. + } = response; + + let (update_tx, update_rx) = mpsc::unbounded(); + let handler = ActiveSessionHandler::new(session_id.clone(), update_tx.clone()); + let session_handler_registration = self.add_dynamic_handler(handler)?; + + Ok(ActiveSession { + session_id, + modes, + meta, + update_rx, + update_tx, + connection: self.clone(), + session_handler_registration, + mcp_handler_registrations, + _responder: PhantomData, + }) + } +} + +/// Session builder for a new session request. +/// Allows you to add MCP servers or set other details for this session. +/// +/// The `BlockState` type parameter tracks whether blocking methods are available: +/// - `NonBlocking` (default): Only [`on_session_start`](Self::on_session_start) is available +/// - `Blocking` (after calling [`block_task`](Self::block_task)): +/// [`run_until`](Self::run_until) and [`start_session`](Self::start_session) become available +#[must_use = "use `start_session`, `run_until`, or `on_session_start` to start the session"] +#[derive(Debug)] +pub struct SessionBuilder< + Counterpart, + Run: RunWithConnectionTo = NullRun, + BlockState: SessionBlockState = NonBlocking, +> where + Counterpart: HasPeer, +{ + connection: ConnectionTo, + request: NewSessionRequest, + dynamic_handler_registrations: Vec>, + run: Run, + block_state: PhantomData, +} + +impl SessionBuilder +where + Counterpart: HasPeer, +{ + fn new(connection: &ConnectionTo, request: NewSessionRequest) -> Self { + SessionBuilder { + connection: connection.clone(), + request, + dynamic_handler_registrations: Vec::default(), + run: NullRun, + block_state: PhantomData, + } + } +} + +impl SessionBuilder +where + Counterpart: HasPeer, + R: RunWithConnectionTo, + BlockState: SessionBlockState, +{ + /// Add the MCP servers from the given registry to this session. + pub fn with_mcp_server( + mut self, + mcp_server: McpServer, + ) -> Result, BlockState>, crate::Error> + where + McpRun: RunWithConnectionTo, + { + let (handler, mcp_run) = mcp_server.into_handler_and_responder(); + self.dynamic_handler_registrations + .push(handler.into_dynamic_handler(&mut self.request, &self.connection)?); + Ok(SessionBuilder { + connection: self.connection, + request: self.request, + dynamic_handler_registrations: self.dynamic_handler_registrations, + run: ChainRun::new(self.run, mcp_run), + block_state: self.block_state, + }) + } + + /// Spawn a task that runs the provided closure once the session starts. + /// + /// Unlike [`start_session`](Self::start_session), this method returns immediately + /// without blocking the current task. The session handshake and closure execution + /// happen in a spawned background task. + /// + /// The closure receives an `ActiveSession<'static, _>` and should return + /// `Result<(), Error>`. If the closure returns an error, it will propagate + /// to the connection's error handling. + /// + /// # Example + /// + /// ``` + /// # use agent_client_protocol_core::{Client, Agent, ConnectTo}; + /// # use agent_client_protocol_core::mcp_server::McpServer; + /// # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + /// # Client.builder().connect_with(transport, async |cx| { + /// # let mcp = McpServer::::builder("tools").build(); + /// cx.build_session_cwd()? + /// .with_mcp_server(mcp)? + /// .on_session_start(async |mut session| { + /// // Do something with the session + /// session.send_prompt("Hello")?; + /// let response = session.read_to_string().await?; + /// Ok(()) + /// })?; + /// // Returns immediately, session runs in background + /// # Ok(()) + /// # }).await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Ordering + /// + /// This callback blocks the dispatch loop until the session starts and your + /// callback completes. See the [`ordering`](crate::concepts::ordering) module for details. + pub fn on_session_start(self, op: F) -> Result<(), crate::Error> + where + R: 'static, + F: FnOnce(ActiveSession<'static, Counterpart>) -> Fut + Send + 'static, + Fut: Future> + Send, + { + let Self { + connection, + request, + dynamic_handler_registrations, + run, + block_state: _, + } = self; + + connection + .send_request_to(Agent, request) + .on_receiving_result({ + let connection = connection.clone(); + async move |result| { + let response = result?; + + connection.spawn(run.run_with_connection_to(connection.clone()))?; + + let active_session = + connection.attach_session(response, dynamic_handler_registrations)?; + + op(active_session).await + } + }) + } + + /// Spawn a proxy session and run a closure with the session ID. + /// + /// A **proxy session** starts the session with the agent and then automatically + /// proxies all session updates (prompts, tool calls, etc.) from the agent back + /// to the client. You don't need to handle any messages yourself - the proxy + /// takes care of forwarding everything. This is useful when you want to inject + /// and/or filter prompts coming from the client but otherwise not be involved + /// in the session. + /// + /// Unlike [`start_session_proxy`](Self::start_session_proxy), this method returns + /// immediately without blocking the current task. The session handshake, client + /// response, and proxy setup all happen in a spawned background task. + /// + /// The closure receives the `SessionId` once the session is established, allowing + /// you to perform any custom work with that ID (e.g., tracking, logging). + /// + /// # Example + /// + /// ``` + /// # use agent_client_protocol_core::{Proxy, Client, Conductor, ConnectTo}; + /// # use agent_client_protocol_core::schema::NewSessionRequest; + /// # use agent_client_protocol_core::mcp_server::McpServer; + /// # async fn example(transport: impl ConnectTo) -> Result<(), agent_client_protocol_core::Error> { + /// Proxy.builder() + /// .on_receive_request_from(Client, async |request: NewSessionRequest, responder, cx| { + /// let mcp = McpServer::::builder("tools").build(); + /// cx.build_session_from(request) + /// .with_mcp_server(mcp)? + /// .on_proxy_session_start(responder, async |session_id| { + /// // Session started + /// Ok(()) + /// }) + /// }, agent_client_protocol_core::on_receive_request!()) + /// .connect_to(transport) + /// .await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Ordering + /// + /// This callback blocks the dispatch loop until the session starts and your + /// callback completes. See the [`ordering`](crate::concepts::ordering) module for details. + pub fn on_proxy_session_start( + self, + responder: Responder, + op: F, + ) -> Result<(), crate::Error> + where + F: FnOnce(SessionId) -> Fut + Send + 'static, + Fut: Future> + Send, + Counterpart: HasPeer, + R: 'static, + { + let Self { + connection, + request, + dynamic_handler_registrations, + run, + block_state: _, + } = self; + + // Spawn off the run and dynamic handlers to run indefinitely + connection.spawn(run.run_with_connection_to(connection.clone()))?; + dynamic_handler_registrations + .into_iter() + .for_each(super::jsonrpc::DynamicHandlerRegistration::run_indefinitely); + + // Send the "new session" request to the agent + connection + .send_request_to(Agent, request) + .on_receiving_result({ + let connection = connection.clone(); + async move |result| { + let response = result?; + + // Extract the session-id from the response and forward + // the response back to the client + let session_id = response.session_id.clone(); + responder.respond(response)?; + + // Install a dynamic handler to proxy messages from this session + connection + .add_dynamic_handler(ProxySessionMessages::new(session_id.clone()))? + .run_indefinitely(); + + op(session_id).await + } + }) + } +} + +impl SessionBuilder +where + Counterpart: HasPeer, + R: RunWithConnectionTo, +{ + /// Mark this session builder as being able to block the current task. + /// + /// After calling this, you can use [`run_until`](Self::run_until) or + /// [`start_session`](Self::start_session) which block the current task. + /// + /// This should not be used from inside a message handler like + /// [`Builder::on_receive_request`](`crate::Builder::on_receive_request`) or [`HandleDispatchFrom`] + /// implementations. + pub fn block_task(self) -> SessionBuilder { + SessionBuilder { + connection: self.connection, + request: self.request, + dynamic_handler_registrations: self.dynamic_handler_registrations, + run: self.run, + block_state: PhantomData, + } + } +} + +impl SessionBuilder +where + Counterpart: HasPeer, + R: RunWithConnectionTo, +{ + /// Run this session synchronously. The current task will be blocked + /// and `op` will be executed with the active session information. + /// This is useful when you have MCP servers that are borrowed from your local + /// stack frame. + /// + /// The `ActiveSession` passed to `op` has a non-`'static` lifetime, which + /// prevents calling [`ActiveSession::proxy_remaining_messages`] (since the + /// responders would terminate when `op` returns). + /// + /// Requires calling [`block_task`](Self::block_task) first. + pub async fn run_until( + self, + op: impl for<'responder> AsyncFnOnce( + ActiveSession<'responder, Counterpart>, + ) -> Result, + ) -> Result { + let Self { + connection, + request, + dynamic_handler_registrations, + run, + block_state: _, + } = self; + + let response = connection + .send_request_to(Agent, request) + .block_task() + .await?; + + let active_session = connection.attach_session(response, dynamic_handler_registrations)?; + + run_until( + run.run_with_connection_to(connection.clone()), + op(active_session), + ) + .await + } + + /// Send the request to create the session and return a handle. + /// This is an alternative to [`Self::run_until`] that avoids rightward + /// drift but at the cost of requiring MCP servers that are `Send` and + /// don't access data from the surrounding scope. + /// + /// Returns an `ActiveSession<'static, _>` because responders are spawned + /// into background tasks that live for the connection lifetime. + /// + /// Requires calling [`block_task`](Self::block_task) first. + pub async fn start_session(self) -> Result, crate::Error> + where + R: 'static, + { + let Self { + connection, + request, + dynamic_handler_registrations, + run, + block_state: _, + } = self; + + let (active_session_tx, active_session_rx) = oneshot::channel(); + + connection.clone().spawn(async move { + let response = connection + .send_request_to(Agent, request) + .block_task() + .await?; + + connection.spawn(run.run_with_connection_to(connection.clone()))?; + + let active_session = + connection.attach_session(response, dynamic_handler_registrations)?; + + active_session_tx + .send(active_session) + .map_err(|_| crate::Error::internal_error())?; + + Ok(()) + })?; + + active_session_rx + .await + .map_err(|_| crate::Error::internal_error()) + } + + /// Start a proxy session that forwards all messages between client and agent. + /// + /// A **proxy session** starts the session with the agent and then automatically + /// proxies all session updates (prompts, tool calls, etc.) from the agent back + /// to the client. You don't need to handle any messages yourself - the proxy + /// takes care of forwarding everything. This is useful when you want to inject + /// and/or filter prompts coming from the client but otherwise not be involved + /// in the session. + /// + /// This is a convenience method that combines [`start_session`](Self::start_session), + /// responding to the client, and [`ActiveSession::proxy_remaining_messages`]. + /// + /// For more control (e.g., to send some messages before proxying), use + /// [`start_session`](Self::start_session) instead and call + /// [`proxy_remaining_messages`](ActiveSession::proxy_remaining_messages) manually. + /// + /// Requires calling [`block_task`](Self::block_task) first. + pub async fn start_session_proxy( + self, + responder: Responder, + ) -> Result + where + Counterpart: HasPeer, + R: 'static, + { + let active_session = self.start_session().await?; + let session_id = active_session.session_id().clone(); + responder.respond(active_session.response())?; + active_session.proxy_remaining_messages()?; + Ok(session_id) + } +} + +/// Active session struct that lets you send prompts and receive updates. +/// +/// The `'responder` lifetime represents the span during which responders +/// (e.g., MCP server handlers) are active. When created via [`SessionBuilder::start_session`], +/// this is `'static` because responders are spawned into background tasks. +/// When created via [`SessionBuilder::run_until`], this is tied to the +/// closure scope, preventing [`Self::proxy_remaining_messages`] from being called +/// (since the responders would die when the closure returns). +#[derive(Debug)] +pub struct ActiveSession<'responder, Link> +where + Link: HasPeer, +{ + session_id: SessionId, + update_rx: mpsc::UnboundedReceiver, + update_tx: mpsc::UnboundedSender, + modes: Option, + meta: Option>, + connection: ConnectionTo, + + /// Registration for the handler that routes session messages to `update_rx`. + /// This is separate from MCP handlers so it can be dropped independently + /// when switching to proxy mode. + session_handler_registration: DynamicHandlerRegistration, + + /// Registrations for MCP server handlers. + /// These will be dropped once the active-session struct is dropped + /// which will cause them to be deregistered. + mcp_handler_registrations: Vec>, + + /// Phantom lifetime representing the responder lifetime. + _responder: PhantomData<&'responder ()>, +} + +/// Incoming message from the agent +#[non_exhaustive] +#[derive(Debug)] +pub enum SessionMessage { + /// Periodic updates with new content, tool requests, etc. + /// Use [`MatchDispatch`] to match on the message type. + SessionMessage(Dispatch), + + /// When a prompt completes, the stop reason. + StopReason(StopReason), +} + +impl ActiveSession<'_, Link> +where + Link: HasPeer, +{ + /// Access the session ID. + pub fn session_id(&self) -> &SessionId { + &self.session_id + } + + /// Access modes available in this session. + pub fn modes(&self) -> &Option { + &self.modes + } + + /// Access meta data from session response. + pub fn meta(&self) -> &Option> { + &self.meta + } + + /// Build a `NewSessionResponse` from the session information. + /// + /// Useful when you need to forward the session response to a client + /// after doing some processing. + pub fn response(&self) -> NewSessionResponse { + NewSessionResponse::new(self.session_id.clone()) + .modes(self.modes.clone()) + .meta(self.meta.clone()) + } + + /// Access the underlying connection context used to communicate with the agent. + pub fn connection(&self) -> ConnectionTo { + self.connection.clone() + } + + /// Send a prompt to the agent. You can then read messages sent in response. + pub fn send_prompt(&mut self, prompt: impl ToString) -> Result<(), crate::Error> { + let update_tx = self.update_tx.clone(); + self.connection + .send_request_to( + Agent, + PromptRequest::new(self.session_id.clone(), vec![prompt.to_string().into()]), + ) + .on_receiving_result(async move |result| { + let PromptResponse { + stop_reason, + meta: _, + .. + } = result?; + + update_tx + .unbounded_send(SessionMessage::StopReason(stop_reason)) + .map_err(crate::util::internal_error)?; + + Ok(()) + }) + } + + /// Read an update from the agent in response to the prompt. + pub async fn read_update(&mut self) -> Result { + use futures::StreamExt; + let message = + self.update_rx.next().await.ok_or_else(|| { + crate::util::internal_error("session channel closed unexpectedly") + })?; + + Ok(message) + } + + /// Read all updates until the end of the turn and create a string. + /// Ignores non-text updates. + pub async fn read_to_string(&mut self) -> Result { + let mut output = String::new(); + loop { + let update = self.read_update().await?; + tracing::trace!(?update, "read_to_string update"); + match update { + SessionMessage::SessionMessage(dispatch) => MatchDispatch::new(dispatch) + .if_notification(async |notif: SessionNotification| match notif.update { + SessionUpdate::AgentMessageChunk(ContentChunk { + content: ContentBlock::Text(text), + meta: _, + .. + }) => { + output.push_str(&text.text); + Ok(()) + } + _ => Ok(()), + }) + .await + .otherwise_ignore()?, + SessionMessage::StopReason(_stop_reason) => break, + } + } + Ok(output) + } +} + +impl ActiveSession<'static, Link> +where + Link: HasPeer, +{ + /// Proxy all remaining messages for this session between client and agent. + /// + /// Use this when you want to inject MCP servers into a session but don't need + /// to actively interact with it after setup. The session messages will be proxied + /// between client and agent automatically. + /// + /// This consumes the `ActiveSession` since you're giving up active control. + /// + /// This method is only available on `ActiveSession<'static, _>` (from + /// [`SessionBuilder::start_session`]) because it requires responders to + /// outlive the method call. + /// + /// # Message Ordering Guarantees + /// + /// This method ensures proper handoff from active session mode to proxy mode + /// without losing or reordering messages: + /// + /// 1. **Stop the session handler** - Drop the registration that routes messages + /// to `update_rx`. After this, no new messages will be queued. + /// 2. **Close the channel** - Drop `update_tx` so we can detect when the channel + /// is fully drained. + /// 3. **Drain queued messages** - Forward any messages that were already queued + /// in `update_rx` to the client, preserving order. + /// 4. **Install proxy handler** - Now that all queued messages are forwarded, + /// install the proxy handler to handle future messages. + /// + /// This sequence prevents the race condition where messages could be delivered + /// out of order or lost during the transition. + pub fn proxy_remaining_messages(self) -> Result<(), crate::Error> + where + Link: HasPeer, + { + // Destructure self to get ownership of all fields + let ActiveSession { + session_id, + mut update_rx, + update_tx, + connection, + session_handler_registration, + mcp_handler_registrations, + // These fields are not needed for proxying + modes: _, + meta: _, + _responder, + } = self; + + // Step 1: Drop the session handler registration. + // This unregisters the handler that was routing messages to update_rx. + // After this point, no new messages will be added to the channel. + drop(session_handler_registration); + + // Step 2: Drop the sender side of the channel. + // This allows us to detect when the channel is fully drained + // (recv will return None when empty and sender is dropped). + drop(update_tx); + + // Step 3: Drain any messages that were already queued and forward to client. + // These messages arrived before we dropped the handler but haven't been + // consumed yet. We must forward them to maintain message ordering. + while let Ok(message) = update_rx.try_recv() { + match message { + SessionMessage::SessionMessage(dispatch) => { + // Forward the message to the client + connection.send_proxied_message_to(Client, dispatch)?; + } + SessionMessage::StopReason(_) => { + // StopReason is internal bookkeeping, not forwarded + } + } + } + + // Step 4: Install the proxy handler for future messages. + // Now that all queued messages have been forwarded, the proxy handler + // can take over. Any new messages will go directly through the proxy. + connection + .add_dynamic_handler(ProxySessionMessages::new(session_id))? + .run_indefinitely(); + + // Keep MCP server handlers alive for the lifetime of the proxy + for registration in mcp_handler_registrations { + registration.run_indefinitely(); + } + + Ok(()) + } +} + +struct ActiveSessionHandler { + session_id: SessionId, + update_tx: mpsc::UnboundedSender, +} + +impl ActiveSessionHandler { + pub fn new(session_id: SessionId, update_tx: mpsc::UnboundedSender) -> Self { + Self { + session_id, + update_tx, + } + } +} + +impl HandleDispatchFrom for ActiveSessionHandler +where + Counterpart: HasPeer, +{ + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + cx: ConnectionTo, + ) -> Result, crate::Error> { + // If this is a message for our session, grab it. + tracing::trace!( + ?message, + handler_session_id = ?self.session_id, + "ActiveSessionHandler::handle_dispatch" + ); + MatchDispatchFrom::new(message, &cx) + .if_message_from(Agent, async |message| { + if let Some(session_id) = message.get_session_id()? { + tracing::trace!( + message_session_id = ?session_id, + handler_session_id = ?self.session_id, + "ActiveSessionHandler::handle_dispatch" + ); + if session_id == self.session_id { + self.update_tx + .unbounded_send(SessionMessage::SessionMessage(message)) + .map_err(crate::util::internal_error)?; + return Ok(Handled::Yes); + } + } + + // Otherwise, pass it through. + Ok(Handled::No { + message, + retry: false, + }) + }) + .await + .done() + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + format!("ActiveSessionHandler({})", self.session_id) + } +} diff --git a/src/agent-client-protocol-core/src/typed.rs b/src/agent-client-protocol-core/src/typed.rs new file mode 100644 index 0000000..6dc70a2 --- /dev/null +++ b/src/agent-client-protocol-core/src/typed.rs @@ -0,0 +1,125 @@ +// Types re-exported from crate root +use jsonrpcmsg::Params; + +use crate::{ + ConnectionTo, Responder, JsonRpcNotification, JsonRpcRequest, UntypedMessage, + util::json_cast, +}; + +/// Utility class for handling untyped requests. +#[must_use] +pub struct TypeRequest { + state: Option, +} + +enum TypeMessageState { + Unhandled(String, Option, Responder), + Handled(Result<(), crate::Error>), +} + +impl TypeRequest { + pub fn new(request: UntypedMessage, responder: Responder) -> Self { + let UntypedMessage { method, params } = request; + let params: Option = json_cast(params).expect("valid params"); + Self { + state: Some(TypeMessageState::Unhandled(method, params, responder)), + } + } + + pub async fn handle_if( + mut self, + op: impl AsyncFnOnce(R, Responder) -> Result<(), crate::Error>, + ) -> Self { + self.state = Some(match self.state.take().expect("valid state") { + TypeMessageState::Unhandled(method, params, responder) => { + match R::parse_message(&method, ¶ms) { + Some(Ok(request)) => { + TypeMessageState::Handled(op(request, responder.cast()).await) + } + + Some(Err(err)) => TypeMessageState::Handled(responder.respond_with_error(err)), + + None => TypeMessageState::Unhandled(method, params, responder), + } + } + + TypeMessageState::Handled(err) => TypeMessageState::Handled(err), + }); + self + } + + pub async fn otherwise( + mut self, + op: impl AsyncFnOnce(UntypedMessage, Responder) -> Result<(), crate::Error>, + ) -> Result<(), crate::Error> { + match self.state.take().expect("valid state") { + TypeMessageState::Unhandled(method, params, responder) => { + match UntypedMessage::new(&method, params) { + Ok(m) => op(m, responder).await, + Err(err) => responder.respond_with_error(err), + } + } + TypeMessageState::Handled(r) => r, + } + } +} + +/// Utility class for handling untyped notifications. +#[must_use] +pub struct TypeNotification { + cx: ConnectionTo, + state: Option, +} + +enum TypeNotificationState { + Unhandled(String, Option), + Handled(Result<(), crate::Error>), +} + +impl TypeNotification { + pub fn new(request: UntypedMessage, cx: &ConnectionTo) -> Self { + let UntypedMessage { method, params } = request; + let params: Option = json_cast(params).expect("valid params"); + Self { + cx: cx.clone(), + state: Some(TypeNotificationState::Unhandled(method, params)), + } + } + + pub async fn handle_if( + mut self, + op: impl AsyncFnOnce(N) -> Result<(), crate::Error>, + ) -> Self { + self.state = Some(match self.state.take().expect("valid state") { + TypeNotificationState::Unhandled(method, params) => { + match N::parse_message(&method, ¶ms) { + Some(Ok(request)) => TypeNotificationState::Handled(op(request).await), + + Some(Err(err)) => { + TypeNotificationState::Handled(self.cx.send_error_notification(err)) + } + + None => TypeNotificationState::Unhandled(method, params), + } + } + + TypeNotificationState::Handled(err) => TypeNotificationState::Handled(err), + }); + self + } + + pub async fn otherwise( + mut self, + op: impl AsyncFnOnce(UntypedMessage) -> Result<(), crate::Error>, + ) -> Result<(), crate::Error> { + match self.state.take().expect("valid state") { + TypeNotificationState::Unhandled(method, params) => { + match UntypedMessage::new(&method, params) { + Ok(m) => op(m).await, + Err(err) => self.cx.send_error_notification(err), + } + } + TypeNotificationState::Handled(r) => r, + } + } +} diff --git a/src/agent-client-protocol-core/src/util.rs b/src/agent-client-protocol-core/src/util.rs new file mode 100644 index 0000000..310b02e --- /dev/null +++ b/src/agent-client-protocol-core/src/util.rs @@ -0,0 +1,184 @@ +// Types re-exported from crate root + +use futures::{ + future::BoxFuture, + stream::{Stream, StreamExt}, +}; + +mod typed; +pub use typed::{MatchDispatch, MatchDispatchFrom, TypeNotification}; + +/// Cast from `N` to `M` by serializing/deserialization to/from JSON. +pub fn json_cast(params: N) -> Result +where + N: serde::Serialize, + M: serde::de::DeserializeOwned, +{ + let json = serde_json::to_value(params).map_err(|e| { + crate::Error::parse_error().data(serde_json::json!({ + "error": e.to_string(), + "phase": "serialization" + })) + })?; + let m = serde_json::from_value(json.clone()).map_err(|e| { + crate::Error::parse_error().data(serde_json::json!({ + "error": e.to_string(), + "json": json, + "phase": "deserialization" + })) + })?; + Ok(m) +} + +/// Creates an internal error with the given message +pub fn internal_error(message: impl ToString) -> crate::Error { + crate::Error::internal_error().data(message.to_string()) +} + +/// Creates a parse error with the given message +pub fn parse_error(message: impl ToString) -> crate::Error { + crate::Error::parse_error().data(message.to_string()) +} + +/// Convert a JSON-RPC id to a serde_json::Value. +pub(crate) fn id_to_json(id: &jsonrpcmsg::Id) -> serde_json::Value { + match id { + jsonrpcmsg::Id::Number(n) => serde_json::Value::Number((*n).into()), + jsonrpcmsg::Id::String(s) => serde_json::Value::String(s.clone()), + jsonrpcmsg::Id::Null => serde_json::Value::Null, + } +} + +pub(crate) fn instrumented_with_connection_name( + name: String, + task: F, +) -> tracing::instrument::Instrumented { + use tracing::Instrument; + + task.instrument(tracing::info_span!("connection", name = name)) +} + +pub(crate) async fn instrument_with_connection_name( + name: Option, + task: impl Future, +) -> R { + if let Some(name) = name { + instrumented_with_connection_name(name.clone(), task).await + } else { + task.await + } +} + +/// Convert a `crate::Error` into a `crate::jsonrpcmsg::Error` +#[must_use] +pub fn into_jsonrpc_error(err: crate::Error) -> crate::jsonrpcmsg::Error { + crate::jsonrpcmsg::Error { + code: err.code.into(), + message: err.message, + data: err.data, + } +} + +/// Run two fallible futures concurrently, returning when both complete successfully +/// or when either fails. +pub async fn both( + a: impl Future>, + b: impl Future>, +) -> Result<(), E> { + let ((), ()) = futures::future::try_join(a, b).await?; + Ok(()) +} + +/// Run `background` until `foreground` completes. +/// +/// Returns the result of `foreground`. If `background` errors before +/// `foreground` completes, the error is propagated. If `background` +/// completes with `Ok(())`, we continue waiting for `foreground`. +pub async fn run_until( + background: impl Future>, + foreground: impl Future>, +) -> Result { + use futures::future::{Either, select}; + use std::pin::pin; + + match select(pin!(background), pin!(foreground)).await { + Either::Left((bg_result, fg_future)) => { + // Background finished first + bg_result?; // propagate error, or if Ok(()), keep waiting + fg_future.await + } + Either::Right((fg_result, _bg_future)) => { + // Foreground finished first, drop background + fg_result + } + } +} + +/// Process items from a stream concurrently. +/// +/// For each item received from `stream`, calls `process_fn` to create a future, +/// then runs all futures concurrently. If any future returns an error, +/// stops processing and returns that error. +/// +/// This is useful for patterns where you receive work items from a channel +/// and want to process them concurrently while respecting backpressure. +pub async fn process_stream_concurrently( + stream: impl Stream, + process_fn: F, + process_fn_hack: impl for<'a> Fn(&'a F, T) -> BoxFuture<'a, Result<(), crate::Error>>, +) -> Result<(), crate::Error> +where + F: AsyncFn(T) -> Result<(), crate::Error>, +{ + use std::pin::pin; + + use futures::stream::{FusedStream, FuturesUnordered}; + use futures_concurrency::future::Race; + + enum Event { + NewItem(Option), + FutureCompleted(Option>), + } + + let mut stream = pin!(stream.fuse()); + let mut futures: FuturesUnordered<_> = FuturesUnordered::new(); + + loop { + // If we have no futures to run, wait until we do. + if futures.is_empty() { + match stream.next().await { + Some(item) => futures.push(process_fn_hack(&process_fn, item)), + None => return Ok(()), + } + continue; + } + + // If there are no more items coming in, just drain our queue and return. + if stream.is_terminated() { + while let Some(result) = futures.next().await { + result?; + } + return Ok(()); + } + + // Otherwise, race between getting a new item and completing a future. + let event = (async { Event::NewItem(stream.next().await) }, async { + Event::FutureCompleted(futures.next().await) + }) + .race() + .await; + + match event { + Event::NewItem(Some(item)) => { + futures.push(process_fn_hack(&process_fn, item)); + } + Event::FutureCompleted(Some(result)) => { + result?; + } + Event::NewItem(None) | Event::FutureCompleted(None) => { + // Stream closed, loop will catch is_terminated + // No futures were pending, shouldn't happen since we checked is_empty + } + } + } +} diff --git a/src/agent-client-protocol-core/src/util/typed.rs b/src/agent-client-protocol-core/src/util/typed.rs new file mode 100644 index 0000000..c94e8dd --- /dev/null +++ b/src/agent-client-protocol-core/src/util/typed.rs @@ -0,0 +1,921 @@ +//! Utilities for pattern matching on untyped JSON-RPC messages. +//! +//! When handling [`UntypedMessage`]s, you can use [`MatchDispatch`] for simple parsing +//! or [`MatchDispatchFrom`] when you need peer-aware transforms (e.g., unwrapping +//! proxy envelopes). +//! +//! # When to use which +//! +//! - **[`MatchDispatchFrom`]**: Preferred over implementing [`HandleDispatchFrom`] directly. +//! Use this in connection handlers when you need to match on message types with +//! proper peer-aware transforms (e.g., unwrapping `SuccessorMessage` envelopes). +//! +//! - **[`MatchDispatch`]**: Use this when you already have an unwrapped message and +//! just need to parse it, such as inside a [`MatchDispatchFrom`] callback or when +//! processing messages that don't need peer transforms. +//! +//! [`HandleDispatchFrom`]: crate::HandleDispatchFrom + +// Types re-exported from crate root +use jsonrpcmsg::Params; + +use crate::{ + ConnectionTo, Dispatch, HandleDispatchFrom, Handled, JsonRpcNotification, JsonRpcRequest, + JsonRpcResponse, Responder, ResponseRouter, UntypedMessage, + role::{HasPeer, Role, handle_incoming_dispatch}, + util::json_cast, +}; + +/// Role-agnostic helper for pattern-matching on untyped JSON-RPC messages. +/// +/// Use this when you already have an unwrapped message and just need to parse it, +/// such as inside a [`MatchDispatchFrom`] callback or when processing messages +/// that don't need peer transforms. +/// +/// For connection handlers where you need proper peer-aware transforms, +/// use [`MatchDispatchFrom`] instead. +/// +/// # Example +/// +/// ``` +/// # use agent_client_protocol_core::Dispatch; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, AgentCapabilities}; +/// # use agent_client_protocol_core::util::MatchDispatch; +/// # async fn example(message: Dispatch) -> Result<(), agent_client_protocol_core::Error> { +/// MatchDispatch::new(message) +/// .if_request(|req: InitializeRequest, responder: agent_client_protocol_core::Responder| async move { +/// let response = InitializeResponse::new(req.protocol_version) +/// .agent_capabilities(AgentCapabilities::new()); +/// responder.respond(response) +/// }) +/// .await +/// .otherwise(|message| async move { +/// match message { +/// Dispatch::Request(_, responder) => { +/// responder.respond_with_error(agent_client_protocol_core::util::internal_error("unknown method")) +/// } +/// Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(()), +/// } +/// }) +/// .await +/// # } +/// ``` +#[must_use] +#[derive(Debug)] +pub struct MatchDispatch { + state: Result, crate::Error>, +} + +impl MatchDispatch { + /// Create a new pattern matcher for the given message. + pub fn new(message: Dispatch) -> Self { + Self { + state: Ok(Handled::No { + message, + retry: false, + }), + } + } + + /// Create a pattern matcher from an existing `Handled` state. + /// + /// This is useful when composing with [`MatchDispatchFrom`] which applies + /// peer transforms before delegating to `MatchDispatch` for parsing. + pub fn from_handled(state: Result, crate::Error>) -> Self { + Self { state } + } + + /// Try to handle the message as a request of type `Req`. + /// + /// If the message can be parsed as `Req`, the handler `op` is called with the parsed + /// request and a typed request context. If parsing fails or the message was already + /// handled by a previous call, this has no effect. + pub async fn if_request( + mut self, + op: impl AsyncFnOnce(Req, Responder) -> Result, + ) -> Self + where + H: crate::IntoHandled<(Req, Responder)>, + { + if let Ok(Handled::No { + message: dispatch, + retry, + }) = self.state + { + self.state = match dispatch { + Dispatch::Request(untyped_request, untyped_responder) => { + if Req::matches_method(untyped_request.method()) { + match Req::parse_message(untyped_request.method(), untyped_request.params()) + { + Ok(typed_request) => { + let typed_responder = untyped_responder.cast(); + match op(typed_request, typed_responder).await { + Ok(result) => match result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: (request, responder), + retry: request_retry, + } => match request.to_untyped_message() { + Ok(untyped) => Ok(Handled::No { + message: Dispatch::Request( + untyped, + responder.erase_to_json(), + ), + retry: retry | request_retry, + }), + Err(err) => Err(err), + }, + }, + Err(err) => Err(err), + } + } + Err(err) => Err(err), + } + } else { + Ok(Handled::No { + message: Dispatch::Request(untyped_request, untyped_responder), + retry, + }) + } + } + Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(Handled::No { + message: dispatch, + retry, + }), + }; + } + self + } + + /// Try to handle the message as a notification of type `N`. + /// + /// If the message can be parsed as `N`, the handler `op` is called with the parsed + /// notification. If parsing fails or the message was already handled, this has no effect. + pub async fn if_notification( + mut self, + op: impl AsyncFnOnce(N) -> Result, + ) -> Self + where + H: crate::IntoHandled, + { + if let Ok(Handled::No { + message: dispatch, + retry, + }) = self.state + { + self.state = match dispatch { + Dispatch::Notification(untyped_notification) => { + if N::matches_method(untyped_notification.method()) { + match N::parse_message( + untyped_notification.method(), + untyped_notification.params(), + ) { + Ok(typed_notification) => match op(typed_notification).await { + Ok(result) => match result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: notification, + retry: notification_retry, + } => match notification.to_untyped_message() { + Ok(untyped) => Ok(Handled::No { + message: Dispatch::Notification(untyped), + retry: retry | notification_retry, + }), + Err(err) => Err(err), + }, + }, + Err(err) => Err(err), + }, + Err(err) => Err(err), + } + } else { + Ok(Handled::No { + message: Dispatch::Notification(untyped_notification), + retry, + }) + } + } + Dispatch::Request(_, _) | Dispatch::Response(_, _) => Ok(Handled::No { + message: dispatch, + retry, + }), + }; + } + self + } + + /// Try to handle the message as a typed `Dispatch`. + /// + /// This attempts to parse the message as either request type `R` or notification type `N`, + /// providing a typed `Dispatch` to the handler if successful. + pub async fn if_message( + mut self, + op: impl AsyncFnOnce(Dispatch) -> Result, + ) -> Self + where + H: crate::IntoHandled>, + { + if let Ok(Handled::No { + message: dispatch, + retry, + }) = self.state + { + self.state = match dispatch.into_typed_dispatch::() { + Ok(Ok(typed_dispatch)) => match op(typed_dispatch).await { + Ok(result) => match result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: typed_dispatch, + retry: message_retry, + } => { + let untyped = match typed_dispatch { + Dispatch::Request(request, responder) => { + match request.to_untyped_message() { + Ok(untyped) => { + Dispatch::Request(untyped, responder.erase_to_json()) + } + Err(err) => return Self { state: Err(err) }, + } + } + Dispatch::Notification(notification) => { + match notification.to_untyped_message() { + Ok(untyped) => Dispatch::Notification(untyped), + Err(err) => return Self { state: Err(err) }, + } + } + Dispatch::Response(result, router) => { + let method = router.method(); + let untyped_result = match result { + Ok(response) => match response.into_json(method) { + Ok(json) => Ok(json), + Err(err) => return Self { state: Err(err) }, + }, + Err(err) => Err(err), + }; + Dispatch::Response(untyped_result, router.erase_to_json()) + } + }; + Ok(Handled::No { + message: untyped, + retry: retry | message_retry, + }) + } + }, + Err(err) => Err(err), + }, + Ok(Err(dispatch)) => Ok(Handled::No { + message: dispatch, + retry, + }), + Err(err) => Err(err), + }; + } + self + } + + /// Try to handle the message as a response to a request of type `Req`. + /// + /// If the message is a `Response` variant and the method matches `Req`, the handler + /// is called with the result (which may be `Ok` or `Err`) and a typed response context. + /// Use this when you need to handle both success and error responses. + /// + /// For handling only successful responses, see [`if_ok_response_to`](Self::if_ok_response_to). + pub async fn if_response_to( + mut self, + op: impl AsyncFnOnce( + Result, + ResponseRouter, + ) -> Result, + ) -> Self + where + H: crate::IntoHandled<( + Result, + ResponseRouter, + )>, + { + if let Ok(Handled::No { + message: dispatch, + retry, + }) = self.state + { + self.state = match dispatch { + Dispatch::Response(result, router) => { + // Check if the request type matches this method + if Req::matches_method(router.method()) { + // Method matches, parse the response + let typed_router: ResponseRouter = router.cast(); + let typed_result = match result { + Ok(value) => Req::Response::from_value(typed_router.method(), value), + Err(err) => Err(err), + }; + + match op(typed_result, typed_router).await { + Ok(handler_result) => match handler_result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: (result, router), + retry: response_retry, + } => { + // Convert typed result back to untyped + let untyped_result = match result { + Ok(response) => response.into_json(router.method()), + Err(err) => Err(err), + }; + Ok(Handled::No { + message: Dispatch::Response( + untyped_result, + router.erase_to_json(), + ), + retry: retry | response_retry, + }) + } + }, + Err(err) => Err(err), + } + } else { + // Method doesn't match, return unhandled + Ok(Handled::No { + message: Dispatch::Response(result, router), + retry, + }) + } + } + Dispatch::Request(_, _) | Dispatch::Notification(_) => Ok(Handled::No { + message: dispatch, + retry, + }), + }; + } + self + } + + /// Try to handle the message as a successful response to a request of type `Req`. + /// + /// If the message is a `Response` variant with an `Ok` result and the method matches `Req`, + /// the handler is called with the parsed response and a typed response context. + /// Error responses are passed through without calling the handler. + /// + /// This is a convenience wrapper around [`if_response_to`](Self::if_response_to) for the + /// common case where you only care about successful responses. + pub async fn if_ok_response_to( + self, + op: impl AsyncFnOnce(Req::Response, ResponseRouter) -> Result, + ) -> Self + where + H: crate::IntoHandled<(Req::Response, ResponseRouter)>, + { + self.if_response_to::(async move |result, router| match result { + Ok(response) => { + let handler_result = op(response, router).await?; + match handler_result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: (resp, router), + retry, + } => Ok(Handled::No { + message: (Ok(resp), router), + retry, + }), + } + } + Err(err) => Ok(Handled::No { + message: (Err(err), router), + retry: false, + }), + }) + .await + } + + /// Complete matching, returning `Handled::No` if no match was found. + pub fn done(self) -> Result, crate::Error> { + self.state + } + + /// Handle messages that didn't match any previous handler. + pub async fn otherwise( + self, + op: impl AsyncFnOnce(Dispatch) -> Result<(), crate::Error>, + ) -> Result<(), crate::Error> { + match self.state { + Ok(Handled::Yes) => Ok(()), + Ok(Handled::No { message, retry: _ }) => op(message).await, + Err(err) => Err(err), + } + } + + /// Handle messages that didn't match any previous handler. + pub fn otherwise_ignore(self) -> Result<(), crate::Error> { + match self.state { + Ok(_) => Ok(()), + Err(err) => Err(err), + } + } +} + +/// Role-aware helper for pattern-matching on untyped JSON-RPC requests. +/// +/// **Prefer this over implementing [`HandleDispatchFrom`] directly.** This provides +/// a more ergonomic API for matching on message types in connection handlers. +/// +/// Use this when you need peer-aware transforms (e.g., unwrapping proxy envelopes) +/// before parsing messages. For simple parsing without peer awareness (e.g., inside +/// a callback), use [`MatchDispatch`] instead. +/// +/// This wraps [`MatchDispatch`] and applies peer-specific message transformations +/// via `remote_style().handle_incoming_dispatch()` before delegating to `MatchDispatch` +/// for the actual parsing. +/// +/// [`HandleDispatchFrom`]: crate::HandleDispatchFrom +/// +/// # Example +/// +/// ``` +/// # use agent_client_protocol_core::Dispatch; +/// # use agent_client_protocol_core::schema::{InitializeRequest, InitializeResponse, PromptRequest, PromptResponse, AgentCapabilities, StopReason}; +/// # use agent_client_protocol_core::util::MatchDispatchFrom; +/// # async fn example(message: Dispatch, cx: &agent_client_protocol_core::ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { +/// MatchDispatchFrom::new(message, cx) +/// .if_request(|req: InitializeRequest, responder: agent_client_protocol_core::Responder| async move { +/// // Handle initialization +/// let response = InitializeResponse::new(req.protocol_version) +/// .agent_capabilities(AgentCapabilities::new()); +/// responder.respond(response) +/// }) +/// .await +/// .if_request(|_req: PromptRequest, responder: agent_client_protocol_core::Responder| async move { +/// // Handle prompts +/// responder.respond(PromptResponse::new(StopReason::EndTurn)) +/// }) +/// .await +/// .otherwise(|message| async move { +/// // Fallback for unrecognized messages +/// match message { +/// Dispatch::Request(_, responder) => responder.respond_with_error(agent_client_protocol_core::util::internal_error("unknown method")), +/// Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(()), +/// } +/// }) +/// .await +/// # } +/// ``` +#[must_use] +#[derive(Debug)] +pub struct MatchDispatchFrom { + state: Result, crate::Error>, + connection: ConnectionTo, +} + +impl MatchDispatchFrom { + /// Create a new pattern matcher for the given untyped request message. + pub fn new(message: Dispatch, cx: &ConnectionTo) -> Self { + Self { + state: Ok(Handled::No { + message, + retry: false, + }), + connection: cx.clone(), + } + } + + /// Try to handle the message as a request of type `Req`. + /// + /// If the message can be parsed as `Req`, the handler `op` is called with the parsed + /// request and a typed request context. If parsing fails or the message was already + /// handled by a previous `handle_if`, this call has no effect. + /// + /// The handler can return either `()` (which becomes `Handled::Yes`) or an explicit + /// `Handled` value to control whether the message should be passed to the next handler. + /// + /// Returns `self` to allow chaining multiple `handle_if` calls. + pub async fn if_request( + self, + op: impl AsyncFnOnce(Req, Responder) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled<(Req, Responder)>, + { + let counterpart = self.connection.counterpart(); + self.if_request_from(counterpart, op).await + } + + /// Try to handle the message as a request of type `Req` from a specific peer. + /// + /// This is similar to [`if_request`](Self::if_request), but first applies peer-specific + /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes when receiving + /// from an agent via a proxy). + /// + /// # Parameters + /// + /// * `peer` - The peer the message is expected to come from + /// * `op` - The handler to call if the message matches + pub async fn if_request_from( + mut self, + peer: Peer, + op: impl AsyncFnOnce(Req, Responder) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled<(Req, Responder)>, + { + if let Ok(Handled::No { message, retry: _ }) = self.state { + self.state = handle_incoming_dispatch( + self.connection.counterpart(), + peer, + message, + self.connection.clone(), + async |dispatch, _connection| { + // Delegate to MatchDispatch for parsing + MatchDispatch::new(dispatch).if_request(op).await.done() + }, + ) + .await; + } + self + } + + /// Try to handle the message as a notification of type `N`. + /// + /// If the message can be parsed as `N`, the handler `op` is called with the parsed + /// notification and connection context. If parsing fails or the message was already + /// handled by a previous `handle_if`, this call has no effect. + /// + /// The handler can return either `()` (which becomes `Handled::Yes`) or an explicit + /// `Handled` value to control whether the message should be passed to the next handler. + /// + /// Returns `self` to allow chaining multiple `handle_if` calls. + pub async fn if_notification( + self, + op: impl AsyncFnOnce(N) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled, + { + let counterpart = self.connection.counterpart(); + self.if_notification_from(counterpart, op).await + } + + /// Try to handle the message as a notification of type `N` from a specific peer. + /// + /// This is similar to [`if_notification`](Self::if_notification), but first applies peer-specific + /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes when receiving + /// from an agent via a proxy). + /// + /// # Parameters + /// + /// * `peer` - The peer the message is expected to come from + /// * `op` - The handler to call if the message matches + pub async fn if_notification_from( + mut self, + peer: Peer, + op: impl AsyncFnOnce(N) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled, + { + if let Ok(Handled::No { message, retry: _ }) = self.state { + self.state = handle_incoming_dispatch( + self.connection.counterpart(), + peer, + message, + self.connection.clone(), + async |dispatch, _connection| { + // Delegate to MatchDispatch for parsing + MatchDispatch::new(dispatch) + .if_notification(op) + .await + .done() + }, + ) + .await; + } + self + } + + /// Try to handle the message as a typed `Dispatch` from a specific peer. + /// + /// This is similar to [`MatchDispatch::if_message`], but first applies peer-specific + /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes). + /// + /// # Parameters + /// + /// * `peer` - The peer the message is expected to come from + /// * `op` - The handler to call if the message matches + pub async fn if_message_from( + mut self, + peer: Peer, + op: impl AsyncFnOnce(Dispatch) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled>, + { + if let Ok(Handled::No { message, retry: _ }) = self.state { + self.state = handle_incoming_dispatch( + self.connection.counterpart(), + peer, + message, + self.connection.clone(), + async |dispatch, _connection| { + // Delegate to MatchDispatch for parsing + MatchDispatch::new(dispatch).if_message(op).await.done() + }, + ) + .await; + } + self + } + + /// Try to handle the message as a response to a request of type `Req`. + /// + /// If the message is a `Response` variant and the method matches `Req`, the handler + /// is called with the result (which may be `Ok` or `Err`) and a typed response context. + /// + /// Unlike requests and notifications, responses don't need peer-specific transforms + /// (they don't have the `SuccessorMessage` envelope structure), so this method + /// delegates directly to [`MatchDispatch::if_response_to`]. + pub async fn if_response_to( + mut self, + op: impl AsyncFnOnce( + Result, + ResponseRouter, + ) -> Result, + ) -> Self + where + H: crate::IntoHandled<( + Result, + ResponseRouter, + )>, + { + if let Ok(Handled::No { message, retry: _ }) = self.state { + self.state = MatchDispatch::new(message) + .if_response_to::(op) + .await + .done(); + } + self + } + + /// Try to handle the message as a successful response to a request of type `Req`. + /// + /// If the message is a `Response` variant with an `Ok` result and the method matches `Req`, + /// the handler is called with the parsed response and a typed response context. + /// Error responses are passed through without calling the handler. + /// + /// This is a convenience wrapper around [`if_response_to`](Self::if_response_to). + pub async fn if_ok_response_to( + self, + op: impl AsyncFnOnce(Req::Response, ResponseRouter) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled<(Req::Response, ResponseRouter)>, + { + let counterpart = self.connection.counterpart(); + self.if_ok_response_to_from::(counterpart, op) + .await + } + + /// Try to handle the message as a response to a request of type `Req` from a specific peer. + /// + /// If the message is a `Response` variant, the method matches `Req`, and the `role_id` + /// matches the expected peer, the handler is called with the result and a typed response context. + /// + /// This is used to filter responses by the peer they came from, which is important + /// in proxy scenarios where responses might arrive from multiple peers. + pub async fn if_response_to_from( + mut self, + peer: Peer, + op: impl AsyncFnOnce( + Result, + ResponseRouter, + ) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled<( + Result, + ResponseRouter, + )>, + { + if let Ok(Handled::No { message, retry: _ }) = self.state { + self.state = handle_incoming_dispatch( + self.connection.counterpart(), + peer, + message, + self.connection.clone(), + async |dispatch, _connection| { + // Delegate to MatchDispatch for parsing + MatchDispatch::new(dispatch) + .if_response_to::(op) + .await + .done() + }, + ) + .await; + } + self + } + + /// Try to handle the message as a successful response to a request of type `Req` from a specific peer. + /// + /// This is a convenience wrapper around [`if_response_to_from`](Self::if_response_to_from) + /// for the common case where you only care about successful responses. + pub async fn if_ok_response_to_from( + self, + peer: Peer, + op: impl AsyncFnOnce(Req::Response, ResponseRouter) -> Result, + ) -> Self + where + Counterpart: HasPeer, + H: crate::IntoHandled<(Req::Response, ResponseRouter)>, + { + self.if_response_to_from::(peer, async move |result, router| match result { + Ok(response) => { + let handler_result = op(response, router).await?; + match handler_result.into_handled() { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message: (resp, router), + retry, + } => Ok(Handled::No { + message: (Ok(resp), router), + retry, + }), + } + } + Err(err) => Ok(Handled::No { + message: (Err(err), router), + retry: false, + }), + }) + .await + } + + /// Complete matching, returning `Handled::No` if no match was found. + pub fn done(self) -> Result, crate::Error> { + match self.state { + Ok(Handled::Yes) => Ok(Handled::Yes), + Ok(Handled::No { message, retry }) => Ok(Handled::No { message, retry }), + Err(err) => Err(err), + } + } + + /// Handle messages that didn't match any previous `handle_if` call. + /// + /// This is the fallback handler that receives the original untyped message if none + /// of the typed handlers matched. You must call this method to complete the pattern + /// matching chain and get the final result. + pub async fn otherwise( + self, + op: impl AsyncFnOnce(Dispatch) -> Result<(), crate::Error>, + ) -> Result<(), crate::Error> { + match self.state { + Ok(Handled::Yes) => Ok(()), + Ok(Handled::No { message, retry: _ }) => op(message).await, + Err(err) => Err(err), + } + } + + /// Handle messages that didn't match any previous `handle_if` call. + /// + /// This is the fallback handler that receives the original untyped message if none + /// of the typed handlers matched. You must call this method to complete the pattern + /// matching chain and get the final result. + pub async fn otherwise_delegate( + self, + mut handler: impl HandleDispatchFrom, + ) -> Result, crate::Error> { + match self.state? { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message, + retry: outer_retry, + } => match handler + .handle_dispatch_from(message, self.connection.clone()) + .await? + { + Handled::Yes => Ok(Handled::Yes), + Handled::No { + message, + retry: inner_retry, + } => Ok(Handled::No { + message, + retry: inner_retry | outer_retry, + }), + }, + } + } +} + +/// Builder for pattern-matching on untyped JSON-RPC notifications. +/// +/// Similar to [`MatchDispatch`] but specifically for notifications (fire-and-forget messages with no response). +/// +/// # Pattern +/// +/// The typical pattern is: +/// 1. Create a `TypeNotification` from an untyped message +/// 2. Chain `.handle_if()` calls for each type you want to try +/// 3. End with `.otherwise()` for messages that don't match any type +/// +/// # Example +/// +/// ``` +/// # use agent_client_protocol_core::{UntypedMessage, ConnectionTo, Agent}; +/// # use agent_client_protocol_core::schema::SessionNotification; +/// # use agent_client_protocol_core::util::TypeNotification; +/// # async fn example(message: UntypedMessage, cx: &ConnectionTo) -> Result<(), agent_client_protocol_core::Error> { +/// TypeNotification::new(message, cx) +/// .handle_if(|notif: SessionNotification| async move { +/// // Handle session notifications +/// println!("Session update: {:?}", notif); +/// Ok(()) +/// }) +/// .await +/// .otherwise(|untyped: UntypedMessage| async move { +/// // Fallback for unrecognized notifications +/// println!("Unknown notification: {}", untyped.method); +/// Ok(()) +/// }) +/// .await +/// # } +/// ``` +/// +/// Since notifications don't expect responses, handlers only receive the parsed +/// notification (not a request context). +#[must_use] +#[derive(Debug)] +pub struct TypeNotification { + cx: ConnectionTo, + state: Option, +} + +#[derive(Debug)] +enum TypeNotificationState { + Unhandled(String, Option), + Handled(Result<(), crate::Error>), +} + +impl TypeNotification { + /// Create a new pattern matcher for the given untyped notification message. + pub fn new(request: UntypedMessage, cx: &ConnectionTo) -> Self { + let UntypedMessage { method, params } = request; + let params: Option = json_cast(params).expect("valid params"); + Self { + cx: cx.clone(), + state: Some(TypeNotificationState::Unhandled(method, params)), + } + } + + /// Try to handle the message as type `N`. + /// + /// If the message can be parsed as `N`, the handler `op` is called with the parsed + /// notification. If parsing fails or the message was already handled by a previous + /// `handle_if`, this call has no effect. + /// + /// Returns `self` to allow chaining multiple `handle_if` calls. + pub async fn handle_if( + mut self, + op: impl AsyncFnOnce(N) -> Result<(), crate::Error>, + ) -> Self { + self.state = Some(match self.state.take().expect("valid state") { + TypeNotificationState::Unhandled(method, params) => { + if N::matches_method(&method) { + match N::parse_message(&method, ¶ms) { + Ok(request) => TypeNotificationState::Handled(op(request).await), + Err(err) => { + TypeNotificationState::Handled(self.cx.send_error_notification(err)) + } + } + } else { + TypeNotificationState::Unhandled(method, params) + } + } + + TypeNotificationState::Handled(err) => TypeNotificationState::Handled(err), + }); + self + } + + /// Handle messages that didn't match any previous `handle_if` call. + /// + /// This is the fallback handler that receives the original untyped message if none + /// of the typed handlers matched. You must call this method to complete the pattern + /// matching chain and get the final result. + pub async fn otherwise( + mut self, + op: impl AsyncFnOnce(UntypedMessage) -> Result<(), crate::Error>, + ) -> Result<(), crate::Error> { + match self.state.take().expect("valid state") { + TypeNotificationState::Unhandled(method, params) => { + match UntypedMessage::new(&method, params) { + Ok(m) => op(m).await, + Err(err) => self.cx.send_error_notification(err), + } + } + TypeNotificationState::Handled(r) => r, + } + } +} diff --git a/src/agent-client-protocol-core/tests/derive_macros.rs b/src/agent-client-protocol-core/tests/derive_macros.rs new file mode 100644 index 0000000..a71799d --- /dev/null +++ b/src/agent-client-protocol-core/tests/derive_macros.rs @@ -0,0 +1,133 @@ +//! Tests for the JsonRpcRequest, JsonRpcNotification, and JsonRpcResponse derive macros. + +use agent_client_protocol_core::{ + JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, +}; +use serde::{Deserialize, Serialize}; + +// ============================================================================ +// Test types using derive macros +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize, JsonRpcRequest)] +#[request(method = "_test/hello", response = HelloResponse)] +struct HelloRequest { + name: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonRpcResponse)] +struct HelloResponse { + greeting: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)] +#[notification(method = "_test/ping")] +struct PingNotification { + timestamp: u64, +} + +// ============================================================================ +// Tests +// ============================================================================ + +#[test] +fn test_jr_request_method() { + let req = HelloRequest { + name: "world".into(), + }; + assert_eq!(req.method(), "_test/hello"); +} + +#[test] +fn test_jr_request_to_untyped() { + let req = HelloRequest { + name: "world".into(), + }; + let untyped = req.to_untyped_message().unwrap(); + assert_eq!(untyped.method(), "_test/hello"); +} + +#[test] +fn test_jr_request_parse_message() { + let original = HelloRequest { + name: "test".into(), + }; + let untyped = original.to_untyped_message().unwrap(); + + // matches_method should return true for correct method + assert!(HelloRequest::matches_method(untyped.method())); + assert!(!HelloRequest::matches_method("wrong/method")); + + // Parse should succeed for matching method + let parsed = HelloRequest::parse_message(untyped.method(), untyped.params()); + assert!(parsed.is_ok()); + let parsed = parsed.unwrap(); + assert_eq!(parsed.name, "test"); + + // Parse should return Err for non-matching method + let wrong_method = HelloRequest::parse_message("wrong/method", untyped.params()); + assert!(wrong_method.is_err()); +} + +#[test] +fn test_jr_request_response_type() { + // This is a compile-time check that the Response type is correctly set + fn assert_response_type>() {} + assert_response_type::(); +} + +#[test] +fn test_jr_notification_method() { + let notif = PingNotification { timestamp: 12345 }; + assert_eq!(notif.method(), "_test/ping"); +} + +#[test] +fn test_jr_notification_to_untyped() { + let notif = PingNotification { timestamp: 12345 }; + let untyped = notif.to_untyped_message().unwrap(); + assert_eq!(untyped.method(), "_test/ping"); +} + +#[test] +fn test_jr_notification_parse_message() { + let original = PingNotification { timestamp: 99999 }; + let untyped = original.to_untyped_message().unwrap(); + + // matches_method should work correctly + assert!(PingNotification::matches_method(untyped.method())); + assert!(!PingNotification::matches_method("wrong/method")); + + let parsed = PingNotification::parse_message(untyped.method(), untyped.params()); + assert!(parsed.is_ok()); + let parsed = parsed.unwrap(); + assert_eq!(parsed.timestamp, 99999); +} + +#[test] +fn test_jr_response_payload_into_json() { + let response = HelloResponse { + greeting: "Hello, world!".into(), + }; + let json = response.into_json("_test/hello").unwrap(); + assert_eq!(json["greeting"], "Hello, world!"); +} + +#[test] +fn test_jr_response_payload_from_value() { + let json = serde_json::json!({ + "greeting": "Hi there!" + }); + let response = HelloResponse::from_value("_test/hello", json).unwrap(); + assert_eq!(response.greeting, "Hi there!"); +} + +// ============================================================================ +// Test that JsonRpcNotification is a marker trait +// ============================================================================ + +#[test] +fn test_jr_notification_is_marker() { + fn assert_notification() {} + assert_notification::(); +} diff --git a/src/agent-client-protocol-core/tests/jsonrpc_advanced.rs b/src/agent-client-protocol-core/tests/jsonrpc_advanced.rs new file mode 100644 index 0000000..eeee00b --- /dev/null +++ b/src/agent-client-protocol-core/tests/jsonrpc_advanced.rs @@ -0,0 +1,368 @@ +//! Advanced feature tests for JSON-RPC layer +//! +//! Tests advanced JSON-RPC capabilities: +//! - Bidirectional communication (both sides can be client+server) +//! - Request ID tracking and matching +//! - Out-of-order response handling + +use agent_client_protocol_core::{ + ConnectionTo, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, Responder, SentRequest, + role::UntypedRole, +}; +use futures::{AsyncRead, AsyncWrite}; +use serde::{Deserialize, Serialize}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Test helper to block and wait for a JSON-RPC response. +async fn recv( + response: SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +/// Helper to set up test streams for testing. +fn setup_test_streams() -> ( + impl AsyncRead, + impl AsyncWrite, + impl AsyncRead, + impl AsyncWrite, +) { + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + (server_reader, server_writer, client_reader, client_writer) +} + +// ============================================================================ +// Test types +// ============================================================================ + +#[derive(Debug, Serialize, Deserialize, Clone)] +struct PingRequest { + value: u32, +} + +impl JsonRpcMessage for PingRequest { + fn matches_method(method: &str) -> bool { + method == "ping" + } + + fn method(&self) -> &'static str { + "ping" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for PingRequest { + type Response = PongResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct PongResponse { + value: u32, +} + +impl JsonRpcResponse for PongResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol_core::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(&value) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct SlowRequest { + delay_ms: u64, + id: u32, +} + +impl JsonRpcMessage for SlowRequest { + fn matches_method(method: &str) -> bool { + method == "slow" + } + + fn method(&self) -> &'static str { + "slow" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for SlowRequest { + type Response = SlowResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct SlowResponse { + id: u32, +} + +impl JsonRpcResponse for SlowResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol_core::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(&value) + } +} + +// ============================================================================ +// Test 1: Bidirectional communication +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_bidirectional_communication() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + // Set up two connections that are symmetric - both can send and receive + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let side_a_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let side_a = UntypedRole.builder().on_receive_request( + async |request: PingRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(PongResponse { + value: request.value + 1, + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let side_b_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + + // Spawn side_a as server + tokio::task::spawn_local(async move { + Box::pin(side_a.connect_to(side_a_transport)).await.ok(); + }); + + // Use side_b as client + let result = UntypedRole + .builder() + .connect_with( + side_b_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + let request = PingRequest { value: 10 }; + let response_future = recv(cx.send_request(request)); + let response: Result = response_future.await; + + assert!(response.is_ok()); + if let Ok(resp) = response { + assert_eq!(resp.value, 11); + } + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 2: Request IDs are properly tracked +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_request_ids() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |request: PingRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(PongResponse { + value: request.value + 1, + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + Box::pin(server.connect_to(server_transport)).await.ok(); + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + // Send multiple requests and verify responses match + let req1 = PingRequest { value: 1 }; + let req2 = PingRequest { value: 2 }; + let req3 = PingRequest { value: 3 }; + + let resp1_future = recv(cx.send_request(req1)); + let resp2_future = recv(cx.send_request(req2)); + let resp3_future = recv(cx.send_request(req3)); + + let resp1: Result = resp1_future.await; + let resp2: Result = resp2_future.await; + let resp3: Result = resp3_future.await; + + // Verify each response corresponds to its request + assert_eq!(resp1.unwrap().value, 2); // 1 + 1 + assert_eq!(resp2.unwrap().value, 3); // 2 + 1 + assert_eq!(resp3.unwrap().value, 4); // 3 + 1 + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 3: Out-of-order responses +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_out_of_order_responses() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |request: SlowRequest, + responder: Responder, + _connection: ConnectionTo| { + // Simulate delay + tokio::time::sleep(tokio::time::Duration::from_millis(request.delay_ms)).await; + responder.respond(SlowResponse { id: request.id }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + Box::pin(server.connect_to(server_transport)).await.ok(); + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + // Send requests with different delays + // Request 1: 100ms delay + // Request 2: 50ms delay + // Request 3: 10ms delay + // Responses should arrive in order: 3, 2, 1 + + let req1 = SlowRequest { + delay_ms: 100, + id: 1, + }; + let req2 = SlowRequest { + delay_ms: 50, + id: 2, + }; + let req3 = SlowRequest { + delay_ms: 10, + id: 3, + }; + + let resp1_future = recv(cx.send_request(req1)); + let resp2_future = recv(cx.send_request(req2)); + let resp3_future = recv(cx.send_request(req3)); + + // Wait for all responses + let resp1: Result = resp1_future.await; + let resp2: Result = resp2_future.await; + let resp3: Result = resp3_future.await; + + // Verify each future got the correct response despite out-of-order arrival + assert_eq!(resp1.unwrap().id, 1); + assert_eq!(resp2.unwrap().id, 2); + assert_eq!(resp3.unwrap().id, 3); + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} diff --git a/src/agent-client-protocol-core/tests/jsonrpc_connection_builder.rs b/src/agent-client-protocol-core/tests/jsonrpc_connection_builder.rs new file mode 100644 index 0000000..a0fb458 --- /dev/null +++ b/src/agent-client-protocol-core/tests/jsonrpc_connection_builder.rs @@ -0,0 +1,749 @@ +//! Integration tests for JSON-RPC connection builder behavior. +//! +//! These tests verify that multiple handlers can be registered on a connection builder +//! and that requests/notifications are routed correctly based on which +//! handler claims them. + +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use agent_client_protocol_core::{ + ConnectTo, ConnectionTo, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, + Responder, SentRequest, role::UntypedRole, util::run_until, +}; +use serde::{Deserialize, Serialize}; +use tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _}; + +/// Test helper to block and wait for a JSON-RPC response. +async fn recv( + response: SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +// ============================================================================ +// Test 1: Multiple handlers with different methods +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct FooRequest { + value: String, +} + +impl JsonRpcMessage for FooRequest { + fn matches_method(method: &str) -> bool { + method == "foo" + } + + fn method(&self) -> &'static str { + "foo" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for FooRequest { + type Response = FooResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct FooResponse { + result: String, +} + +impl JsonRpcResponse for FooResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol_core::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(&value) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct BarRequest { + value: String, +} + +impl JsonRpcMessage for BarRequest { + fn matches_method(method: &str) -> bool { + method == "bar" + } + + fn method(&self) -> &'static str { + "bar" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for BarRequest { + type Response = BarResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct BarResponse { + result: String, +} + +impl JsonRpcResponse for BarResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol_core::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(&value) + } +} + +#[tokio::test(flavor = "current_thread")] +async fn test_multiple_handlers_different_methods() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + // Chain both handlers + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole + .builder() + .on_receive_request( + async |request: FooRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(FooResponse { + result: format!("foo: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_request( + async |request: BarRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(BarResponse { + result: format!("bar: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + // Test foo request + let foo_response = recv(cx.send_request(FooRequest { + value: "test1".to_string(), + })) + .await + .map_err(|e| -> agent_client_protocol_core::Error { + agent_client_protocol_core::util::internal_error(format!( + "Foo request failed: {e:?}" + )) + })?; + assert_eq!(foo_response.result, "foo: test1"); + + // Test bar request + let bar_response = recv(cx.send_request(BarRequest { + value: "test2".to_string(), + })) + .await + .map_err(|e| -> agent_client_protocol_core::Error { + agent_client_protocol_core::util::internal_error(format!( + "Bar request failed: {e:?}" + )) + })?; + assert_eq!(bar_response.result, "bar: test2"); + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 2: Handler priority/ordering (first handler gets first chance) +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct TrackRequest { + value: String, +} + +impl JsonRpcMessage for TrackRequest { + fn matches_method(method: &str) -> bool { + method == "track" + } + + fn method(&self) -> &'static str { + "track" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for TrackRequest { + type Response = FooResponse; +} + +#[tokio::test(flavor = "current_thread")] +async fn test_handler_priority_ordering() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let handled = Arc::new(Mutex::new(Vec::new())); + + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + // First handler in chain should get first chance + let handled_clone1 = handled.clone(); + let handled_clone2 = handled.clone(); + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole + .builder() + .on_receive_request( + async move |request: TrackRequest, + responder: Responder, + _connection: ConnectionTo| { + handled_clone1.lock().unwrap().push("handler1".to_string()); + responder.respond(FooResponse { + result: format!("handler1: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_request( + async move |request: TrackRequest, + responder: Responder, + _connection: ConnectionTo| { + handled_clone2.lock().unwrap().push("handler2".to_string()); + responder.respond(FooResponse { + result: format!("handler2: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + let response = recv(cx.send_request(TrackRequest { + value: "test".to_string(), + })) + .await + .map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Track request failed: {e:?}" + )) + })?; + + // First handler should have handled it + assert_eq!(response.result, "handler1: test"); + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + + // Verify only handler1 was invoked + let handled_by = handled.lock().unwrap(); + assert_eq!(handled_by.len(), 1); + assert_eq!(handled_by[0], "handler1"); + })) + .await; +} + +// ============================================================================ +// Test 3: Fallthrough behavior (handler passes to next) +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Method1Request { + value: String, +} + +impl JsonRpcMessage for Method1Request { + fn matches_method(method: &str) -> bool { + method == "method1" + } + + fn method(&self) -> &'static str { + "method1" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for Method1Request { + type Response = FooResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Method2Request { + value: String, +} + +impl JsonRpcMessage for Method2Request { + fn matches_method(method: &str) -> bool { + method == "method2" + } + + fn method(&self) -> &'static str { + "method2" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for Method2Request { + type Response = FooResponse; +} + +#[tokio::test(flavor = "current_thread")] +async fn test_fallthrough_behavior() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let handled = Arc::new(Mutex::new(Vec::new())); + + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + // Handler1 only handles "method1", Handler2 only handles "method2" + let handled_clone1 = handled.clone(); + let handled_clone2 = handled.clone(); + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole + .builder() + .on_receive_request( + async move |request: Method1Request, + responder: Responder, + _connection: ConnectionTo| { + handled_clone1.lock().unwrap().push("method1".to_string()); + responder.respond(FooResponse { + result: format!("method1: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_request( + async move |request: Method2Request, + responder: Responder, + _connection: ConnectionTo| { + handled_clone2.lock().unwrap().push("method2".to_string()); + responder.respond(FooResponse { + result: format!("method2: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + // Send method2 - should fallthrough handler1 to handler2 + let response = recv(cx.send_request(Method2Request { + value: "fallthrough".to_string(), + })) + .await + .map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Method2 request failed: {e:?}" + )) + })?; + + assert_eq!(response.result, "method2: fallthrough"); + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + + // Verify only method2 was handled (handler1 passed through) + let handled_methods = handled.lock().unwrap(); + assert_eq!(handled_methods.len(), 1); + assert_eq!(handled_methods[0], "method2"); + })) + .await; +} + +// ============================================================================ +// Test 4: No handler claims request +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_no_handler_claims() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + // Handler that only handles "foo" + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |request: FooRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(FooResponse { + result: format!("foo: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + // Send "bar" request which no handler claims + let response_result = recv(cx.send_request(BarRequest { + value: "unclaimed".to_string(), + })) + .await; + + // Should get an error (method not found) + assert!(response_result.is_err()); + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 5: Handler can claim notifications +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct EventNotification { + event: String, +} + +impl JsonRpcMessage for EventNotification { + fn matches_method(method: &str) -> bool { + method == "event" + } + + fn method(&self) -> &'static str { + "event" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcNotification for EventNotification {} + +#[tokio::test(flavor = "current_thread")] +async fn test_handler_claims_notification() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let events = Arc::new(Mutex::new(Vec::new())); + + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + // EventHandler claims notifications + let events_clone = events.clone(); + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_notification( + async move |notification: EventNotification, _connection: ConnectionTo| { + events_clone.lock().unwrap().push(notification.event); + Ok(()) + }, + agent_client_protocol_core::on_receive_notification!(), + ); + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + cx.send_notification(EventNotification { + event: "test_event".to_string(), + }) + .map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Failed to send notification: {e:?}" + )) + })?; + + // Give server time to process + tokio::time::sleep(Duration::from_millis(100)).await; + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + + let received_events = events.lock().unwrap(); + assert_eq!(received_events.len(), 1); + assert_eq!(received_events[0], "test_event"); + })) + .await; +} + +// ============================================================================ +// Test 6: Builder implements Component +// ============================================================================ + +#[tokio::test] +async fn test_connection_builder_as_component() -> Result<(), agent_client_protocol_core::Error> { + // Create duplex streams + let (server_stream, client_stream) = tokio::io::duplex(8192); + let (server_read, server_write) = tokio::io::split(server_stream); + let (client_read, client_write) = tokio::io::split(client_stream); + + // Create a connection builder (server side) + let server_builder = UntypedRole.builder().on_receive_request( + async |request: FooRequest, + responder: Responder, + _cx: ConnectionTo| { + responder.respond(FooResponse { + result: format!("component: {}", request.value), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + // Create ByteStreams for both sides + let server_transport = agent_client_protocol_core::ByteStreams::new( + server_write.compat_write(), + server_read.compat(), + ); + let client_transport = agent_client_protocol_core::ByteStreams::new( + client_write.compat_write(), + client_read.compat(), + ); + + // Use Builder as a Component via run_until + Box::pin(run_until( + // This uses Component::serve on Builder + ConnectTo::::connect_to(server_builder, server_transport), + async move { + // Client side + UntypedRole + .builder() + .connect_with(client_transport, async |cx| { + let response = recv(cx.send_request(FooRequest { + value: "test".to_string(), + })) + .await?; + + assert_eq!(response.result, "component: test"); + Ok(()) + }) + .await + }, + )) + .await +} diff --git a/src/agent-client-protocol-core/tests/jsonrpc_edge_cases.rs b/src/agent-client-protocol-core/tests/jsonrpc_edge_cases.rs new file mode 100644 index 0000000..052722d --- /dev/null +++ b/src/agent-client-protocol-core/tests/jsonrpc_edge_cases.rs @@ -0,0 +1,376 @@ +//! Edge case tests for JSON-RPC layer +//! +//! Tests various edge cases and boundary conditions: +//! - Empty requests +//! - Null parameters +//! - Server shutdown scenarios +//! - Client disconnect handling + +use agent_client_protocol_core::{ + ConnectionTo, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, Responder, SentRequest, + role::UntypedRole, +}; +use futures::{AsyncRead, AsyncWrite}; +use serde::{Deserialize, Serialize}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Test helper to block and wait for a JSON-RPC response. +async fn recv( + response: SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +/// Helper to set up test streams. +fn setup_test_streams() -> ( + impl AsyncRead, + impl AsyncWrite, + impl AsyncRead, + impl AsyncWrite, +) { + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + (server_reader, server_writer, client_reader, client_writer) +} + +// ============================================================================ +// Test types +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct EmptyRequest; + +impl JsonRpcMessage for EmptyRequest { + fn matches_method(method: &str) -> bool { + method == "empty_method" + } + + fn method(&self) -> &'static str { + "empty_method" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + _params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + Ok(EmptyRequest) + } +} + +impl JsonRpcRequest for EmptyRequest { + type Response = SimpleResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct OptionalParamsRequest { + #[serde(skip_serializing_if = "Option::is_none")] + value: Option, +} + +impl JsonRpcMessage for OptionalParamsRequest { + fn matches_method(method: &str) -> bool { + method == "optional_params_method" + } + + fn method(&self) -> &'static str { + "optional_params_method" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for OptionalParamsRequest { + type Response = SimpleResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct SimpleResponse { + result: String, +} + +impl JsonRpcResponse for SimpleResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol_core::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(&value) + } +} + +// ============================================================================ +// Test 1: Empty request (no parameters) +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_empty_request() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |_request: EmptyRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(SimpleResponse { + result: "Got empty request".to_string(), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + Box::pin(server.connect_to(server_transport)).await.ok(); + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + let request = EmptyRequest; + + let result: Result = recv(cx.send_request(request)).await; + + // Should succeed + assert!(result.is_ok()); + if let Ok(response) = result { + assert_eq!(response.result, "Got empty request"); + } + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 2: Null parameters +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_null_params() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |_request: OptionalParamsRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(SimpleResponse { + result: "Has params: true".to_string(), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + Box::pin(server.connect_to(server_transport)).await.ok(); + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + let request = OptionalParamsRequest { value: None }; + + let result: Result = recv(cx.send_request(request)).await; + + // Should succeed - handler should handle null/missing params + assert!(result.is_ok()); + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 3: Server shutdown with pending requests +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_server_shutdown() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + local + .run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |_request: EmptyRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(SimpleResponse { + result: "Got empty request".to_string(), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + let server_handle = tokio::task::spawn_local(async move { + Box::pin(server.connect_to(server_transport)).await.ok(); + }); + + let client_result = tokio::task::spawn_local(async move { + client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + let request = EmptyRequest; + + // Send request and get future for response + let response_future = recv(cx.send_request(request)); + + // Give the request time to be sent over the wire + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + + // Try to get response (server should still be running briefly) + let _result: Result = response_future.await; + + // Could succeed or fail depending on timing + // The important thing is that it doesn't hang + Ok(()) + }, + ) + .await + }); + + // Let the client send its request + tokio::time::sleep(tokio::time::Duration::from_millis(5)).await; + + // Abort the server + server_handle.abort(); + + // Wait for client to finish + let result = client_result.await; + assert!(result.is_ok(), "Test failed: {result:?}"); + }) + .await; +} + +// ============================================================================ +// Test 4: Client disconnect mid-request +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_client_disconnect() { + use tokio::io::AsyncWriteExt; + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + local + .run_until(async { + let (mut client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, _client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |_request: EmptyRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(SimpleResponse { + result: "Got empty request".to_string(), + }) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + tokio::task::spawn_local(async move { + drop(Box::pin(server.connect_to(server_transport)).await); + }); + + // Send partial request and then disconnect + let partial_request = b"{\"jsonrpc\":\"2.0\",\"method\":\"empty_method\",\"id\":1"; + client_writer.write_all(partial_request).await.unwrap(); + client_writer.flush().await.unwrap(); + + // Drop the writer to disconnect + drop(client_writer); + + // Give server time to process the disconnect + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + // Server should handle this gracefully and terminate + // (We can't really assert much here except that the test completes) + }) + .await; +} diff --git a/src/agent-client-protocol-core/tests/jsonrpc_error_handling.rs b/src/agent-client-protocol-core/tests/jsonrpc_error_handling.rs new file mode 100644 index 0000000..8f9154b --- /dev/null +++ b/src/agent-client-protocol-core/tests/jsonrpc_error_handling.rs @@ -0,0 +1,445 @@ +//! Error handling tests for JSON-RPC layer +//! +//! Tests various error conditions: +//! - Invalid JSON +//! - Unknown methods +//! - Handler-returned errors +//! - Serialization failures +//! - Missing/invalid parameters + +use agent_client_protocol_core::{ + ConnectionTo, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, Responder, SentRequest, + role::UntypedRole, +}; +use expect_test::expect; +use futures::{AsyncRead, AsyncWrite}; +use serde::{Deserialize, Serialize}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Test helper to block and wait for a JSON-RPC response. +async fn recv( + response: SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +/// Helper to set up test streams. +fn setup_test_streams() -> ( + impl AsyncRead, + impl AsyncWrite, + impl AsyncRead, + impl AsyncWrite, +) { + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + (server_reader, server_writer, client_reader, client_writer) +} + +// ============================================================================ +// Test types +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct SimpleRequest { + message: String, +} + +impl JsonRpcMessage for SimpleRequest { + fn matches_method(method: &str) -> bool { + method == "simple_method" + } + + fn method(&self) -> &'static str { + "simple_method" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for SimpleRequest { + type Response = SimpleResponse; +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct SimpleResponse { + result: String, +} + +impl JsonRpcResponse for SimpleResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol_core::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(&value) + } +} + +// ============================================================================ +// Test 1: Invalid JSON (complete line with parse error) +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_invalid_json() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + local + .run_until(async { + // Create duplex streams for bidirectional communication + let (mut client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, mut client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + + // No handlers - all requests will return errors + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder(); + + // Spawn server + tokio::task::spawn_local(async move { + drop(server.connect_to(server_transport).await); + }); + + // Send invalid JSON + let invalid_json = b"{\"method\": \"test\", \"id\": 1, INVALID}\n"; + client_writer.write_all(invalid_json).await.unwrap(); + client_writer.flush().await.unwrap(); + + // Read response + let mut buffer = vec![0u8; 1024]; + let n = client_reader.read(&mut buffer).await.unwrap(); + let response_str = String::from_utf8_lossy(&buffer[..n]); + + // Parse as JSON and verify structure + let response: serde_json::Value = + serde_json::from_str(response_str.trim()).expect("Response should be valid JSON"); + + // Use expect_test to verify the exact structure + expect![[r#" + { + "error": { + "code": -32700, + "data": { + "line": "{\"method\": \"test\", \"id\": 1, INVALID}" + }, + "message": "Parse error" + }, + "jsonrpc": "2.0" + }"#]] + .assert_eq(&serde_json::to_string_pretty(&response).unwrap()); + }) + .await; +} + +// ============================================================================ +// Test 1b: Incomplete line (EOF mid-message) +// ============================================================================ + +#[tokio::test] +#[ignore = "hangs indefinitely - see https://github.com/agentclientprotocol/rust-sdk/issues/64"] +async fn test_incomplete_line() { + use futures::io::Cursor; + + // Incomplete JSON input - no newline, simulates client disconnect + let incomplete_json = b"{\"method\": \"test\", \"id\": 1"; + let input = Cursor::new(incomplete_json.to_vec()); + let output = Cursor::new(Vec::new()); + + // No handlers needed for EOF test + let transport = agent_client_protocol_core::ByteStreams::new(output, input); + let connection = UntypedRole.builder(); + + // The server should handle EOF mid-message gracefully + let result = connection.connect_to(transport).await; + + // Server should terminate cleanly when hitting EOF + assert!(result.is_ok() || result.is_err()); +} + +// ============================================================================ +// Test 2: Unknown method (no handler claims) +// ============================================================================ + +#[tokio::test(flavor = "current_thread")] +async fn test_unknown_method() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + // No handlers - all requests will be "method not found" + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder(); + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + // Spawn server + tokio::task::spawn_local(async move { + server.connect_to(server_transport).await.ok(); + }); + + // Send request from client + let result = client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + let request = SimpleRequest { + message: "test".to_string(), + }; + + let result: Result = recv(cx.send_request(request)).await; + + // Should get an error because no handler claims the method + assert!(result.is_err()); + if let Err(err) = result { + // Should be "method not found" or similar error + assert!(matches!( + err.code, + agent_client_protocol_core::ErrorCode::MethodNotFound + )); + } + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 3: Handler returns error +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct ErrorRequest { + value: String, +} + +impl JsonRpcMessage for ErrorRequest { + fn matches_method(method: &str) -> bool { + method == "error_method" + } + + fn method(&self) -> &'static str { + "error_method" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for ErrorRequest { + type Response = SimpleResponse; +} + +#[tokio::test(flavor = "current_thread")] +async fn test_handler_returns_error() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |_request: ErrorRequest, + responder: Responder, + _connection: ConnectionTo| { + // Explicitly return an error + responder.respond_with_error(agent_client_protocol_core::Error::internal_error()) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + Box::pin(server.connect_to(server_transport)).await.ok(); + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + let request = ErrorRequest { + value: "trigger error".to_string(), + }; + + let result: Result = recv(cx.send_request(request)).await; + + // Should get the error the handler returned + assert!(result.is_err()); + if let Err(err) = result { + assert!(matches!( + err.code, + agent_client_protocol_core::ErrorCode::InternalError + )); + } + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +// ============================================================================ +// Test 4: Request without required params +// ============================================================================ + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct EmptyRequest; + +impl JsonRpcMessage for EmptyRequest { + fn matches_method(method: &str) -> bool { + method == "strict_method" + } + + fn method(&self) -> &'static str { + "strict_method" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + _params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + Ok(EmptyRequest) + } +} + +impl JsonRpcRequest for EmptyRequest { + type Response = SimpleResponse; +} + +#[tokio::test(flavor = "current_thread")] +async fn test_missing_required_params() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + // Handler that validates params - since EmptyRequest has no params but we're checking + // against SimpleRequest which requires a message field, this will fail + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |_request: EmptyRequest, + responder: Responder, + _connection: ConnectionTo| { + // This will be called, but EmptyRequest parsing already succeeded + // The test is actually checking if EmptyRequest (no params) fails to parse as SimpleRequest + // But with the new API, EmptyRequest parses successfully since it expects no params + // We need to manually check - but actually the parse_request for EmptyRequest + // accepts anything for "strict_method", so the error must come from somewhere else + responder.respond_with_error(agent_client_protocol_core::Error::invalid_params()) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + Box::pin(server.connect_to(server_transport)).await.ok(); + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol_core::Error> { + // Send request with no params (EmptyRequest has no fields) + let request = EmptyRequest; + + let result: Result = recv(cx.send_request(request)).await; + + // Should get invalid_params error + assert!(result.is_err()); + if let Err(err) = result { + assert!(matches!( + err.code, + agent_client_protocol_core::ErrorCode::InvalidParams + )); // JSONRPC_INVALID_PARAMS + } + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} diff --git a/src/agent-client-protocol-core/tests/jsonrpc_hello.rs b/src/agent-client-protocol-core/tests/jsonrpc_hello.rs new file mode 100644 index 0000000..4f8d333 --- /dev/null +++ b/src/agent-client-protocol-core/tests/jsonrpc_hello.rs @@ -0,0 +1,408 @@ +//! Integration test for basic JSON-RPC communication. +//! +//! This test sets up two JSON-RPC connections and verifies they can +//! exchange simple "hello world" messages. + +use agent_client_protocol_core::role::UntypedRole; +use agent_client_protocol_core::{ + ConnectionTo, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Responder, + SentRequest, +}; +use futures::{AsyncRead, AsyncWrite}; +use serde::{Deserialize, Serialize}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Test helper to block and wait for a JSON-RPC response. +async fn recv( + response: SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +/// Helper to set up a client-server pair for testing. +/// Returns (server_reader, server_writer, client_reader, client_writer) for manual setup. +fn setup_test_streams() -> ( + impl AsyncRead, + impl AsyncWrite, + impl AsyncRead, + impl AsyncWrite, +) { + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + + let server_reader = server_reader.compat(); + let server_writer = server_writer.compat_write(); + let client_reader = client_reader.compat(); + let client_writer = client_writer.compat_write(); + + (server_reader, server_writer, client_reader, client_writer) +} + +/// A simple "ping" request. +#[derive(Debug, Clone, Serialize, Deserialize)] +struct PingRequest { + message: String, +} + +impl JsonRpcMessage for PingRequest { + fn matches_method(method: &str) -> bool { + method == "ping" + } + + fn method(&self) -> &'static str { + "ping" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcRequest for PingRequest { + type Response = PongResponse; +} + +/// A simple "pong" response. +#[derive(Debug, Clone, Serialize, Deserialize)] +struct PongResponse { + echo: String, +} + +impl JsonRpcResponse for PongResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + serde_json::to_value(self).map_err(agent_client_protocol_core::Error::into_internal_error) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(&value) + } +} + +#[tokio::test(flavor = "current_thread")] +async fn test_hello_world() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async move |request: PingRequest, + responder: Responder, + _connection: ConnectionTo| { + let pong = PongResponse { + echo: format!("pong: {}", request.message), + }; + responder.respond(pong) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + // Spawn the server in the background + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + // Use the client to send a ping and wait for a pong + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + let request = PingRequest { + message: "hello world".to_string(), + }; + + let response = recv(cx.send_request(request)).await.map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Request failed: {e:?}" + )) + })?; + + assert_eq!(response.echo, "pong: hello world"); + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +/// A simple notification message +#[derive(Debug, Clone, Serialize, Deserialize)] +struct LogNotification { + message: String, +} + +impl JsonRpcMessage for LogNotification { + fn matches_method(method: &str) -> bool { + method == "log" + } + + fn method(&self) -> &'static str { + "log" + } + + fn to_untyped_message( + &self, + ) -> Result { + agent_client_protocol_core::UntypedMessage::new(self.method(), self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !Self::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcNotification for LogNotification {} + +#[tokio::test(flavor = "current_thread")] +async fn test_notification() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let logs = Arc::new(Mutex::new(Vec::new())); + let logs_clone = logs.clone(); + + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_notification( + { + let logs = logs_clone.clone(); + async move |notification: LogNotification, _cx: ConnectionTo| { + logs.lock().unwrap().push(notification.message); + Ok(()) + } + }, + agent_client_protocol_core::on_receive_notification!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + // Send a notification (no response expected) + cx.send_notification(LogNotification { + message: "test log 1".to_string(), + }) + .map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Failed to send notification: {e:?}" + )) + })?; + + cx.send_notification(LogNotification { + message: "test log 2".to_string(), + }) + .map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Failed to send notification: {e:?}" + )) + })?; + + // Give the server time to process notifications + tokio::time::sleep(Duration::from_millis(100)).await; + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + + let received_logs = logs.lock().unwrap(); + assert_eq!(received_logs.len(), 2); + assert_eq!(received_logs[0], "test log 1"); + assert_eq!(received_logs[1], "test log 2"); + })) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn test_multiple_sequential_requests() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |request: PingRequest, + responder: Responder, + _connection: ConnectionTo| { + let pong = PongResponse { + echo: format!("pong: {}", request.message), + }; + responder.respond(pong) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + // Send multiple requests sequentially + for i in 1..=5 { + let request = PingRequest { + message: format!("message {i}"), + }; + + let response = recv(cx.send_request(request)).await.map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Request {i} failed: {e:?}" + )) + })?; + + assert_eq!(response.echo, format!("pong: message {i}")); + } + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn test_concurrent_requests() { + use tokio::task::LocalSet; + + let local = LocalSet::new(); + + Box::pin(local.run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol_core::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |request: PingRequest, + responder: Responder, + _connection: ConnectionTo| { + let pong = PongResponse { + echo: format!("pong: {}", request.message), + }; + responder.respond(pong) + }, + agent_client_protocol_core::on_receive_request!(), + ); + + let client_transport = + agent_client_protocol_core::ByteStreams::new(client_writer, client_reader); + let client = UntypedRole.builder(); + + tokio::task::spawn_local(async move { + if let Err(e) = Box::pin(server.connect_to(server_transport)).await { + eprintln!("Server error: {e:?}"); + } + }); + + let result = client + .connect_with( + client_transport, + async |cx| -> std::result::Result<(), agent_client_protocol_core::Error> { + // Send multiple requests concurrently + let mut responses = Vec::new(); + + for i in 1..=5 { + let request = PingRequest { + message: format!("concurrent message {i}"), + }; + + // Start all requests without awaiting + responses.push((i, cx.send_request(request))); + } + + // Now await all responses + for (i, response_future) in responses { + let response = recv(response_future).await.map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Request {i} failed: {e:?}" + )) + })?; + + assert_eq!(response.echo, format!("pong: concurrent message {i}")); + } + + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + })) + .await; +} diff --git a/src/agent-client-protocol-core/tests/match_dispatch.rs b/src/agent-client-protocol-core/tests/match_dispatch.rs new file mode 100644 index 0000000..b11b73c --- /dev/null +++ b/src/agent-client-protocol-core/tests/match_dispatch.rs @@ -0,0 +1,275 @@ +use agent_client_protocol_core::Role; +use agent_client_protocol_core::role::UntypedRole; +use agent_client_protocol_core::util::MatchDispatch; +use agent_client_protocol_core::{ + ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, JsonRpcMessage, JsonRpcRequest, + JsonRpcResponse, Responder, util::MatchDispatchFrom, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct EchoRequestResponse { + text: Vec, +} + +impl JsonRpcMessage for EchoRequestResponse { + fn matches_method(method: &str) -> bool { + method == "echo" + } + + fn method(&self) -> &'static str { + "echo" + } + + fn to_untyped_message( + &self, + ) -> Result { + Ok(agent_client_protocol_core::UntypedMessage { + method: self.method().to_string(), + params: agent_client_protocol_core::util::json_cast(self)?, + }) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if !::matches_method(method) { + return Err(agent_client_protocol_core::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } +} + +impl JsonRpcResponse for EchoRequestResponse { + fn into_json( + self, + _method: &str, + ) -> Result { + agent_client_protocol_core::util::json_cast(self) + } + + fn from_value( + _method: &str, + value: serde_json::Value, + ) -> Result { + agent_client_protocol_core::util::json_cast(value) + } +} + +impl JsonRpcRequest for EchoRequestResponse { + type Response = EchoRequestResponse; +} + +struct EchoHandler; + +impl HandleDispatchFrom for EchoHandler { + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + _connection: ConnectionTo, + ) -> Result, agent_client_protocol_core::Error> { + MatchDispatch::new(message) + .if_request(async move |request: EchoRequestResponse, responder| { + responder.respond(request) + }) + .await + .done() + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + "TestHandler" + } +} + +#[tokio::test] +async fn modify_message_en_route() -> Result<(), agent_client_protocol_core::Error> { + // Demonstrate a case where we modify a message + // using a `HandleDispatchFrom` invoked from `MatchDispatch` + + struct TestComponent; + + impl ConnectTo for TestComponent { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + UntypedRole + .builder() + .with_handler(PushHandler { + message: "b".to_string(), + }) + .with_handler(EchoHandler) + .connect_to(client) + .await + } + } + + struct PushHandler { + message: String, + } + + impl HandleDispatchFrom for PushHandler { + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + cx: ConnectionTo, + ) -> Result, agent_client_protocol_core::Error> { + MatchDispatchFrom::new(message, &cx) + .if_request(async move |mut request: EchoRequestResponse, responder| { + request.text.push(self.message.clone()); + Ok(Handled::No { + message: (request, responder), + retry: false, + }) + }) + .await + .done() + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + "TestHandler" + } + } + + UntypedRole + .builder() + .connect_with(TestComponent, async |cx| { + let result = cx + .send_request(EchoRequestResponse { + text: vec!["a".to_string()], + }) + .block_task() + .await?; + + expect_test::expect![[r#" + EchoRequestResponse { + text: [ + "a", + "b", + ], + } + "#]] + .assert_debug_eq(&result); + Ok(()) + }) + .await +} + +#[tokio::test] +async fn modify_message_en_route_inline() -> Result<(), agent_client_protocol_core::Error> { + // Demonstrate a case where we modify a message en route using an `on_receive_request` call + + struct TestComponent; + + impl ConnectTo for TestComponent { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + UntypedRole + .builder() + .on_receive_request( + async move |mut request: EchoRequestResponse, + responder: Responder, + _connection: ConnectionTo| { + request.text.push("b".to_string()); + Ok(Handled::No { + message: (request, responder), + retry: false, + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .with_handler(EchoHandler) + .connect_to(client) + .await + } + } + + UntypedRole + .builder() + .connect_with(TestComponent, async |cx| { + let result = cx + .send_request(EchoRequestResponse { + text: vec!["a".to_string()], + }) + .block_task() + .await?; + + expect_test::expect![[r#" + EchoRequestResponse { + text: [ + "a", + "b", + ], + } + "#]] + .assert_debug_eq(&result); + Ok(()) + }) + .await +} + +#[tokio::test] +async fn modify_message_and_stop() -> Result<(), agent_client_protocol_core::Error> { + // Demonstrate a case where we have an async handler that just returns `()` + // in front (and hence we never see the `'b`). + + struct TestComponent; + + impl ConnectTo for TestComponent { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + UntypedRole + .builder() + .on_receive_request( + async move |request: EchoRequestResponse, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(request) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_request( + async move |mut request: EchoRequestResponse, + responder: Responder, + _connection: ConnectionTo| { + request.text.push("b".to_string()); + Ok(Handled::No { + message: (request, responder), + retry: false, + }) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .with_handler(EchoHandler) + .connect_to(client) + .await + } + } + + UntypedRole + .builder() + .connect_with(TestComponent, async |cx| { + let result = cx + .send_request(EchoRequestResponse { + text: vec!["a".to_string()], + }) + .block_task() + .await?; + + expect_test::expect![[r#" + EchoRequestResponse { + text: [ + "a", + ], + } + "#]] + .assert_debug_eq(&result); + Ok(()) + }) + .await +} diff --git a/src/agent-client-protocol-derive/CHANGELOG.md b/src/agent-client-protocol-derive/CHANGELOG.md new file mode 100644 index 0000000..1d013ff --- /dev/null +++ b/src/agent-client-protocol-derive/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-derive/Cargo.toml b/src/agent-client-protocol-derive/Cargo.toml new file mode 100644 index 0000000..597e553 --- /dev/null +++ b/src/agent-client-protocol-derive/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "agent-client-protocol-derive" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "Derive macros for Agent Client Protocol JSON-RPC traits" + +[lib] +proc-macro = true + +[dependencies] +proc-macro2.workspace = true +quote.workspace = true +syn.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-derive/src/lib.rs b/src/agent-client-protocol-derive/src/lib.rs new file mode 100644 index 0000000..58814df --- /dev/null +++ b/src/agent-client-protocol-derive/src/lib.rs @@ -0,0 +1,340 @@ +//! Derive macros for Agent Client Protocol JSON-RPC traits. +//! +//! This crate provides derive macros to reduce boilerplate when implementing +//! custom JSON-RPC requests, notifications, and response types. +//! +//! # Example +//! +//! ```ignore +//! use agent_client_protocol::{JsonRpcRequest, JsonRpcNotification, JsonRpcResponse}; +//! +//! #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcRequest)] +//! #[request(method = "_hello", response = HelloResponse)] +//! struct HelloRequest { +//! name: String, +//! } +//! +//! #[derive(Debug, Serialize, Deserialize, JsonRpcResponse)] +//! #[response(method = "_hello")] +//! struct HelloResponse { +//! greeting: String, +//! } +//! +//! #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)] +//! #[notification(method = "_ping")] +//! struct PingNotification { +//! timestamp: u64, +//! } +//! ``` +//! +//! # Using within the `agent_client_protocol` crate +//! +//! When using these derives within the `agent_client_protocol` crate itself, add `crate = crate`: +//! +//! ```ignore +//! #[derive(JsonRpcRequest)] +//! #[request(method = "_foo", response = FooResponse, crate = crate)] +//! struct FooRequest { ... } +//! ``` + +use proc_macro::TokenStream; +use quote::quote; +use syn::{DeriveInput, Expr, Lit, Path, Type, parse_macro_input}; + +/// Derive macro for implementing `JsonRpcRequest` and `JsonRpcMessage` traits. +/// +/// # Attributes +/// +/// - `#[request(method = "method_name", response = ResponseType)]` +/// - `#[request(method = "method_name", response = ResponseType, crate = crate)]` - for use within the `agent_client_protocol_core` crate +/// +/// # Example +/// +/// ```ignore +/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcRequest)] +/// #[request(method = "_hello", response = HelloResponse)] +/// struct HelloRequest { +/// name: String, +/// } +/// ``` +#[proc_macro_derive(JsonRpcRequest, attributes(request))] +pub fn derive_json_rpc_request(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + let name = &input.ident; + + // Parse attributes + let (method, response_type, krate) = match parse_request_attrs(&input) { + Ok(attrs) => attrs, + Err(e) => return e.to_compile_error().into(), + }; + + let expanded = quote! { + impl #krate::JsonRpcMessage for #name { + fn matches_method(method: &str) -> bool { + method == #method + } + + fn method(&self) -> &str { + #method + } + + fn to_untyped_message(&self) -> Result<#krate::UntypedMessage, #krate::Error> { + #krate::UntypedMessage::new(#method, self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if method != #method { + return Err(#krate::Error::method_not_found()); + } + #krate::util::json_cast(params) + } + } + + impl #krate::JsonRpcRequest for #name { + type Response = #response_type; + } + }; + + TokenStream::from(expanded) +} + +/// Derive macro for implementing `JsonRpcNotification` and `JsonRpcMessage` traits. +/// +/// # Attributes +/// +/// - `#[notification(method = "method_name")]` +/// - `#[notification(method = "method_name", crate = crate)]` - for use within the `agent_client_protocol_core` crate +/// +/// # Example +/// +/// ```ignore +/// #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)] +/// #[notification(method = "_ping")] +/// struct PingNotification { +/// timestamp: u64, +/// } +/// ``` +#[proc_macro_derive(JsonRpcNotification, attributes(notification))] +pub fn derive_json_rpc_notification(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + let name = &input.ident; + + // Parse attributes + let (method, krate) = match parse_notification_attrs(&input) { + Ok(attrs) => attrs, + Err(e) => return e.to_compile_error().into(), + }; + + let expanded = quote! { + impl #krate::JsonRpcMessage for #name { + fn matches_method(method: &str) -> bool { + method == #method + } + + fn method(&self) -> &str { + #method + } + + fn to_untyped_message(&self) -> Result<#krate::UntypedMessage, #krate::Error> { + #krate::UntypedMessage::new(#method, self) + } + + fn parse_message( + method: &str, + params: &impl serde::Serialize, + ) -> Result { + if method != #method { + return Err(#krate::Error::method_not_found()); + } + #krate::util::json_cast(params) + } + } + + impl #krate::JsonRpcNotification for #name {} + }; + + TokenStream::from(expanded) +} + +/// Derive macro for implementing `JsonRpcResponse` trait. +/// +/// # Attributes +/// +/// - `#[response(crate = crate)]` - for use within the `agent_client_protocol_core` crate +/// +/// # Example +/// +/// ```ignore +/// #[derive(Debug, Serialize, Deserialize, JsonRpcResponse)] +/// struct HelloResponse { +/// greeting: String, +/// } +/// ``` +#[proc_macro_derive(JsonRpcResponse, attributes(response))] +pub fn derive_json_rpc_response_payload(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + let name = &input.ident; + + let krate = match parse_response_attrs(&input) { + Ok(attrs) => attrs, + Err(e) => return e.to_compile_error().into(), + }; + + let expanded = quote! { + impl #krate::JsonRpcResponse for #name { + fn into_json(self, _method: &str) -> Result { + serde_json::to_value(self).map_err(#krate::Error::into_internal_error) + } + + fn from_value(_method: &str, value: serde_json::Value) -> Result { + #krate::util::json_cast(value) + } + } + }; + + TokenStream::from(expanded) +} + +fn default_crate_path() -> Path { + syn::parse_quote!(agent_client_protocol_core) +} + +fn parse_request_attrs(input: &DeriveInput) -> syn::Result<(String, Type, Path)> { + let mut method: Option = None; + let mut response_type: Option = None; + let mut krate: Option = None; + + for attr in &input.attrs { + if !attr.path().is_ident("request") { + continue; + } + + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("method") { + let value: Expr = meta.value()?.parse()?; + if let Expr::Lit(expr_lit) = value + && let Lit::Str(lit_str) = expr_lit.lit + { + method = Some(lit_str.value()); + return Ok(()); + } + return Err(meta.error("expected string literal for method")); + } + + if meta.path.is_ident("response") { + let value: Expr = meta.value()?.parse()?; + if let Expr::Path(expr_path) = value { + response_type = Some(Type::Path(syn::TypePath { + qself: None, + path: expr_path.path, + })); + return Ok(()); + } + return Err(meta.error("expected type for response")); + } + + if meta.path.is_ident("crate") { + let value: Expr = meta.value()?.parse()?; + if let Expr::Path(expr_path) = value { + krate = Some(expr_path.path); + return Ok(()); + } + return Err(meta.error("expected path for crate")); + } + + Err(meta.error("unknown attribute")) + })?; + } + + let method = method.ok_or_else(|| { + syn::Error::new_spanned( + &input.ident, + "missing required attribute: #[request(method = \"...\")]", + ) + })?; + + let response_type = response_type.ok_or_else(|| { + syn::Error::new_spanned( + &input.ident, + "missing required attribute: #[request(response = ...)]", + ) + })?; + + Ok(( + method, + response_type, + krate.unwrap_or_else(default_crate_path), + )) +} + +fn parse_notification_attrs(input: &DeriveInput) -> syn::Result<(String, Path)> { + let mut method: Option = None; + let mut krate: Option = None; + + for attr in &input.attrs { + if !attr.path().is_ident("notification") { + continue; + } + + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("method") { + let value: Expr = meta.value()?.parse()?; + if let Expr::Lit(expr_lit) = value + && let Lit::Str(lit_str) = expr_lit.lit + { + method = Some(lit_str.value()); + return Ok(()); + } + return Err(meta.error("expected string literal for method")); + } + + if meta.path.is_ident("crate") { + let value: Expr = meta.value()?.parse()?; + if let Expr::Path(expr_path) = value { + krate = Some(expr_path.path); + return Ok(()); + } + return Err(meta.error("expected path for crate")); + } + + Err(meta.error("unknown attribute")) + })?; + } + + let method = method.ok_or_else(|| { + syn::Error::new_spanned( + &input.ident, + "missing required attribute: #[notification(method = \"...\")]", + ) + })?; + + Ok((method, krate.unwrap_or_else(default_crate_path))) +} + +fn parse_response_attrs(input: &DeriveInput) -> syn::Result { + let mut krate: Option = None; + + for attr in &input.attrs { + if !attr.path().is_ident("response") { + continue; + } + + attr.parse_nested_meta(|meta| { + if meta.path.is_ident("crate") { + let value: Expr = meta.value()?.parse()?; + if let Expr::Path(expr_path) = value { + krate = Some(expr_path.path); + return Ok(()); + } + return Err(meta.error("expected path for crate")); + } + + Err(meta.error("unknown attribute")) + })?; + } + + Ok(krate.unwrap_or_else(default_crate_path)) +} diff --git a/src/agent-client-protocol-rmcp/CHANGELOG.md b/src/agent-client-protocol-rmcp/CHANGELOG.md new file mode 100644 index 0000000..1d013ff --- /dev/null +++ b/src/agent-client-protocol-rmcp/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-rmcp/Cargo.toml b/src/agent-client-protocol-rmcp/Cargo.toml new file mode 100644 index 0000000..48c9efe --- /dev/null +++ b/src/agent-client-protocol-rmcp/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "agent-client-protocol-rmcp" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "rmcp integration for Agent Client Protocol MCP servers" +keywords = ["acp", "agent", "proxy", "mcp", "rmcp"] +categories = ["development-tools"] + +[dependencies] +agent-client-protocol-core.workspace = true +futures.workspace = true +futures-concurrency.workspace = true +rmcp.workspace = true +tokio.workspace = true +tokio-util.workspace = true + +[dev-dependencies] +schemars.workspace = true +serde.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-rmcp/README.md b/src/agent-client-protocol-rmcp/README.md new file mode 100644 index 0000000..c4fb9c5 --- /dev/null +++ b/src/agent-client-protocol-rmcp/README.md @@ -0,0 +1,41 @@ +# agent-client-protocol-rmcp + +[rmcp](https://docs.rs/rmcp) integration for [Agent Client Protocol](https://agentclientprotocol.com/) MCP servers. + +## Overview + +This crate bridges [rmcp](https://docs.rs/rmcp)-based MCP server implementations with the ACP MCP server framework from `agent-client-protocol-core`. It lets you use any rmcp service as an MCP server in an ACP proxy. + +## Usage + +Use the `McpServerExt` trait to create an MCP server from an rmcp service: + +```rust +use agent_client_protocol_core::mcp_server::McpServer; +use agent_client_protocol_rmcp::McpServerExt; + +let server = McpServer::from_rmcp("my-server", MyRmcpService::new); + +// Use as a handler in a proxy +Proxy.builder() + .with_mcp_server(server) + .connect_to(transport) + .await?; +``` + +## Why a Separate Crate? + +This crate is separate from `agent-client-protocol-core` to avoid coupling the core protocol crate to the `rmcp` dependency. This allows: + +- `agent-client-protocol-core` to remain focused on the ACP protocol +- `agent-client-protocol-rmcp` to track `rmcp` updates independently +- Breaking changes in `rmcp` only require updating this crate + +## Related Crates + +- **[agent-client-protocol-core](../agent-client-protocol-core/)** — Core ACP protocol types and traits +- **[agent-client-protocol-tokio](../agent-client-protocol-tokio/)** — Tokio utilities for spawning agent processes + +## License + +Apache-2.0 diff --git a/src/agent-client-protocol-rmcp/examples/with_mcp_server.rs b/src/agent-client-protocol-rmcp/examples/with_mcp_server.rs new file mode 100644 index 0000000..171d36a --- /dev/null +++ b/src/agent-client-protocol-rmcp/examples/with_mcp_server.rs @@ -0,0 +1,103 @@ +//! Proxy with MCP server example +//! +//! This proxy provides a simple MCP server with an "echo" tool. +//! Demonstrates how to add MCP tools to any agent through a proxy. +//! +//! Run with: +//! ```bash +//! cargo run --example with_mcp_server +//! ``` + +use agent_client_protocol_core::{Proxy, mcp_server::McpServer}; +use agent_client_protocol_rmcp::McpServerExt; +use rmcp::{ + ErrorData as McpError, ServerHandler, + handler::server::{router::tool::ToolRouter, wrapper::Parameters}, + model::*, + tool, tool_handler, tool_router, +}; +use serde::{Deserialize, Serialize}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Parameters for the echo tool +#[derive(Debug, Serialize, Deserialize, schemars::JsonSchema)] +struct EchoParams { + /// The message to echo back + message: String, +} + +/// Simple MCP server with an echo tool +#[derive(Clone, Debug)] +pub struct ExampleMcpServer { + tool_router: ToolRouter, +} + +impl ExampleMcpServer { + #[must_use] + pub fn new() -> Self { + Self { + tool_router: Self::tool_router(), + } + } +} + +impl Default for ExampleMcpServer { + fn default() -> Self { + Self::new() + } +} + +#[tool_router] +impl ExampleMcpServer { + /// Echo tool - returns the input message + #[tool(description = "Echoes back the input message")] + async fn echo( + &self, + Parameters(params): Parameters, + ) -> Result { + Ok(CallToolResult::success(vec![Content::text(format!( + "Echo: {}", + params.message + ))])) + } +} + +#[tool_handler] +impl ServerHandler for ExampleMcpServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("example-mcp-server", "0.1.0")) + .with_protocol_version(ProtocolVersion::V_2024_11_05) + .with_instructions("A simple example MCP server with an echo tool") + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + // Initialize tracing for debugging + tracing_subscriber::fmt() + .with_target(true) + .with_writer(std::io::stderr) + .init(); + + tracing::info!("MCP server proxy starting"); + + // Create an MCP server from the rmcp service + let mcp_server = McpServer::from_rmcp("example", ExampleMcpServer::new); + + // Set up the proxy connection with our MCP server + // ProxyToConductor already has proxy behavior built into its default_message_handler + let proxy = Proxy + .builder() + .name("mcp-server-proxy") + // Register the MCP server as a handler + .with_mcp_server(mcp_server) + // Start serving + .connect_to(agent_client_protocol_core::ByteStreams::new( + tokio::io::stdout().compat_write(), + tokio::io::stdin().compat(), + )); + Box::pin(proxy).await?; + + Ok(()) +} diff --git a/src/agent-client-protocol-rmcp/src/lib.rs b/src/agent-client-protocol-rmcp/src/lib.rs new file mode 100644 index 0000000..4b436c9 --- /dev/null +++ b/src/agent-client-protocol-rmcp/src/lib.rs @@ -0,0 +1,133 @@ +//! # agent-client-protocol-rmcp +//! +//! This crate provides integration between [rmcp](https://docs.rs/rmcp) MCP servers +//! and the Agent Client Protocol MCP server framework. +//! +//! ## Usage +//! +//! Create an MCP server from an rmcp service using the extension trait: +//! +//! ```ignore +//! use agent_client_protocol_core::mcp_server::McpServer; +//! use agent_client_protocol_rmcp::McpServerExt; +//! +//! let server = McpServer::from_rmcp("my-server", MyRmcpService::new); +//! +//! // Use as a handler +//! Proxy.builder() +//! .with_handler(server) +//! .serve(client) +//! .await?; +//! ``` + +use agent_client_protocol_core::mcp_server::{McpConnectionTo, McpServer, McpServerConnect}; +use agent_client_protocol_core::role::{self, HasPeer}; +use agent_client_protocol_core::{Agent, ByteStreams, ConnectTo, DynConnectTo, NullRun, Role}; +use futures_concurrency::future::TryJoin as _; +use rmcp::ServiceExt; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +pub trait McpServerExt +where + Counterpart: HasPeer, +{ + /// Create an MCP server from something that implements the [`McpServerConnect`] trait. + /// + /// # See also + /// + /// See [`McpServer::builder`] to construct MCP servers from Rust code. + fn from_rmcp( + name: impl ToString, + new_fn: impl Fn() -> S + Send + Sync + 'static, + ) -> McpServer + where + S: rmcp::Service, + { + struct RmcpServer { + name: String, + new_fn: F, + } + + impl McpServerConnect for RmcpServer + where + Counterpart: Role, + F: Fn() -> S + Send + Sync + 'static, + S: rmcp::Service, + { + fn name(&self) -> String { + self.name.clone() + } + + fn connect( + &self, + _cx: McpConnectionTo, + ) -> DynConnectTo { + let service = (self.new_fn)(); + DynConnectTo::new(RmcpServerComponent { service }) + } + } + + McpServer::new( + RmcpServer { + name: name.to_string(), + new_fn, + }, + NullRun, + ) + } +} + +impl McpServerExt for McpServer where + Counterpart: HasPeer +{ +} + +/// Component wrapper for rmcp services. +struct RmcpServerComponent { + service: S, +} + +impl ConnectTo for RmcpServerComponent +where + S: rmcp::Service, +{ + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + // Create tokio byte streams that rmcp expects + let (mcp_server_stream, mcp_client_stream) = tokio::io::duplex(8192); + let (mcp_server_read, mcp_server_write) = tokio::io::split(mcp_server_stream); + let (mcp_client_read, mcp_client_write) = tokio::io::split(mcp_client_stream); + + let bytes_to_acp = async { + // Create ByteStreams component for the client side + let byte_streams = + ByteStreams::new(mcp_client_write.compat_write(), mcp_client_read.compat()); + + // Spawn task to connect byte_streams to the provided client + drop(ConnectTo::::connect_to(byte_streams, client).await); + + Ok(()) + }; + + let bytes_to_rmcp = async { + // Run the rmcp server with the server side of the duplex stream + let running_server = self + .service + .serve((mcp_server_read, mcp_server_write)) + .await + .map_err(agent_client_protocol_core::Error::into_internal_error)?; + + // Wait for the server to finish + running_server + .waiting() + .await + .map(|_quit_reason| ()) + .map_err(agent_client_protocol_core::Error::into_internal_error) + }; + + (bytes_to_acp, bytes_to_rmcp).try_join().await?; + Ok(()) + } +} diff --git a/src/agent-client-protocol-test/CHANGELOG.md b/src/agent-client-protocol-test/CHANGELOG.md new file mode 100644 index 0000000..1d013ff --- /dev/null +++ b/src/agent-client-protocol-test/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-test/Cargo.toml b/src/agent-client-protocol-test/Cargo.toml new file mode 100644 index 0000000..e404df6 --- /dev/null +++ b/src/agent-client-protocol-test/Cargo.toml @@ -0,0 +1,37 @@ +[package] +name = "agent-client-protocol-test" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "Test utilities and mock implementations for the Agent Client Protocol" +publish = false + +[[bin]] +name = "mcp-echo-server" +path = "src/bin/mcp_echo_server.rs" + +[[bin]] +name = "testy" +path = "src/bin/testy.rs" + +[dependencies] +agent-client-protocol-core.workspace = true +agent-client-protocol-tokio.workspace = true +yopo.workspace = true +anyhow.workspace = true +futures.workspace = true +rmcp = { workspace = true, features = ["server", "client", "transport-child-process", "transport-streamable-http-client-reqwest"] } +schemars.workspace = true +serde.workspace = true +serde_json.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true +uuid.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-test/examples/arrow_proxy.rs b/src/agent-client-protocol-test/examples/arrow_proxy.rs new file mode 100644 index 0000000..73fc6ab --- /dev/null +++ b/src/agent-client-protocol-test/examples/arrow_proxy.rs @@ -0,0 +1,29 @@ +//! Arrow proxy example that can be run as a standalone process. +//! +//! This adds '>' prefix to session updates for testing conductor proxy chains. + +use agent_client_protocol_test::arrow_proxy::run_arrow_proxy; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +#[tokio::main] +async fn main() -> Result<(), Box> { + // Initialize tracing for debugging + tracing_subscriber::fmt() + .with_target(true) + .with_writer(std::io::stderr) + .init(); + + tracing::info!("Arrow proxy starting"); + + // Create connection over stdio with compat layer for tokio/futures interop + let stdin = tokio::io::stdin().compat(); + let stdout = tokio::io::stdout().compat_write(); + + // Run the arrow proxy + Box::pin(run_arrow_proxy( + agent_client_protocol_core::ByteStreams::new(stdout, stdin), + )) + .await?; + + Ok(()) +} diff --git a/src/agent-client-protocol-test/src/arrow_proxy.rs b/src/agent-client-protocol-test/src/arrow_proxy.rs new file mode 100644 index 0000000..5a95767 --- /dev/null +++ b/src/agent-client-protocol-test/src/arrow_proxy.rs @@ -0,0 +1,56 @@ +//! A simple test proxy that adds `>` prefix to session update messages. +//! +//! This proxy demonstrates basic proxy functionality by intercepting +//! `session/update` notifications and prepending `>` to the content. + +use agent_client_protocol_core::schema::{ + ContentBlock, ContentChunk, SessionNotification, SessionUpdate, +}; +use agent_client_protocol_core::{Agent, Client, ConnectTo, Proxy}; + +/// Run the arrow proxy that adds `>` to each session update. +/// +/// # Arguments +/// +/// * `transport` - Component to the predecessor (conductor or another proxy) +pub async fn run_arrow_proxy( + transport: impl ConnectTo + 'static, +) -> Result<(), agent_client_protocol_core::Error> { + Proxy + .builder() + .name("arrow-proxy") + // Intercept session notifications from successor (agent) and modify them. + // Using on_receive_notification_from(Agent, ...) automatically unwraps + // SuccessorMessage envelopes. + .on_receive_notification_from( + Agent, + async |mut notification: SessionNotification, cx| { + // Modify the content by adding > prefix + if let SessionUpdate::AgentMessageChunk(ContentChunk { content, .. }) = + &mut notification.update + // Add > prefix to text content + && let ContentBlock::Text(text_content) = content + { + text_content.text = format!(">{}", text_content.text); + } else { + // Don't modify other update types + } + + // Forward modified notification to predecessor (client) + cx.send_notification_to(Client, notification)?; + Ok(()) + }, + agent_client_protocol_core::on_receive_notification!(), + ) + .connect_to(transport) + .await +} + +#[cfg(test)] +mod tests { + #[test] + fn test_arrow_proxy_compiles() { + // Basic smoke test that the arrow proxy module compiles + // Full integration tests with conductor will be in agent-client-protocol-conductor tests + } +} diff --git a/src/agent-client-protocol-test/src/bin/mcp_echo_server.rs b/src/agent-client-protocol-test/src/bin/mcp_echo_server.rs new file mode 100644 index 0000000..88ebe78 --- /dev/null +++ b/src/agent-client-protocol-test/src/bin/mcp_echo_server.rs @@ -0,0 +1,80 @@ +//! Simple MCP echo server for testing +//! +//! This server provides a basic "echo" tool that returns the input message. +//! It runs over stdio and can be used for testing MCP integrations. + +use rmcp::{ + ErrorData as McpError, ServerHandler, ServiceExt, + handler::server::{router::tool::ToolRouter, wrapper::Parameters}, + model::*, + tool, tool_handler, tool_router, +}; +use serde::{Deserialize, Serialize}; + +/// Parameters for the echo tool +#[derive(Debug, Serialize, Deserialize, schemars::JsonSchema)] +struct EchoParams { + /// The message to echo back + message: String, +} + +/// Simple MCP server with an echo tool +#[derive(Clone)] +struct EchoServer { + tool_router: ToolRouter, +} + +impl EchoServer { + fn new() -> Self { + Self { + tool_router: Self::tool_router(), + } + } +} + +#[tool_router] +impl EchoServer { + /// Echo tool - returns the input message + #[tool(description = "Echoes back the input message")] + async fn echo( + &self, + Parameters(params): Parameters, + ) -> Result { + Ok(CallToolResult::success(vec![Content::text(format!( + "Echo: {}", + params.message + ))])) + } +} + +#[tool_handler] +impl ServerHandler for EchoServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + .with_server_info(Implementation::new("mcp-echo-server", "1.0.0")) + .with_protocol_version(ProtocolVersion::V_2024_11_05) + .with_instructions("A simple MCP server with an echo tool for testing") + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + // Initialize tracing + tracing_subscriber::fmt() + .with_writer(std::io::stderr) + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")), + ) + .init(); + + tracing::info!("Starting MCP echo server"); + + // Run the server on stdio + let service = EchoServer::new().serve(rmcp::transport::stdio()).await?; + + // Keep the service running indefinitely + service.waiting().await?; + + Ok(()) +} diff --git a/src/agent-client-protocol-test/src/bin/testy.rs b/src/agent-client-protocol-test/src/bin/testy.rs new file mode 100644 index 0000000..3a03c44 --- /dev/null +++ b/src/agent-client-protocol-test/src/bin/testy.rs @@ -0,0 +1,11 @@ +use agent_client_protocol_core::ConnectTo; +use agent_client_protocol_test::testy::Testy; + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + tracing_subscriber::fmt() + .with_writer(std::io::stderr) + .init(); + Box::pin(Testy::new().connect_to(agent_client_protocol_tokio::Stdio::new())).await?; + Ok(()) +} diff --git a/src/agent-client-protocol-test/src/lib.rs b/src/agent-client-protocol-test/src/lib.rs new file mode 100644 index 0000000..843db92 --- /dev/null +++ b/src/agent-client-protocol-test/src/lib.rs @@ -0,0 +1,226 @@ +use agent_client_protocol_core::*; +use serde::{Deserialize, Serialize}; + +pub mod arrow_proxy; +pub mod test_binaries; +pub mod testy; + +/// A mock transport for doctests that panics if actually used. +/// This is only for documentation examples that don't actually run. +#[derive(Debug)] +pub struct MockTransport; + +impl ConnectTo for MockTransport { + async fn connect_to(self, _client: impl ConnectTo) -> Result<(), Error> { + panic!("MockTransport should never be used in running code - it's only for doctests") + } +} + +// Mock request/response types +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MyRequest {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MyResponse { + pub status: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProcessRequest { + pub data: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProcessResponse { + pub result: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProcessStarted {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AnalyzeRequest { + pub data: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AnalysisStarted { + pub job_id: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryRequest { + pub id: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryResponse { + pub data: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ValidateRequest { + pub data: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ValidateResponse { + pub is_valid: bool, + pub error: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecuteRequest {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExecuteResponse { + pub result: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OtherRequest {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OtherResponse { + pub value: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProxyRequest { + pub inner_request: UntypedMessage, +} + +// Mock notification types +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SessionUpdate {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StatusUpdate { + pub message: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ProcessComplete { + pub result: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AnalysisComplete { + pub result: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct QueryComplete {} + +// Implement JsonRpcMessage for all types +macro_rules! impl_jr_message { + ($type:ty, $method:expr) => { + impl JsonRpcMessage for $type { + fn matches_method(method: &str) -> bool { + method == $method + } + fn method(&self) -> &str { + $method + } + fn to_untyped_message(&self) -> Result { + UntypedMessage::new($method, self) + } + fn parse_message(method: &str, params: &impl Serialize) -> Result { + if !Self::matches_method(method) { + return Err(crate::Error::method_not_found()); + } + agent_client_protocol_core::util::json_cast(params) + } + } + }; +} + +// Implement JsonRpcRequest for request types +macro_rules! impl_jr_request { + ($req:ty, $resp:ty, $method:expr) => { + impl_jr_message!($req, $method); + impl JsonRpcRequest for $req { + type Response = $resp; + } + }; +} + +// Implement JsonRpcNotification for notification types +macro_rules! impl_jr_notification { + ($type:ty, $method:expr) => { + impl_jr_message!($type, $method); + impl JsonRpcNotification for $type {} + }; +} + +impl_jr_request!(MyRequest, MyResponse, "myRequest"); +impl_jr_request!(ProcessRequest, ProcessResponse, "processRequest"); +impl_jr_request!(AnalyzeRequest, AnalysisStarted, "analyzeRequest"); +impl_jr_request!(QueryRequest, QueryResponse, "queryRequest"); +impl_jr_request!(ValidateRequest, ValidateResponse, "validateRequest"); +impl_jr_request!(ExecuteRequest, ExecuteResponse, "executeRequest"); +impl_jr_request!(OtherRequest, OtherResponse, "otherRequest"); +impl_jr_request!(ProxyRequest, serde_json::Value, "proxyRequest"); + +impl_jr_notification!(SessionUpdate, "sessionUpdate"); +impl_jr_notification!(StatusUpdate, "statusUpdate"); +impl_jr_notification!(ProcessComplete, "processComplete"); +impl_jr_notification!(AnalysisComplete, "analysisComplete"); +impl_jr_notification!(QueryComplete, "queryComplete"); +impl_jr_notification!(ProcessStarted, "processStarted"); + +// Implement JsonRpcResponse for response types +macro_rules! impl_jr_response_payload { + ($type:ty, $method:expr) => { + impl JsonRpcResponse for $type { + fn into_json(self, _method: &str) -> Result { + Ok(serde_json::to_value(self)?) + } + fn from_value(_method: &str, value: serde_json::Value) -> Result { + Ok(serde_json::from_value(value)?) + } + } + }; +} + +impl_jr_response_payload!(MyResponse, "myRequest"); +impl_jr_response_payload!(ProcessResponse, "processRequest"); +impl_jr_response_payload!(AnalysisStarted, "analyzeRequest"); +impl_jr_response_payload!(QueryResponse, "queryRequest"); +impl_jr_response_payload!(ValidateResponse, "validateRequest"); +impl_jr_response_payload!(ExecuteResponse, "executeRequest"); +impl_jr_response_payload!(OtherResponse, "otherRequest"); + +// Mock async functions +#[expect(clippy::unused_async)] +pub async fn expensive_analysis(_data: &str) -> Result { + Ok("analysis result".into()) +} + +#[expect(clippy::unused_async)] +pub async fn expensive_operation(_data: &str) -> Result { + Ok("operation result".into()) +} + +pub fn update_session_state(_update: &SessionUpdate) -> Result<(), crate::Error> { + Ok(()) +} + +pub fn process(data: &str) -> Result { + Ok(data.to_string()) +} + +// Helper to create a mock connection for examples +pub fn mock_connection() -> Builder { + Client.builder() +} + +pub trait Make { + fn make() -> Self; +} + +impl Make for T { + fn make() -> Self { + panic!() + } +} diff --git a/src/agent-client-protocol-test/src/test_binaries.rs b/src/agent-client-protocol-test/src/test_binaries.rs new file mode 100644 index 0000000..0cc4082 --- /dev/null +++ b/src/agent-client-protocol-test/src/test_binaries.rs @@ -0,0 +1,86 @@ +//! Utilities for locating pre-built test binaries. +//! +//! Integration tests that spawn subprocesses should use pre-built binaries +//! rather than `cargo run` to avoid recursive cargo invocations during +//! `cargo test --all`. +//! +//! Run `just prep-tests` before running tests to build all required binaries. + +use std::path::{Path, PathBuf}; + +/// Returns the workspace root directory. +fn workspace_root() -> PathBuf { + let manifest_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + manifest_dir + .parent() + .expect("agent-client-protocol-test should be in src/") + .parent() + .expect("src/ should be in workspace root") + .to_path_buf() +} + +/// Returns the path to a binary in the target/debug directory. +#[must_use] +pub fn debug_binary(name: &str) -> PathBuf { + workspace_root().join("target/debug").join(name) +} + +/// Returns the path to an example binary in the target/debug/examples directory. +#[must_use] +pub fn debug_example(name: &str) -> PathBuf { + workspace_root().join("target/debug/examples").join(name) +} + +/// Asserts that a binary exists, panicking with a helpful message if not. +/// +/// # Panics +/// +/// Panics if the binary does not exist, with a message instructing the user +/// to run `just prep-tests`. +pub fn require_binary(path: &Path) { + assert!( + path.exists(), + "Binary not found at {}.\n\ + Run `just prep-tests` before running these tests.", + path.display(), + ); +} + +/// Returns the path to the agent-client-protocol-conductor binary, asserting it exists. +#[must_use] +pub fn conductor_binary() -> PathBuf { + let path = debug_binary("agent-client-protocol-conductor"); + require_binary(&path); + path +} + +/// Returns the path to the test-agent binary, asserting it exists. +#[must_use] +pub fn testy_binary() -> PathBuf { + let path = debug_binary("testy"); + require_binary(&path); + path +} + +/// Returns an AcpAgent configured for the test agent. +#[must_use] +pub fn testy() -> agent_client_protocol_tokio::AcpAgent { + agent_client_protocol_tokio::AcpAgent::from_args([testy_binary().to_string_lossy().to_string()]) + .expect("failed to create test agent") +} + +/// Returns the path to the mcp-echo-server binary, asserting it exists. +#[must_use] +pub fn mcp_echo_server_binary() -> PathBuf { + let path = debug_binary("mcp-echo-server"); + require_binary(&path); + path +} + +/// Returns the path to the arrow_proxy example, asserting it exists. +#[must_use] +pub fn arrow_proxy_example() -> PathBuf { + let path = debug_example("arrow_proxy"); + require_binary(&path); + path +} diff --git a/src/agent-client-protocol-test/src/testy.rs b/src/agent-client-protocol-test/src/testy.rs new file mode 100644 index 0000000..3e3d5ea --- /dev/null +++ b/src/agent-client-protocol-test/src/testy.rs @@ -0,0 +1,301 @@ +//! testy: you friendly neighborhood ACP test agent with typed JSON commands. +//! +//! The agent accepts JSON-serialized [`TestyCommand`] values as prompt text. + +use agent_client_protocol_core::schema::{ + AgentCapabilities, ContentBlock, ContentChunk, InitializeRequest, InitializeResponse, + McpServer, NewSessionRequest, NewSessionResponse, PromptRequest, PromptResponse, SessionId, + SessionNotification, SessionUpdate, StopReason, TextContent, +}; +use agent_client_protocol_core::{Agent, Client, ConnectTo, ConnectionTo, Responder}; +use anyhow::Result; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +/// Commands that can be sent as prompt text (serialized as JSON) to the [`Testy`]. +/// +/// Tests construct these as typed values and serialize to JSON via [`TestyCommand::to_prompt`]. +/// The agent deserializes the prompt text and dispatches accordingly. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "command", rename_all = "snake_case")] +pub enum TestyCommand { + /// Responds with `"Hello, world!"`. + Greet, + + /// Echoes the given message back as the response. + Echo { message: String }, + + /// Invokes an MCP tool and returns the result. + /// The agent must have been given MCP servers in the `NewSessionRequest`. + CallTool { + server: String, + tool: String, + #[serde(default)] + params: serde_json::Value, + }, + + /// Lists tools from the named MCP server. + ListTools { server: String }, +} + +impl TestyCommand { + /// Serialize this command to a JSON string suitable for use as prompt text. + #[must_use] + pub fn to_prompt(&self) -> String { + serde_json::to_string(self).expect("TestyCommand serialization should not fail") + } +} + +/// Session data for each active session. +#[derive(Clone, Debug)] +struct SessionData { + mcp_servers: Vec, +} + +/// A minimal ACP test agent. +/// +/// Implements `ConnectTo` and handles `InitializeRequest`, `NewSessionRequest`, +/// and `PromptRequest`. Prompt text is parsed as a JSON [`TestyCommand`]; if parsing fails, +/// the agent responds with `"Hello, world!"` (equivalent to [`TestyCommand::Greet`]). +#[derive(Clone, Debug)] +pub struct Testy { + sessions: Arc>>, +} + +impl Testy { + #[must_use] + pub fn new() -> Self { + Self { + sessions: Arc::new(Mutex::new(HashMap::new())), + } + } + + fn create_session(&self, session_id: &SessionId, mcp_servers: Vec) { + let mut sessions = self.sessions.lock().unwrap(); + sessions.insert(session_id.clone(), SessionData { mcp_servers }); + } + + fn get_mcp_servers(&self, session_id: &SessionId) -> Option> { + let sessions = self.sessions.lock().unwrap(); + sessions + .get(session_id) + .map(|session| session.mcp_servers.clone()) + } + + async fn process_prompt( + &self, + request: PromptRequest, + responder: Responder, + connection: ConnectionTo, + ) -> Result<(), agent_client_protocol_core::Error> { + let session_id = request.session_id.clone(); + let input_text = extract_text_from_prompt(&request.prompt); + + let command: TestyCommand = + serde_json::from_str(&input_text).unwrap_or(TestyCommand::Greet); + + let response_text = match command { + TestyCommand::Greet => "Hello, world!".to_string(), + + TestyCommand::Echo { message } => message, + + TestyCommand::CallTool { + server, + tool, + params, + } => match self + .execute_tool_call(&session_id, &server, &tool, params) + .await + { + Ok(result) => format!("OK: {result}"), + Err(e) => format!("ERROR: {e}"), + }, + + TestyCommand::ListTools { server } => { + match self.list_tools(&session_id, &server).await { + Ok(tools) => format!("Available tools:\n{tools}"), + Err(e) => format!("ERROR: {e}"), + } + } + }; + + connection.send_notification(SessionNotification::new( + session_id, + SessionUpdate::AgentMessageChunk(ContentChunk::new(response_text.into())), + ))?; + + responder.respond(PromptResponse::new(StopReason::EndTurn)) + } + + /// Helper to execute an operation with a spawned MCP client. + async fn with_mcp_client( + &self, + session_id: &SessionId, + server_name: &str, + operation: F, + ) -> Result + where + F: FnOnce(rmcp::service::RunningService) -> Fut, + Fut: std::future::Future>, + { + use rmcp::{ + ServiceExt, + transport::{ConfigureCommandExt, TokioChildProcess}, + }; + use tokio::process::Command; + + let mcp_servers = self + .get_mcp_servers(session_id) + .ok_or_else(|| anyhow::anyhow!("Session not found"))?; + + let mcp_server = mcp_servers + .iter() + .find(|server| match server { + McpServer::Stdio(stdio) => stdio.name == server_name, + McpServer::Http(http) => http.name == server_name, + McpServer::Sse(sse) => sse.name == server_name, + _ => false, + }) + .ok_or_else(|| anyhow::anyhow!("MCP server '{server_name}' not found"))?; + + match mcp_server { + McpServer::Stdio(stdio) => { + let mcp_client = () + .serve(TokioChildProcess::new( + Command::new(&stdio.command).configure(|cmd| { + cmd.args(&stdio.args); + for env_var in &stdio.env { + cmd.env(&env_var.name, &env_var.value); + } + }), + )?) + .await?; + + operation(mcp_client).await + } + McpServer::Http(http) => { + use rmcp::transport::StreamableHttpClientTransport; + + let mcp_client = + ().serve(StreamableHttpClientTransport::from_uri(http.url.as_str())) + .await?; + + operation(mcp_client).await + } + McpServer::Sse(_) => Err(anyhow::anyhow!("SSE MCP servers not yet supported")), + _ => Err(anyhow::anyhow!("Unknown MCP server type")), + } + } + + async fn list_tools(&self, session_id: &SessionId, server_name: &str) -> Result { + self.with_mcp_client(session_id, server_name, async move |mcp_client| { + let tools_result = mcp_client.list_tools(None).await?; + mcp_client.cancel().await?; + + let tools_list = tools_result + .tools + .iter() + .map(|tool| { + format!( + " - {}: {}", + tool.name, + tool.description.as_deref().unwrap_or("No description") + ) + }) + .collect::>() + .join("\n"); + + Ok(tools_list) + }) + .await + } + + async fn execute_tool_call( + &self, + session_id: &SessionId, + server_name: &str, + tool_name: &str, + params: serde_json::Value, + ) -> Result { + use rmcp::model::CallToolRequestParams; + + let params_obj = params.as_object().cloned().unwrap_or_default(); + let tool_name = tool_name.to_string(); + + self.with_mcp_client(session_id, server_name, async move |mcp_client| { + let tool_result = mcp_client + .call_tool(CallToolRequestParams::new(tool_name).with_arguments(params_obj)) + .await?; + + mcp_client.cancel().await?; + + Ok(format!("{tool_result:?}")) + }) + .await + } +} + +impl Default for Testy { + fn default() -> Self { + Self::new() + } +} + +/// Extract text content from prompt blocks. +fn extract_text_from_prompt(blocks: &[ContentBlock]) -> String { + blocks + .iter() + .filter_map(|block| match block { + ContentBlock::Text(TextContent { text, .. }) => Some(text.clone()), + _ => None, + }) + .collect::>() + .join(" ") +} + +impl ConnectTo for Testy { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + Agent + .builder() + .name("test-agent") + .on_receive_request( + async |initialize: InitializeRequest, responder, _cx| { + responder.respond( + InitializeResponse::new(initialize.protocol_version) + .agent_capabilities(AgentCapabilities::new()), + ) + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_request( + { + let agent = self.clone(); + async move |request: NewSessionRequest, responder, _cx| { + let session_id = SessionId::new(uuid::Uuid::new_v4().to_string()); + agent.create_session(&session_id, request.mcp_servers); + responder.respond(NewSessionResponse::new(session_id)) + } + }, + agent_client_protocol_core::on_receive_request!(), + ) + .on_receive_request( + { + let agent = self.clone(); + async move |request: PromptRequest, responder, cx| { + let cx_clone = cx.clone(); + cx.spawn({ + let agent = agent.clone(); + async move { agent.process_prompt(request, responder, cx_clone).await } + }) + } + }, + agent_client_protocol_core::on_receive_request!(), + ) + .connect_to(client) + .await + } +} diff --git a/src/agent-client-protocol-tokio/CHANGELOG.md b/src/agent-client-protocol-tokio/CHANGELOG.md new file mode 100644 index 0000000..1d013ff --- /dev/null +++ b/src/agent-client-protocol-tokio/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-tokio/Cargo.toml b/src/agent-client-protocol-tokio/Cargo.toml new file mode 100644 index 0000000..be67f97 --- /dev/null +++ b/src/agent-client-protocol-tokio/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "agent-client-protocol-tokio" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "Tokio-based utilities for the Agent Client Protocol" +keywords = ["acp", "agent", "protocol", "ai", "tokio"] +categories = ["development-tools"] + +[dependencies] +agent-client-protocol-core.workspace = true +futures.workspace = true +serde.workspace = true +serde_json.workspace = true +shell-words.workspace = true +tokio.workspace = true +tokio-util.workspace = true + +[dev-dependencies] +agent-client-protocol-test.workspace = true +expect-test.workspace = true +tokio.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-tokio/README.md b/src/agent-client-protocol-tokio/README.md new file mode 100644 index 0000000..146f230 --- /dev/null +++ b/src/agent-client-protocol-tokio/README.md @@ -0,0 +1,60 @@ +# agent-client-protocol-tokio + +Tokio-based utilities for working with [ACP](https://agentclientprotocol.com/) agents. + +## What's in this crate? + +This crate provides helpers for spawning and connecting to ACP agents using the Tokio async runtime: + +- **`AcpAgent`** — Configuration for spawning agent processes, parseable from command strings or JSON +- **`Stdio`** — A transport that connects over stdin/stdout with optional debug logging + +## Usage + +The main use case is spawning an agent process and connecting to it: + +```rust +use agent_client_protocol_core::{Client, ConnectTo}; +use agent_client_protocol_tokio::AcpAgent; +use std::str::FromStr; + +let agent = AcpAgent::from_str("python my_agent.py")?; + +// The agent process is spawned automatically when connected +Client.builder() + .name("my-client") + .connect_to(agent) + .await?; +``` + +You can also add debug logging to inspect the wire protocol: + +```rust +use agent_client_protocol_tokio::{AcpAgent, LineDirection}; + +let agent = AcpAgent::from_str("python my_agent.py")? + .with_debug(|line, direction| { + eprintln!("{direction:?}: {line}"); + }); +``` + +## When to use this crate + +Use `agent-client-protocol-tokio` when you need to: + +- Spawn agent processes from your code +- Test agents by programmatically launching them +- Build tools that orchestrate multiple agents + +If you're implementing an agent that listens on stdin/stdout, you only need the core +[`agent-client-protocol-core`](../agent-client-protocol-core/) crate. + +## Related Crates + +- **[agent-client-protocol-core](../agent-client-protocol-core/)** — Core ACP protocol types and traits +- **[agent-client-protocol-derive](../agent-client-protocol-derive/)** — Derive macros for JSON-RPC traits +- **[agent-client-protocol-trace-viewer](../agent-client-protocol-trace-viewer/)** — Interactive trace visualization + +## License + +Apache-2.0 diff --git a/src/agent-client-protocol-tokio/src/acp_agent.rs b/src/agent-client-protocol-tokio/src/acp_agent.rs new file mode 100644 index 0000000..1139ae2 --- /dev/null +++ b/src/agent-client-protocol-tokio/src/acp_agent.rs @@ -0,0 +1,589 @@ +//! Utilities for connecting to ACP agents and proxies. +//! +//! This module provides [`AcpAgent`], a convenient wrapper around [`agent_client_protocol_core::schema::McpServer`] +//! that can be parsed from either a command string or JSON configuration. + +use std::path::PathBuf; +use std::str::FromStr; +use std::sync::Arc; + +use agent_client_protocol_core::{Client, Conductor, Role}; +use tokio::process::Child; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +/// Direction of a line being sent or received. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LineDirection { + /// Line being sent to the agent (stdin) + Stdin, + /// Line being received from the agent (stdout) + Stdout, + /// Line being received from the agent (stderr) + Stderr, +} + +/// A component representing an external ACP agent running in a separate process. +/// +/// `AcpAgent` implements the [`agent_client_protocol_core::ConnectTo`] trait for spawning and communicating with +/// external agents or proxies via stdio. It handles process spawning, stream setup, and +/// byte stream serialization automatically. This is the primary way to connect to agents +/// that run as separate executables. +/// +/// This is a wrapper around [`agent_client_protocol_core::schema::McpServer`] that provides convenient parsing +/// from command-line strings or JSON configurations. +/// +/// # Use Cases +/// +/// - **External agents**: Connect to agents written in any language (Python, Node.js, Rust, etc.) +/// - **Proxy chains**: Spawn intermediate proxies that transform or intercept messages +/// - **Conductor components**: Use with [`agent_client_protocol_conductor::Conductor`] to build proxy chains +/// - **Subprocess isolation**: Run potentially untrusted code in a separate process +/// +/// # Examples +/// +/// Parse from a command string: +/// ``` +/// # use agent_client_protocol_tokio::AcpAgent; +/// # use std::str::FromStr; +/// let agent = AcpAgent::from_str("python my_agent.py --verbose").unwrap(); +/// ``` +/// +/// Parse from JSON: +/// ``` +/// # use agent_client_protocol_tokio::AcpAgent; +/// # use std::str::FromStr; +/// let agent = AcpAgent::from_str(r#"{"type": "stdio", "name": "my-agent", "command": "python", "args": ["my_agent.py"], "env": []}"#).unwrap(); +/// ``` +/// +/// Use as a component to connect to an external agent: +/// ```ignore +/// use agent_client_protocol_core::{Client, Builder}; +/// use agent_client_protocol_tokio::AcpAgent; +/// use std::str::FromStr; +/// +/// # async fn example() -> Result<(), Box> { +/// let agent = AcpAgent::from_str("python my_agent.py")?; +/// +/// // The agent process will be spawned automatically when connected +/// Client.builder() +/// .connect_to(agent) +/// .await? +/// .connect_with(|cx| async move { +/// // Use the connection to communicate with the agent process +/// Ok(()) +/// }) +/// .await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// [`agent_client_protocol_conductor::Conductor`]: https://docs.rs/agent-client-protocol-conductor/latest/agent_client_protocol_conductor/struct.Conductor.html +pub struct AcpAgent { + server: agent_client_protocol_core::schema::McpServer, + debug_callback: Option>, +} + +impl std::fmt::Debug for AcpAgent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AcpAgent") + .field("server", &self.server) + .field( + "debug_callback", + &self.debug_callback.as_ref().map(|_| "..."), + ) + .finish() + } +} + +impl AcpAgent { + /// Create a new `AcpAgent` from an [`agent_client_protocol_core::schema::McpServer`] configuration. + #[must_use] + pub fn new(server: agent_client_protocol_core::schema::McpServer) -> Self { + Self { + server, + debug_callback: None, + } + } + + /// Create an ACP agent for Zed Industries' Claude Code tool. + /// Just runs `npx -y @zed-industries/claude-code-acp@latest`. + #[must_use] + pub fn zed_claude_code() -> Self { + Self::from_str("npx -y @zed-industries/claude-code-acp@latest").expect("valid bash command") + } + + /// Create an ACP agent for Zed Industries' Codex tool. + /// Just runs `npx -y @zed-industries/codex-acp@latest`. + #[must_use] + pub fn zed_codex() -> Self { + Self::from_str("npx -y @zed-industries/codex-acp@latest").expect("valid bash command") + } + + /// Create an ACP agent for Google's Gemini CLI. + /// Just runs `npx -y -- @google/gemini-cli@latest --experimental-acp`. + #[must_use] + pub fn google_gemini() -> Self { + Self::from_str("npx -y -- @google/gemini-cli@latest --experimental-acp") + .expect("valid bash command") + } + + /// Get the underlying [`agent_client_protocol_core::schema::McpServer`] configuration. + #[must_use] + pub fn server(&self) -> &agent_client_protocol_core::schema::McpServer { + &self.server + } + + /// Convert into the underlying [`agent_client_protocol_core::schema::McpServer`] configuration. + #[must_use] + pub fn into_server(self) -> agent_client_protocol_core::schema::McpServer { + self.server + } + + /// Add a debug callback that will be invoked for each line sent/received. + /// + /// The callback receives the line content and the direction (stdin/stdout/stderr). + /// This is useful for logging, debugging, or monitoring agent communication. + /// + /// # Example + /// + /// ```no_run + /// # use agent_client_protocol_tokio::{AcpAgent, LineDirection}; + /// # use std::str::FromStr; + /// let agent = AcpAgent::from_str("python my_agent.py") + /// .unwrap() + /// .with_debug(|line, direction| { + /// eprintln!("{:?}: {}", direction, line); + /// }); + /// ``` + #[must_use] + pub fn with_debug(mut self, callback: F) -> Self + where + F: Fn(&str, LineDirection) + Send + Sync + 'static, + { + self.debug_callback = Some(Arc::new(callback)); + self + } + + /// Spawn the process and get stdio streams. + /// Used internally by the Component trait implementation. + pub fn spawn_process( + &self, + ) -> Result< + ( + tokio::process::ChildStdin, + tokio::process::ChildStdout, + tokio::process::ChildStderr, + Child, + ), + agent_client_protocol_core::Error, + > { + match &self.server { + agent_client_protocol_core::schema::McpServer::Stdio(stdio) => { + let mut cmd = tokio::process::Command::new(&stdio.command); + cmd.args(&stdio.args); + for env_var in &stdio.env { + cmd.env(&env_var.name, &env_var.value); + } + cmd.stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()); + + let mut child = cmd + .spawn() + .map_err(agent_client_protocol_core::Error::into_internal_error)?; + + let child_stdin = child.stdin.take().ok_or_else(|| { + agent_client_protocol_core::util::internal_error("Failed to open stdin") + })?; + let child_stdout = child.stdout.take().ok_or_else(|| { + agent_client_protocol_core::util::internal_error("Failed to open stdout") + })?; + let child_stderr = child.stderr.take().ok_or_else(|| { + agent_client_protocol_core::util::internal_error("Failed to open stderr") + })?; + + Ok((child_stdin, child_stdout, child_stderr, child)) + } + agent_client_protocol_core::schema::McpServer::Http(_) => { + Err(agent_client_protocol_core::util::internal_error( + "HTTP transport not yet supported by AcpAgent", + )) + } + agent_client_protocol_core::schema::McpServer::Sse(_) => { + Err(agent_client_protocol_core::util::internal_error( + "SSE transport not yet supported by AcpAgent", + )) + } + _ => Err(agent_client_protocol_core::util::internal_error( + "Unknown MCP server transport type", + )), + } + } +} + +/// A wrapper around Child that kills the process when dropped. +struct ChildGuard(Child); + +impl ChildGuard { + async fn wait(&mut self) -> std::io::Result { + self.0.wait().await + } +} + +impl Drop for ChildGuard { + fn drop(&mut self) { + drop(self.0.start_kill()); + } +} + +/// Waits for a child process and returns an error if it exits with non-zero status. +/// +/// The error message includes any stderr output collected by the background task. +/// When dropped, the child process is killed. +async fn monitor_child( + child: Child, + stderr_rx: tokio::sync::oneshot::Receiver, +) -> Result<(), agent_client_protocol_core::Error> { + let mut guard = ChildGuard(child); + + // Wait for the child to exit + let status = guard.wait().await.map_err(|e| { + agent_client_protocol_core::util::internal_error(format!("Failed to wait for process: {e}")) + })?; + + if status.success() { + Ok(()) + } else { + // Get stderr content if available + let stderr = stderr_rx.await.unwrap_or_default(); + + let message = if stderr.is_empty() { + format!("Process exited with {status}") + } else { + format!("Process exited with {status}: {stderr}") + }; + + Err(agent_client_protocol_core::util::internal_error(message)) + } +} + +/// Roles that an ACP agent executable can potentially serve. +pub trait AcpAgentCounterpartRole: Role {} + +impl AcpAgentCounterpartRole for Client {} + +impl AcpAgentCounterpartRole for Conductor {} + +impl agent_client_protocol_core::ConnectTo + for AcpAgent +{ + async fn connect_to( + self, + client: impl agent_client_protocol_core::ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + use futures::AsyncBufReadExt; + use futures::AsyncWriteExt; + use futures::StreamExt; + use futures::io::BufReader; + + let (child_stdin, child_stdout, child_stderr, child) = self.spawn_process()?; + + // Create a channel to collect stderr for error reporting + let (stderr_tx, stderr_rx) = tokio::sync::oneshot::channel::(); + + // Spawn a task to read stderr, optionally calling the debug callback + let debug_callback = self.debug_callback.clone(); + tokio::spawn(async move { + let stderr_reader = BufReader::new(child_stderr.compat()); + let mut stderr_lines = stderr_reader.lines(); + let mut collected = String::new(); + while let Some(line_result) = stderr_lines.next().await { + if let Ok(line) = line_result { + // Call debug callback if present + if let Some(ref callback) = debug_callback { + callback(&line, LineDirection::Stderr); + } + // Always collect for error reporting + if !collected.is_empty() { + collected.push('\n'); + } + collected.push_str(&line); + } + } + drop(stderr_tx.send(collected)); + }); + + // Create a future that monitors the child process for early exit + let child_monitor = monitor_child(child, stderr_rx); + + // Convert stdio to line streams with optional debug inspection + let incoming_lines = if let Some(callback) = self.debug_callback.clone() { + Box::pin( + BufReader::new(child_stdout.compat()) + .lines() + .inspect(move |result| { + if let Ok(line) = result { + callback(line, LineDirection::Stdout); + } + }), + ) + as std::pin::Pin> + Send>> + } else { + Box::pin(BufReader::new(child_stdout.compat()).lines()) + }; + + // Create a sink that writes lines (with newlines) to stdin with optional debug logging + let outgoing_sink = if let Some(callback) = self.debug_callback.clone() { + Box::pin(futures::sink::unfold( + (child_stdin.compat_write(), callback), + async move |(mut writer, callback), line: String| { + callback(&line, LineDirection::Stdin); + let mut bytes = line.into_bytes(); + bytes.push(b'\n'); + writer.write_all(&bytes).await?; + Ok::<_, std::io::Error>((writer, callback)) + }, + )) + as std::pin::Pin + Send>> + } else { + Box::pin(futures::sink::unfold( + child_stdin.compat_write(), + async move |mut writer, line: String| { + let mut bytes = line.into_bytes(); + bytes.push(b'\n'); + writer.write_all(&bytes).await?; + Ok::<_, std::io::Error>(writer) + }, + )) + }; + + // Race the protocol against child process exit + // If the child exits early (e.g., with an error), we return that error + let protocol_future = agent_client_protocol_core::ConnectTo::::connect_to( + agent_client_protocol_core::Lines::new(outgoing_sink, incoming_lines), + client, + ); + + tokio::select! { + result = protocol_future => result, + result = child_monitor => result, + } + } +} + +impl AcpAgent { + /// Create an `AcpAgent` from an iterator of command-line arguments. + /// + /// Leading arguments of the form `NAME=value` are parsed as environment variables. + /// The first non-env argument is the command, and the rest are arguments. + /// + /// # Example + /// + /// ``` + /// # use agent_client_protocol_tokio::AcpAgent; + /// let agent = AcpAgent::from_args([ + /// "RUST_LOG=debug", + /// "cargo", + /// "run", + /// "-p", + /// "my-crate", + /// ]).unwrap(); + /// ``` + pub fn from_args(args: I) -> Result + where + I: IntoIterator, + T: ToString, + { + let args: Vec = args.into_iter().map(|s| s.to_string()).collect(); + + if args.is_empty() { + return Err(agent_client_protocol_core::util::internal_error( + "Arguments cannot be empty", + )); + } + + let mut env = vec![]; + let mut command_idx = 0; + + // Parse leading FOO=bar arguments as environment variables + for (i, arg) in args.iter().enumerate() { + if let Some((name, value)) = parse_env_var(arg) { + env.push(agent_client_protocol_core::schema::EnvVariable::new( + name, value, + )); + command_idx = i + 1; + } else { + break; + } + } + + if command_idx >= args.len() { + return Err(agent_client_protocol_core::util::internal_error( + "No command found (only environment variables provided)", + )); + } + + let command = PathBuf::from(&args[command_idx]); + let cmd_args = args[command_idx + 1..].to_vec(); + + // Generate a name from the command + let name = command + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("agent") + .to_string(); + + Ok(AcpAgent { + server: agent_client_protocol_core::schema::McpServer::Stdio( + agent_client_protocol_core::schema::McpServerStdio::new(name, command) + .args(cmd_args) + .env(env), + ), + debug_callback: None, + }) + } +} + +/// Parse a string as an environment variable assignment (NAME=value). +/// Returns None if it doesn't match the pattern. +fn parse_env_var(s: &str) -> Option<(String, String)> { + // Must contain '=' and the part before must be a valid env var name + let eq_pos = s.find('=')?; + if eq_pos == 0 { + return None; + } + + let name = &s[..eq_pos]; + let value = &s[eq_pos + 1..]; + + // Env var names must start with a letter or underscore, and contain only + // alphanumeric characters and underscores + let mut chars = name.chars(); + let first = chars.next()?; + if !first.is_ascii_alphabetic() && first != '_' { + return None; + } + if !chars.all(|c| c.is_ascii_alphanumeric() || c == '_') { + return None; + } + + Some((name.to_string(), value.to_string())) +} + +impl FromStr for AcpAgent { + type Err = agent_client_protocol_core::Error; + + fn from_str(s: &str) -> Result { + let trimmed = s.trim(); + + // If it starts with '{', try to parse as JSON + if trimmed.starts_with('{') { + let server: agent_client_protocol_core::schema::McpServer = + serde_json::from_str(trimmed).map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Failed to parse JSON: {e}" + )) + })?; + return Ok(Self { + server, + debug_callback: None, + }); + } + + // Otherwise, parse as a command string + let parts = shell_words::split(trimmed).map_err(|e| { + agent_client_protocol_core::util::internal_error(format!( + "Failed to parse command: {e}" + )) + })?; + + Self::from_args(parts) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_simple_command() { + let agent = AcpAgent::from_str("python agent.py").unwrap(); + match agent.server { + agent_client_protocol_core::schema::McpServer::Stdio(stdio) => { + assert_eq!(stdio.name, "python"); + assert_eq!(stdio.command, PathBuf::from("python")); + assert_eq!(stdio.args, vec!["agent.py"]); + assert!(stdio.env.is_empty()); + } + _ => panic!("Expected Stdio variant"), + } + } + + #[test] + fn test_parse_command_with_args() { + let agent = AcpAgent::from_str("node server.js --port 8080 --verbose").unwrap(); + match agent.server { + agent_client_protocol_core::schema::McpServer::Stdio(stdio) => { + assert_eq!(stdio.name, "node"); + assert_eq!(stdio.command, PathBuf::from("node")); + assert_eq!(stdio.args, vec!["server.js", "--port", "8080", "--verbose"]); + assert!(stdio.env.is_empty()); + } + _ => panic!("Expected Stdio variant"), + } + } + + #[test] + fn test_parse_command_with_quotes() { + let agent = AcpAgent::from_str(r#"python "my agent.py" --name "Test Agent""#).unwrap(); + match agent.server { + agent_client_protocol_core::schema::McpServer::Stdio(stdio) => { + assert_eq!(stdio.name, "python"); + assert_eq!(stdio.command, PathBuf::from("python")); + assert_eq!(stdio.args, vec!["my agent.py", "--name", "Test Agent"]); + assert!(stdio.env.is_empty()); + } + _ => panic!("Expected Stdio variant"), + } + } + + #[test] + fn test_parse_json_stdio() { + let json = r#"{ + "type": "stdio", + "name": "my-agent", + "command": "/usr/bin/python", + "args": ["agent.py", "--verbose"], + "env": [] + }"#; + let agent = AcpAgent::from_str(json).unwrap(); + match agent.server { + agent_client_protocol_core::schema::McpServer::Stdio(stdio) => { + assert_eq!(stdio.name, "my-agent"); + assert_eq!(stdio.command, PathBuf::from("/usr/bin/python")); + assert_eq!(stdio.args, vec!["agent.py", "--verbose"]); + assert!(stdio.env.is_empty()); + } + _ => panic!("Expected Stdio variant"), + } + } + + #[test] + fn test_parse_json_http() { + let json = r#"{ + "type": "http", + "name": "remote-agent", + "url": "https://example.com/agent", + "headers": [] + }"#; + let agent = AcpAgent::from_str(json).unwrap(); + match agent.server { + agent_client_protocol_core::schema::McpServer::Http(http) => { + assert_eq!(http.name, "remote-agent"); + assert_eq!(http.url, "https://example.com/agent"); + assert!(http.headers.is_empty()); + } + _ => panic!("Expected Http variant"), + } + } +} diff --git a/src/agent-client-protocol-tokio/src/lib.rs b/src/agent-client-protocol-tokio/src/lib.rs new file mode 100644 index 0000000..5936e0c --- /dev/null +++ b/src/agent-client-protocol-tokio/src/lib.rs @@ -0,0 +1,104 @@ +//! Tokio-based utilities for the Agent Client Protocol +//! +//! This crate provides higher-level functionality for working with ACP +//! that requires the Tokio async runtime, such as spawning agent processes +//! and creating connections. + +mod acp_agent; + +pub use acp_agent::{AcpAgent, LineDirection}; +use agent_client_protocol_core::{ByteStreams, ConnectTo, Role}; +use std::sync::Arc; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +pub struct Stdio { + debug_callback: Option>, +} + +impl std::fmt::Debug for Stdio { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Stdio").finish_non_exhaustive() + } +} + +impl Stdio { + #[must_use] + pub fn new() -> Self { + Self { + debug_callback: None, + } + } + + #[must_use] + pub fn with_debug(mut self, callback: F) -> Self + where + F: Fn(&str, LineDirection) + Send + Sync + 'static, + { + self.debug_callback = Some(Arc::new(callback)); + self + } +} + +impl Default for Stdio { + fn default() -> Self { + Self::new() + } +} + +impl ConnectTo for Stdio { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol_core::Error> { + if let Some(callback) = self.debug_callback { + use futures::AsyncBufReadExt; + use futures::AsyncWriteExt; + use futures::StreamExt; + use futures::io::BufReader; + + // With debug: use Lines with interception + let stdin = tokio::io::stdin(); + let stdout = tokio::io::stdout(); + + // Convert stdio to line streams with debug inspection + let incoming_callback = callback.clone(); + let incoming_lines = Box::pin(BufReader::new(stdin.compat()).lines().inspect( + move |result| { + if let Ok(line) = result { + incoming_callback(line, LineDirection::Stdin); + } + }, + )) + as std::pin::Pin> + Send>>; + + // Create a sink that writes lines with debug logging + let outgoing_sink = Box::pin(futures::sink::unfold( + (stdout.compat_write(), callback), + async move |(mut writer, callback), line: String| { + callback(&line, LineDirection::Stdout); + let mut bytes = line.into_bytes(); + bytes.push(b'\n'); + writer.write_all(&bytes).await?; + Ok::<_, std::io::Error>((writer, callback)) + }, + )) + as std::pin::Pin + Send>>; + + ConnectTo::::connect_to( + agent_client_protocol_core::Lines::new(outgoing_sink, incoming_lines), + client, + ) + .await + } else { + // Without debug: use simple ByteStreams + ConnectTo::::connect_to( + ByteStreams::new( + tokio::io::stdout().compat_write(), + tokio::io::stdin().compat(), + ), + client, + ) + .await + } + } +} diff --git a/src/agent-client-protocol-tokio/tests/debug_logging.rs b/src/agent-client-protocol-tokio/tests/debug_logging.rs new file mode 100644 index 0000000..a56490b --- /dev/null +++ b/src/agent-client-protocol-tokio/tests/debug_logging.rs @@ -0,0 +1,128 @@ +//! Integration test for AcpAgent debug logging + +use agent_client_protocol_core::schema::InitializeRequest; +use agent_client_protocol_core::{Client, ConnectTo}; +use agent_client_protocol_test::test_binaries::testy; +use agent_client_protocol_tokio::LineDirection; +use std::sync::{Arc, Mutex}; + +/// Test helper to receive a JSON-RPC response +async fn recv( + response: agent_client_protocol_core::SentRequest, +) -> Result { + let (tx, rx) = tokio::sync::oneshot::channel(); + response.on_receiving_result(async move |result| { + tx.send(result) + .map_err(|_| agent_client_protocol_core::Error::internal_error()) + })?; + rx.await + .map_err(|_| agent_client_protocol_core::Error::internal_error())? +} + +#[tokio::test] +async fn test_acp_agent_debug_callback() -> Result<(), Box> { + use tokio::io::duplex; + use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + + // Collect debug output + #[derive(Debug, Clone, Default)] + struct DebugLog { + lines: Arc>>, + } + + impl DebugLog { + fn log(&self, line: &str, direction: LineDirection) { + self.lines + .lock() + .unwrap() + .push((line.to_string(), direction)); + } + + fn get_lines(&self) -> Vec<(String, LineDirection)> { + self.lines.lock().unwrap().clone() + } + } + + let debug_log = DebugLog::default(); + + // Create an agent that runs test-agent + let agent = testy().with_debug({ + let debug_log = debug_log.clone(); + move |line, direction| { + debug_log.log(line, direction); + } + }); + + // Set up client <-> agent communication + let (client_out, agent_in) = duplex(1024); + let (agent_out, client_in) = duplex(1024); + + let transport = + agent_client_protocol_core::ByteStreams::new(client_out.compat_write(), client_in.compat()); + + let client = Client + .builder() + .name("test-client") + .with_spawned(|_cx| async move { + ConnectTo::::connect_to( + agent, + agent_client_protocol_core::ByteStreams::new( + agent_out.compat_write(), + agent_in.compat(), + ), + ) + .await + }) + .connect_with(transport, async |connection_to_client| { + // Send an initialize request + let _init_response = recv(connection_to_client.send_request(InitializeRequest::new( + agent_client_protocol_core::schema::ProtocolVersion::LATEST, + ))) + .await?; + + Ok(()) + }); + Box::pin(client).await?; + + // Verify debug output was captured + let logged_lines = debug_log.get_lines(); + + // Should have at least some stdin and stdout lines + let stdin_count = logged_lines + .iter() + .filter(|(_, dir)| *dir == LineDirection::Stdin) + .count(); + let stdout_count = logged_lines + .iter() + .filter(|(_, dir)| *dir == LineDirection::Stdout) + .count(); + + assert!( + stdin_count > 0, + "Expected at least one stdin line, got {stdin_count}" + ); + assert!( + stdout_count > 0, + "Expected at least one stdout line, got {stdout_count}" + ); + + // Check that we logged the initialize request (contains "initialize" method) + let has_initialize_request = logged_lines.iter().any(|(line, dir)| { + *dir == LineDirection::Stdin && line.contains("\"method\":\"initialize\"") + }); + assert!( + has_initialize_request, + "Expected to find initialize request in debug log" + ); + + // Check that we logged the initialize response (contains result field) + let has_initialize_response = logged_lines + .iter() + .any(|(line, dir)| *dir == LineDirection::Stdout && line.contains("\"result\"")); + assert!( + has_initialize_response, + "Expected to find initialize response in debug log" + ); + + Ok(()) +} diff --git a/src/agent-client-protocol-trace-viewer/CHANGELOG.md b/src/agent-client-protocol-trace-viewer/CHANGELOG.md new file mode 100644 index 0000000..1d013ff --- /dev/null +++ b/src/agent-client-protocol-trace-viewer/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/agent-client-protocol-trace-viewer/Cargo.toml b/src/agent-client-protocol-trace-viewer/Cargo.toml new file mode 100644 index 0000000..9b2c23a --- /dev/null +++ b/src/agent-client-protocol-trace-viewer/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "agent-client-protocol-trace-viewer" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "Interactive sequence diagram viewer for Agent Client Protocol trace files" +keywords = ["acp", "trace", "debugging", "visualization"] +categories = ["development-tools::debugging"] + +[lib] +name = "agent_client_protocol_trace_viewer" +path = "src/lib.rs" + +[[bin]] +name = "agent-client-protocol-trace-viewer" +path = "src/main.rs" + +[dependencies] +anyhow.workspace = true +axum.workspace = true +clap.workspace = true +open.workspace = true +serde.workspace = true +serde_json.workspace = true +tokio.workspace = true +tower.workspace = true +tower-http.workspace = true + +[lints] +workspace = true diff --git a/src/agent-client-protocol-trace-viewer/src/lib.rs b/src/agent-client-protocol-trace-viewer/src/lib.rs new file mode 100644 index 0000000..a07fee0 --- /dev/null +++ b/src/agent-client-protocol-trace-viewer/src/lib.rs @@ -0,0 +1,204 @@ +//! Agent Client Protocol Trace Viewer Library +//! +//! Provides an interactive sequence diagram viewer for ACP trace events. +//! Can serve events from memory (for live viewing) or from a file. + +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; + +use axum::{ + Router, + extract::State, + http::StatusCode, + response::{Html, IntoResponse, Response}, + routing::get, +}; +use tokio::net::TcpListener; + +/// The HTML viewer page (embedded at compile time). +pub const VIEWER_HTML: &str = include_str!("viewer.html"); + +/// Source of trace events for the viewer. +#[derive(Debug, Clone)] +pub enum TraceSource { + /// Read events from a file (re-reads on each request for live updates). + File(PathBuf), + /// Read events from shared memory. + Memory(Arc>>), +} + +/// Handle to push events when using memory-backed trace source. +#[derive(Debug, Clone)] +pub struct TraceHandle { + events: Arc>>, +} + +impl TraceHandle { + /// Push a new event to the trace. + pub fn push(&self, event: serde_json::Value) { + self.events.lock().unwrap().push(event); + } + + /// Get the current number of events. + #[must_use] + pub fn len(&self) -> usize { + self.events.lock().unwrap().len() + } + + /// Check if empty. + #[must_use] + pub fn is_empty(&self) -> bool { + self.events.lock().unwrap().is_empty() + } +} + +struct AppState { + source: TraceSource, +} + +/// Configuration for the trace viewer server. +#[derive(Debug)] +pub struct TraceViewerConfig { + /// Port to serve on (0 = auto-select). + pub port: u16, + /// Whether to open browser automatically. + pub open_browser: bool, +} + +impl Default for TraceViewerConfig { + fn default() -> Self { + Self { + port: 0, + open_browser: true, + } + } +} + +/// Start the trace viewer server with a memory-backed event source. +/// +/// Returns a handle to push events and a future that runs the server. +/// The server will poll for new events automatically. +/// +/// # Example +/// +/// ```no_run +/// # async fn example() -> anyhow::Result<()> { +/// let (handle, server) = agent_client_protocol_trace_viewer::serve_memory(Default::default())?; +/// +/// // Push events from your application +/// handle.push(serde_json::json!({"type": "request", "method": "test"})); +/// +/// // Run the server (or spawn it) +/// server.await?; +/// # Ok(()) +/// # } +/// ``` +pub fn serve_memory( + config: TraceViewerConfig, +) -> anyhow::Result<( + TraceHandle, + impl std::future::Future>, +)> { + let events = Arc::new(Mutex::new(Vec::new())); + let handle = TraceHandle { + events: events.clone(), + }; + let source = TraceSource::Memory(events); + + let server = serve_impl(source, config); + Ok((handle, server)) +} + +/// Start the trace viewer server with a file-backed event source. +/// +/// The file is re-read on each request, allowing live updates as the file grows. +/// +/// # Example +/// +/// ```no_run +/// # use std::path::PathBuf; +/// # async fn example() -> anyhow::Result<()> { +/// agent_client_protocol_trace_viewer::serve_file(PathBuf::from("trace.jsons"), Default::default()).await?; +/// # Ok(()) +/// # } +/// ``` +pub async fn serve_file(path: PathBuf, config: TraceViewerConfig) -> anyhow::Result<()> { + serve_impl(TraceSource::File(path), config).await +} + +async fn serve_impl(source: TraceSource, config: TraceViewerConfig) -> anyhow::Result<()> { + let state = Arc::new(AppState { source }); + + let app = Router::new() + .route("/", get(serve_viewer)) + .route("/events", get(serve_events)) + .with_state(state); + + let listener = TcpListener::bind(format!("127.0.0.1:{}", config.port)).await?; + let addr = listener.local_addr()?; + + eprintln!("Trace viewer at http://{addr}"); + + if config.open_browser { + let url = format!("http://{addr}"); + if let Err(e) = open::that(&url) { + eprintln!("Failed to open browser: {e}. Open {url} manually."); + } + } + + axum::serve(listener, app).await?; + Ok(()) +} + +/// Serve the main viewer HTML page. +async fn serve_viewer() -> Html<&'static str> { + Html(VIEWER_HTML) +} + +/// Serve the trace events as JSON. +async fn serve_events(State(state): State>) -> Response { + match &state.source { + TraceSource::File(path) => serve_events_from_file(path).await, + TraceSource::Memory(events) => serve_events_from_memory(events), + } +} + +async fn serve_events_from_file(path: &PathBuf) -> Response { + match tokio::fs::read_to_string(path).await { + Ok(content) => { + let events: Vec = content + .lines() + .filter(|line| !line.trim().is_empty()) + .filter_map(|line| serde_json::from_str(line).ok()) + .collect(); + + match serde_json::to_string(&events) { + Ok(json) => { + (StatusCode::OK, [("content-type", "application/json")], json).into_response() + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to serialize events: {e}"), + ) + .into_response(), + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to read trace file: {e}"), + ) + .into_response(), + } +} + +fn serve_events_from_memory(events: &Arc>>) -> Response { + let events = events.lock().unwrap(); + match serde_json::to_string(&*events) { + Ok(json) => (StatusCode::OK, [("content-type", "application/json")], json).into_response(), + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to serialize events: {e}"), + ) + .into_response(), + } +} diff --git a/src/agent-client-protocol-trace-viewer/src/main.rs b/src/agent-client-protocol-trace-viewer/src/main.rs new file mode 100644 index 0000000..00d67b3 --- /dev/null +++ b/src/agent-client-protocol-trace-viewer/src/main.rs @@ -0,0 +1,51 @@ +//! Agent Client Protocol Trace Viewer +//! +//! Interactive sequence diagram viewer for ACP trace files. +//! +//! Usage: +//! ```bash +//! agent-client-protocol-trace-viewer ./trace.jsons +//! ``` +//! +//! This starts a local HTTP server and opens a browser to view the trace +//! as an interactive sequence diagram. + +use std::path::PathBuf; + +use clap::Parser; + +#[derive(Parser, Debug)] +#[command( + author, + version, + about = "Interactive sequence diagram viewer for ACP trace files" +)] +struct Args { + /// Path to the trace file (.jsons) + trace_file: PathBuf, + + /// Port to serve on (default: auto-select) + #[arg(short, long)] + port: Option, + + /// Don't open browser automatically + #[arg(long)] + no_open: bool, +} + +#[tokio::main] +async fn main() -> anyhow::Result<()> { + let args = Args::parse(); + + // Verify trace file exists + if !args.trace_file.exists() { + anyhow::bail!("Trace file not found: {}", args.trace_file.display()); + } + + let config = agent_client_protocol_trace_viewer::TraceViewerConfig { + port: args.port.unwrap_or(0), + open_browser: !args.no_open, + }; + + agent_client_protocol_trace_viewer::serve_file(args.trace_file, config).await +} diff --git a/src/agent-client-protocol-trace-viewer/src/viewer.html b/src/agent-client-protocol-trace-viewer/src/viewer.html new file mode 100644 index 0000000..206cd55 --- /dev/null +++ b/src/agent-client-protocol-trace-viewer/src/viewer.html @@ -0,0 +1,767 @@ + + + + + + SACP Trace Viewer + + + +
+

SACP Trace Viewer

+
+ + + + +
+
+ +
+
+
Loading trace...
+
+ +
+
+
+

Message Details

+ +
+
+

+                
+
+
+ + + + diff --git a/src/yopo/CHANGELOG.md b/src/yopo/CHANGELOG.md new file mode 100644 index 0000000..1d013ff --- /dev/null +++ b/src/yopo/CHANGELOG.md @@ -0,0 +1,6 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). diff --git a/src/yopo/Cargo.toml b/src/yopo/Cargo.toml new file mode 100644 index 0000000..f028526 --- /dev/null +++ b/src/yopo/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "agent-client-protocol-yopo" +version = "0.1.0" +edition.workspace = true +authors.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +description = "YOPO (You Only Prompt Once) - A simple ACP client for one-shot prompts" + +[lib] +name = "yopo" +path = "src/lib.rs" + +[[bin]] +name = "yopo" +path = "src/main.rs" + +[dependencies] +agent-client-protocol-core.workspace = true +agent-client-protocol-tokio.workspace = true +clap.workspace = true +tokio.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true + +[lints] +workspace = true diff --git a/src/yopo/src/lib.rs b/src/yopo/src/lib.rs new file mode 100644 index 0000000..46505cd --- /dev/null +++ b/src/yopo/src/lib.rs @@ -0,0 +1,224 @@ +//! YOPO (You Only Prompt Once) - A simple library for testing ACP agents +//! +//! Provides a convenient API for running one-shot prompts against ACP components. + +use agent_client_protocol_core::schema::{ + AudioContent, ContentBlock, EmbeddedResourceResource, ImageContent, InitializeRequest, + ProtocolVersion, RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse, + SelectedPermissionOutcome, SessionNotification, TextContent, +}; +use agent_client_protocol_core::util::MatchDispatch; +use agent_client_protocol_core::{Agent, Client, ConnectTo, Dispatch, Handled, UntypedMessage}; +use std::path::PathBuf; + +/// Converts a `ContentBlock` to its string representation. +/// +/// This function provides standard string conversions for different content types: +/// - `Text`: Returns the text content +/// - `Image`: Returns a placeholder like `[Image: image/png]` +/// - `Audio`: Returns a placeholder like `[Audio: audio/wav]` +/// - `ResourceLink`: Returns the URI +/// - `Resource`: Returns the URI +/// +/// # Example +/// +/// ```no_run +/// use yopo::content_block_to_string; +/// use agent_client_protocol_core::schema::{ContentBlock, TextContent}; +/// +/// let block = ContentBlock::Text(TextContent::new("Hello".to_string())); +/// assert_eq!(content_block_to_string(&block), "Hello"); +/// ``` +#[must_use] +pub fn content_block_to_string(block: &ContentBlock) -> String { + match block { + ContentBlock::Text(TextContent { text, .. }) => text.clone(), + ContentBlock::Image(ImageContent { mime_type, .. }) => { + format!("[Image: {mime_type}]") + } + ContentBlock::Audio(AudioContent { mime_type, .. }) => { + format!("[Audio: {mime_type}]") + } + ContentBlock::ResourceLink(link) => link.uri.clone(), + ContentBlock::Resource(resource) => match &resource.resource { + EmbeddedResourceResource::TextResourceContents(text) => text.uri.clone(), + EmbeddedResourceResource::BlobResourceContents(blob) => blob.uri.clone(), + _ => "[Unknown resource type]".to_string(), + }, + _ => "[Unknown content type]".to_string(), + } +} + +/// Runs a single prompt against a component with a callback for each content block. +/// +/// This function: +/// - Spawns the component +/// - Initializes the agent +/// - Creates a new session +/// - Sends the prompt +/// - Auto-approves all permission requests +/// - Calls the callback with each `ContentBlock` from agent messages +/// - Returns when the prompt completes +/// +/// The callback receives each `ContentBlock` as it arrives and can process it +/// asynchronously (e.g., print it, accumulate it, etc.). +/// +/// # Example +/// +/// ```ignore +/// use yopo::{prompt_with_callback, content_block_to_string}; +/// use agent_client_protocol_tokio::AcpAgent; +/// use std::str::FromStr; +/// +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// let agent = AcpAgent::from_str("python agent.py")?; +/// prompt_with_callback(agent, "What is 2+2?", async |block| { +/// print!("{}", content_block_to_string(&block)); +/// }).await?; +/// # Ok(()) +/// # } +/// ``` +pub async fn prompt_with_callback( + component: impl ConnectTo, + prompt_text: impl ToString, + mut callback: impl AsyncFnMut(ContentBlock) + Send, +) -> Result<(), agent_client_protocol_core::Error> { + // Convert prompt to String + let prompt_text = prompt_text.to_string(); + + // Run the client + Client + .builder() + .on_receive_dispatch( + async |message: Dispatch, _cx| { + tracing::trace!("received: {:?}", message.message()); + Ok(Handled::No { + message, + retry: false, + }) + }, + agent_client_protocol_core::on_receive_dispatch!(), + ) + .connect_with(component, |cx: agent_client_protocol_core::ConnectionTo| async move { + // Initialize the agent + let _init_response = cx + .send_request(InitializeRequest::new(ProtocolVersion::LATEST)) + .block_task() + .await?; + + let mut session = cx + .build_session(PathBuf::from(".")) + .block_task() + .start_session() + .await?; + + session.send_prompt(prompt_text)?; + + loop { + let update = session.read_update().await?; + match update { + agent_client_protocol_core::SessionMessage::SessionMessage(message) => { + MatchDispatch::new(message) + .if_notification(async |notification: SessionNotification| { + tracing::debug!( + ?notification, + "yopo: received SessionNotification" + ); + // Call the callback for each agent message chunk + if let agent_client_protocol_core::schema::SessionUpdate::AgentMessageChunk( + content_chunk, + ) = notification.update + { + callback(content_chunk.content).await; + } + Ok(()) + }) + .await + .if_request(async |request: RequestPermissionRequest, responder| { + // Auto-approve all permission requests by selecting the first option + // that looks "allow-ish" + let outcome = request + .options + .iter() + .find(|option| match option.kind { + agent_client_protocol_core::schema::PermissionOptionKind::AllowOnce + | agent_client_protocol_core::schema::PermissionOptionKind::AllowAlways => true, + agent_client_protocol_core::schema::PermissionOptionKind::RejectOnce + | agent_client_protocol_core::schema::PermissionOptionKind::RejectAlways + | _ => false, + }) + .map_or(RequestPermissionOutcome::Cancelled, |option| { + RequestPermissionOutcome::Selected( + SelectedPermissionOutcome::new( + option.option_id.clone(), + ), + ) + }); + + responder.respond(RequestPermissionResponse::new(outcome))?; + + Ok(()) + }) + .await + .otherwise(async |_msg| Ok(())) + .await?; + } + agent_client_protocol_core::SessionMessage::StopReason(stop_reason) => match stop_reason { + agent_client_protocol_core::schema::StopReason::EndTurn => break, + agent_client_protocol_core::schema::StopReason::MaxTokens => todo!(), + agent_client_protocol_core::schema::StopReason::MaxTurnRequests => todo!(), + agent_client_protocol_core::schema::StopReason::Refusal => todo!(), + agent_client_protocol_core::schema::StopReason::Cancelled => todo!(), + _ => todo!(), + }, + _ => todo!(), + } + } + + Ok(()) + }) + .await?; + + Ok(()) +} + +/// Runs a single prompt against a component and returns the accumulated text response. +/// +/// This function: +/// - Spawns the component +/// - Initializes the agent +/// - Creates a new session +/// - Sends the prompt +/// - Auto-approves all permission requests +/// - Accumulates all content from agent messages using [`content_block_to_string`] +/// - Returns the complete response as a String +/// +/// This is a convenience wrapper around [`prompt_with_callback`] that accumulates +/// all content blocks into a single string. +/// +/// # Example +/// +/// ```ignore +/// use yopo::prompt; +/// use agent_client_protocol_tokio::AcpAgent; +/// use std::str::FromStr; +/// +/// # async fn example() -> Result<(), agent_client_protocol_core::Error> { +/// let agent = AcpAgent::from_str("python agent.py")?; +/// let response = prompt(agent, "What is 2+2?").await?; +/// assert!(response.contains("4")); +/// # Ok(()) +/// # } +/// ``` +pub async fn prompt( + component: impl ConnectTo, + prompt_text: impl ToString, +) -> Result { + let mut accumulated_text = String::new(); + prompt_with_callback(component, prompt_text, async |block| { + let text = content_block_to_string(&block); + accumulated_text.push_str(&text); + }) + .await?; + Ok(accumulated_text) +} diff --git a/src/yopo/src/main.rs b/src/yopo/src/main.rs new file mode 100644 index 0000000..7c92916 --- /dev/null +++ b/src/yopo/src/main.rs @@ -0,0 +1,84 @@ +//! YOPO (You Only Prompt Once) - A simple ACP client for one-shot prompts +//! +//! This client: +//! - Takes a prompt and agent command as arguments +//! - Spawns the agent +//! - Sends the prompt +//! - Auto-approves all permission requests +//! - Prints content progressively as it arrives +//! - Runs until the agent completes +//! +//! # Usage +//! +//! With a command (arguments are concatenated): +//! ```bash +//! yopo "What is 2+2?" python my_agent.py +//! yopo "Hello!" cargo run --release +//! ``` +//! +//! With JSON config: +//! ```bash +//! yopo "Hello!" '{"type":"stdio","name":"my-agent","command":"python","args":["agent.py"],"env":[]}' +//! ``` + +use agent_client_protocol_tokio::AcpAgent; +use clap::Parser; +use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; + +#[derive(Parser, Debug)] +#[command(author, version, about = "YOPO - You Only Prompt Once", long_about = None)] +struct Args { + /// The prompt to send to the agent + prompt: String, + + /// Agent command (multiple arguments are joined with spaces) or JSON config + #[arg(required = true, num_args = 1..)] + agent_args: Vec, + + /// Set logging level (trace, debug, info, warn, error) + #[arg(short, long)] + log: Option, +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let args = Args::parse(); + + // Initialize tracing to stderr + let env_filter = if let Some(level) = args.log { + EnvFilter::new(format!("yopo={level}")) + } else { + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("yopo=info")) + }; + + tracing_subscriber::registry() + .with(env_filter) + .with( + tracing_subscriber::fmt::layer() + .with_target(true) + .with_writer(std::io::stderr), + ) + .init(); + + let prompt = &args.prompt; + + // Parse the agent configuration from args + let agent = AcpAgent::from_args(args.agent_args)?; + + eprintln!("🚀 Spawning agent and running prompt..."); + + // Use the library function with callback to print progressively + Box::pin(yopo::prompt_with_callback( + agent, + prompt.as_str(), + |block| async move { + print!("{}", yopo::content_block_to_string(&block)); + }, + )) + .await?; + + println!(); // Final newline + eprintln!("✅ Agent completed!"); + + Ok(()) +}