diff --git a/.github/ci-path-filters.yml b/.github/ci-path-filters.yml index a589613fe..3b783a62b 100644 --- a/.github/ci-path-filters.yml +++ b/.github/ci-path-filters.yml @@ -163,12 +163,14 @@ rust: - 'Cargo.toml' - 'crates/**/Cargo.toml' - 'crates/**/*.rs' - # The embedded pi extension is `.ts`/`.json` under the CLI crate. Changes must run - # the Rust tests that verify the embedded asset inventory and package contents. - 'crates/cli/assets/**' + # The managed pi extension is `.ts`/`.json` under the CLI daemon tree. Changes must run + # the Rust tests that verify the embedded asset inventory and package contents. + - 'crates/cli/src/daemon/managed/pi_extension/**' - 'crates/worker-proto/proto/**' - 'crates/ffi/cbindgen.toml' - 'integrations/coding-agents/**' + - 'scripts/latency_benchmark/daemon_transport/**' go: - 'crates/ffi/Cargo.toml' @@ -211,6 +213,8 @@ pi: # suite and symlinks the rest, so an edit to the extension itself lands here and would # otherwise never reach the suite that tests it. - 'crates/cli/assets/pi-extension/**' + - 'crates/cli/src/daemon/managed/pi_extension/**' + - 'crates/cli/tests/managed_pi_extension_tests.mjs' - 'integrations/pi/**' python: diff --git a/.github/workflows/ci_rust.yml b/.github/workflows/ci_rust.yml index b4be6602a..e035b70e1 100644 --- a/.github/workflows/ci_rust.yml +++ b/.github/workflows/ci_rust.yml @@ -111,6 +111,7 @@ jobs: with: cache: false toolchain: ${{ steps.ci-config.outputs.rust_version }} + components: clippy, rustfmt - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 with: @@ -169,6 +170,16 @@ jobs: name: rust-${{ matrix.platform }} verbose: true + - name: Check daemon transport benchmark + if: ${{ matrix.platform == 'linux-amd64' }} + working-directory: ${{ env.NEMO_RELAY_CI_WORKSPACE }} + run: just check-daemon-transport-benchmark + + - name: Run daemon transport benchmark smoke check + if: ${{ matrix.platform == 'linux-amd64' }} + working-directory: ${{ env.NEMO_RELAY_CI_WORKSPACE }} + run: just daemon-transport-benchmark-smoke + Package: name: Package (${{ matrix.platform }}) if: ${{ inputs.run_package }} diff --git a/.gitignore b/.gitignore index 2fa985d81..6f005ee19 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ # Rust /target +/scripts/latency_benchmark/daemon_transport/target # Python /.venv diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index affb9e07e..03fbf9b6a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -159,6 +159,13 @@ repos: types: [rust] pass_filenames: false + - id: daemon-transport-benchmark-rust + name: daemon transport benchmark fmt and clippy + entry: just check-daemon-transport-benchmark + language: system + files: '^scripts/latency_benchmark/daemon_transport/.*\.(rs|toml)$' + pass_filenames: false + - id: cargo-deny name: cargo deny entry: cargo deny check diff --git a/ATTRIBUTIONS-Rust.md b/ATTRIBUTIONS-Rust.md index 33da74e3d..cb75afb05 100644 --- a/ATTRIBUTIONS-Rust.md +++ b/ATTRIBUTIONS-Rust.md @@ -10,6 +10,7 @@ This project uses the following third-party libraries. Each library is open-sour This file is automatically generated. Please do not edit it directly. Regenerate with `./scripts/generate_attributions.sh rust`. ## ahash - 0.8.12 + **Repository URL**: https://github.com/tkaitchuck/ahash **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -218,6 +219,7 @@ limitations under the License. ``` ## aho-corasick - 1.1.4 + **Repository URL**: https://github.com/BurntSushi/aho-corasick **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -246,6 +248,7 @@ THE SOFTWARE. ``` ## allocator-api2 - 0.2.21 + **Repository URL**: https://github.com/zakarumych/allocator-api2 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -326,6 +329,7 @@ limitations under the License. ``` ## android_system_properties - 0.1.5 + **Repository URL**: https://github.com/nical/android_system_properties **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -406,6 +410,7 @@ limitations under the License. ``` ## anstream - 1.0.0 + **Repository URL**: https://github.com/rust-cli/anstyle.git **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -615,6 +620,7 @@ limitations under the License. ``` ## anstyle - 1.0.14 + **Repository URL**: https://github.com/rust-cli/anstyle.git **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -824,6 +830,7 @@ limitations under the License. ``` ## anstyle-parse - 1.0.0 + **Repository URL**: https://github.com/rust-cli/anstyle.git **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -1033,6 +1040,7 @@ limitations under the License. ``` ## anstyle-query - 1.1.5 + **Repository URL**: https://github.com/rust-cli/anstyle.git **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -1242,6 +1250,7 @@ limitations under the License. ``` ## anstyle-wincon - 3.0.11 + **Repository URL**: https://github.com/rust-cli/anstyle.git **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -1451,6 +1460,7 @@ limitations under the License. ``` ## anyhow - 1.0.104 + **Repository URL**: https://github.com/dtolnay/anyhow **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -1531,6 +1541,7 @@ limitations under the License. ``` ## arc-swap - 1.9.1 + **Repository URL**: https://github.com/vorner/arc-swap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -1739,6 +1750,7 @@ limitations under the License. ``` ## arcstr - 1.2.0 + **Repository URL**: https://github.com/thomcc/arcstr **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -1947,6 +1959,7 @@ limitations under the License. ``` ## async-lock - 3.4.2 + **Repository URL**: https://github.com/smol-rs/async-lock **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -2155,6 +2168,7 @@ limitations under the License. ``` ## async-stream - 0.3.6 + **Repository URL**: https://github.com/tokio-rs/async-stream **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -2180,6 +2194,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## async-stream-impl - 0.3.6 + **Repository URL**: https://github.com/tokio-rs/async-stream **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -2205,6 +2220,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## async-trait - 0.1.89 + **Repository URL**: https://github.com/dtolnay/async-trait **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -2285,6 +2301,7 @@ limitations under the License. ``` ## atomic - 0.6.1 + **Repository URL**: https://github.com/Amanieu/atomic-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -2493,6 +2510,7 @@ limitations under the License. ``` ## atomic-waker - 1.1.2 + **Repository URL**: https://github.com/smol-rs/atomic-waker **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -2701,6 +2719,7 @@ limitations under the License. ``` ## autocfg - 1.5.0 + **Repository URL**: https://github.com/cuviper/autocfg **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -2909,6 +2928,7 @@ limitations under the License. ``` ## aws-lc-rs - 1.17.3 + **Repository URL**: https://github.com/aws/aws-lc-rs **License Type(s)**: ISC ### License: https://spdx.org/licenses/ISC.html @@ -2924,6 +2944,7 @@ THE SOFTWARE IS PROVIDED "AS IS" AND ISC DISCLAIMS ALL WARRANTIES WITH REGARD TO ``` ## aws-lc-sys - 0.43.0 + **Repository URL**: https://github.com/aws/aws-lc-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -3004,6 +3025,7 @@ limitations under the License. ``` ## aws-lc-sys - 0.43.0 + **Repository URL**: https://github.com/aws/aws-lc-rs **License Type(s)**: BSD-3-Clause ### License: https://spdx.org/licenses/BSD-3-Clause.html @@ -3022,6 +3044,7 @@ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ``` ## aws-lc-sys - 0.43.0 + **Repository URL**: https://github.com/aws/aws-lc-rs **License Type(s)**: ISC ### License: https://spdx.org/licenses/ISC.html @@ -3037,6 +3060,7 @@ THE SOFTWARE IS PROVIDED "AS IS" AND ISC DISCLAIMS ALL WARRANTIES WITH REGARD TO ``` ## aws-lc-sys - 0.43.0 + **Repository URL**: https://github.com/aws/aws-lc-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -3066,6 +3090,7 @@ SOFTWARE. ``` ## axum - 0.8.9 + **Repository URL**: https://github.com/tokio-rs/axum **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -3098,6 +3123,7 @@ DEALINGS IN THE SOFTWARE. ``` ## axum-core - 0.5.6 + **Repository URL**: https://github.com/tokio-rs/axum **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -3132,6 +3158,7 @@ DEALINGS IN THE SOFTWARE. ``` ## backon - 1.6.0 + **Repository URL**: https://github.com/Xuanwo/backon **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -3340,6 +3367,7 @@ DEALINGS IN THE SOFTWARE. ``` ## base64 - 0.22.1 + **Repository URL**: https://github.com/marshallpierce/rust-base64 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -3548,6 +3576,7 @@ limitations under the License. ``` ## bit-set - 0.8.0 + **Repository URL**: https://github.com/contain-rs/bit-set **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -3756,6 +3785,7 @@ limitations under the License. ``` ## bit-vec - 0.8.0 + **Repository URL**: https://github.com/contain-rs/bit-vec **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -3964,6 +3994,7 @@ limitations under the License. ``` ## bitflags - 2.11.0 + **Repository URL**: https://github.com/bitflags/bitflags **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -4172,6 +4203,7 @@ limitations under the License. ``` ## block-buffer - 0.10.4 + **Repository URL**: https://github.com/RustCrypto/utils **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -4380,6 +4412,7 @@ limitations under the License. ``` ## block-buffer - 0.12.0 + **Repository URL**: https://github.com/RustCrypto/utils **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -4588,6 +4621,7 @@ limitations under the License. ``` ## borrow-or-share - 0.2.4 + **Repository URL**: https://github.com/yescallop/borrow-or-share **License Type(s)**: MIT-0 ### License: https://spdx.org/licenses/MIT-0.html @@ -4613,6 +4647,7 @@ SOFTWARE. ``` ## bumpalo - 3.20.2 + **Repository URL**: https://github.com/fitzgen/bumpalo **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -4821,6 +4856,7 @@ limitations under the License. ``` ## bytecount - 0.6.9 + **Repository URL**: https://github.com/llogiq/bytecount **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -5029,6 +5065,7 @@ limitations under the License. ``` ## bytemuck - 1.25.1 + **Repository URL**: https://github.com/Lokathor/bytemuck **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -5097,6 +5134,7 @@ limitations under the License. ``` ## bytemuck_derive - 1.11.0 + **Repository URL**: https://github.com/Lokathor/bytemuck **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -5165,6 +5203,7 @@ limitations under the License. ``` ## byteorder - 1.5.0 + **Repository URL**: https://github.com/BurntSushi/byteorder **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -5193,6 +5232,7 @@ THE SOFTWARE. ``` ## bytes - 1.11.1 + **Repository URL**: https://github.com/tokio-rs/bytes **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -5225,6 +5265,7 @@ DEALINGS IN THE SOFTWARE. ``` ## cbindgen - 0.29.2 + **Repository URL**: https://github.com/mozilla/cbindgen **License Type(s)**: MPL-2.0 ### License: https://spdx.org/licenses/MPL-2.0.html @@ -5605,6 +5646,7 @@ Exhibit B - "Incompatible With Secondary Licenses" Notice ``` ## cc - 1.2.60 + **Repository URL**: https://github.com/rust-lang/cc-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -5813,6 +5855,7 @@ limitations under the License. ``` ## cfg-if - 1.0.4 + **Repository URL**: https://github.com/rust-lang/cfg-if **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -6021,6 +6064,7 @@ limitations under the License. ``` ## cfg_aliases - 0.2.1 + **Repository URL**: https://github.com/katharostech/cfg_aliases **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -6038,6 +6082,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI ``` ## chacha20 - 0.10.2 + **Repository URL**: https://github.com/RustCrypto/stream-ciphers **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -6246,6 +6291,7 @@ limitations under the License. ``` ## chrono - 0.4.44 + **Repository URL**: https://github.com/chronotope/chrono **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -6493,6 +6539,7 @@ limitations under the License. ``` ## clap - 4.6.0 + **Repository URL**: https://github.com/clap-rs/clap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -6702,6 +6749,7 @@ limitations under the License. ``` ## clap_builder - 4.6.0 + **Repository URL**: https://github.com/clap-rs/clap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -6911,6 +6959,7 @@ limitations under the License. ``` ## clap_complete - 4.6.5 + **Repository URL**: https://github.com/clap-rs/clap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -7120,6 +7169,7 @@ limitations under the License. ``` ## clap_derive - 4.6.0 + **Repository URL**: https://github.com/clap-rs/clap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -7329,6 +7379,7 @@ limitations under the License. ``` ## clap_lex - 1.1.0 + **Repository URL**: https://github.com/clap-rs/clap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -7538,6 +7589,7 @@ limitations under the License. ``` ## cmake - 0.1.58 + **Repository URL**: https://github.com/rust-lang/cmake-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -7746,6 +7798,7 @@ limitations under the License. ``` ## colorchoice - 1.0.5 + **Repository URL**: https://github.com/rust-cli/anstyle.git **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -7955,6 +8008,7 @@ limitations under the License. ``` ## combine - 4.6.7 + **Repository URL**: https://github.com/Marwes/combine **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -7984,6 +8038,7 @@ THE SOFTWARE. ``` ## console - 0.15.11 + **Repository URL**: https://github.com/console-rs/console **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -8013,6 +8068,7 @@ SOFTWARE. ``` ## console - 0.16.3 + **Repository URL**: https://github.com/console-rs/console **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -8042,6 +8098,7 @@ SOFTWARE. ``` ## const-oid - 0.10.2 + **Repository URL**: https://github.com/RustCrypto/formats **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -8250,6 +8307,7 @@ limitations under the License. ``` ## convert_case - 0.6.0 + **Repository URL**: https://github.com/rutrum/convert-case **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -8278,6 +8336,7 @@ SOFTWARE. ``` ## core-foundation - 0.10.1 + **Repository URL**: https://github.com/servo/core-foundation-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -8486,6 +8545,7 @@ limitations under the License. ``` ## core-foundation-sys - 0.8.7 + **Repository URL**: https://github.com/servo/core-foundation-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -8694,6 +8754,7 @@ limitations under the License. ``` ## cpufeatures - 0.2.17 + **Repository URL**: https://github.com/RustCrypto/utils **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -8902,6 +8963,7 @@ limitations under the License. ``` ## cpufeatures - 0.3.0 + **Repository URL**: https://github.com/RustCrypto/utils **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -9110,6 +9172,7 @@ limitations under the License. ``` ## crc-fast - 1.10.0 + **Repository URL**: https://github.com/awesomized/crc-fast-rust **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -9318,6 +9381,7 @@ limitations under the License. ``` ## crossbeam - 0.8.4 + **Repository URL**: https://github.com/crossbeam-rs/crossbeam **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -9526,6 +9590,7 @@ limitations under the License. ``` ## crossbeam-channel - 0.5.16 + **Repository URL**: https://github.com/crossbeam-rs/crossbeam **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -9734,6 +9799,7 @@ limitations under the License. ``` ## crossbeam-deque - 0.8.7 + **Repository URL**: https://github.com/crossbeam-rs/crossbeam **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -9942,6 +10008,7 @@ limitations under the License. ``` ## crossbeam-epoch - 0.9.20 + **Repository URL**: https://github.com/crossbeam-rs/crossbeam **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -10150,6 +10217,7 @@ limitations under the License. ``` ## crossbeam-queue - 0.3.13 + **Repository URL**: https://github.com/crossbeam-rs/crossbeam **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -10358,6 +10426,7 @@ limitations under the License. ``` ## crossbeam-utils - 0.8.21 + **Repository URL**: https://github.com/crossbeam-rs/crossbeam **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -10566,6 +10635,7 @@ limitations under the License. ``` ## crypto-common - 0.1.7 + **Repository URL**: https://github.com/RustCrypto/traits **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -10774,6 +10844,7 @@ limitations under the License. ``` ## crypto-common - 0.2.1 + **Repository URL**: https://github.com/RustCrypto/traits **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -10982,6 +11053,7 @@ limitations under the License. ``` ## ctor - 0.2.9 + **Repository URL**: https://github.com/mmastrac/rust-ctor **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -11190,6 +11262,7 @@ Apache License ``` ## data-encoding - 2.11.0 + **Repository URL**: https://github.com/ia0/data-encoding **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -11219,6 +11292,7 @@ SOFTWARE. ``` ## deranged - 0.5.8 + **Repository URL**: https://github.com/jhpratt/deranged **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -11428,6 +11502,7 @@ SOFTWARE. ``` ## dialoguer - 0.11.0 + **Repository URL**: https://github.com/console-rs/dialoguer **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -11457,6 +11532,7 @@ SOFTWARE. ``` ## digest - 0.10.7 + **Repository URL**: https://github.com/RustCrypto/traits **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -11665,6 +11741,7 @@ limitations under the License. ``` ## digest - 0.11.2 + **Repository URL**: https://github.com/RustCrypto/traits **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -11873,6 +11950,7 @@ limitations under the License. ``` ## displaydoc - 0.2.5 + **Repository URL**: https://github.com/yaahc/displaydoc **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -12081,6 +12159,7 @@ limitations under the License. ``` ## dunce - 1.0.5 + **Repository URL**: https://gitlab.com/kornelski/dunce **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -12161,6 +12240,7 @@ limitations under the License. ``` ## dyn-clone - 1.0.20 + **Repository URL**: https://github.com/dtolnay/dyn-clone **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -12241,6 +12321,7 @@ limitations under the License. ``` ## either - 1.15.0 + **Repository URL**: https://github.com/rayon-rs/either **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -12449,6 +12530,7 @@ limitations under the License. ``` ## email_address - 0.2.9 + **Repository URL**: https://github.com/johnstonskj/rust-email_address.git **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -12477,6 +12559,7 @@ SOFTWARE. ``` ## encode_unicode - 1.0.0 + **Repository URL**: https://github.com/tormol/encode_unicode **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -12686,6 +12769,7 @@ SOFTWARE. ``` ## encoding_rs - 0.8.35 + **Repository URL**: https://github.com/hsivonen/encoding_rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -12895,6 +12979,7 @@ SOFTWARE. ``` ## encoding_rs - 0.8.35 + **Repository URL**: https://github.com/hsivonen/encoding_rs **License Type(s)**: BSD-3-Clause ### License: https://spdx.org/licenses/BSD-3-Clause.html @@ -12928,6 +13013,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ``` ## env_filter - 0.1.4 + **Repository URL**: https://github.com/rust-cli/env_logger **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -13137,6 +13223,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ``` ## equivalent - 1.0.2 + **Repository URL**: https://github.com/indexmap-rs/equivalent **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -13345,6 +13432,7 @@ limitations under the License. ``` ## erased-serde - 0.4.10 + **Repository URL**: https://github.com/dtolnay/erased-serde **License Type(s)**: MIT OR Apache-2.0 ### License: https://spdx.org/licenses/ @@ -13556,6 +13644,7 @@ DEALINGS IN THE SOFTWARE. ``` ## errno - 0.3.14 + **Repository URL**: https://github.com/lambda-fairy/rust-errno **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -13764,6 +13853,7 @@ limitations under the License. ``` ## event-listener - 5.4.2 + **Repository URL**: https://github.com/smol-rs/event-listener **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -13972,6 +14062,7 @@ limitations under the License. ``` ## event-listener-strategy - 0.5.4 + **Repository URL**: https://github.com/smol-rs/event-listener-strategy **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -14180,6 +14271,7 @@ limitations under the License. ``` ## fancy-regex - 0.18.0 + **Repository URL**: https://github.com/fancy-regex/fancy-regex **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -14208,6 +14300,7 @@ THE SOFTWARE. ``` ## fastrand - 2.4.1 + **Repository URL**: https://github.com/smol-rs/fastrand **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -14416,6 +14509,7 @@ limitations under the License. ``` ## find-msvc-tools - 0.1.9 + **Repository URL**: https://github.com/rust-lang/cc-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -14624,6 +14718,7 @@ limitations under the License. ``` ## fixedbitset - 0.5.7 + **Repository URL**: https://github.com/petgraph/fixedbitset **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -14832,6 +14927,7 @@ limitations under the License. ``` ## fluent-uri - 0.4.1 + **Repository URL**: https://github.com/yescallop/fluent-uri-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -14860,6 +14956,7 @@ SOFTWARE. ``` ## fnv - 1.0.7 + **Repository URL**: https://github.com/servo/rust-fnv **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -15068,6 +15165,7 @@ limitations under the License. ``` ## foldhash - 0.1.5 + **Repository URL**: https://github.com/orlp/foldhash **License Type(s)**: Zlib ### License: https://spdx.org/licenses/Zlib.html @@ -15094,6 +15192,7 @@ the following restrictions: ``` ## foldhash - 0.2.0 + **Repository URL**: https://github.com/orlp/foldhash **License Type(s)**: Zlib ### License: https://spdx.org/licenses/Zlib.html @@ -15120,6 +15219,7 @@ the following restrictions: ``` ## form_urlencoded - 1.2.2 + **Repository URL**: https://github.com/servo/rust-url **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -15328,6 +15428,7 @@ limitations under the License. ``` ## fraction - 0.15.4 + **Repository URL**: https://github.com/dnsl48/fraction.git **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -15536,6 +15637,7 @@ limitations under the License. ``` ## fs2 - 0.4.3 + **Repository URL**: https://github.com/danburkert/fs2-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -15744,6 +15846,7 @@ limitations under the License. ``` ## fs_extra - 1.3.0 + **Repository URL**: https://github.com/webdesus/fs_extra **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -15772,6 +15875,7 @@ SOFTWARE. ``` ## futures - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -15981,6 +16085,7 @@ limitations under the License. ``` ## futures-channel - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -16190,6 +16295,7 @@ limitations under the License. ``` ## futures-core - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -16399,6 +16505,7 @@ limitations under the License. ``` ## futures-executor - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -16608,6 +16715,7 @@ limitations under the License. ``` ## futures-io - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -16817,6 +16925,7 @@ limitations under the License. ``` ## futures-macro - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -17026,6 +17135,7 @@ limitations under the License. ``` ## futures-sink - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -17235,6 +17345,7 @@ limitations under the License. ``` ## futures-task - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -17444,6 +17555,7 @@ limitations under the License. ``` ## futures-util - 0.3.32 + **Repository URL**: https://github.com/rust-lang/futures-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -17653,6 +17765,7 @@ limitations under the License. ``` ## generic-array - 0.14.7 + **Repository URL**: https://github.com/fizyk20/generic-array.git **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -17681,6 +17794,7 @@ SOFTWARE. ``` ## getrandom - 0.2.17 + **Repository URL**: https://github.com/rust-random/getrandom **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -17889,6 +18003,7 @@ limitations under the License. ``` ## getrandom - 0.3.4 + **Repository URL**: https://github.com/rust-random/getrandom **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -18097,6 +18212,7 @@ limitations under the License. ``` ## getrandom - 0.4.2 + **Repository URL**: https://github.com/rust-random/getrandom **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -18304,7 +18420,8 @@ See the License for the specific language governing permissions and limitations under the License. ``` -## h2 - 0.4.16 +## h2 - 0.4.19 + **Repository URL**: https://github.com/hyperium/h2 **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -18337,6 +18454,7 @@ DEALINGS IN THE SOFTWARE. ``` ## hashbrown - 0.15.5 + **Repository URL**: https://github.com/rust-lang/hashbrown **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -18545,6 +18663,7 @@ limitations under the License. ``` ## hashbrown - 0.16.1 + **Repository URL**: https://github.com/rust-lang/hashbrown **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -18753,6 +18872,7 @@ limitations under the License. ``` ## hashbrown - 0.17.0 + **Repository URL**: https://github.com/rust-lang/hashbrown **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -18961,6 +19081,7 @@ limitations under the License. ``` ## heck - 0.5.0 + **Repository URL**: https://github.com/withoutboats/heck **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -19169,6 +19290,7 @@ limitations under the License. ``` ## http - 1.4.0 + **Repository URL**: https://github.com/hyperium/http **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -19377,6 +19499,7 @@ limitations under the License. ``` ## http-body - 1.0.1 + **Repository URL**: https://github.com/hyperium/http-body **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -19409,6 +19532,7 @@ DEALINGS IN THE SOFTWARE. ``` ## http-body-util - 0.1.3 + **Repository URL**: https://github.com/hyperium/http-body **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -19441,6 +19565,7 @@ DEALINGS IN THE SOFTWARE. ``` ## httparse - 1.10.1 + **Repository URL**: https://github.com/seanmonstar/httparse **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -19649,6 +19774,7 @@ limitations under the License. ``` ## httpdate - 1.0.3 + **Repository URL**: https://github.com/pyfisch/httpdate **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -19857,6 +19983,7 @@ limitations under the License. ``` ## humantime - 2.3.0 + **Repository URL**: https://github.com/chronotope/humantime **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -20066,6 +20193,7 @@ limitations under the License. ``` ## hybrid-array - 0.4.10 + **Repository URL**: https://github.com/RustCrypto/hybrid-array **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -20273,7 +20401,8 @@ See the License for the specific language governing permissions and limitations under the License. ``` -## hyper - 1.9.0 +## hyper - 1.11.1 + **Repository URL**: https://github.com/hyperium/hyper **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -20300,6 +20429,7 @@ THE SOFTWARE. ``` ## hyper-rustls - 0.27.9 + **Repository URL**: https://github.com/rustls/hyper-rustls **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -20508,6 +20638,7 @@ limitations under the License. ``` ## hyper-timeout - 0.5.2 + **Repository URL**: https://github.com/hjr3/hyper-timeout **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -20716,6 +20847,7 @@ limitations under the License. ``` ## hyper-util - 0.1.20 + **Repository URL**: https://github.com/hyperium/hyper-util **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -20742,6 +20874,7 @@ THE SOFTWARE. ``` ## iana-time-zone - 0.1.65 + **Repository URL**: https://github.com/strawlab/iana-time-zone **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -20950,6 +21083,7 @@ limitations under the License. ``` ## iana-time-zone-haiku - 0.1.2 + **Repository URL**: https://github.com/strawlab/iana-time-zone **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -21158,6 +21292,7 @@ limitations under the License. ``` ## icu_collections - 2.2.0 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -21211,6 +21346,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## icu_locale_core - 2.2.0 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -21264,6 +21400,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## icu_normalizer - 2.2.0 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -21317,6 +21454,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## icu_normalizer_data - 2.2.0 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -21370,6 +21508,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## icu_properties - 2.2.0 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -21423,6 +21562,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## icu_properties_data - 2.2.0 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -21476,6 +21616,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## icu_provider - 2.2.0 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -21529,6 +21670,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## id-arena - 2.3.0 + **Repository URL**: https://github.com/fitzgen/id-arena **License Type(s)**: MIT/Apache-2.0 ### License: https://spdx.org/licenses/ @@ -21767,6 +21909,7 @@ DEALINGS IN THE SOFTWARE. ``` ## idna - 1.1.0 + **Repository URL**: https://github.com/servo/rust-url/ **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -21975,6 +22118,7 @@ limitations under the License. ``` ## idna_adapter - 1.2.1 + **Repository URL**: https://github.com/hsivonen/idna_adapter **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -22183,6 +22327,7 @@ limitations under the License. ``` ## indexmap - 2.14.0 + **Repository URL**: https://github.com/indexmap-rs/indexmap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -22391,6 +22536,7 @@ limitations under the License. ``` ## ipnet - 2.12.0 + **Repository URL**: https://github.com/krisprice/ipnet **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -22599,6 +22745,7 @@ limitations under the License. ``` ## iri-string - 0.7.12 + **Repository URL**: https://github.com/lo48576/iri-string **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -22808,6 +22955,7 @@ limitations under the License. ``` ## is_terminal_polyfill - 1.70.2 + **Repository URL**: https://github.com/polyfill-rs/is_terminal_polyfill **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23017,6 +23165,7 @@ limitations under the License. ``` ## itertools - 0.14.0 + **Repository URL**: https://github.com/rust-itertools/itertools **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23225,6 +23374,7 @@ limitations under the License. ``` ## itertools - 0.15.0 + **Repository URL**: https://github.com/rust-itertools/itertools **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23433,6 +23583,7 @@ limitations under the License. ``` ## itoa - 1.0.18 + **Repository URL**: https://github.com/dtolnay/itoa **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23513,6 +23664,7 @@ limitations under the License. ``` ## jni - 0.22.4 + **Repository URL**: https://github.com/jni-rs/jni-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23593,6 +23745,7 @@ limitations under the License. ``` ## jni-macros - 0.22.4 + **Repository URL**: https://github.com/jni-rs/jni-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23673,6 +23826,7 @@ limitations under the License. ``` ## jni-sys - 0.4.1 + **Repository URL**: https://github.com/jni-rs/jni-sys **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23882,6 +24036,7 @@ limitations under the License. ``` ## jni-sys-macros - 0.4.1 + **Repository URL**: https://github.com/jni-rs/jni-sys **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -23962,6 +24117,7 @@ limitations under the License. ``` ## jobserver - 0.1.35 + **Repository URL**: https://github.com/rust-lang/jobserver-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -24170,6 +24326,7 @@ limitations under the License. ``` ## js-sys - 0.3.95 + **Repository URL**: https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/js-sys **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -24378,6 +24535,7 @@ limitations under the License. ``` ## jsonschema - 0.46.8 + **Repository URL**: https://github.com/Stranger6667/jsonschema **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -24406,6 +24564,7 @@ SOFTWARE. ``` ## lazy_static - 1.5.0 + **Repository URL**: https://github.com/rust-lang-nursery/lazy-static.rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -24614,6 +24773,7 @@ limitations under the License. ``` ## leb128fmt - 0.1.0 + **Repository URL**: https://github.com/bluk/leb128fmt **License Type(s)**: MIT OR Apache-2.0 ### License: https://spdx.org/licenses/ @@ -24850,6 +25010,7 @@ DEALINGS IN THE SOFTWARE. ``` ## libc - 0.2.185 + **Repository URL**: https://github.com/rust-lang/libc **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -24930,6 +25091,7 @@ limitations under the License. ``` ## libloading - 0.8.9 + **Repository URL**: https://github.com/nagisa/rust_libloading/ **License Type(s)**: ISC ### License: https://spdx.org/licenses/ISC.html @@ -24949,6 +25111,7 @@ THIS SOFTWARE. ``` ## linux-raw-sys - 0.12.1 + **Repository URL**: https://github.com/sunfishcode/linux-raw-sys **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -25157,6 +25320,7 @@ limitations under the License. ``` ## listeners - 0.4.0 + **Repository URL**: https://github.com/GyulyVGC/listeners **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -25185,6 +25349,7 @@ SOFTWARE. ``` ## litemap - 0.8.2 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -25238,6 +25403,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## lock_api - 0.4.14 + **Repository URL**: https://github.com/Amanieu/parking_lot **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -25446,6 +25612,7 @@ limitations under the License. ``` ## log - 0.4.29 + **Repository URL**: https://github.com/rust-lang/log **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -25654,6 +25821,7 @@ limitations under the License. ``` ## lru-slab - 0.1.2 + **Repository URL**: https://github.com/Ralith/lru-slab **License Type(s)**: MIT OR Apache-2.0 OR Zlib ### License: https://spdx.org/licenses/ @@ -25898,6 +26066,7 @@ the following restrictions: ``` ## matchit - 0.8.4 + **Repository URL**: https://github.com/ibraheemdev/matchit **License Type(s)**: BSD-3-Clause ### License: https://spdx.org/licenses/BSD-3-Clause.html @@ -25934,6 +26103,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ``` ## matchit - 0.8.4 + **Repository URL**: https://github.com/ibraheemdev/matchit **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -25962,6 +26132,7 @@ SOFTWARE. ``` ## md-5 - 0.11.0 + **Repository URL**: https://github.com/RustCrypto/hashes **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -26170,6 +26341,7 @@ limitations under the License. ``` ## memchr - 2.8.0 + **Repository URL**: https://github.com/BurntSushi/memchr **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26198,6 +26370,7 @@ THE SOFTWARE. ``` ## micromap - 0.3.0 + **Repository URL**: https://github.com/yegor256/micromap **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26224,6 +26397,7 @@ SOFTWARE. ``` ## mime - 0.3.17 + **Repository URL**: https://github.com/hyperium/mime **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -26432,6 +26606,7 @@ limitations under the License. ``` ## mio - 1.2.0 + **Repository URL**: https://github.com/tokio-rs/mio **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26458,6 +26633,7 @@ THE SOFTWARE. ``` ## multimap - 0.10.1 + **Repository URL**: https://github.com/havarnov/multimap **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -26666,6 +26842,7 @@ limitations under the License. ``` ## napi - 2.16.17 + **Repository URL**: https://github.com/napi-rs/napi-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26691,6 +26868,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## napi-build - 2.3.1 + **Repository URL**: https://github.com/napi-rs/napi-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26716,6 +26894,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## napi-derive - 2.16.13 + **Repository URL**: https://github.com/napi-rs/napi-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26741,6 +26920,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## napi-derive-backend - 1.0.75 + **Repository URL**: https://github.com/napi-rs/napi-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26766,6 +26946,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## napi-sys - 2.4.0 + **Repository URL**: https://github.com/napi-rs/napi-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26791,6 +26972,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## nom - 8.0.0 + **Repository URL**: https://github.com/rust-bakery/nom **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -26818,6 +27000,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## ntapi - 0.4.3 + **Repository URL**: https://github.com/MSxDOS/ntapi **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -27027,6 +27210,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## num - 0.4.3 + **Repository URL**: https://github.com/rust-num/num **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -27235,6 +27419,7 @@ limitations under the License. ``` ## num-bigint - 0.4.6 + **Repository URL**: https://github.com/rust-num/num-bigint **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -27443,6 +27628,7 @@ limitations under the License. ``` ## num-cmp - 0.1.0 + **Repository URL**: https://github.com/lifthrasiir/num-cmp **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -27523,6 +27709,7 @@ limitations under the License. ``` ## num-complex - 0.4.6 + **Repository URL**: https://github.com/rust-num/num-complex **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -27731,6 +27918,7 @@ limitations under the License. ``` ## num-conv - 0.2.2 + **Repository URL**: https://github.com/jhpratt/num-conv **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -27811,6 +27999,7 @@ limitations under the License. ``` ## num-integer - 0.1.46 + **Repository URL**: https://github.com/rust-num/num-integer **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -28019,6 +28208,7 @@ limitations under the License. ``` ## num-iter - 0.1.45 + **Repository URL**: https://github.com/rust-num/num-iter **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -28227,6 +28417,7 @@ limitations under the License. ``` ## num-rational - 0.4.2 + **Repository URL**: https://github.com/rust-num/num-rational **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -28435,6 +28626,7 @@ limitations under the License. ``` ## num-traits - 0.2.19 + **Repository URL**: https://github.com/rust-num/num-traits **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -28643,6 +28835,7 @@ limitations under the License. ``` ## objc2-core-foundation - 0.3.1 + **Repository URL**: https://github.com/madsmtm/objc2 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -28723,6 +28916,7 @@ limitations under the License. ``` ## objc2-io-kit - 0.3.1 + **Repository URL**: https://github.com/madsmtm/objc2 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -28803,6 +28997,7 @@ limitations under the License. ``` ## object_store - 0.14.1 + **Repository URL**: https://github.com/apache/arrow-rs-object-store **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -29014,6 +29209,7 @@ limitations under the License. ``` ## once_cell - 1.21.4 + **Repository URL**: https://github.com/matklad/once_cell **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -29222,6 +29418,7 @@ limitations under the License. ``` ## once_cell_polyfill - 1.70.2 + **Repository URL**: https://github.com/polyfill-rs/once_cell_polyfill **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -29431,6 +29628,7 @@ limitations under the License. ``` ## openssl-probe - 0.2.1 + **Repository URL**: https://github.com/rustls/openssl-probe **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -29639,6 +29837,7 @@ limitations under the License. ``` ## opentelemetry - 0.32.0 + **Repository URL**: https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -29719,6 +29918,7 @@ limitations under the License. ``` ## opentelemetry-http - 0.32.0 + **Repository URL**: https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-http **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -29799,6 +29999,7 @@ limitations under the License. ``` ## opentelemetry-otlp - 0.32.0 + **Repository URL**: https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-otlp **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -29879,6 +30080,7 @@ limitations under the License. ``` ## opentelemetry-proto - 0.32.0 + **Repository URL**: https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-proto **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -30087,6 +30289,7 @@ limitations under the License. ``` ## opentelemetry-semantic-conventions - 0.32.1 + **Repository URL**: https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-semantic-conventions **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -30167,6 +30370,7 @@ limitations under the License. ``` ## opentelemetry_sdk - 0.32.1 + **Repository URL**: https://github.com/open-telemetry/opentelemetry-rust/tree/main/opentelemetry-sdk **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -30247,6 +30451,7 @@ limitations under the License. ``` ## ordered-float - 2.10.1 + **Repository URL**: https://github.com/reem/rust-ordered-float **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -30279,6 +30484,7 @@ DEALINGS IN THE SOFTWARE. ``` ## outref - 0.5.2 + **Repository URL**: https://github.com/Nugine/outref **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -30307,6 +30513,7 @@ SOFTWARE. ``` ## parking - 2.2.1 + **Repository URL**: https://github.com/smol-rs/parking **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -30515,6 +30722,7 @@ limitations under the License. ``` ## parking_lot - 0.12.5 + **Repository URL**: https://github.com/Amanieu/parking_lot **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -30723,6 +30931,7 @@ limitations under the License. ``` ## parking_lot_core - 0.9.12 + **Repository URL**: https://github.com/Amanieu/parking_lot **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -30931,6 +31140,7 @@ limitations under the License. ``` ## pem - 3.0.6 + **Repository URL**: https://github.com/jcreekmore/pem-rs.git **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -30959,6 +31169,7 @@ SOFTWARE. ``` ## percent-encoding - 2.3.2 + **Repository URL**: https://github.com/servo/rust-url/ **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -31167,6 +31378,7 @@ limitations under the License. ``` ## petgraph - 0.8.3 + **Repository URL**: https://github.com/petgraph/petgraph **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -31375,6 +31587,7 @@ limitations under the License. ``` ## pin-project - 1.1.11 + **Repository URL**: https://github.com/taiki-e/pin-project **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -31455,6 +31668,7 @@ limitations under the License. ``` ## pin-project-internal - 1.1.11 + **Repository URL**: https://github.com/taiki-e/pin-project **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -31535,6 +31749,7 @@ limitations under the License. ``` ## pin-project-lite - 0.2.17 + **Repository URL**: https://github.com/taiki-e/pin-project-lite **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -31615,6 +31830,7 @@ limitations under the License. ``` ## pkg-config - 0.3.33 + **Repository URL**: https://github.com/rust-lang/pkg-config-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -31823,6 +32039,7 @@ limitations under the License. ``` ## portable-atomic - 1.13.1 + **Repository URL**: https://github.com/taiki-e/portable-atomic **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -31903,6 +32120,7 @@ limitations under the License. ``` ## potential_utf - 0.1.5 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -31956,6 +32174,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## powerfmt - 0.2.0 + **Repository URL**: https://github.com/jhpratt/powerfmt **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -32165,6 +32384,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## ppv-lite86 - 0.2.21 + **Repository URL**: https://github.com/cryptocorrosion/cryptocorrosion **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -32373,6 +32593,7 @@ limitations under the License. ``` ## prettyplease - 0.2.37 + **Repository URL**: https://github.com/dtolnay/prettyplease **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -32453,6 +32674,7 @@ limitations under the License. ``` ## proc-macro2 - 1.0.106 + **Repository URL**: https://github.com/dtolnay/proc-macro2 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -32533,6 +32755,7 @@ limitations under the License. ``` ## prost - 0.14.3 + **Repository URL**: https://github.com/tokio-rs/prost **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -32741,6 +32964,7 @@ limitations under the License. ``` ## prost-build - 0.14.3 + **Repository URL**: https://github.com/tokio-rs/prost **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -32949,6 +33173,7 @@ limitations under the License. ``` ## prost-derive - 0.14.3 + **Repository URL**: https://github.com/tokio-rs/prost **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -33157,6 +33382,7 @@ limitations under the License. ``` ## prost-types - 0.14.3 + **Repository URL**: https://github.com/tokio-rs/prost **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -33365,6 +33591,7 @@ limitations under the License. ``` ## protoc-bin-vendored - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33390,6 +33617,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-linux-aarch_64 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33415,6 +33643,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-linux-ppcle_64 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33440,6 +33669,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-linux-s390_64 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33465,6 +33695,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-linux-x86_32 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33490,6 +33721,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-linux-x86_64 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33515,6 +33747,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-macos-aarch_64 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33540,6 +33773,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-macos-x86_64 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33565,6 +33799,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## protoc-bin-vendored-win32 - 3.2.0 + **Repository URL**: https://github.com/stepancheg/rust-protoc-bin-vendored/ **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33590,6 +33825,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## pulldown-cmark - 0.13.4 + **Repository URL**: https://github.com/raphlinus/pulldown-cmark **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -33618,6 +33854,7 @@ THE SOFTWARE. ``` ## pulldown-cmark-to-cmark - 22.0.0 + **Repository URL**: https://github.com/Byron/pulldown-cmark-to-cmark **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -33826,6 +34063,7 @@ THE SOFTWARE. ``` ## pyo3 - 0.29.0 + **Repository URL**: https://github.com/pyo3/pyo3 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -33906,6 +34144,7 @@ limitations under the License. ``` ## pyo3-async-runtimes - 0.29.0 + **Repository URL**: https://github.com/PyO3/pyo3-async-runtimes **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -34102,6 +34341,7 @@ limitations under the License. ``` ## pyo3-build-config - 0.29.0 + **Repository URL**: https://github.com/pyo3/pyo3 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -34182,6 +34422,7 @@ limitations under the License. ``` ## pyo3-ffi - 0.29.0 + **Repository URL**: https://github.com/pyo3/pyo3 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -34262,6 +34503,7 @@ limitations under the License. ``` ## pyo3-macros - 0.29.0 + **Repository URL**: https://github.com/pyo3/pyo3 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -34342,6 +34584,7 @@ limitations under the License. ``` ## pyo3-macros-backend - 0.29.0 + **Repository URL**: https://github.com/pyo3/pyo3 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -34422,6 +34665,7 @@ limitations under the License. ``` ## pythonize - 0.29.0 + **Repository URL**: https://github.com/davidhewitt/pythonize **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -34448,6 +34692,7 @@ SOFTWARE. ``` ## quick-xml - 0.41.0 + **Repository URL**: https://github.com/tafia/quick-xml **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -34478,6 +34723,7 @@ THE SOFTWARE. ``` ## quinn - 0.11.11 + **Repository URL**: https://github.com/quinn-rs/quinn **License Type(s)**: MIT OR Apache-2.0 ### License: https://spdx.org/licenses/ @@ -34698,6 +34944,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI ``` ## quinn-proto - 0.11.15 + **Repository URL**: https://github.com/quinn-rs/quinn **License Type(s)**: MIT OR Apache-2.0 ### License: https://spdx.org/licenses/ @@ -34918,6 +35165,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI ``` ## quinn-udp - 0.5.14 + **Repository URL**: https://github.com/quinn-rs/quinn **License Type(s)**: MIT OR Apache-2.0 ### License: https://spdx.org/licenses/ @@ -35138,6 +35386,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI ``` ## quote - 1.0.45 + **Repository URL**: https://github.com/dtolnay/quote **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -35218,6 +35467,7 @@ limitations under the License. ``` ## r-efi - 5.3.0 + **Repository URL**: https://github.com/r-efi/r-efi **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -35298,6 +35548,7 @@ limitations under the License. ``` ## r-efi - 6.0.0 + **Repository URL**: https://github.com/r-efi/r-efi **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -35378,6 +35629,7 @@ limitations under the License. ``` ## rand - 0.10.1 + **Repository URL**: https://github.com/rust-random/rand **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -35458,6 +35710,7 @@ limitations under the License. ``` ## rand - 0.9.3 + **Repository URL**: https://github.com/rust-random/rand **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -35538,6 +35791,7 @@ limitations under the License. ``` ## rand_chacha - 0.9.0 + **Repository URL**: https://github.com/rust-random/rand **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -35618,6 +35872,7 @@ limitations under the License. ``` ## rand_core - 0.10.1 + **Repository URL**: https://github.com/rust-random/rand_core **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -35812,6 +36067,7 @@ APPENDIX: How to apply the Apache License to your work. ``` ## rand_core - 0.9.5 + **Repository URL**: https://github.com/rust-random/rand **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -36006,6 +36262,7 @@ APPENDIX: How to apply the Apache License to your work. ``` ## rcgen - 0.13.2 + **Repository URL**: https://github.com/rustls/rcgen **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -36086,6 +36343,7 @@ limitations under the License. ``` ## redis - 1.2.0 + **Repository URL**: https://github.com/redis-rs/redis-rs **License Type(s)**: BSD-3-Clause ### License: https://spdx.org/licenses/BSD-3-Clause.html @@ -36126,6 +36384,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ``` ## redox_syscall - 0.5.18 + **Repository URL**: https://gitlab.redox-os.org/redox-os/syscall **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -36155,6 +36414,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## ref-cast - 1.0.25 + **Repository URL**: https://github.com/dtolnay/ref-cast **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -36235,6 +36495,7 @@ limitations under the License. ``` ## ref-cast-impl - 1.0.25 + **Repository URL**: https://github.com/dtolnay/ref-cast **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -36315,6 +36576,7 @@ limitations under the License. ``` ## referencing - 0.46.8 + **Repository URL**: https://github.com/Stranger6667/jsonschema **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -36343,6 +36605,7 @@ SOFTWARE. ``` ## regex - 1.12.3 + **Repository URL**: https://github.com/rust-lang/regex **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -36551,6 +36814,7 @@ limitations under the License. ``` ## regex-automata - 0.4.14 + **Repository URL**: https://github.com/rust-lang/regex **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -36759,6 +37023,7 @@ limitations under the License. ``` ## regex-syntax - 0.8.10 + **Repository URL**: https://github.com/rust-lang/regex **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -36967,6 +37232,7 @@ limitations under the License. ``` ## reqwest - 0.12.28 + **Repository URL**: https://github.com/seanmonstar/reqwest **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -37175,6 +37441,7 @@ limitations under the License. ``` ## reqwest - 0.13.4 + **Repository URL**: https://github.com/seanmonstar/reqwest **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -37383,6 +37650,7 @@ limitations under the License. ``` ## ring - 0.17.14 + **Repository URL**: https://github.com/briansmith/ring **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -37591,6 +37859,7 @@ limitations under the License. ``` ## ring - 0.17.14 + **Repository URL**: https://github.com/briansmith/ring **License Type(s)**: ISC ### License: https://spdx.org/licenses/ISC.html @@ -37611,6 +37880,7 @@ CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. ``` ## rustc-hash - 2.1.2 + **Repository URL**: https://github.com/rust-lang/rustc-hash **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -37822,6 +38092,7 @@ DEALINGS IN THE SOFTWARE. ``` ## rustc_version - 0.4.1 + **Repository URL**: https://github.com/djc/rustc-version-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -38030,6 +38301,7 @@ limitations under the License. ``` ## rustix - 1.1.4 + **Repository URL**: https://github.com/bytecodealliance/rustix **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -38238,6 +38510,7 @@ limitations under the License. ``` ## rustls - 0.23.40 + **Repository URL**: https://github.com/rustls/rustls **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -38446,6 +38719,7 @@ limitations under the License. ``` ## rustls-native-certs - 0.8.3 + **Repository URL**: https://github.com/rustls/rustls-native-certs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -38654,6 +38928,7 @@ limitations under the License. ``` ## rustls-pki-types - 1.14.1 + **Repository URL**: https://github.com/rustls/pki-types **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -38862,6 +39137,7 @@ limitations under the License. ``` ## rustls-platform-verifier - 0.7.0 + **Repository URL**: https://github.com/rustls/rustls-platform-verifier **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -39070,6 +39346,7 @@ limitations under the License. ``` ## rustls-platform-verifier-android - 0.1.1 + **Repository URL**: https://github.com/rustls/rustls-platform-verifier **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -39150,6 +39427,7 @@ limitations under the License. ``` ## rustls-webpki - 0.103.13 + **Repository URL**: https://github.com/rustls/webpki **License Type(s)**: ISC ### License: https://spdx.org/licenses/ISC.html @@ -39176,6 +39454,7 @@ third-party/chromium/LICENSE. ``` ## rustversion - 1.0.22 + **Repository URL**: https://github.com/dtolnay/rustversion **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -39256,6 +39535,7 @@ limitations under the License. ``` ## ryu - 1.0.23 + **Repository URL**: https://github.com/dtolnay/ryu **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -39336,6 +39616,7 @@ limitations under the License. ``` ## ryu-js - 1.0.2 + **Repository URL**: https://github.com/boa-dev/ryu-js **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -39416,6 +39697,7 @@ limitations under the License. ``` ## same-file - 1.0.6 + **Repository URL**: https://github.com/BurntSushi/same-file **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -39444,6 +39726,7 @@ THE SOFTWARE. ``` ## schannel - 0.1.29 + **Repository URL**: https://github.com/steffengy/schannel-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -39458,6 +39741,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI ``` ## schemars - 0.8.22 + **Repository URL**: https://github.com/GREsau/schemars **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -39486,6 +39770,7 @@ SOFTWARE. ``` ## schemars_derive - 0.8.22 + **Repository URL**: https://github.com/GREsau/schemars **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -39514,6 +39799,7 @@ SOFTWARE. ``` ## scopeguard - 1.2.0 + **Repository URL**: https://github.com/bluss/scopeguard **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -39722,6 +40008,7 @@ limitations under the License. ``` ## security-framework - 3.7.0 + **Repository URL**: https://github.com/kornelski/rust-security-framework **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -39930,6 +40217,7 @@ limitations under the License. ``` ## security-framework-sys - 2.17.0 + **Repository URL**: https://github.com/kornelski/rust-security-framework **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -40138,6 +40426,7 @@ limitations under the License. ``` ## semver - 1.0.28 + **Repository URL**: https://github.com/dtolnay/semver **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -40218,6 +40507,7 @@ limitations under the License. ``` ## serde - 1.0.228 + **Repository URL**: https://github.com/serde-rs/serde **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -40298,6 +40588,7 @@ limitations under the License. ``` ## serde_buf - 0.1.2 + **Repository URL**: https://github.com/KodrAus/serde_buf.git **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -40532,6 +40823,7 @@ SOFTWARE. ``` ## serde_core - 1.0.228 + **Repository URL**: https://github.com/serde-rs/serde **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -40612,6 +40904,7 @@ limitations under the License. ``` ## serde_derive - 1.0.228 + **Repository URL**: https://github.com/serde-rs/serde **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -40692,6 +40985,7 @@ limitations under the License. ``` ## serde_derive_internals - 0.29.1 + **Repository URL**: https://github.com/serde-rs/serde **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -40772,6 +41066,7 @@ limitations under the License. ``` ## serde_fmt - 1.1.0 + **Repository URL**: https://github.com/KodrAus/serde_fmt.git **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -41006,6 +41301,7 @@ SOFTWARE. ``` ## serde_json - 1.0.149 + **Repository URL**: https://github.com/serde-rs/json **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -41086,6 +41382,7 @@ limitations under the License. ``` ## serde_json_canonicalizer - 0.3.2 + **Repository URL**: https://github.com/evik42/serde-json-canonicalizer **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -41114,6 +41411,7 @@ SOFTWARE. ``` ## serde_path_to_error - 0.1.20 + **Repository URL**: https://github.com/dtolnay/path-to-error **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -41194,6 +41492,7 @@ limitations under the License. ``` ## serde_spanned - 1.1.1 + **Repository URL**: https://github.com/toml-rs/toml **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -41403,6 +41702,7 @@ limitations under the License. ``` ## serde_urlencoded - 0.7.1 + **Repository URL**: https://github.com/nox/serde_urlencoded **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -41483,6 +41783,7 @@ limitations under the License. ``` ## serde_yaml - 0.9.34+deprecated + **Repository URL**: https://github.com/dtolnay/serde-yaml **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -41563,6 +41864,7 @@ limitations under the License. ``` ## sha1 - 0.10.6 + **Repository URL**: https://github.com/RustCrypto/hashes **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -41771,6 +42073,7 @@ limitations under the License. ``` ## sha1_smol - 1.0.1 + **Repository URL**: https://github.com/mitsuhiko/sha1-smol **License Type(s)**: BSD-3-Clause ### License: https://spdx.org/licenses/BSD-3-Clause.html @@ -41789,6 +42092,7 @@ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ``` ## sha2 - 0.11.0 + **Repository URL**: https://github.com/RustCrypto/hashes **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -41997,6 +42301,7 @@ limitations under the License. ``` ## shell-words - 1.1.1 + **Repository URL**: https://github.com/tmiasko/shell-words **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -42205,6 +42510,7 @@ limitations under the License. ``` ## shlex - 1.3.0 + **Repository URL**: https://github.com/comex/rust-shlex **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -42285,6 +42591,7 @@ limitations under the License. ``` ## signal-hook-registry - 1.4.8 + **Repository URL**: https://github.com/vorner/signal-hook **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -42493,6 +42800,7 @@ limitations under the License. ``` ## simd_cesu8 - 1.2.0 + **Repository URL**: https://github.com/seancroach/simd_cesu8 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -42701,6 +43009,7 @@ limitations under the License. ``` ## simdutf8 - 0.1.5 + **Repository URL**: https://github.com/rusticstuff/simdutf8 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -42781,6 +43090,7 @@ limitations under the License. ``` ## slab - 0.4.12 + **Repository URL**: https://github.com/tokio-rs/slab **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -42813,6 +43123,7 @@ DEALINGS IN THE SOFTWARE. ``` ## smallvec - 1.15.1 + **Repository URL**: https://github.com/servo/rust-smallvec **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -43021,6 +43332,7 @@ limitations under the License. ``` ## socket2 - 0.6.3 + **Repository URL**: https://github.com/rust-lang/socket2 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -43229,6 +43541,7 @@ limitations under the License. ``` ## spdlog-internal - 0.2.1 + **Repository URL**: https://github.com/SpriteOvO/spdlog-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -43309,6 +43622,7 @@ limitations under the License. ``` ## spdlog-macros - 0.3.1 + **Repository URL**: https://github.com/SpriteOvO/spdlog-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -43389,6 +43703,7 @@ limitations under the License. ``` ## spdlog-rs - 0.5.3 + **Repository URL**: https://github.com/SpriteOvO/spdlog-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -43597,6 +43912,7 @@ Apache License ``` ## spin - 0.10.1 + **Repository URL**: https://github.com/mvdnes/spin-rs.git **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -43625,6 +43941,7 @@ SOFTWARE. ``` ## stable_deref_trait - 1.2.1 + **Repository URL**: https://github.com/storyyeller/stable_deref_trait **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -43833,6 +44150,7 @@ limitations under the License. ``` ## strsim - 0.11.1 + **Repository URL**: https://github.com/rapidfuzz/strsim-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -43863,6 +44181,7 @@ SOFTWARE. ``` ## strum - 0.27.2 + **Repository URL**: https://github.com/Peternator7/strum **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -43891,6 +44210,7 @@ SOFTWARE. ``` ## strum_macros - 0.27.2 + **Repository URL**: https://github.com/Peternator7/strum **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -43919,6 +44239,7 @@ SOFTWARE. ``` ## subtle - 2.6.1 + **Repository URL**: https://github.com/dalek-cryptography/subtle **License Type(s)**: BSD-3-Clause ### License: https://spdx.org/licenses/BSD-3-Clause.html @@ -43955,6 +44276,7 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ``` ## sval - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -44189,6 +44511,7 @@ SOFTWARE. ``` ## sval_buffer - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -44423,6 +44746,7 @@ SOFTWARE. ``` ## sval_dynamic - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -44657,6 +44981,7 @@ SOFTWARE. ``` ## sval_fmt - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -44891,6 +45216,7 @@ SOFTWARE. ``` ## sval_json - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -45125,6 +45451,7 @@ SOFTWARE. ``` ## sval_nested - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -45139,6 +45466,7 @@ SOFTWARE. ``` ## sval_ref - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -45373,6 +45701,7 @@ SOFTWARE. ``` ## sval_serde - 2.20.0 + **Repository URL**: https://github.com/sval-rs/sval **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -45607,6 +45936,7 @@ SOFTWARE. ``` ## syn - 2.0.117 + **Repository URL**: https://github.com/dtolnay/syn **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -45687,6 +46017,7 @@ limitations under the License. ``` ## sync_wrapper - 1.0.2 + **Repository URL**: https://github.com/Actyx/sync_wrapper **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -45767,6 +46098,7 @@ limitations under the License. ``` ## synstructure - 0.13.2 + **Repository URL**: https://github.com/mystor/synstructure **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -45781,6 +46113,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI ``` ## sysinfo - 0.38.1 + **Repository URL**: https://github.com/GuillaumeGomez/sysinfo **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -45810,6 +46143,7 @@ SOFTWARE. ``` ## target-lexicon - 0.13.5 + **Repository URL**: https://github.com/bytecodealliance/target-lexicon **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46037,6 +46371,7 @@ Software. ``` ## tdigest - 0.2.3 + **Repository URL**: https://github.com/MnO2/t-digest **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46117,6 +46452,7 @@ limitations under the License. ``` ## tempfile - 3.27.0 + **Repository URL**: https://github.com/Stebalien/tempfile **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46325,6 +46661,7 @@ limitations under the License. ``` ## thiserror - 1.0.69 + **Repository URL**: https://github.com/dtolnay/thiserror **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46405,6 +46742,7 @@ limitations under the License. ``` ## thiserror - 2.0.18 + **Repository URL**: https://github.com/dtolnay/thiserror **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46485,6 +46823,7 @@ limitations under the License. ``` ## thiserror-impl - 1.0.69 + **Repository URL**: https://github.com/dtolnay/thiserror **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46565,6 +46904,7 @@ limitations under the License. ``` ## thiserror-impl - 2.0.18 + **Repository URL**: https://github.com/dtolnay/thiserror **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46645,6 +46985,7 @@ limitations under the License. ``` ## time - 0.3.53 + **Repository URL**: https://github.com/time-rs/time **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46725,6 +47066,7 @@ limitations under the License. ``` ## time-core - 0.1.9 + **Repository URL**: https://github.com/time-rs/time **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -46805,6 +47147,7 @@ limitations under the License. ``` ## tinystr - 0.8.3 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -46858,6 +47201,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## tinyvec - 1.11.0 + **Repository URL**: https://github.com/Lokathor/tinyvec **License Type(s)**: Zlib OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -47092,6 +47436,7 @@ Permission is granted to anyone to use this software for any purpose, including ``` ## tinyvec_macros - 0.1.1 + **Repository URL**: https://github.com/Soveu/tinyvec_macros **License Type(s)**: MIT OR Apache-2.0 OR Zlib ### License: https://spdx.org/licenses/ @@ -47351,6 +47696,7 @@ freely, subject to the following restrictions: ``` ## tokio - 1.51.1 + **Repository URL**: https://github.com/tokio-rs/tokio **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -47379,6 +47725,7 @@ SOFTWARE. ``` ## tokio-macros - 2.7.0 + **Repository URL**: https://github.com/tokio-rs/tokio **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -47408,6 +47755,7 @@ SOFTWARE. ``` ## tokio-rustls - 0.26.4 + **Repository URL**: https://github.com/rustls/tokio-rustls **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -47616,6 +47964,7 @@ limitations under the License. ``` ## tokio-stream - 0.1.18 + **Repository URL**: https://github.com/tokio-rs/tokio **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -47644,6 +47993,7 @@ SOFTWARE. ``` ## tokio-tungstenite - 0.27.0 + **Repository URL**: https://github.com/snapview/tokio-tungstenite **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -47671,6 +48021,7 @@ THE SOFTWARE. ``` ## tokio-util - 0.7.18 + **Repository URL**: https://github.com/tokio-rs/tokio **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -47699,6 +48050,7 @@ SOFTWARE. ``` ## toml - 0.9.12+spec-1.1.0 + **Repository URL**: https://github.com/toml-rs/toml **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -47908,6 +48260,7 @@ SOFTWARE. ``` ## toml_datetime - 0.7.5+spec-1.1.0 + **Repository URL**: https://github.com/toml-rs/toml **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -48117,6 +48470,7 @@ SOFTWARE. ``` ## toml_edit - 0.23.10+spec-1.0.0 + **Repository URL**: https://github.com/toml-rs/toml **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -48326,6 +48680,7 @@ SOFTWARE. ``` ## toml_parser - 1.1.2+spec-1.1.0 + **Repository URL**: https://github.com/toml-rs/toml **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -48535,6 +48890,7 @@ SOFTWARE. ``` ## toml_writer - 1.1.1+spec-1.1.0 + **Repository URL**: https://github.com/toml-rs/toml **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -48744,6 +49100,7 @@ SOFTWARE. ``` ## tonic - 0.14.5 + **Repository URL**: https://github.com/hyperium/tonic **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48770,6 +49127,7 @@ THE SOFTWARE. ``` ## tonic-build - 0.14.6 + **Repository URL**: https://github.com/hyperium/tonic **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48796,6 +49154,7 @@ THE SOFTWARE. ``` ## tonic-prost - 0.14.5 + **Repository URL**: https://github.com/hyperium/tonic **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48821,6 +49180,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## tonic-prost-build - 0.14.6 + **Repository URL**: https://github.com/hyperium/tonic **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48846,6 +49206,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## tonic-types - 0.14.5 + **Repository URL**: https://github.com/hyperium/tonic **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48872,6 +49233,7 @@ THE SOFTWARE. ``` ## tower - 0.5.3 + **Repository URL**: https://github.com/tower-rs/tower **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48904,6 +49266,7 @@ DEALINGS IN THE SOFTWARE. ``` ## tower-http - 0.6.8 + **Repository URL**: https://github.com/tower-rs/tower-http **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48936,6 +49299,7 @@ DEALINGS IN THE SOFTWARE. ``` ## tower-layer - 0.3.3 + **Repository URL**: https://github.com/tower-rs/tower **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -48968,6 +49332,7 @@ DEALINGS IN THE SOFTWARE. ``` ## tower-service - 0.3.3 + **Repository URL**: https://github.com/tower-rs/tower **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -49000,6 +49365,7 @@ DEALINGS IN THE SOFTWARE. ``` ## tracing - 0.1.44 + **Repository URL**: https://github.com/tokio-rs/tracing **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -49032,6 +49398,7 @@ DEALINGS IN THE SOFTWARE. ``` ## tracing-attributes - 0.1.31 + **Repository URL**: https://github.com/tokio-rs/tracing **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -49064,6 +49431,7 @@ DEALINGS IN THE SOFTWARE. ``` ## tracing-core - 0.1.36 + **Repository URL**: https://github.com/tokio-rs/tracing **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -49096,6 +49464,7 @@ DEALINGS IN THE SOFTWARE. ``` ## try-lock - 0.2.5 + **Repository URL**: https://github.com/seanmonstar/try-lock **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -49124,6 +49493,7 @@ THE SOFTWARE. ``` ## tungstenite - 0.27.0 + **Repository URL**: https://github.com/snapview/tungstenite-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -49332,6 +49702,7 @@ limitations under the License. ``` ## typed-builder - 0.23.2 + **Repository URL**: https://github.com/idanarye/rust-typed-builder **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -49540,6 +49911,7 @@ limitations under the License. ``` ## typed-builder-macro - 0.23.2 + **Repository URL**: https://github.com/idanarye/rust-typed-builder **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -49748,6 +50120,7 @@ limitations under the License. ``` ## typeid - 1.0.3 + **Repository URL**: https://github.com/dtolnay/typeid **License Type(s)**: MIT OR Apache-2.0 ### License: https://spdx.org/licenses/ @@ -49959,6 +50332,7 @@ DEALINGS IN THE SOFTWARE. ``` ## typenum - 1.19.0 + **Repository URL**: https://github.com/paholg/typenum **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -50167,6 +50541,7 @@ limitations under the License. ``` ## unicase - 2.9.0 + **Repository URL**: https://github.com/seanmonstar/unicase **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -50375,6 +50750,7 @@ limitations under the License. ``` ## unicode-general-category - 1.1.0 + **Repository URL**: https://github.com/yeslogic/unicode-general-category **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -50583,6 +50959,7 @@ limitations under the License. ``` ## unicode-ident - 1.0.24 + **Repository URL**: https://github.com/dtolnay/unicode-ident **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -50663,6 +51040,7 @@ limitations under the License. ``` ## unicode-ident - 1.0.24 + **Repository URL**: https://github.com/dtolnay/unicode-ident **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -50709,6 +51087,7 @@ authorization of the copyright holder. ``` ## unicode-segmentation - 1.13.2 + **Repository URL**: https://github.com/unicode-rs/unicode-segmentation **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -50917,6 +51296,7 @@ limitations under the License. ``` ## unicode-width - 0.2.2 + **Repository URL**: https://github.com/unicode-rs/unicode-width **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -51125,6 +51505,7 @@ limitations under the License. ``` ## unicode-xid - 0.2.6 + **Repository URL**: https://github.com/unicode-rs/unicode-xid **License Type(s)**: MIT OR Apache-2.0 ### License: https://spdx.org/licenses/ @@ -51363,6 +51744,7 @@ DEALINGS IN THE SOFTWARE. ``` ## unsafe-libyaml - 0.2.11 + **Repository URL**: https://github.com/dtolnay/unsafe-libyaml **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -51393,6 +51775,7 @@ DEALINGS IN THE SOFTWARE. ``` ## untrusted - 0.9.0 + **Repository URL**: https://github.com/briansmith/untrusted **License Type(s)**: ISC ### License: https://spdx.org/licenses/ISC.html @@ -51413,6 +51796,7 @@ DEALINGS IN THE SOFTWARE. ``` ## url - 2.5.8 + **Repository URL**: https://github.com/servo/rust-url **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -51621,6 +52005,7 @@ limitations under the License. ``` ## utf-8 - 0.7.6 + **Repository URL**: https://github.com/SimonSapin/rust-utf8 **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -51701,6 +52086,7 @@ limitations under the License. ``` ## utf8_iter - 1.0.4 + **Repository URL**: https://github.com/hsivonen/utf8_iter **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -51910,6 +52296,7 @@ limitations under the License. ``` ## utf8parse - 0.2.2 + **Repository URL**: https://github.com/alacritty/vte **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -51990,6 +52377,7 @@ limitations under the License. ``` ## uuid - 1.18.1 + **Repository URL**: https://github.com/uuid-rs/uuid **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -52198,6 +52586,7 @@ limitations under the License. ``` ## uuid-simd - 0.8.0 + **Repository URL**: https://github.com/Nugine/simd **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -52223,6 +52612,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## value-bag - 1.13.0 + **Repository URL**: https://github.com/sval-rs/value-bag **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -52431,6 +52821,7 @@ limitations under the License. ``` ## value-bag-serde1 - 1.13.0 + **Repository URL**: https://crates.io/crates/value-bag-serde1 **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -52665,6 +53056,7 @@ SOFTWARE. ``` ## value-bag-sval2 - 1.13.0 + **Repository URL**: https://crates.io/crates/value-bag-sval2 **License Type(s)**: Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -52899,6 +53291,7 @@ SOFTWARE. ``` ## version_check - 0.9.5 + **Repository URL**: https://github.com/SergioBenitez/version_check **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -53107,6 +53500,7 @@ limitations under the License. ``` ## vsimd - 0.8.0 + **Repository URL**: https://github.com/Nugine/simd **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -53132,6 +53526,7 @@ USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## walkdir - 2.5.0 + **Repository URL**: https://github.com/BurntSushi/walkdir **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -53160,6 +53555,7 @@ THE SOFTWARE. ``` ## want - 0.3.1 + **Repository URL**: https://github.com/seanmonstar/want **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -53187,6 +53583,7 @@ THE SOFTWARE. ``` ## wasi - 0.11.1+wasi-snapshot-preview1 + **Repository URL**: https://github.com/bytecodealliance/wasi **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -53395,6 +53792,7 @@ limitations under the License. ``` ## wasip2 - 1.0.2+wasi-0.2.9 + **Repository URL**: https://github.com/bytecodealliance/wasi-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -53475,6 +53873,7 @@ limitations under the License. ``` ## wasip3 - 0.4.0+wasi-0.3.0-rc-2026-01-06 + **Repository URL**: https://github.com/bytecodealliance/wasi-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -53555,6 +53954,7 @@ limitations under the License. ``` ## wasm-bindgen - 0.2.118 + **Repository URL**: https://github.com/wasm-bindgen/wasm-bindgen **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -53763,6 +54163,7 @@ limitations under the License. ``` ## wasm-bindgen-futures - 0.4.68 + **Repository URL**: https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/futures **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -53971,6 +54372,7 @@ limitations under the License. ``` ## wasm-bindgen-macro - 0.2.118 + **Repository URL**: https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/macro **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -54179,6 +54581,7 @@ limitations under the License. ``` ## wasm-bindgen-macro-support - 0.2.118 + **Repository URL**: https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/macro-support **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -54387,6 +54790,7 @@ limitations under the License. ``` ## wasm-bindgen-shared - 0.2.118 + **Repository URL**: https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/shared **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -54595,6 +54999,7 @@ limitations under the License. ``` ## wasm-encoder - 0.244.0 + **Repository URL**: https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wasm-encoder **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -54604,6 +55009,7 @@ No package license file was found in the downloaded crate archive. ``` ## wasm-metadata - 0.244.0 + **Repository URL**: https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wasm-metadata **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -54613,6 +55019,7 @@ No package license file was found in the downloaded crate archive. ``` ## wasm-streams - 0.4.2 + **Repository URL**: https://github.com/MattiasBuelens/wasm-streams/ **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -54693,6 +55100,7 @@ limitations under the License. ``` ## wasm-streams - 0.5.0 + **Repository URL**: https://github.com/MattiasBuelens/wasm-streams/ **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -54773,6 +55181,7 @@ limitations under the License. ``` ## wasmparser - 0.244.0 + **Repository URL**: https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wasmparser **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -54782,6 +55191,7 @@ No package license file was found in the downloaded crate archive. ``` ## web-sys - 0.3.95 + **Repository URL**: https://github.com/wasm-bindgen/wasm-bindgen/tree/master/crates/web-sys **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -54990,6 +55400,7 @@ limitations under the License. ``` ## web-time - 1.1.0 + **Repository URL**: https://github.com/daxpedda/web-time **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -55198,6 +55609,7 @@ limitations under the License. ``` ## webpki-root-certs - 1.0.9 + **Repository URL**: https://github.com/rustls/webpki-roots **License Type(s)**: CDLA-Permissive-2.0 ### License: https://spdx.org/licenses/CDLA-Permissive-2.0.html @@ -55266,6 +55678,7 @@ insights. ``` ## winapi - 0.3.9 + **Repository URL**: https://github.com/retep998/winapi-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -55474,6 +55887,7 @@ insights. ``` ## winapi-i686-pc-windows-gnu - 0.4.0 + **Repository URL**: https://github.com/retep998/winapi-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -55554,6 +55968,7 @@ limitations under the License. ``` ## winapi-util - 0.1.11 + **Repository URL**: https://github.com/BurntSushi/winapi-util **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -55582,6 +55997,7 @@ THE SOFTWARE. ``` ## winapi-x86_64-pc-windows-gnu - 0.4.0 + **Repository URL**: https://github.com/retep998/winapi-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -55662,6 +56078,7 @@ limitations under the License. ``` ## windows - 0.62.2 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -55870,6 +56287,7 @@ limitations under the License. ``` ## windows-collections - 0.3.2 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -56078,6 +56496,7 @@ limitations under the License. ``` ## windows-core - 0.62.2 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -56286,6 +56705,7 @@ limitations under the License. ``` ## windows-future - 0.3.2 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -56494,6 +56914,7 @@ limitations under the License. ``` ## windows-implement - 0.60.2 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -56702,6 +57123,7 @@ limitations under the License. ``` ## windows-interface - 0.59.3 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -56910,6 +57332,7 @@ limitations under the License. ``` ## windows-link - 0.2.1 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -57118,6 +57541,7 @@ limitations under the License. ``` ## windows-numerics - 0.3.1 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -57326,6 +57750,7 @@ limitations under the License. ``` ## windows-result - 0.4.1 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -57534,6 +57959,7 @@ limitations under the License. ``` ## windows-strings - 0.5.1 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -57742,6 +58168,7 @@ limitations under the License. ``` ## windows-sys - 0.52.0 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -57950,6 +58377,7 @@ limitations under the License. ``` ## windows-sys - 0.59.0 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -58158,6 +58586,7 @@ limitations under the License. ``` ## windows-sys - 0.61.2 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -58366,6 +58795,7 @@ limitations under the License. ``` ## windows-targets - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -58574,6 +59004,7 @@ limitations under the License. ``` ## windows-threading - 0.2.1 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -58782,6 +59213,7 @@ limitations under the License. ``` ## windows_aarch64_gnullvm - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -58990,6 +59422,7 @@ limitations under the License. ``` ## windows_aarch64_msvc - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -59198,6 +59631,7 @@ limitations under the License. ``` ## windows_i686_gnu - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -59406,6 +59840,7 @@ limitations under the License. ``` ## windows_i686_gnullvm - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -59614,6 +60049,7 @@ limitations under the License. ``` ## windows_i686_msvc - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -59822,6 +60258,7 @@ limitations under the License. ``` ## windows_x86_64_gnu - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -60030,6 +60467,7 @@ limitations under the License. ``` ## windows_x86_64_gnullvm - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -60238,6 +60676,7 @@ limitations under the License. ``` ## windows_x86_64_msvc - 0.52.6 + **Repository URL**: https://github.com/microsoft/windows-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -60446,6 +60885,7 @@ limitations under the License. ``` ## winnow - 0.7.15 + **Repository URL**: https://github.com/winnow-rs/winnow **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -60471,6 +60911,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## winnow - 1.0.1 + **Repository URL**: https://github.com/winnow-rs/winnow **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -60496,6 +60937,7 @@ WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ``` ## wit-bindgen - 0.51.0 + **Repository URL**: https://github.com/bytecodealliance/wit-bindgen **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -60704,6 +61146,7 @@ limitations under the License. ``` ## wit-bindgen-core - 0.51.0 + **Repository URL**: https://github.com/bytecodealliance/wit-bindgen **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -61164,6 +61607,7 @@ DEALINGS IN THE SOFTWARE. ``` ## wit-bindgen-rust - 0.51.0 + **Repository URL**: https://github.com/bytecodealliance/wit-bindgen **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -61624,6 +62068,7 @@ DEALINGS IN THE SOFTWARE. ``` ## wit-bindgen-rust-macro - 0.51.0 + **Repository URL**: https://github.com/bytecodealliance/wit-bindgen **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -62084,6 +62529,7 @@ DEALINGS IN THE SOFTWARE. ``` ## wit-component - 0.244.0 + **Repository URL**: https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wit-component **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -62093,6 +62539,7 @@ No package license file was found in the downloaded crate archive. ``` ## wit-parser - 0.244.0 + **Repository URL**: https://github.com/bytecodealliance/wasm-tools/tree/main/crates/wit-parser **License Type(s)**: Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -62102,6 +62549,7 @@ No package license file was found in the downloaded crate archive. ``` ## writeable - 0.6.3 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -62155,6 +62603,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## xxhash-rust - 0.8.15 + **Repository URL**: https://github.com/DoumanAsh/xxhash-rust **License Type(s)**: BSL-1.0 ### License: https://spdx.org/licenses/BSL-1.0.html @@ -62185,6 +62634,7 @@ DEALINGS IN THE SOFTWARE. ``` ## yasna - 0.5.2 + **Repository URL**: https://github.com/qnighy/yasna.rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -62265,6 +62715,7 @@ limitations under the License. ``` ## yoke - 0.8.2 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -62318,6 +62769,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## yoke-derive - 0.8.2 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -62371,6 +62823,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zerocopy - 0.8.48 + **Repository URL**: https://github.com/google/zerocopy **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -62580,6 +63033,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zerocopy-derive - 0.8.48 + **Repository URL**: https://github.com/google/zerocopy **License Type(s)**: BSD-2-Clause OR Apache-2.0 OR MIT ### License: https://spdx.org/licenses/ @@ -62848,6 +63302,7 @@ DEALINGS IN THE SOFTWARE. ``` ## zerofrom - 0.1.7 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -62901,6 +63356,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zerofrom-derive - 0.1.7 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -62954,6 +63410,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zeroize - 1.8.2 + **Repository URL**: https://github.com/RustCrypto/utils **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -63163,6 +63620,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zerotrie - 0.2.4 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -63216,6 +63674,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zerovec - 0.11.6 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -63269,6 +63728,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zerovec-derive - 0.11.3 + **Repository URL**: https://github.com/unicode-org/icu4x **License Type(s)**: Unicode-3.0 ### License: https://spdx.org/licenses/Unicode-3.0.html @@ -63322,6 +63782,7 @@ ICU 1.8.1 to ICU 57.1 © 1995-2016 International Business Machines Corporation a ``` ## zmij - 1.0.21 + **Repository URL**: https://github.com/dtolnay/zmij **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -63352,6 +63813,7 @@ DEALINGS IN THE SOFTWARE. ``` ## zstd - 0.13.3 + **Repository URL**: https://github.com/gyscos/zstd-rs **License Type(s)**: MIT ### License: https://spdx.org/licenses/MIT.html @@ -63367,6 +63829,7 @@ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLI ``` ## zstd-safe - 7.2.4 + **Repository URL**: https://github.com/gyscos/zstd-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html @@ -63447,6 +63910,7 @@ limitations under the License. ``` ## zstd-sys - 2.0.16+zstd.1.5.7 + **Repository URL**: https://github.com/gyscos/zstd-rs **License Type(s)**: Apache-2.0 ### License: https://spdx.org/licenses/Apache-2.0.html diff --git a/Cargo.lock b/Cargo.lock index 0ee67ac77..22ab6587c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -805,7 +805,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1062,9 +1062,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.16" +version = "0.4.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" +checksum = "ef8e5e5a340588f4452631496976cf8636d4a7ecf600239fdc27615d2530bc16" dependencies = [ "atomic-waker", "bytes", @@ -1173,9 +1173,9 @@ dependencies = [ [[package]] name = "hyper" -version = "1.9.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6299f016b246a94207e63da54dbe807655bf9e00044f73ded42c3ac5305fbcca" +checksum = "27b501faa50e7a26c3d3560ca625132f4078a17771f4810baf70475ae48cbe43" dependencies = [ "atomic-waker", "bytes", @@ -1807,6 +1807,7 @@ dependencies = [ "http", "http-body-util", "hyper", + "hyper-rustls", "hyper-util", "jsonschema", "libc", @@ -2980,7 +2981,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3039,7 +3040,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -3609,7 +3610,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4405,7 +4406,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 124c5b7ae..6a8d792c9 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -42,8 +42,9 @@ futures-util = "0.3" fs2 = "0.4" http = "1" http-body-util = "0.1" -hyper = "1" -hyper-util = { version = "0.1", features = ["tokio"] } +hyper = { version = "1.11.1", features = ["client", "server", "http1", "http2"] } +hyper-rustls = { version = "0.27", default-features = false, features = ["http1", "http2", "native-tokio", "ring", "tls12"] } +hyper-util = { version = "0.1", features = ["client-legacy", "http1", "http2", "server-auto", "service", "tokio"] } dialoguer = { version = "0.11", default-features = false, features = ["password"] } jsonschema = { version = "0.46.6", default-features = false } listeners = "0.4" diff --git a/crates/cli/src/bootstrap/state.rs b/crates/cli/src/bootstrap/state.rs index 865b32ca9..46ea67a9a 100644 --- a/crates/cli/src/bootstrap/state.rs +++ b/crates/cli/src/bootstrap/state.rs @@ -21,7 +21,7 @@ use crate::gateway::client::{RelayHealth, probe, request_shutdown}; pub(crate) const BOOTSTRAP_STATE_DIR_ENV: &str = "NEMO_RELAY_BOOTSTRAP_STATE_DIR"; pub(crate) const BOOTSTRAP_SHUTDOWN_TOKEN_ENV: &str = "NEMO_RELAY_BOOTSTRAP_SHUTDOWN_TOKEN"; const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); -const UNHEALTHY_GATEWAY_TERMINATION_TIMEOUT: Duration = Duration::from_secs(1); +const UNHEALTHY_GATEWAY_TERMINATION_TIMEOUT: Duration = Duration::from_secs(2); #[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)] pub(super) struct OwnerRecord { diff --git a/crates/cli/src/commands/configure/mod.rs b/crates/cli/src/commands/configure/mod.rs index 974f2cb2d..51883c08e 100644 --- a/crates/cli/src/commands/configure/mod.rs +++ b/crates/cli/src/commands/configure/mod.rs @@ -69,3 +69,7 @@ pub(super) async fn execute( } Ok(ExitCode::SUCCESS) } + +#[cfg(test)] +#[path = "../../../tests/coverage/commands/configure_command_tests.rs"] +mod tests; diff --git a/crates/cli/src/commands/daemon.rs b/crates/cli/src/commands/daemon.rs new file mode 100644 index 000000000..3ee3333da --- /dev/null +++ b/crates/cli/src/commands/daemon.rs @@ -0,0 +1,282 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::net::Ipv4Addr; +use std::path::PathBuf; +use std::process::ExitCode; + +use clap::{Args, Subcommand, ValueEnum}; + +use super::root::AgentArg; +use crate::daemon; +use crate::daemon::common::address::{ + DEFAULT_DAEMON_BIND, DEFAULT_DAEMON_PORT, DEFAULT_WORKER_BIND, +}; +use crate::error::CliError; + +/// Run or connect to the multi-user NeMo Relay daemon. +#[derive(Debug, Clone, Args)] +#[command(args_conflicts_with_subcommands = true)] +pub(crate) struct DaemonCommand { + /// Address on which the daemon accepts public requests. + #[arg(long, default_value_t = DEFAULT_DAEMON_BIND, value_parser = parse_bind_address)] + pub(crate) bind: Ipv4Addr, + /// Port on which the daemon accepts public requests. + #[arg(long, default_value_t = DEFAULT_DAEMON_PORT, value_parser = parse_nonzero_port)] + pub(crate) port: u16, + /// Concrete URL at which clients can reach a daemon bound to 0.0.0.0. + #[arg(long)] + pub(crate) advertise_address: Option, + /// PEM certificate chain for a native TLS daemon listener. + #[arg(long, requires = "tls_key")] + pub(crate) tls_cert: Option, + /// PKCS#8 PEM private key for a native TLS daemon listener. + #[arg(long, requires = "tls_cert")] + pub(crate) tls_key: Option, + /// Route directly to configured providers and never activate a worker. + #[arg(long)] + pub(crate) pass_through: bool, + #[command(subcommand)] + pub(crate) command: Option, +} + +#[derive(Debug, Clone, Subcommand)] +pub(crate) enum DaemonSubcommand { + /// Register this MCP process with an explicitly selected daemon. + Mcp(DaemonMcpCommand), + /// Forward a managed coding-agent hook to an explicitly selected daemon. + Hook(DaemonHookCommand), + /// Run a worker activated and controlled by an explicitly selected daemon. + Worker(DaemonWorkerCommand), + /// Create an immutable administrator-managed integration bundle. + ManagedBundle(DaemonManagedBundleCommand), +} + +#[derive(Debug, Clone, Args)] +pub(crate) struct DaemonMcpCommand { + /// Absolute daemon URL, including scheme, host, and port. + #[arg(long, value_parser = parse_daemon_address)] + pub(crate) daemon_address: String, +} + +#[derive(Debug, Clone, Args)] +pub(crate) struct DaemonHookCommand { + /// Coding agent whose native hook payload is read from standard input. + #[arg(value_enum)] + pub(crate) agent: AgentArg, + /// Absolute daemon URL, including scheme, host, and port. + #[arg(long, value_parser = parse_daemon_address)] + pub(crate) daemon_address: String, + /// Allow the coding agent to continue when hook delivery fails. + #[arg(long, conflicts_with = "fail_closed")] + pub(crate) fail_open: bool, + /// Return a failure when the hook cannot be delivered or is rejected. + #[arg(long, conflicts_with = "fail_open")] + pub(crate) fail_closed: bool, +} + +#[derive(Debug, Clone, Args)] +pub(crate) struct DaemonWorkerCommand { + /// Absolute daemon URL, including scheme, host, and port. + #[arg(long, value_parser = parse_daemon_address)] + pub(crate) daemon_address: String, + /// Address on which the worker accepts daemon requests. + #[arg(long, default_value_t = DEFAULT_WORKER_BIND, value_parser = parse_bind_address)] + pub(crate) bind: Ipv4Addr, + /// Prescribed worker port. Omit to let the operating system select a port. + #[arg(long, value_parser = parse_nonzero_port)] + pub(crate) port: Option, + /// Concrete daemon-reachable host or IP for a worker bound to 0.0.0.0. + #[arg(long)] + pub(crate) advertise_address: Option, +} + +#[derive(Debug, Clone, Args)] +#[command( + long_about = "Create a new immutable administrator-managed integration bundle. This is separate from personal `nemo-relay install`: artifacts contain only fixed deployment values and an existing bundle is never rewritten with different bytes. The command prints the canonical bundle SHA-256 to stdout for separate administrator provisioning.", + after_help = "On success, stdout contains only the canonical bundle SHA-256 for separate administrator provisioning. The dispatcher is checked lexically for the target platform. It must be an absolute stable system path outside known user and temporary directories. Filesystem ownership is not checked while building because the bundle may be created on a different operating system; deployment tooling must install the dispatcher with administrator-controlled ownership and permissions." +)] +pub(crate) struct DaemonManagedBundleCommand { + /// New bundle directory. An existing byte-identical bundle is left untouched. + #[arg(long)] + pub(crate) output: PathBuf, + /// Fixed absolute daemon URL embedded identically for every managed user. + #[arg(long, value_parser = parse_daemon_address)] + pub(crate) daemon_address: String, + /// Absolute, stable administrator dispatcher path embedded in every artifact. + /// + /// The path is validated lexically for the selected target platform. It must be outside + /// known user and temporary directories. Ownership is enforced at deployment time because + /// cross-platform bundles may be built on a different operating system. + #[arg(long, value_name = "ABSOLUTE-PATH")] + pub(crate) dispatcher_command: String, + /// Operating system on which the managed artifacts will be deployed. + #[arg(long, value_enum)] + pub(crate) platform: ManagedPlatformArg, + /// Managed coding agent to include. Repeat this option to include multiple agents. + #[arg(long = "agent", value_enum, required = true)] + pub(crate) agents: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +#[value(rename_all = "lower")] +pub(crate) enum ManagedPlatformArg { + Linux, + Macos, + Windows, +} + +impl From for daemon::managed::ManagedPlatform { + fn from(platform: ManagedPlatformArg) -> Self { + match platform { + ManagedPlatformArg::Linux => Self::Linux, + ManagedPlatformArg::Macos => Self::Macos, + ManagedPlatformArg::Windows => Self::Windows, + } + } +} + +pub(crate) async fn execute( + command: DaemonCommand, + server: &crate::commands::serve::ServerArgs, +) -> Result { + match command.command { + None => { + if command.bind == Ipv4Addr::UNSPECIFIED && command.advertise_address.is_none() { + return Err(CliError::Config( + "a daemon bound to 0.0.0.0 requires --advertise-address".into(), + )); + } + daemon::serve(daemon::ServerOptions { + bind: command.bind, + port: command.port, + advertise_address: command.advertise_address, + pass_through: command.pass_through, + gateway: server.to_runtime(), + tls_cert: command.tls_cert, + tls_key: command.tls_key, + }) + .await?; + } + Some(DaemonSubcommand::Mcp(command)) => { + daemon::mcp::run(daemon::mcp::Options { + daemon_address: command.daemon_address, + }) + .await?; + } + Some(DaemonSubcommand::Hook(command)) => { + daemon::hook::run(daemon::hook::Options { + agent: command.agent.into(), + daemon_address: command.daemon_address, + failure_policy: if command.fail_closed { + crate::hooks::HookFailurePolicy::FailClosed + } else if command.fail_open { + crate::hooks::HookFailurePolicy::FailOpen + } else { + crate::hooks::HookFailurePolicy::Default + }, + }) + .await?; + } + Some(DaemonSubcommand::Worker(command)) => { + if command.bind == Ipv4Addr::UNSPECIFIED && command.advertise_address.is_none() { + return Err(CliError::Config( + "a worker bound to 0.0.0.0 requires --advertise-address".into(), + )); + } + daemon::worker::run(daemon::worker::Options { + daemon_address: command.daemon_address, + bind: command.bind, + port: command.port, + advertise_address: command.advertise_address, + }) + .await?; + } + Some(DaemonSubcommand::ManagedBundle(command)) => { + let agents = command.agents.into_iter().map(|agent| match agent { + AgentArg::Codex => daemon::managed::ManagedAgent::Codex, + AgentArg::Claude => daemon::managed::ManagedAgent::ClaudeCode, + AgentArg::Pi => daemon::managed::ManagedAgent::Pi, + }); + let spec = daemon::managed::ManagedBundleSpec::new( + command.daemon_address, + command.dispatcher_command, + command.platform.into(), + agents, + )?; + let sha256 = daemon::managed::write_new_bundle(&command.output, &spec)?; + println!("{sha256}"); + } + } + + Ok(ExitCode::SUCCESS) +} + +fn parse_bind_address(value: &str) -> Result { + let address = value + .parse::() + .map_err(|_| "bind address must be 127.0.0.1 or 0.0.0.0".to_string())?; + if matches!(address, Ipv4Addr::LOCALHOST | Ipv4Addr::UNSPECIFIED) { + Ok(address) + } else { + Err("bind address must be 127.0.0.1 or 0.0.0.0".into()) + } +} + +fn parse_nonzero_port(value: &str) -> Result { + match value.parse::() { + Ok(0) => Err("an explicitly supplied port must be between 1 and 65535".into()), + Ok(port) => Ok(port), + Err(_) => Err("port must be between 1 and 65535".into()), + } +} + +fn parse_daemon_address(value: &str) -> Result { + let uri = value + .parse::() + .map_err(|_| "daemon address must be an absolute HTTP or HTTPS URL".to_string())?; + let scheme = uri + .scheme_str() + .filter(|scheme| matches!(*scheme, "http" | "https")) + .ok_or_else(|| "daemon address must use http or https".to_string())?; + let authority = uri + .authority() + .ok_or_else(|| "daemon address must include a host and explicit port".to_string())?; + if authority.port_u16().is_none() { + return Err("daemon address must include an explicit port".into()); + } + let url = reqwest::Url::parse(value) + .map_err(|_| "daemon address must be an absolute HTTP or HTTPS URL".to_string())?; + if !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + { + return Err( + "daemon address cannot contain credentials, a non-root path, query, or fragment".into(), + ); + } + let host = url + .host_str() + .ok_or_else(|| "daemon address must include a host".to_string())?; + let host = host.trim_matches(['[', ']']); + if host + .parse::() + .is_ok_and(|address| address.is_unspecified()) + { + return Err("0.0.0.0 is a bind address and cannot be a daemon target".into()); + } + let loopback = host.eq_ignore_ascii_case("localhost") + || host + .parse::() + .is_ok_and(|address| address.is_loopback()); + if scheme == "http" && !loopback { + return Err("non-loopback daemon addresses must use https".into()); + } + Ok(value.trim_end_matches('/').to_string()) +} + +#[cfg(test)] +#[path = "../../tests/coverage/commands/daemon_tests.rs"] +mod tests; diff --git a/crates/cli/src/commands/diagnostics.rs b/crates/cli/src/commands/diagnostics.rs index 396f5ebbd..189f64ef5 100644 --- a/crates/cli/src/commands/diagnostics.rs +++ b/crates/cli/src/commands/diagnostics.rs @@ -19,6 +19,19 @@ pub(crate) struct DoctorCommand { pub(crate) plugin: Option, #[arg(long, requires = "plugin")] pub(crate) install_dir: Option, + /// Validate only an administrator-managed daemon bundle and its managed environment. + #[arg( + long, + requires = "managed_bundle_sha256", + conflicts_with_all = ["agent", "plugin", "install_dir", "offline"] + )] + pub(crate) managed_bundle: Option, + /// Trusted canonical bundle SHA-256 emitted when the managed bundle was created. + /// + /// This value must be provisioned separately by an administrator. Digest values stored inside + /// the bundle are never used as its trust root. + #[arg(long, requires = "managed_bundle")] + pub(crate) managed_bundle_sha256: Option, #[arg(long)] pub(crate) json: bool, #[arg( @@ -39,6 +52,16 @@ pub(super) async fn execute( server: &super::serve::ServerArgs, logging_fallback_error: Option<&CliError>, ) -> Result { + if let Some(bundle) = command.managed_bundle { + let expected_sha256 = command + .managed_bundle_sha256 + .expect("clap requires a managed bundle SHA-256"); + return crate::diagnostics::run_managed_bundle_doctor( + &bundle, + &expected_sha256, + command.json, + ); + } if let Some(plugin) = command.plugin { return execute_plugin_doctor(plugin, command.install_dir, command.json); } diff --git a/crates/cli/src/commands/integrations.rs b/crates/cli/src/commands/integrations.rs index e41457416..f0da8cefc 100644 --- a/crates/cli/src/commands/integrations.rs +++ b/crates/cli/src/commands/integrations.rs @@ -41,6 +41,13 @@ pub(crate) fn execute(command: IntegrationsCommand) -> Result Result { + refresh_with_installer(command, crate::agents::install_integration) +} + +fn refresh_with_installer( + command: RefreshCommand, + mut install: impl FnMut(CodingAgent, InstallRequest) -> Result, +) -> Result { let targets = refresh_targets(command.install_dir.as_deref())?; let managed_targets = targets .iter() @@ -79,7 +86,7 @@ fn refresh(command: RefreshCommand) -> Result { dry_run: command.dry_run, skip_doctor: false, }; - let result = match crate::agents::install_integration(agent, request) { + let result = match install(agent, request) { Ok(status) if status == ExitCode::SUCCESS => Ok(()), Ok(_) => Err(format!( "{} at {} returned a nonzero status", @@ -171,3 +178,7 @@ fn refresh_targets( Ok(targets) } + +#[cfg(test)] +#[path = "../../tests/coverage/commands/integrations_tests.rs"] +mod tests; diff --git a/crates/cli/src/commands/mod.rs b/crates/cli/src/commands/mod.rs index 33268fad7..51d85f951 100644 --- a/crates/cli/src/commands/mod.rs +++ b/crates/cli/src/commands/mod.rs @@ -5,6 +5,7 @@ mod completions; mod configure; +mod daemon; mod diagnostics; mod gateway; mod hook_forward; @@ -23,7 +24,6 @@ use std::process::ExitCode; use clap::Parser; #[cfg(test)] -use self::completions::CompletionsCommand; #[cfg(test)] use self::model_pricing::PricingCommand; #[cfg(test)] @@ -92,6 +92,14 @@ fn configure_logging(cli: &Cli) -> Result { cli.logging.resolve_without_ambient_config() } Some(Command::Mcp) => cli.logging.resolve(None), + Some(Command::Daemon(command)) + if matches!( + command.command.as_ref(), + Some(daemon::DaemonSubcommand::Mcp(_) | daemon::DaemonSubcommand::Hook(_)) + ) => + { + cli.logging.resolve_without_ambient_config() + } Some(Command::Run(command)) => cli .logging .resolve(command.config.as_deref().or(cli.server.config.as_deref())), @@ -189,6 +197,7 @@ async fn run_command( bootstrap_shutdown_token: Option, ) -> Result { match command { + Command::Daemon(command) => daemon::execute(command, server).await, Command::HookForward(command) => { hook_forward::execute(command).await?; Ok(ExitCode::SUCCESS) @@ -284,11 +293,6 @@ async fn serve_gateway( Ok(ExitCode::SUCCESS) } -#[cfg(test)] -fn run_completions(command: CompletionsCommand) -> Result { - completions::execute(command) -} - #[cfg(test)] fn run_plugins(command: PluginsCommand, server: &ServerArgs) -> Result { plugins::execute(command, server) diff --git a/crates/cli/src/commands/root.rs b/crates/cli/src/commands/root.rs index 1f6ae4b09..5b7a7e906 100644 --- a/crates/cli/src/commands/root.rs +++ b/crates/cli/src/commands/root.rs @@ -5,6 +5,7 @@ use clap::{Parser, Subcommand, ValueEnum}; use super::completions::CompletionsCommand; use super::configure::ConfigCommand; +use super::daemon::DaemonCommand; use super::diagnostics::{AgentsCommand, DoctorCommand}; use super::gateway::GatewayCommand; use super::hook_forward::HookForwardCommand; @@ -51,6 +52,8 @@ pub(crate) struct Cli { #[derive(Debug, Clone, Subcommand)] pub(crate) enum Command { + /// Run or connect to the multi-user NeMo Relay daemon. + Daemon(DaemonCommand), /// Run Claude Code with observability (setup on first use) #[command( long_about = "Run Anthropic's `claude` CLI under an ephemeral NeMo Relay gateway. \ @@ -146,6 +149,7 @@ pub(crate) enum Command { impl Command { pub(crate) fn log_name(&self) -> &'static str { match self { + Self::Daemon(_) => "daemon", Self::Claude(_) => "claude", Self::Codex(_) => "codex", Self::Pi(_) => "pi", @@ -165,12 +169,17 @@ impl Command { } } - /// Configuration-editing commands remain available even when operational logging settings are - /// invalid, so users can repair their configuration. + /// Administrative and configuration-editing commands must not depend on ambient personal + /// logging state. This keeps repair and managed deployment workflows deterministic. pub(crate) fn skips_logging(&self) -> bool { matches!(self, Self::Config(_)) || matches!(self, Self::Gateway(command) if command.is_stop()) || matches!(self, Self::Plugins(command) if command.is_edit()) + || matches!(self, Self::Doctor(command) if command.managed_bundle.is_some()) + || matches!(self, Self::Daemon(command) if matches!( + command.command.as_ref(), + Some(super::daemon::DaemonSubcommand::ManagedBundle(_)) + )) || matches!(self, Self::HookForward(command) if transparent_hook_is_inert(command)) } } diff --git a/crates/cli/src/configuration/mod.rs b/crates/cli/src/configuration/mod.rs index 77c46b7f7..6d8a814b7 100644 --- a/crates/cli/src/configuration/mod.rs +++ b/crates/cli/src/configuration/mod.rs @@ -113,6 +113,49 @@ pub(crate) fn resolve_server_config(args: &GatewayOverrides) -> Result Result { + let system_directory = system_config_dir(); + resolve_managed_worker_config_from_paths( + system_directory.join("config.toml"), + system_directory.join(PLUGINS_TOML), + ) +} + +fn resolve_managed_worker_config_from_paths( + config_path: PathBuf, + plugin_config_path: PathBuf, +) -> Result { + let resolved = load_config_from_paths( + vec![(config_path, false)], + vec![plugin_config_path.clone()], + apply_managed_worker_env_config, + )?; + enforce_required_dynamic_plugin_startup(Some(&plugin_config_path), &resolved)?; + log::info!( + target: "nemo_relay.configuration", + event = "configuration_resolved", + mode = "managed_worker", + dynamic_plugin_count = resolved.dynamic_plugins.len(); + "Managed worker configuration resolved" + ); + Ok(ManagedWorkerConfig { + resolved, + plugin_config_path, + }) +} + /// Resolves only operational logging from the normal config discovery scope. /// /// This intentionally avoids plugin discovery and activation so logging can be initialized before @@ -1114,10 +1157,28 @@ pub(crate) const PLUGINS_TOML: &str = "plugins.toml"; fn load_shared_config( explicit: Option<&PathBuf>, plugin_config_path: Option<&PathBuf>, +) -> Result { + let config_paths = config_paths(explicit) + .into_iter() + .map(|path| { + let required = explicit == Some(&path); + (path, required) + }) + .collect(); + load_config_from_paths( + config_paths, + plugin_config_paths(explicit, plugin_config_path), + apply_env_config, + ) +} + +fn load_config_from_paths( + config_paths: Vec<(PathBuf, bool)>, + plugin_config_paths: Vec, + apply_environment: fn(&mut GatewayConfig) -> Result<(), CliError>, ) -> Result { let mut merged = toml::Value::Table(toml::map::Map::new()); - for path in config_paths(explicit) { - let required = explicit == Some(&path); + for (path, required) in config_paths { let Some(raw) = read_config_file(&path, required, "configuration")? else { continue; }; @@ -1144,14 +1205,14 @@ fn load_shared_config( } merge_gateway_config_toml(&mut merged, parsed); } - let plugin_toml = load_plugin_toml_config(explicit, plugin_config_path)?; + let plugin_toml = load_plugin_toml_config_from_paths(plugin_config_paths)?; let mut resolved = ResolvedConfig { gateway: GatewayConfig::default(), ..ResolvedConfig::default() }; apply_file_config(&mut resolved, merged)?; apply_plugin_toml_config(&mut resolved, plugin_toml); - apply_env_config(&mut resolved.gateway)?; + apply_environment(&mut resolved.gateway)?; Ok(resolved) } @@ -1386,13 +1447,6 @@ struct FileDynamicPluginConfig { config: Option>, } -fn load_plugin_toml_config( - explicit: Option<&PathBuf>, - plugin_config_path: Option<&PathBuf>, -) -> Result, CliError> { - load_plugin_toml_config_from_paths(plugin_config_paths(explicit, plugin_config_path)) -} - /// Returns the plugin configuration paths selected by the same rules as runtime resolution. /// /// Diagnostics use this so they report the same explicit-or-user and system layers as @@ -1673,6 +1727,25 @@ fn apply_env_config(config: &mut GatewayConfig) -> Result<(), CliError> { Ok(()) } +// Managed workers accept provider secrets from their inherited environment but never allow a +// user's environment to change administrator-selected endpoints, limits, or listener settings. +// `OPENAI_API_KEY` and `ANTHROPIC_API_KEY` are intentionally read later at request forwarding time. +fn apply_managed_worker_env_config(config: &mut GatewayConfig) -> Result<(), CliError> { + if let Ok(value) = std::env::var("NEMO_RELAY_OPENAI_AUTH_HEADER") { + config.openai_auth_header = Some(validate_auth_header( + "NEMO_RELAY_OPENAI_AUTH_HEADER", + value, + )?); + } + if let Ok(value) = std::env::var("NEMO_RELAY_ANTHROPIC_AUTH_HEADER") { + config.anthropic_auth_header = Some(validate_auth_header( + "NEMO_RELAY_ANTHROPIC_AUTH_HEADER", + value, + )?); + } + Ok(()) +} + fn replace_upstream_base_url( base_url: &mut String, auth_header: &mut Option, diff --git a/crates/cli/src/daemon/broker/lifecycle.rs b/crates/cli/src/daemon/broker/lifecycle.rs new file mode 100644 index 000000000..10fed39e7 --- /dev/null +++ b/crates/cli/src/daemon/broker/lifecycle.rs @@ -0,0 +1,245 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use crate::daemon::common::identity::Fingerprint; +use crate::daemon::common::protocol::{SensitiveString, WorkerLaunch}; +use crate::daemon::common::transport::PooledClient; +#[cfg(test)] +use crate::daemon::common::transport::pooled_client; + +/// Stable identity for one connected MCP process. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct McpSessionId(String); + +impl McpSessionId { + /// Constructs a non-empty MCP session identifier. + pub(crate) fn new(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() { + return Err(LifecycleError::EmptyIdentifier); + } + Ok(Self(value)) + } + + /// Returns the wire representation. + pub(crate) fn as_str(&self) -> &str { + &self.0 + } +} + +/// An immutable destination for one registered worker generation. +pub(crate) struct WorkerTarget { + worker_id: String, + endpoint: String, + session_token: SensitiveString, + client: Arc, + in_flight: AtomicUsize, +} + +impl WorkerTarget { + /// Creates an authenticated worker target. + #[cfg(test)] + pub(crate) fn new( + worker_id: impl Into, + endpoint: impl Into, + session_token: SensitiveString, + ) -> Result { + let client = pooled_client().map_err(|_| LifecycleError::TransportInitialization)?; + Self::with_client(worker_id, endpoint, session_token, client) + } + + /// Creates a target with a standalone pool for deterministic transport tests. + #[cfg(test)] + pub(crate) fn with_client( + worker_id: impl Into, + endpoint: impl Into, + session_token: SensitiveString, + client: PooledClient, + ) -> Result { + Self::with_shared_client(worker_id, endpoint, session_token, Arc::new(client)) + } + + /// Creates a target that retains a handle to the daemon's process-wide worker pool service. + pub(crate) fn with_shared_client( + worker_id: impl Into, + endpoint: impl Into, + session_token: SensitiveString, + client: Arc, + ) -> Result { + let worker_id = worker_id.into(); + let endpoint = endpoint.into(); + if worker_id.is_empty() || endpoint.is_empty() { + return Err(LifecycleError::EmptyIdentifier); + } + Ok(Self { + worker_id, + endpoint, + session_token, + client, + in_flight: AtomicUsize::new(0), + }) + } + + /// Returns the worker generation identifier. + pub(crate) fn worker_id(&self) -> &str { + &self.worker_id + } + + /// Returns the daemon-reachable worker endpoint. + pub(crate) fn endpoint(&self) -> &str { + &self.endpoint + } + + /// Returns the internal daemon-to-worker credential. + pub(crate) fn session_token(&self) -> &str { + self.session_token.expose() + } + + /// Returns the shared pool selected for this worker's transport trust identity. + pub(crate) fn client(&self) -> &PooledClient { + self.client.as_ref() + } + + /// Returns the number of requests accepted by the broker and not yet dropped. + pub(crate) fn in_flight(&self) -> usize { + self.in_flight.load(Ordering::Acquire) + } + + pub(super) fn acquire(self: &Arc, fingerprint: Fingerprint) -> WorkerRequest { + self.in_flight.fetch_add(1, Ordering::Relaxed); + WorkerRequest { + fingerprint, + target: Arc::clone(self), + } + } +} + +impl fmt::Debug for WorkerTarget { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("WorkerTarget") + .field("worker_id", &self.worker_id) + .field("endpoint", &self.endpoint) + .field("session_token", &self.session_token) + .field("in_flight", &self.in_flight()) + .finish() + } +} + +/// An accepted request's ownership of one worker target. +pub(crate) struct WorkerRequest { + fingerprint: Fingerprint, + target: Arc, +} + +impl WorkerRequest { + /// Returns the stable route identity that selected this worker. + pub(crate) const fn fingerprint(&self) -> Fingerprint { + self.fingerprint + } + + /// Returns the immutable worker target. + pub(crate) fn target(&self) -> &Arc { + &self.target + } + + /// Returns the internal daemon-to-worker credential. + pub(crate) fn session_token(&self) -> &str { + self.target.session_token() + } +} + +impl fmt::Debug for WorkerRequest { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("WorkerRequest") + .field("fingerprint", &self.fingerprint) + .field("target", &self.target) + .finish() + } +} + +impl Drop for WorkerRequest { + fn drop(&mut self) { + self.target.in_flight.fetch_sub(1, Ordering::Release); + } +} + +/// A route destination resolved before any request-body frame is polled. +#[derive(Debug)] +pub(crate) enum ResolvedTarget { + Worker(WorkerRequest), + PassThrough, +} + +/// The externally useful category of a route's lifecycle state. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RouteStateKind { + Empty, + Activating, + Ready, + Draining, + PassThrough, + Recovering, +} + +/// Internal state for one fingerprint route. +#[derive(Debug)] +pub(crate) enum RouteState { + Empty, + Activating { + owner: McpSessionId, + launch: WorkerLaunch, + }, + Ready { + target: Arc, + }, + Draining { + target: Arc, + deadline_unix_ms: u64, + }, + PassThrough { + permanent: bool, + }, + Recovering { + target: Option>, + owner: Option, + deadline_unix_ms: u64, + }, +} + +impl RouteState { + /// Returns the state category without exposing credentials or mutable internals. + pub(crate) const fn kind(&self) -> RouteStateKind { + match self { + Self::Empty => RouteStateKind::Empty, + Self::Activating { .. } => RouteStateKind::Activating, + Self::Ready { .. } => RouteStateKind::Ready, + Self::Draining { .. } => RouteStateKind::Draining, + Self::PassThrough { .. } => RouteStateKind::PassThrough, + Self::Recovering { .. } => RouteStateKind::Recovering, + } + } +} + +/// Validation failures for strongly typed lifecycle identifiers and targets. +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum LifecycleError { + #[error("a lifecycle identifier or endpoint cannot be empty")] + EmptyIdentifier, + #[error("failed to construct a worker transport pool")] + #[cfg(test)] + TransportInitialization, +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/lifecycle_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/broker/mod.rs b/crates/cli/src/daemon/broker/mod.rs new file mode 100644 index 000000000..a5282082e --- /dev/null +++ b/crates/cli/src/daemon/broker/mod.rs @@ -0,0 +1,8 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-user-machine route lifecycle and lock-bounded broker registry. + +pub(crate) mod lifecycle; +pub(crate) mod registry; +pub(crate) mod server; diff --git a/crates/cli/src/daemon/broker/registry.rs b/crates/cli/src/daemon/broker/registry.rs new file mode 100644 index 000000000..ec73fdf83 --- /dev/null +++ b/crates/cli/src/daemon/broker/registry.rs @@ -0,0 +1,1124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{BTreeMap, HashMap}; +use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; + +use thiserror::Error; + +use super::lifecycle::{McpSessionId, ResolvedTarget, RouteState, RouteStateKind, WorkerTarget}; +use crate::daemon::common::identity::{Fingerprint, TokenDigest}; +use crate::daemon::common::protocol::{BrokerDirective, WorkerLaunch}; + +const DEFAULT_RETRY_AFTER_MS: u64 = 100; +const MAX_ROUTE_BINDINGS: usize = 4_096; +const MAX_MCP_REFERENCES_PER_ROUTE: usize = 1_024; + +/// An authenticated MCP registration applied idempotently by session ID. +#[derive(Debug, Clone)] +pub(crate) struct McpRegistration { + pub(crate) fingerprint: Fingerprint, + pub(crate) token_digest: TokenDigest, + pub(crate) session_id: McpSessionId, + pub(crate) lease_expires_at_unix_ms: u64, +} + +/// A lock-bounded broker registry keyed by stable user-machine fingerprint. +pub(crate) struct Registry { + global_pass_through: bool, + retry_after_ms: u64, + route_capacity: usize, + inner: RwLock, +} + +impl Registry { + /// Authorizes a worker recovery without mutating route state. + /// + /// The returned permit captures the exact route generation and must be presented again when + /// publishing the worker after its authenticated readiness probe. + pub(crate) fn authorize_worker_recovery( + &self, + fingerprint: Fingerprint, + worker_id: &str, + ) -> Result { + let inner = self.read(); + let route = inner + .routes + .get(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + if route.refs.is_empty() { + return Err(RegistryError::NoLiveMcpReferences); + } + match &route.state { + RouteState::Activating { launch, .. } => Ok(RecoveryPermit::Activating { + activation_id: launch.activation_id.clone(), + }), + RouteState::Ready { target } if target.worker_id() == worker_id => { + Ok(RecoveryPermit::ExistingWorker { + worker_id: worker_id.to_owned(), + recovering: false, + }) + } + RouteState::Recovering { + target: Some(target), + .. + } if target.worker_id() == worker_id => Ok(RecoveryPermit::ExistingWorker { + worker_id: worker_id.to_owned(), + recovering: true, + }), + _ => Err(RegistryError::RecoveryNotAuthorized), + } + } + + /// Publishes a recovered worker only if the preflighted route generation is unchanged. + pub(crate) fn publish_recovered_worker( + &self, + fingerprint: Fingerprint, + permit: &RecoveryPermit, + target: Arc, + ) -> Result, RegistryError> { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + if route.refs.is_empty() { + return Err(RegistryError::NoLiveMcpReferences); + } + let authorized = match (&route.state, permit) { + ( + RouteState::Activating { launch, .. }, + RecoveryPermit::Activating { activation_id }, + ) => launch.activation_id == *activation_id, + ( + RouteState::Ready { target }, + RecoveryPermit::ExistingWorker { + worker_id, + recovering: false, + }, + ) => target.worker_id() == worker_id, + ( + RouteState::Recovering { + target: Some(target), + .. + }, + RecoveryPermit::ExistingWorker { + worker_id, + recovering: true, + }, + ) => target.worker_id() == worker_id, + _ => false, + }; + if !authorized { + return Err(RegistryError::RecoveryGenerationChanged); + } + let canceled_activation = match &route.state { + RouteState::Activating { launch, .. } => Some(launch.activation_id.clone()), + _ => None, + }; + route.state = RouteState::Ready { target }; + Ok(canceled_activation) + } + + /// Creates an empty registry. + pub(crate) fn new(global_pass_through: bool) -> Self { + Self { + global_pass_through, + retry_after_ms: DEFAULT_RETRY_AFTER_MS, + route_capacity: MAX_ROUTE_BINDINGS, + inner: RwLock::new(RegistryInner::default()), + } + } + + #[cfg(test)] + pub(crate) fn with_route_capacity(mut self, route_capacity: usize) -> Self { + self.route_capacity = route_capacity; + self + } + + /// Changes the retry hint used by wait directives. + #[cfg(test)] + pub(crate) fn with_retry_after_ms(mut self, retry_after_ms: u64) -> Self { + self.retry_after_ms = retry_after_ms; + self + } + + /// Restores a persisted token/fingerprint binding without creating a live reference. + #[cfg(test)] + pub(crate) fn restore_binding( + &self, + fingerprint: Fingerprint, + token_digest: TokenDigest, + ) -> Result<(), RegistryError> { + let mut inner = self.write(); + evict_inactive_routes_at_capacity(&mut inner, fingerprint, self.route_capacity); + validate_binding(&inner, fingerprint, token_digest)?; + validate_capacity(&inner, fingerprint, self.route_capacity)?; + inner.tokens.insert(token_digest, fingerprint); + inner + .routes + .entry(fingerprint) + .or_insert_with(|| RouteEntry::new(token_digest, self.global_pass_through)); + Ok(()) + } + + /// Registers or renews an MCP and returns the daemon's authoritative directive. + /// + /// The launch plan is used only when this call wins the empty-route singleflight. + pub(crate) fn register_mcp( + &self, + registration: McpRegistration, + launch: WorkerLaunch, + ) -> Result { + let mut inner = self.write(); + evict_inactive_routes_at_capacity( + &mut inner, + registration.fingerprint, + self.route_capacity, + ); + validate_binding(&inner, registration.fingerprint, registration.token_digest)?; + validate_capacity(&inner, registration.fingerprint, self.route_capacity)?; + inner + .tokens + .insert(registration.token_digest, registration.fingerprint); + let route = inner + .routes + .entry(registration.fingerprint) + .or_insert_with(|| { + RouteEntry::new(registration.token_digest, self.global_pass_through) + }); + if !route.refs.contains_key(®istration.session_id) + && route.refs.len() >= MAX_MCP_REFERENCES_PER_ROUTE + { + return Err(RegistryError::McpReferenceCapacityReached); + } + route.refs.insert( + registration.session_id.clone(), + registration.lease_expires_at_unix_ms, + ); + Ok(route.directive_for(®istration.session_id, launch, self.retry_after_ms)) + } + + /// Renews an existing MCP reference without changing its lifecycle state. + pub(crate) fn renew_mcp( + &self, + fingerprint: Fingerprint, + session_id: &McpSessionId, + lease_expires_at_unix_ms: u64, + ) -> Result<(), RegistryError> { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + let expiry = route + .refs + .get_mut(session_id) + .ok_or(RegistryError::UnknownMcpSession)?; + *expiry = lease_expires_at_unix_ms; + Ok(()) + } + + /// Returns the current broker decision for an existing MCP reference without creating a new + /// activation. This lets an activation owner poll readiness over its authenticated session. + pub(crate) fn current_directive( + &self, + fingerprint: Fingerprint, + session_id: &McpSessionId, + ) -> Result { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + if !route.refs.contains_key(session_id) { + return Err(RegistryError::UnknownMcpSession); + } + Ok(route.current_directive(session_id, self.retry_after_ms)) + } + + /// Releases an MCP reference and begins teardown when the final reference leaves. + pub(crate) fn release_mcp( + &self, + fingerprint: Fingerprint, + session_id: &McpSessionId, + drain_deadline_unix_ms: u64, + ) -> Result { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + if route.refs.remove(session_id).is_none() { + return Ok(ReleaseAction::NoChange); + } + Ok(route.after_reference_removed(session_id, drain_deadline_unix_ms)) + } + + /// Expires dead MCP leases and returns any resulting teardown or ownership actions. + pub(crate) fn expire_mcp_leases( + &self, + now_unix_ms: u64, + drain_deadline_unix_ms: u64, + ) -> Vec<(Fingerprint, ReleaseAction)> { + let mut inner = self.write(); + let mut actions = Vec::new(); + for (fingerprint, route) in &mut inner.routes { + let expired: Vec<_> = route + .refs + .iter() + .filter_map(|(session, expiry)| (*expiry <= now_unix_ms).then_some(session.clone())) + .collect(); + let removed_owner = match &route.state { + RouteState::Activating { owner, .. } + if expired.iter().any(|session| session == owner) => + { + Some(owner.clone()) + } + RouteState::Recovering { + owner: Some(owner), .. + } if expired.iter().any(|session| session == owner) => Some(owner.clone()), + _ => None, + }; + let Some(removed_session) = removed_owner.or_else(|| expired.first().cloned()) else { + continue; + }; + for session in expired { + route.refs.remove(&session); + } + let action = route.after_reference_removed(&removed_session, drain_deadline_unix_ms); + if !matches!(action, ReleaseAction::NoChange) { + actions.push((*fingerprint, action)); + } + } + actions + } + + /// Publishes a worker only when its one-time activation ID matches the active generation. + pub(crate) fn mark_worker_ready( + &self, + fingerprint: Fingerprint, + activation_id: &str, + target: Arc, + ) -> Result<(), RegistryError> { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + if route.refs.is_empty() { + return Err(RegistryError::NoLiveMcpReferences); + } + match &route.state { + RouteState::Activating { launch, .. } if launch.activation_id == activation_id => { + route.state = RouteState::Ready { target }; + Ok(()) + } + RouteState::Activating { .. } => Err(RegistryError::ActivationMismatch), + state => Err(RegistryError::InvalidState { + expected: RouteStateKind::Activating, + actual: state.kind(), + }), + } + } + + /// Converts a failed authenticated activation into shared transient pass-through. + pub(crate) fn mark_activation_failed( + &self, + fingerprint: Fingerprint, + activation_id: &str, + ) -> Result<(), RegistryError> { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + match &route.state { + RouteState::Activating { launch, .. } if launch.activation_id == activation_id => { + route.state = if route.refs.is_empty() { + RouteState::Empty + } else { + RouteState::PassThrough { permanent: false } + }; + Ok(()) + } + RouteState::Activating { .. } => Err(RegistryError::ActivationMismatch), + state => Err(RegistryError::InvalidState { + expected: RouteStateKind::Activating, + actual: state.kind(), + }), + } + } + + /// Expires activation grants and moves every still-referenced route to transient pass-through. + /// + /// This is deliberately separate from expiring the server's secret-bearing grant table: the + /// broker lifecycle must never remain `Activating` after its signed launch deadline passes. + pub(crate) fn expire_activations(&self, now_unix_ms: u64) -> Vec { + let mut inner = self.write(); + let mut expired = Vec::new(); + for (fingerprint, route) in &mut inner.routes { + let activation_id = match &route.state { + RouteState::Activating { launch, .. } if launch.deadline_unix_ms <= now_unix_ms => { + Some(launch.activation_id.clone()) + } + _ => None, + }; + let Some(activation_id) = activation_id else { + continue; + }; + route.state = if route.refs.is_empty() { + RouteState::Empty + } else { + RouteState::PassThrough { permanent: false } + }; + expired.push(ExpiredActivation { + fingerprint: *fingerprint, + activation_id, + }); + } + expired + } + + /// Converts an authenticated worker communication failure into route-wide pass-through. + /// + /// The worker ID prevents a delayed failure from an old stream from displacing a newer ready + /// generation. An activation that raced the failed request is returned so its grant can be + /// revoked by the control plane. + pub(crate) fn mark_worker_communication_failed( + &self, + fingerprint: Fingerprint, + worker_id: &str, + ) -> Result, RegistryError> { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + let state = std::mem::replace(&mut route.state, RouteState::Empty); + let canceled_activation = match state { + RouteState::Ready { target } if target.worker_id() == worker_id => None, + RouteState::Draining { + target, + deadline_unix_ms, + } if target.worker_id() == worker_id => { + route.state = RouteState::Draining { + target, + deadline_unix_ms, + }; + return Err(RegistryError::InvalidState { + expected: RouteStateKind::Ready, + actual: RouteStateKind::Draining, + }); + } + RouteState::Recovering { target, .. } + if target + .as_ref() + .is_none_or(|target| target.worker_id() == worker_id) => + { + None + } + RouteState::Activating { launch, .. } => Some(launch.activation_id), + RouteState::PassThrough { permanent } => { + route.state = RouteState::PassThrough { permanent }; + return Ok(None); + } + RouteState::Empty => return Ok(None), + RouteState::Ready { target } => { + route.state = RouteState::Ready { target }; + return Err(RegistryError::WorkerMismatch); + } + RouteState::Draining { + target, + deadline_unix_ms, + } => { + route.state = RouteState::Draining { + target, + deadline_unix_ms, + }; + return Err(RegistryError::WorkerMismatch); + } + RouteState::Recovering { + target, + owner, + deadline_unix_ms, + } => { + route.state = RouteState::Recovering { + target, + owner, + deadline_unix_ms, + }; + return Err(RegistryError::WorkerMismatch); + } + }; + route.state = if route.refs.is_empty() { + RouteState::Empty + } else { + RouteState::PassThrough { permanent: false } + }; + Ok(canceled_activation) + } + + /// Forces an authenticated route into transient pass-through after activation setup fails. + pub(crate) fn mark_route_pass_through( + &self, + fingerprint: Fingerprint, + ) -> Result, RegistryError> { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + if matches!(&route.state, RouteState::PassThrough { permanent: true }) { + return Ok(None); + } + let canceled_activation = match &route.state { + RouteState::Activating { launch, .. } => Some(launch.activation_id.clone()), + _ => None, + }; + route.state = if route.refs.is_empty() { + RouteState::Empty + } else { + RouteState::PassThrough { permanent: false } + }; + Ok(canceled_activation) + } + + /// Records a ready worker failure and nominates one live MCP to relaunch it. + pub(crate) fn worker_failed( + &self, + fingerprint: Fingerprint, + worker_id: &str, + recovery_deadline_unix_ms: u64, + ) -> Result { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + let state = std::mem::replace(&mut route.state, RouteState::Empty); + match state { + RouteState::Ready { target } if target.worker_id() == worker_id => { + let owner = route.refs.keys().next().cloned(); + if let Some(owner) = owner { + route.state = RouteState::Recovering { + target: None, + owner: Some(owner.clone()), + deadline_unix_ms: recovery_deadline_unix_ms, + }; + Ok(WorkerFailureAction::NominateMcp { session_id: owner }) + } else { + Ok(WorkerFailureAction::RouteEmpty) + } + } + RouteState::Ready { target } => { + route.state = RouteState::Ready { target }; + Err(RegistryError::WorkerMismatch) + } + other => { + let actual = other.kind(); + route.state = other; + Err(RegistryError::InvalidState { + expected: RouteStateKind::Ready, + actual, + }) + } + } + } + + /// Installs a fresh activation plan after the broker nominates a replacement owner. + pub(crate) fn begin_relaunch( + &self, + fingerprint: Fingerprint, + session_id: &McpSessionId, + launch: WorkerLaunch, + ) -> Result { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + match &route.state { + RouteState::Recovering { + owner: Some(owner), .. + } if owner == session_id => { + route.state = RouteState::Activating { + owner: session_id.clone(), + launch: launch.clone(), + }; + Ok(launch.into_directive()) + } + RouteState::Recovering { .. } => Err(RegistryError::NotLaunchOwner), + state => Err(RegistryError::InvalidState { + expected: RouteStateKind::Recovering, + actual: state.kind(), + }), + } + } + + /// Places a restored route into bounded daemon-restart recovery. + #[cfg(test)] + pub(crate) fn begin_recovery( + &self, + fingerprint: Fingerprint, + target: Option>, + deadline_unix_ms: u64, + ) -> Result<(), RegistryError> { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + let owner = route.refs.keys().next().cloned(); + route.state = RouteState::Recovering { + target, + owner, + deadline_unix_ms, + }; + Ok(()) + } + + /// Ends restart recovery at its deadline and returns the required next action. + #[cfg(test)] + pub(crate) fn finish_recovery( + &self, + fingerprint: Fingerprint, + now_unix_ms: u64, + drain_deadline_unix_ms: u64, + ) -> Result { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + let state = std::mem::replace(&mut route.state, RouteState::Empty); + match state { + RouteState::Recovering { + target, + deadline_unix_ms, + .. + } if now_unix_ms >= deadline_unix_ms => match (target, route.refs.is_empty()) { + (Some(target), false) => { + route.state = RouteState::Ready { target }; + Ok(RecoveryAction::WorkerRecovered) + } + (Some(target), true) => { + route.state = RouteState::Draining { + target: Arc::clone(&target), + deadline_unix_ms: drain_deadline_unix_ms, + }; + Ok(RecoveryAction::BeginDrain { + target, + deadline_unix_ms: drain_deadline_unix_ms, + }) + } + (None, false) => { + let session_id = route + .refs + .keys() + .next() + .expect("route has live references") + .clone(); + route.state = RouteState::Recovering { + target: None, + owner: Some(session_id.clone()), + deadline_unix_ms, + }; + Ok(RecoveryAction::NominateMcp { session_id }) + } + (None, true) => Ok(RecoveryAction::RouteEmpty), + }, + RouteState::Recovering { + target, + owner, + deadline_unix_ms, + } => { + route.state = RouteState::Recovering { + target, + owner, + deadline_unix_ms, + }; + Err(RegistryError::RecoveryInProgress) + } + other => { + let actual = other.kind(); + route.state = other; + Err(RegistryError::InvalidState { + expected: RouteStateKind::Recovering, + actual, + }) + } + } + } + + /// Completes a drained worker after all requests finish or the deadline elapses. + pub(crate) fn finish_draining( + &self, + fingerprint: Fingerprint, + now_unix_ms: u64, + ) -> Result { + let mut inner = self.write(); + let route = inner + .routes + .get_mut(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + match &route.state { + RouteState::Draining { + target, + deadline_unix_ms, + } if target.in_flight() != 0 && now_unix_ms < *deadline_unix_ms => { + return Err(RegistryError::DrainInProgress); + } + RouteState::Draining { .. } => {} + state => { + return Err(RegistryError::InvalidState { + expected: RouteStateKind::Draining, + actual: state.kind(), + }); + } + } + route.state = RouteState::Empty; + Ok(route + .refs + .keys() + .next() + .cloned() + .map_or(DrainCompletion::RouteEmpty, |session_id| { + DrainCompletion::ActivationRequired { session_id } + })) + } + + /// Resolves and acquires a route from the request header without polling its body. + pub(crate) fn resolve_target( + &self, + token_digest: &TokenDigest, + ) -> Result { + let inner = self.read(); + let fingerprint = inner + .tokens + .get(token_digest) + .ok_or(ResolveError::UnknownToken)?; + let route = inner + .routes + .get(fingerprint) + .ok_or(ResolveError::UnknownToken)?; + match &route.state { + RouteState::Ready { target } => { + Ok(ResolvedTarget::Worker(target.acquire(*fingerprint))) + } + RouteState::PassThrough { .. } if !route.refs.is_empty() => { + Ok(ResolvedTarget::PassThrough) + } + RouteState::PassThrough { .. } => { + Err(ResolveError::Unavailable(RouteStateKind::PassThrough)) + } + state => Err(ResolveError::Unavailable(state.kind())), + } + } + + /// Returns a credential-free route snapshot for status and tests. + #[cfg(test)] + pub(crate) fn snapshot( + &self, + fingerprint: Fingerprint, + ) -> Result { + let inner = self.read(); + let route = inner + .routes + .get(&fingerprint) + .ok_or(RegistryError::UnknownRoute)?; + let (launch_owner, endpoint, in_flight) = match &route.state { + RouteState::Activating { owner, .. } => (Some(owner.clone()), None, 0), + RouteState::Ready { target } | RouteState::Draining { target, .. } => { + (None, Some(target.endpoint().to_owned()), target.in_flight()) + } + RouteState::Recovering { target, owner, .. } => ( + owner.clone(), + target.as_ref().map(|target| target.endpoint().to_owned()), + target.as_ref().map_or(0, |target| target.in_flight()), + ), + RouteState::Empty | RouteState::PassThrough { .. } => (None, None, 0), + }; + Ok(RouteSnapshot { + state: route.state.kind(), + reference_count: route.refs.len(), + launch_owner, + endpoint, + in_flight, + }) + } + + fn read(&self) -> RwLockReadGuard<'_, RegistryInner> { + self.inner.read().unwrap_or_else(|error| error.into_inner()) + } + + fn write(&self) -> RwLockWriteGuard<'_, RegistryInner> { + self.inner + .write() + .unwrap_or_else(|error| error.into_inner()) + } +} + +fn evict_inactive_routes_at_capacity( + inner: &mut RegistryInner, + incoming: Fingerprint, + capacity: usize, +) { + if inner.routes.contains_key(&incoming) || inner.routes.len() < capacity { + return; + } + let removable = inner.routes.iter().find_map(|(fingerprint, route)| { + (route.refs.is_empty() + && matches!( + route.state, + RouteState::Empty | RouteState::PassThrough { permanent: true } + )) + .then_some((*fingerprint, route.token_digest)) + }); + if let Some((fingerprint, token_digest)) = removable { + inner.routes.remove(&fingerprint); + inner.tokens.remove(&token_digest); + } +} + +#[derive(Default)] +struct RegistryInner { + routes: HashMap, + tokens: HashMap, +} + +struct RouteEntry { + token_digest: TokenDigest, + refs: BTreeMap, + state: RouteState, +} + +impl RouteEntry { + fn new(token_digest: TokenDigest, global_pass_through: bool) -> Self { + Self { + token_digest, + refs: BTreeMap::new(), + state: if global_pass_through { + RouteState::PassThrough { permanent: true } + } else { + RouteState::Empty + }, + } + } + + fn directive_for( + &mut self, + session_id: &McpSessionId, + launch: WorkerLaunch, + retry_after_ms: u64, + ) -> BrokerDirective { + match &self.state { + RouteState::Empty => { + self.state = RouteState::Activating { + owner: session_id.clone(), + launch: launch.clone(), + }; + launch.into_directive() + } + RouteState::Activating { + owner, + launch: active_launch, + } if owner == session_id => active_launch.clone().into_directive(), + RouteState::Activating { .. } + | RouteState::Draining { .. } + | RouteState::Recovering { target: None, .. } => { + BrokerDirective::WaitForWorker { retry_after_ms } + } + RouteState::Ready { target } => BrokerDirective::ReuseWorker { + endpoint: target.endpoint().to_owned(), + }, + RouteState::PassThrough { .. } => BrokerDirective::UsePassThrough, + RouteState::Recovering { + target: Some(target), + .. + } => { + let endpoint = target.endpoint().to_owned(); + self.state = RouteState::Ready { + target: Arc::clone(target), + }; + BrokerDirective::ReuseWorker { endpoint } + } + } + } + + fn current_directive( + &mut self, + session_id: &McpSessionId, + retry_after_ms: u64, + ) -> BrokerDirective { + match &self.state { + RouteState::Empty => BrokerDirective::WaitForWorker { retry_after_ms }, + RouteState::Activating { + owner, + launch: active_launch, + } if owner == session_id => active_launch.clone().into_directive(), + RouteState::Activating { .. } + | RouteState::Draining { .. } + | RouteState::Recovering { target: None, .. } => { + BrokerDirective::WaitForWorker { retry_after_ms } + } + RouteState::Ready { target } => BrokerDirective::ReuseWorker { + endpoint: target.endpoint().to_owned(), + }, + RouteState::PassThrough { .. } => BrokerDirective::UsePassThrough, + RouteState::Recovering { + target: Some(target), + .. + } => { + let endpoint = target.endpoint().to_owned(); + self.state = RouteState::Ready { + target: Arc::clone(target), + }; + BrokerDirective::ReuseWorker { endpoint } + } + } + } + + fn after_reference_removed( + &mut self, + removed_session: &McpSessionId, + drain_deadline_unix_ms: u64, + ) -> ReleaseAction { + if !self.refs.is_empty() { + return self.transfer_owner_if_needed(removed_session); + } + let state = std::mem::replace(&mut self.state, RouteState::Empty); + match state { + RouteState::Activating { launch, .. } => ReleaseAction::CancelActivation { + activation_id: launch.activation_id, + }, + RouteState::Ready { target } => { + self.state = RouteState::Draining { + target: Arc::clone(&target), + deadline_unix_ms: drain_deadline_unix_ms, + }; + ReleaseAction::BeginDrain { + target, + deadline_unix_ms: drain_deadline_unix_ms, + } + } + RouteState::Recovering { + target: Some(target), + .. + } => { + self.state = RouteState::Draining { + target: Arc::clone(&target), + deadline_unix_ms: drain_deadline_unix_ms, + }; + ReleaseAction::BeginDrain { + target, + deadline_unix_ms: drain_deadline_unix_ms, + } + } + RouteState::PassThrough { permanent: true } => { + self.state = RouteState::PassThrough { permanent: true }; + ReleaseAction::NoChange + } + RouteState::Draining { + target, + deadline_unix_ms, + } => { + self.state = RouteState::Draining { + target, + deadline_unix_ms, + }; + ReleaseAction::NoChange + } + RouteState::Empty + | RouteState::PassThrough { permanent: false } + | RouteState::Recovering { target: None, .. } => ReleaseAction::NoChange, + } + } + + fn transfer_owner_if_needed(&mut self, removed_session: &McpSessionId) -> ReleaseAction { + let replacement = self + .refs + .keys() + .next() + .expect("route has live references") + .clone(); + match &mut self.state { + RouteState::Activating { owner, launch } if owner == removed_session => { + *owner = replacement.clone(); + ReleaseAction::TransferActivation { + session_id: replacement, + directive: launch.clone().into_directive(), + } + } + RouteState::Recovering { owner, .. } + if owner.as_ref().is_some_and(|owner| owner == removed_session) => + { + *owner = Some(replacement.clone()); + ReleaseAction::NominateMcp { + session_id: replacement, + } + } + _ => ReleaseAction::NoChange, + } + } +} + +fn validate_binding( + inner: &RegistryInner, + fingerprint: Fingerprint, + token_digest: TokenDigest, +) -> Result<(), RegistryError> { + if inner + .tokens + .get(&token_digest) + .is_some_and(|existing| *existing != fingerprint) + { + return Err(RegistryError::TokenAlreadyBound); + } + if inner + .routes + .get(&fingerprint) + .is_some_and(|existing| !existing.token_digest.matches(&token_digest)) + { + return Err(RegistryError::FingerprintTokenMismatch); + } + Ok(()) +} + +fn validate_capacity( + inner: &RegistryInner, + fingerprint: Fingerprint, + route_capacity: usize, +) -> Result<(), RegistryError> { + if !inner.routes.contains_key(&fingerprint) && inner.routes.len() >= route_capacity { + return Err(RegistryError::RouteCapacityReached); + } + Ok(()) +} + +/// A credential-free lifecycle snapshot. +#[cfg(test)] +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct RouteSnapshot { + pub(crate) state: RouteStateKind, + pub(crate) reference_count: usize, + pub(crate) launch_owner: Option, + pub(crate) endpoint: Option, + pub(crate) in_flight: usize, +} + +/// An activation grant whose signed launch deadline elapsed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ExpiredActivation { + pub(crate) fingerprint: Fingerprint, + pub(crate) activation_id: String, +} + +/// Work required after releasing or expiring an MCP reference. +#[derive(Debug)] +pub(crate) enum ReleaseAction { + NoChange, + CancelActivation { + activation_id: String, + }, + BeginDrain { + target: Arc, + deadline_unix_ms: u64, + }, + TransferActivation { + session_id: McpSessionId, + directive: BrokerDirective, + }, + NominateMcp { + session_id: McpSessionId, + }, +} + +/// Work required after a ready worker disconnects. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum WorkerFailureAction { + NominateMcp { session_id: McpSessionId }, + RouteEmpty, +} + +/// Exact route generation authorized to attempt worker recovery. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum RecoveryPermit { + Activating { activation_id: String }, + ExistingWorker { worker_id: String, recovering: bool }, +} + +/// Work required when bounded daemon-restart recovery ends. +#[cfg(test)] +#[derive(Debug)] +#[allow(dead_code)] +pub(crate) enum RecoveryAction { + WorkerRecovered, + BeginDrain { + target: Arc, + deadline_unix_ms: u64, + }, + NominateMcp { + session_id: McpSessionId, + }, + RouteEmpty, +} + +/// Result of completing a worker drain. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum DrainCompletion { + RouteEmpty, + ActivationRequired { session_id: McpSessionId }, +} + +/// Registry mutation failures. +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum RegistryError { + #[error("the route token is already bound to a different user-machine fingerprint")] + TokenAlreadyBound, + #[error("the user-machine fingerprint is already bound to a different route token")] + FingerprintTokenMismatch, + #[error("the broker route does not exist")] + UnknownRoute, + #[error("the broker route binding capacity has been reached")] + RouteCapacityReached, + #[error("the MCP session is not registered on this route")] + UnknownMcpSession, + #[error("the route MCP reference capacity has been reached")] + McpReferenceCapacityReached, + #[error("the route no longer has a live MCP reference")] + NoLiveMcpReferences, + #[error("the worker activation ID does not match the active generation")] + ActivationMismatch, + #[error("the worker ID does not match the active generation")] + WorkerMismatch, + #[error("this MCP session is not the nominated launch owner")] + NotLaunchOwner, + #[error("this route generation is not eligible for worker recovery")] + RecoveryNotAuthorized, + #[error("the route generation changed during worker recovery")] + RecoveryGenerationChanged, + #[error("the route is {actual:?}; expected {expected:?}")] + InvalidState { + expected: RouteStateKind, + actual: RouteStateKind, + }, + #[error("worker recovery is still within its grace period")] + #[cfg(test)] + RecoveryInProgress, + #[error("the worker still has in-flight requests before its drain deadline")] + DrainInProgress, +} + +/// Request-route resolution failures. +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum ResolveError { + #[error("the route token is unknown")] + UnknownToken, + #[error("the authenticated route is not ready: {0:?}")] + Unavailable(RouteStateKind), +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/registry_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/broker/server.rs b/crates/cli/src/daemon/broker/server.rs new file mode 100644 index 000000000..4fb97174b --- /dev/null +++ b/crates/cli/src/daemon/broker/server.rs @@ -0,0 +1,2466 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Public daemon listener, authenticated broker control plane, and streaming data plane. + +use std::collections::HashMap; +use std::error::Error as StdError; +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::path::Path; +use std::pin::Pin; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use axum::body::Body; +use axum::extract::{ConnectInfo, DefaultBodyLimit, State}; +use axum::http::header::{AUTHORIZATION, CONTENT_TYPE, RETRY_AFTER}; +use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, Request, Response, StatusCode, Uri}; +use axum::middleware::{Next, from_fn_with_state}; +use axum::response::IntoResponse; +use axum::routing::post; +use axum::serve::ListenerExt; +use axum::{Json, Router}; +use base64::Engine; +use bytes::Bytes; +use hyper::body::Body as HttpBody; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use hyper_util::server::conn::auto::Builder as ConnectionBuilder; +use hyper_util::service::TowerToHyperService; +use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer}; +use serde::Serialize; +use serde_json::json; +use tokio::net::TcpListener; +use tokio::sync::Semaphore; + +use super::lifecycle::{McpSessionId, ResolvedTarget, WorkerRequest, WorkerTarget}; +use super::registry::{ + ExpiredActivation, McpRegistration, RecoveryPermit, Registry, RegistryError, ReleaseAction, + ResolveError, WorkerFailureAction, +}; +use crate::configuration::GatewayConfig; +use crate::daemon::ServerOptions; +use crate::daemon::common::address::{daemon_url, validate_bind_ip}; +use crate::daemon::common::control::{ + ACTIVATION_LIFETIME_MS, ActivationFailedPayload, CHALLENGE_LIFETIME_MS, CHALLENGE_PATH, + CLIENT_TOKEN_HEADER, ChallengeRequest, ChallengeResponse, DRAIN_LIFETIME_MS, EmptyPayload, + MAX_CONTROL_BODY_BYTES, MCP_ACTIVATION_FAILED_PATH, MCP_HEARTBEAT_INTERVAL_MS, + MCP_HEARTBEAT_PATH, MCP_LEASE_MS, MCP_REGISTER_PATH, MCP_RELEASE_PATH, McpHeartbeatResponse, + McpRegisterRequest, McpRegisterResponse, RECOVERY_LIFETIME_MS, SessionRequest, + WORKER_DRAIN_PATH, WORKER_HEARTBEAT_INTERVAL_MS, WORKER_HEARTBEAT_PATH, WORKER_LEASE_MS, + WORKER_PROBE_PATH, WORKER_READY_PATH, WORKER_RECOVER_PATH, WORKER_REGISTER_PATH, + WORKER_ROUTE_FAILURE_HEADER, WORKER_TOKEN_HEADER, WorkerDrainRequest, WorkerGenerationGrant, + WorkerHeartbeatPayload, WorkerNetworkHint, WorkerReadyPayload, WorkerRecoverRequest, + WorkerRegisterRequest, WorkerRegisterResponse, now_unix_ms, random_secret, +}; +use crate::daemon::common::identity::{ + ChallengeId, ChallengeRecord, Fingerprint, MachineIdentity, TokenDigest, +}; +use crate::daemon::common::protocol::{ + BrokerDirective, Capabilities, ComponentRole, HandshakeProof, SensitiveString, WorkerLaunch, +}; +use crate::daemon::common::routes::{ProviderRoute, PublicRoute}; +use crate::daemon::common::state::{ + ActiveWorkerGenerations, RouteCredential, load_or_create_daemon_identity, +}; +use crate::daemon::common::transport::{ + PooledClient, RelayBody, box_body, hold_body, pooled_client, prepare_forward_request, + prepare_forward_response, +}; +use crate::daemon::common::worker_tls::WorkerClientPool; +use crate::error::CliError; + +const RESPONSE_HEAD_TIMEOUT: Duration = Duration::from_secs(60); +const MAINTENANCE_INTERVAL: Duration = Duration::from_secs(1); +const MAX_PENDING_CHALLENGES: usize = 512; +const MAX_PENDING_MCP_CHALLENGES: usize = 384; +const MAX_PENDING_WORKER_CHALLENGES: usize = 128; +const MAX_STAGED_WORKER_SESSIONS: usize = 4_096; +const MAX_MCP_CONTROL_SESSIONS: usize = 8_192; +const MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256; +const TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); +// A single peer cannot consume either role's challenge pool, even across a window boundary. +const CHALLENGES_PER_PEER_WINDOW: u32 = 16; +const MAX_CHALLENGE_PEERS: usize = 1_024; +type ChallengePeers = Arc>>; + +struct PendingChallenge { + request: ChallengeRequest, + record: ChallengeRecord, +} + +struct Activation { + fingerprint: Fingerprint, + secret_digest: TokenDigest, + deadline_unix_ms: u64, + consumed: bool, + bind_ip: Ipv4Addr, + port: u16, + advertise_address: Option, +} + +struct McpControlSession { + fingerprint: Fingerprint, + token_digest: TokenDigest, + secret: SensitiveString, + secret_digest: TokenDigest, + lease_expires_at_unix_ms: u64, + last_sequence: u64, + last_request_id: String, + last_heartbeat: Option, + worker_network: WorkerNetworkHint, + released: bool, +} + +#[derive(Clone)] +struct CachedHeartbeat { + sequence: u64, + request_id: String, + response: McpHeartbeatResponse, +} + +struct WorkerControlSession { + fingerprint: Fingerprint, + worker_id: String, + secret: SensitiveString, + secret_digest: TokenDigest, + last_sequence: u64, + last_request_id: String, + next_daemon_sequence: u64, + lease_expires_at_unix_ms: u64, + pending_target: Arc, + publication: WorkerPublication, + published: bool, + generation_grant: WorkerGenerationGrant, +} + +#[derive(Clone)] +enum WorkerPublication { + Activation { activation_id: String }, + Recovery { permit: RecoveryPermit }, +} + +struct DaemonState { + registry: Registry, + identity: MachineIdentity, + descriptor: crate::daemon::common::protocol::ComponentDescriptor, + instance_id: String, + public_origin: String, + config: GatewayConfig, + upstream: PooledClient, + worker_clients: WorkerClientPool, + challenges: Mutex>, + activations: Mutex>, + mcp_sessions: Mutex>, + mcp_heartbeat_serialization: Mutex<()>, + worker_sessions: Mutex>, + pending_directives: Mutex>, + active_worker_generations: ActiveWorkerGenerations, + worker_generation_publication: Mutex<()>, +} + +pub(crate) async fn serve(options: ServerOptions) -> Result<(), CliError> { + validate_bind_ip(options.bind, "daemon")?; + let bind = SocketAddr::new(IpAddr::V4(options.bind), options.port); + let listener = TcpListener::bind(bind).await.map_err(|error| { + CliError::Launch(format!("failed to bind daemon listener {bind}: {error}")) + })?; + let local = listener.local_addr()?; + let public_origin = daemon_origin(&options, local)?; + let resolved = crate::configuration::resolve_server_config(&options.gateway)?; + let state = Arc::new(DaemonState { + registry: Registry::new(options.pass_through), + identity: load_or_create_daemon_identity()?, + descriptor: crate::daemon::common::control::descriptor(ComponentRole::Daemon), + instance_id: uuid::Uuid::now_v7().to_string(), + public_origin, + config: resolved.gateway, + upstream: pooled_client().map_err(|error| CliError::Launch(error.to_string()))?, + worker_clients: WorkerClientPool::new()?, + challenges: Mutex::new(HashMap::new()), + activations: Mutex::new(HashMap::new()), + mcp_sessions: Mutex::new(HashMap::new()), + mcp_heartbeat_serialization: Mutex::new(()), + worker_sessions: Mutex::new(HashMap::new()), + pending_directives: Mutex::new(HashMap::new()), + active_worker_generations: ActiveWorkerGenerations::load()?, + worker_generation_publication: Mutex::new(()), + }); + spawn_maintenance(Arc::clone(&state)); + let app = router(Arc::clone(&state)); + let address = local.to_string(); + log::info!( + target: "nemo_relay.daemon", + event = "daemon_listening", + address = address.as_str(), + public_origin = state.public_origin.as_str(), + pass_through = options.pass_through; + "NeMo Relay daemon is listening" + ); + match (&options.tls_cert, &options.tls_key) { + (Some(certificate), Some(key)) => { + let tls = load_tls_config(certificate, key)?; + serve_tls(listener, app, tls).await + } + (None, None) => axum::serve( + listener.tap_io(|stream| { + let _ = stream.set_nodelay(true); + }), + app.into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(shutdown_signal()) + .await + .map_err(CliError::Io), + _ => Err(CliError::Config( + "--tls-cert and --tls-key must be supplied together".into(), + )), + } +} + +fn router(state: Arc) -> Router { + let peers: ChallengePeers = Arc::new(Mutex::new(HashMap::new())); + let control = Router::new() + .route( + CHALLENGE_PATH, + post(issue_challenge).route_layer(from_fn_with_state(peers, limit_challenges)), + ) + .route(MCP_REGISTER_PATH, post(register_mcp)) + .route(MCP_HEARTBEAT_PATH, post(heartbeat_mcp)) + .route(MCP_RELEASE_PATH, post(release_mcp)) + .route(MCP_ACTIVATION_FAILED_PATH, post(activation_failed)) + .route(WORKER_REGISTER_PATH, post(register_worker)) + .route(WORKER_RECOVER_PATH, post(recover_worker)) + .route(WORKER_READY_PATH, post(ready_worker)) + .route(WORKER_HEARTBEAT_PATH, post(heartbeat_worker)) + .layer(DefaultBodyLimit::max(MAX_CONTROL_BODY_BYTES)); + Router::new() + .merge(control) + .fallback(public_proxy) + .with_state(state) +} + +async fn limit_challenges( + State(peers): State, + request: Request, + next: Next, +) -> Response { + // Use the transport peer, never caller-controlled forwarding headers. In-process services + // without connection metadata share a conservative bucket rather than bypassing admission. + let peer = request + .extensions() + .get::>() + .map_or(IpAddr::V4(Ipv4Addr::UNSPECIFIED), |peer| peer.0.ip()); + let admitted = { + let mut peers = lock(&peers); + let now = tokio::time::Instant::now(); + let window = Duration::from_millis(CHALLENGE_LIFETIME_MS); + peers.retain(|_, (start, _)| now.duration_since(*start) < window); + if !peers.contains_key(&peer) && peers.len() >= MAX_CHALLENGE_PEERS { + false + } else { + let (_, count) = peers.entry(peer).or_insert((now, 0)); + if *count >= CHALLENGES_PER_PEER_WINDOW { + false + } else { + *count += 1; + true + } + } + }; + if !admitted { + let mut response = control_message( + StatusCode::TOO_MANY_REQUESTS, + "authentication challenge rate limit exceeded", + ); + response + .headers_mut() + .insert(RETRY_AFTER, HeaderValue::from_static("15")); + return response; + } + next.run(request).await +} + +async fn issue_challenge( + State(state): State>, + Json(request): Json, +) -> Response { + if let Err(error) = request.initiator.validate() { + return control_error(StatusCode::UNAUTHORIZED, error); + } + if !has_required_transport_capabilities(&request.initiator) { + return control_message( + StatusCode::UPGRADE_REQUIRED, + "component lacks required lossless streaming and trailer capabilities", + ); + } + if request.initiator.role == ComponentRole::Daemon + || request.initiator_public_identity.fingerprint() != request.initiator_fingerprint + || request.initiator_instance_id.is_empty() + || request.initiator_instance_id.len() > 256 + { + return control_message(StatusCode::UNAUTHORIZED, "invalid component identity"); + } + let now = now_unix_ms(); + let record = match ChallengeRecord::generate(now, CHALLENGE_LIFETIME_MS) { + Ok(record) => record, + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + }; + let challenge = record.challenge(); + let signed_request = request.clone(); + let initiator_role = request.initiator.role; + let pending = PendingChallenge { request, record }; + let mut challenges = lock(&state.challenges); + if !reserve_challenge_slot(&mut challenges, now, initiator_role) { + return control_message( + StatusCode::TOO_MANY_REQUESTS, + "too many pending authentication challenges", + ); + } + challenges.insert(challenge.id, pending); + drop(challenges); + let mut response = ChallengeResponse { + daemon: state.descriptor.clone(), + daemon_instance_id: state.instance_id.clone(), + daemon_public_identity: state.identity.public_identity(), + daemon_fingerprint: state.identity.fingerprint(), + challenge, + daemon_challenge_proof: state.identity.sign(b"pending-daemon-challenge"), + }; + let canonical = + match crate::daemon::common::control::daemon_challenge_bytes(&signed_request, &response) { + Ok(canonical) => canonical, + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + }; + response.daemon_challenge_proof = state.identity.sign(&canonical); + Json(response).into_response() +} + +fn reserve_challenge_slot( + challenges: &mut HashMap, + now_unix_ms: u64, + role: ComponentRole, +) -> bool { + challenges.retain(|_, pending| pending.record.challenge().expires_at_unix_ms > now_unix_ms); + let role_limit = match role { + ComponentRole::Mcp => MAX_PENDING_MCP_CHALLENGES, + ComponentRole::Worker => MAX_PENDING_WORKER_CHALLENGES, + ComponentRole::Daemon => return false, + }; + challenges.len() < MAX_PENDING_CHALLENGES + && challenges + .values() + .filter(|pending| pending.request.initiator.role == role) + .count() + < role_limit +} + +fn reserve_worker_session_slot( + sessions: &mut HashMap, + now_unix_ms: u64, + worker_id: &str, + capacity: usize, +) -> bool { + sessions + .retain(|_, session| session.published || session.lease_expires_at_unix_ms > now_unix_ms); + if sessions.contains_key(worker_id) { + return false; + } + sessions + .values() + .filter(|session| !session.published) + .count() + < capacity +} + +async fn register_mcp( + State(state): State>, + headers: HeaderMap, + Json(request): Json, +) -> Response { + let credential = match public_credential(&headers) { + Ok(credential) => credential, + Err(response) => return response, + }; + let transcript = &request.proof.transcript; + if transcript.initiator.role != ComponentRole::Mcp + || transcript.route_token_digest != Some(credential.digest()) + { + return control_message(StatusCode::UNAUTHORIZED, "route credential proof mismatch"); + } + let daemon_proof = match validate_registration(&state, &request.proof) { + Ok(proof) => proof, + Err(response) => return response, + }; + if request + .worker_network + .verify( + &transcript.daemon_target, + &transcript.initiator_instance_id, + &transcript.challenge_id, + &transcript.initiator_fingerprint, + &transcript.initiator_public_identity, + ) + .is_err() + { + return control_message( + StatusCode::UNAUTHORIZED, + "invalid worker network hint proof", + ); + } + let session_id = match McpSessionId::new(transcript.initiator_instance_id.clone()) { + Ok(session_id) => session_id, + Err(error) => return control_error(StatusCode::BAD_REQUEST, error), + }; + let now = now_unix_ms(); + expire_activation_routes(&state, now); + let launch = match fresh_launch(request.worker_network.hint.clone()) { + Ok(launch) => launch, + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + }; + let fresh_session_token = match random_secret(32).and_then(|secret| { + SensitiveString::new(secret).map_err(|error| CliError::Launch(error.to_string())) + }) { + Ok(token) => token, + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + }; + let lease_expires_at_unix_ms = now.saturating_add(MCP_LEASE_MS); + let mut sessions = lock(&state.mcp_sessions); + sessions.retain(|_, session| session.lease_expires_at_unix_ms > now); + if !sessions.contains_key(session_id.as_str()) && sessions.len() >= MAX_MCP_CONTROL_SESSIONS { + return control_message(StatusCode::TOO_MANY_REQUESTS, "too many live MCP sessions"); + } + let (session_token, reuse_session) = match select_mcp_session_token( + &sessions, + session_id.as_str(), + transcript.initiator_fingerprint, + credential.digest(), + request.worker_network.hint.clone(), + now, + fresh_session_token, + ) { + Ok(selection) => selection, + Err(response) => return response, + }; + // Enrollment is open to reachable clients with a valid identity proof. The registry binds + // this credential digest to that fingerprint and rejects attempts to rebind either side. + let directive = match state.registry.register_mcp( + McpRegistration { + fingerprint: transcript.initiator_fingerprint, + token_digest: credential.digest(), + session_id: session_id.clone(), + lease_expires_at_unix_ms, + }, + launch, + ) { + Ok(directive) => directive, + Err(error) => return registry_error(error), + }; + if reuse_session { + let session = sessions + .get_mut(session_id.as_str()) + .expect("a selected reusable MCP session must still exist while locked"); + session.lease_expires_at_unix_ms = lease_expires_at_unix_ms; + } else { + sessions.insert( + session_id.as_str().to_owned(), + McpControlSession { + fingerprint: transcript.initiator_fingerprint, + token_digest: credential.digest(), + secret_digest: TokenDigest::from_token(session_token.expose().as_bytes()), + secret: session_token.clone(), + lease_expires_at_unix_ms, + last_sequence: 0, + last_request_id: String::new(), + last_heartbeat: None, + worker_network: request.worker_network.hint, + released: false, + }, + ); + } + drop(sessions); + if !reuse_session { + lock(&state.pending_directives).remove(session_id.as_str()); + } + remember_activation(&state, transcript.initiator_fingerprint, &directive); + Json(McpRegisterResponse { + daemon_proof, + session_token, + heartbeat_interval_ms: MCP_HEARTBEAT_INTERVAL_MS, + directive, + }) + .into_response() +} + +#[allow(clippy::result_large_err)] +fn select_mcp_session_token( + sessions: &HashMap, + session_id: &str, + fingerprint: Fingerprint, + token_digest: TokenDigest, + worker_network: WorkerNetworkHint, + now_unix_ms: u64, + fresh: SensitiveString, +) -> Result<(SensitiveString, bool), Response> { + let reusable = sessions + .get(session_id) + .filter(|session| session.lease_expires_at_unix_ms > now_unix_ms && !session.released); + if reusable.is_some_and(|session| { + session.fingerprint != fingerprint + || !session.token_digest.matches(&token_digest) + || session.worker_network != worker_network + }) { + return Err(control_message( + StatusCode::UNAUTHORIZED, + "MCP session ID is already bound to another authenticated route", + )); + } + Ok(reusable.map_or((fresh, false), |session| (session.secret.clone(), true))) +} + +async fn heartbeat_mcp( + State(state): State>, + Json(request): Json>, +) -> Response { + // The critical section contains no I/O. Serializing it closes the small window in which a + // concurrent lost-response retry could observe the accepted sequence before its response was + // cached, while leaving the request data plane entirely lock-free. + let _heartbeat_serialization = lock(&state.mcp_heartbeat_serialization); + let lease_expires_at_unix_ms = now_unix_ms().saturating_add(MCP_LEASE_MS); + let authenticated = match authenticate_mcp(&state, &request, Some(lease_expires_at_unix_ms)) { + Ok(authenticated) => authenticated, + Err(response) => return response, + }; + if authenticated.released { + return control_message(StatusCode::UNAUTHORIZED, "MCP session was already released"); + } + if authenticated.duplicate { + return authenticated.cached_heartbeat.map_or_else( + || { + control_message( + StatusCode::CONFLICT, + "duplicate request does not match the cached heartbeat response", + ) + }, + |response| Json(response).into_response(), + ); + } + if let Err(error) = state.registry.renew_mcp( + authenticated.fingerprint, + &authenticated.session_id, + lease_expires_at_unix_ms, + ) { + lock(&state.mcp_sessions).remove(authenticated.session_id.as_str()); + lock(&state.pending_directives).remove(authenticated.session_id.as_str()); + return registry_error(error); + } + let directive = lock(&state.pending_directives) + .remove(authenticated.session_id.as_str()) + .map_or_else( + || { + state + .registry + .current_directive(authenticated.fingerprint, &authenticated.session_id) + .map(Some) + }, + |directive| Ok(Some(directive)), + ); + let response = match directive { + Ok(directive) => McpHeartbeatResponse { directive }, + Err(error) => return registry_error(error), + }; + if let Some(session) = lock(&state.mcp_sessions).get_mut(authenticated.session_id.as_str()) { + session.last_heartbeat = Some(CachedHeartbeat { + sequence: request.sequence, + request_id: request.request_id, + response: response.clone(), + }); + } + Json(response).into_response() +} + +async fn release_mcp( + State(state): State>, + Json(request): Json>, +) -> Response { + let authenticated = match authenticate_mcp(&state, &request, None) { + Ok(authenticated) => authenticated, + Err(response) => return response, + }; + if authenticated.released { + return if authenticated.duplicate { + StatusCode::NO_CONTENT.into_response() + } else { + control_message(StatusCode::UNAUTHORIZED, "MCP session was already released") + }; + } + if authenticated.duplicate { + return StatusCode::NO_CONTENT.into_response(); + } + if let Some(session) = lock(&state.mcp_sessions).get_mut(authenticated.session_id.as_str()) { + session.released = true; + } + lock(&state.pending_directives).remove(authenticated.session_id.as_str()); + let action = match state.registry.release_mcp( + authenticated.fingerprint, + &authenticated.session_id, + now_unix_ms().saturating_add(DRAIN_LIFETIME_MS), + ) { + Ok(action) => action, + Err(RegistryError::UnknownRoute | RegistryError::UnknownMcpSession) => { + return StatusCode::NO_CONTENT.into_response(); + } + Err(error) => return registry_error(error), + }; + handle_release_action(Arc::clone(&state), authenticated.fingerprint, action); + StatusCode::NO_CONTENT.into_response() +} + +async fn activation_failed( + State(state): State>, + Json(request): Json>, +) -> Response { + if request.payload.activation_id.len() > 128 || request.payload.reason.len() > 2_048 { + return control_message( + StatusCode::BAD_REQUEST, + "activation failure payload is too large", + ); + } + let authenticated = match authenticate_mcp(&state, &request, None) { + Ok(authenticated) => authenticated, + Err(response) => return response, + }; + if authenticated.released { + return control_message(StatusCode::UNAUTHORIZED, "MCP session was already released"); + } + if authenticated.duplicate { + return StatusCode::NO_CONTENT.into_response(); + } + match state + .registry + .mark_activation_failed(authenticated.fingerprint, &request.payload.activation_id) + { + Ok(()) => { + revoke_activation(&state, &request.payload.activation_id); + let fingerprint = authenticated.fingerprint.to_string(); + log::error!( + target: "nemo_relay.daemon", + event = "worker_activation_failed", + fingerprint = fingerprint.as_str(), + reason = request.payload.reason.as_str(); + "Worker activation failed; route changed to pass-through" + ); + StatusCode::NO_CONTENT.into_response() + } + Err(error) => registry_error(error), + } +} + +async fn register_worker( + State(state): State>, + Json(request): Json, +) -> Response { + if request.proof.transcript.initiator.role != ComponentRole::Worker + || request.proof.transcript.route_token_digest.is_some() + { + return control_message(StatusCode::UNAUTHORIZED, "invalid worker proof"); + } + let daemon_proof = match validate_registration(&state, &request.proof) { + Ok(proof) => proof, + Err(response) => return response, + }; + let now = now_unix_ms(); + let (activation_fingerprint, replay) = { + let mut activations = lock(&state.activations); + let Some(activation) = activations.get_mut(&request.activation_id) else { + return control_message( + StatusCode::UNAUTHORIZED, + "unknown or consumed activation grant", + ); + }; + if activation.deadline_unix_ms <= now + || activation.fingerprint != request.proof.transcript.initiator_fingerprint + || !activation.secret_digest.matches(&TokenDigest::from_token( + request.activation_token.expose().as_bytes(), + )) + { + return control_message(StatusCode::UNAUTHORIZED, "invalid activation grant"); + } + if !activation_endpoint_matches(&request.endpoint, activation) { + let fingerprint = activation.fingerprint; + drop(activations); + let _ = state + .registry + .mark_activation_failed(fingerprint, &request.activation_id); + revoke_activation(&state, &request.activation_id); + return control_message( + StatusCode::BAD_REQUEST, + "worker endpoint does not match the signed activation policy", + ); + } + if activation.consumed { + (activation.fingerprint, true) + } else { + activation.consumed = true; + (activation.fingerprint, false) + } + }; + if replay { + return replay_worker_registration( + &state, + activation_fingerprint, + &request.worker_id, + &request.endpoint, + request.tls_root_certificate.as_deref(), + Some(&request.activation_id), + None, + daemon_proof, + ) + .unwrap_or_else(|| { + control_message( + StatusCode::UNAUTHORIZED, + "activation grant was consumed by another worker registration", + ) + }); + } + let publication = WorkerPublication::Activation { + activation_id: request.activation_id.clone(), + }; + let response = stage_worker( + &state, + activation_fingerprint, + request.worker_id, + request.endpoint, + request.tls_root_certificate, + None, + publication.clone(), + daemon_proof, + ); + if !response.status().is_success() { + fail_worker_publication(&state, activation_fingerprint, &publication); + } + response +} + +async fn recover_worker( + State(state): State>, + Json(request): Json, +) -> Response { + if request.proof.transcript.initiator.role != ComponentRole::Worker + || request.proof.transcript.route_token_digest.is_some() + { + return control_message(StatusCode::UNAUTHORIZED, "invalid worker proof"); + } + let daemon_proof = match validate_registration(&state, &request.proof) { + Ok(proof) => proof, + Err(response) => return response, + }; + let fingerprint = request.proof.transcript.initiator_fingerprint; + if request + .generation_grant + .verify( + &request.worker_id, + fingerprint, + &request.endpoint, + request.tls_root_certificate.as_deref(), + &state.identity.public_identity(), + ) + .is_err() + { + return control_message( + StatusCode::UNAUTHORIZED, + "invalid worker recovery generation", + ); + } + tokio::task::spawn_blocking(move || { + recover_worker_after_validation(state, request, daemon_proof, fingerprint) + }) + .await + .unwrap_or_else(|error| { + control_error( + StatusCode::INTERNAL_SERVER_ERROR, + CliError::Launch(format!("worker recovery task failed: {error}")), + ) + }) +} + +fn recover_worker_after_validation( + state: Arc, + request: WorkerRecoverRequest, + daemon_proof: HandshakeProof, + fingerprint: Fingerprint, +) -> Response { + match state + .active_worker_generations + .matches(fingerprint, &request.generation_grant.generation_id) + { + Ok(true) => {} + Ok(false) => { + return control_message( + StatusCode::UNAUTHORIZED, + "invalid worker recovery generation", + ); + } + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + } + if let Some(response) = replay_worker_registration( + &state, + fingerprint, + &request.worker_id, + &request.endpoint, + request.tls_root_certificate.as_deref(), + None, + Some(&request.generation_grant.generation_id), + daemon_proof.clone(), + ) { + return response; + } + let permit = match state + .registry + .authorize_worker_recovery(fingerprint, &request.worker_id) + { + Ok(permit) => permit, + Err(error) => return registry_error(error), + }; + let publication = WorkerPublication::Recovery { permit }; + let generation_id = request.generation_grant.generation_id.clone(); + let response = stage_worker( + &state, + fingerprint, + request.worker_id, + request.endpoint, + request.tls_root_certificate, + Some(request.generation_grant), + publication.clone(), + daemon_proof, + ); + if !response.status().is_success() + && revoke_active_worker_generation(&state, fingerprint, &generation_id) + { + fail_worker_publication(&state, fingerprint, &publication); + } + response +} + +#[allow(clippy::too_many_arguments)] +fn replay_worker_registration( + state: &DaemonState, + fingerprint: Fingerprint, + worker_id: &str, + endpoint: &str, + tls_root_certificate: Option<&str>, + activation_id: Option<&str>, + generation_id: Option<&str>, + daemon_proof: HandshakeProof, +) -> Option> { + let sessions = lock(&state.worker_sessions); + let session = sessions.get(worker_id)?; + let publication_matches = match (&session.publication, activation_id, generation_id) { + ( + WorkerPublication::Activation { + activation_id: staged, + }, + Some(expected), + None, + ) => staged == expected, + (_, None, Some(expected)) => session.generation_grant.generation_id == expected, + _ => false, + }; + if session.fingerprint != fingerprint + || session.pending_target.endpoint() != endpoint + || !publication_matches + || session + .generation_grant + .verify( + worker_id, + fingerprint, + endpoint, + tls_root_certificate, + &state.identity.public_identity(), + ) + .is_err() + { + return None; + } + let data_token = + SensitiveString::new(session.pending_target.session_token().to_owned()).ok()?; + Some( + Json(WorkerRegisterResponse { + daemon_proof, + session_token: session.secret.clone(), + data_token, + heartbeat_interval_ms: WORKER_HEARTBEAT_INTERVAL_MS, + generation_grant: session.generation_grant.clone(), + }) + .into_response(), + ) +} + +#[allow(clippy::too_many_arguments)] +fn stage_worker( + state: &Arc, + fingerprint: Fingerprint, + worker_id: String, + endpoint: String, + tls_root_certificate: Option, + generation_grant: Option, + publication: WorkerPublication, + daemon_proof: HandshakeProof, +) -> Response { + if worker_id.is_empty() + || worker_id.len() > 256 + || endpoint.len() > 2_048 + || validate_worker_endpoint(&endpoint, tls_root_certificate.as_deref()).is_err() + { + return control_message(StatusCode::BAD_REQUEST, "invalid worker endpoint"); + } + let generation_grant = match generation_grant { + Some(grant) => grant, + None => match WorkerGenerationGrant::issue( + &worker_id, + fingerprint, + &endpoint, + tls_root_certificate.as_deref(), + &state.identity, + ) { + Ok(grant) => grant, + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + }, + }; + let control_secret = match random_secret(32).and_then(|secret| { + SensitiveString::new(secret).map_err(|error| CliError::Launch(error.to_string())) + }) { + Ok(secret) => secret, + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + }; + let data_secret = match random_secret(32).and_then(|secret| { + SensitiveString::new(secret).map_err(|error| CliError::Launch(error.to_string())) + }) { + Ok(secret) => secret, + Err(error) => return control_error(StatusCode::INTERNAL_SERVER_ERROR, error), + }; + let mut worker_sessions = lock(&state.worker_sessions); + if !reserve_worker_session_slot( + &mut worker_sessions, + now_unix_ms(), + &worker_id, + MAX_STAGED_WORKER_SESSIONS, + ) { + return control_message( + StatusCode::TOO_MANY_REQUESTS, + "too many staged worker sessions", + ); + } + // Select or construct the process-wide pool only after this worker has won admission. Keeping + // the session lock through the synchronous client construction makes the collision check and + // insertion atomic, without holding it across connection acquisition or network I/O. + let worker_client = match state.worker_clients.client(tls_root_certificate.as_deref()) { + Ok(client) => client, + Err(error) => return control_error(StatusCode::BAD_REQUEST, error), + }; + let target = match WorkerTarget::with_shared_client( + worker_id.clone(), + endpoint, + data_secret.clone(), + worker_client, + ) { + Ok(target) => Arc::new(target), + Err(error) => return control_error(StatusCode::BAD_REQUEST, error), + }; + worker_sessions.insert( + worker_id.clone(), + WorkerControlSession { + fingerprint, + worker_id, + secret_digest: TokenDigest::from_token(control_secret.expose().as_bytes()), + secret: control_secret.clone(), + last_sequence: 0, + last_request_id: String::new(), + next_daemon_sequence: 0, + lease_expires_at_unix_ms: now_unix_ms().saturating_add(WORKER_LEASE_MS), + pending_target: target, + publication, + published: false, + generation_grant: generation_grant.clone(), + }, + ); + drop(worker_sessions); + Json(WorkerRegisterResponse { + daemon_proof, + session_token: control_secret, + data_token: data_secret, + heartbeat_interval_ms: WORKER_HEARTBEAT_INTERVAL_MS, + generation_grant, + }) + .into_response() +} + +async fn ready_worker( + State(state): State>, + Json(request): Json>, +) -> Response { + let candidate = { + let mut sessions = lock(&state.worker_sessions); + let Some(session) = sessions.get_mut(&request.session_id) else { + return control_message(StatusCode::UNAUTHORIZED, "unknown worker session"); + }; + if request.payload.worker_id != session.worker_id { + return control_message(StatusCode::UNAUTHORIZED, "worker identity mismatch"); + } + if let Err(response) = authenticate_sequence( + session.secret_digest, + &mut session.last_sequence, + &mut session.last_request_id, + &request, + ) { + return response; + } + if session.published { + return StatusCode::NO_CONTENT.into_response(); + } + ( + session.fingerprint, + Arc::clone(&session.pending_target), + session.publication.clone(), + session.generation_grant.generation_id.clone(), + ) + }; + let (fingerprint, target, publication, generation_id) = candidate; + if let Err(error) = probe_worker(&target).await { + let fail_route = match &publication { + WorkerPublication::Activation { .. } => true, + WorkerPublication::Recovery { .. } => { + let revoke_state = Arc::clone(&state); + tokio::task::spawn_blocking(move || { + revoke_active_worker_generation(&revoke_state, fingerprint, &generation_id) + }) + .await + .unwrap_or_else(|error| { + log::error!( + target: "nemo_relay.daemon", + event = "worker_generation_revocation_join_failed", + error_kind = if error.is_panic() { "panic" } else { "cancelled" }; + "Worker generation revocation task failed after readiness probe; continuing cleanup" + ); + true + }) + } + }; + if fail_route { + fail_worker_publication(&state, fingerprint, &publication); + } + lock(&state.worker_sessions).remove(target.worker_id()); + return control_error(StatusCode::BAD_GATEWAY, error); + } + tokio::task::spawn_blocking(move || { + publish_ready_worker(state, fingerprint, target, publication, generation_id) + }) + .await + .unwrap_or_else(|error| { + control_error( + StatusCode::INTERNAL_SERVER_ERROR, + CliError::Launch(format!("worker readiness publication task failed: {error}")), + ) + }) +} + +fn publish_ready_worker( + state: Arc, + fingerprint: Fingerprint, + target: Arc, + publication: WorkerPublication, + generation_id: String, +) -> Response { + // This serializes the durable generation update with broker publication so restart recovery + // cannot race a replacement readiness probe. The complete file-backed transaction runs on a + // blocking thread and this section performs no network I/O. + let _generation_publication = lock(&state.worker_generation_publication); + let previous_generation = match &publication { + WorkerPublication::Activation { activation_id } => { + let current = lock(&state.activations) + .get(activation_id) + .is_some_and(|activation| { + activation.fingerprint == fingerprint && activation.consumed + }); + if !current { + lock(&state.worker_sessions).remove(target.worker_id()); + return control_message( + StatusCode::UNAUTHORIZED, + "worker activation is no longer current", + ); + } + match state + .active_worker_generations + .publish(fingerprint, &generation_id) + { + Ok(previous) => previous, + Err(error) => { + fail_worker_publication(&state, fingerprint, &publication); + lock(&state.worker_sessions).remove(target.worker_id()); + return control_error(StatusCode::INTERNAL_SERVER_ERROR, error); + } + } + } + WorkerPublication::Recovery { .. } => { + match state + .active_worker_generations + .matches(fingerprint, &generation_id) + { + Ok(true) => {} + Ok(false) => { + lock(&state.worker_sessions).remove(target.worker_id()); + return control_message( + StatusCode::UNAUTHORIZED, + "worker recovery generation was revoked before readiness", + ); + } + Err(error) => { + fail_worker_publication(&state, fingerprint, &publication); + lock(&state.worker_sessions).remove(target.worker_id()); + return control_error(StatusCode::INTERNAL_SERVER_ERROR, error); + } + } + None + } + }; + let publication_result = match &publication { + WorkerPublication::Activation { activation_id } => state + .registry + .mark_worker_ready(fingerprint, activation_id, Arc::clone(&target)) + .map(|()| Some(activation_id.clone())), + WorkerPublication::Recovery { permit } => { + state + .registry + .publish_recovered_worker(fingerprint, permit, Arc::clone(&target)) + } + }; + let canceled_activation = match publication_result { + Ok(canceled_activation) => canceled_activation, + Err(error) => { + if matches!(&publication, WorkerPublication::Activation { .. }) + && let Err(restore_error) = state.active_worker_generations.restore_if_matches( + fingerprint, + &generation_id, + previous_generation.as_deref(), + ) + { + log::error!( + target: "nemo_relay.daemon", + event = "worker_generation_restore_failed", + error_kind = restore_error.log_kind(); + "Failed to restore durable worker generation after publication race" + ); + } + lock(&state.worker_sessions).remove(target.worker_id()); + return registry_error(error); + } + }; + if let Some(activation_id) = canceled_activation { + revoke_activation(&state, &activation_id); + } + if let Some(session) = lock(&state.worker_sessions).get_mut(target.worker_id()) { + session.published = true; + session.lease_expires_at_unix_ms = now_unix_ms().saturating_add(WORKER_LEASE_MS); + } + StatusCode::NO_CONTENT.into_response() +} + +fn fail_worker_publication( + state: &DaemonState, + fingerprint: Fingerprint, + publication: &WorkerPublication, +) { + let canceled_activation = match publication { + WorkerPublication::Activation { activation_id } => state + .registry + .mark_activation_failed(fingerprint, activation_id) + .ok() + .map(|()| activation_id.clone()), + WorkerPublication::Recovery { .. } => state + .registry + .mark_route_pass_through(fingerprint) + .ok() + .flatten(), + }; + if let Some(activation_id) = canceled_activation { + revoke_activation(state, &activation_id); + } +} + +async fn probe_worker(target: &Arc) -> Result<(), CliError> { + let uri = format!( + "{}{}", + target.endpoint().trim_end_matches('/'), + WORKER_PROBE_PATH + ) + .parse::() + .map_err(|error| CliError::Launch(format!("invalid worker readiness endpoint: {error}")))?; + let request = Request::get(uri) + .header(WORKER_TOKEN_HEADER, target.session_token()) + .body(box_body(http_body_util::Empty::::new()))?; + let response = tokio::time::timeout(Duration::from_secs(2), target.client().request(request)) + .await + .map_err(|_| CliError::Launch("worker readiness probe timed out".into()))? + .map_err(|error| CliError::Launch(format!("worker readiness probe failed: {error}")))?; + if response.status() != StatusCode::NO_CONTENT { + return Err(CliError::Launch(format!( + "worker readiness probe returned HTTP {}", + response.status() + ))); + } + Ok(()) +} + +async fn heartbeat_worker( + State(state): State>, + Json(request): Json>, +) -> Response { + let mut sessions = lock(&state.worker_sessions); + let Some(session) = sessions.get_mut(&request.session_id) else { + return control_message(StatusCode::UNAUTHORIZED, "unknown worker session"); + }; + if request.payload.worker_id != session.worker_id { + return control_message(StatusCode::UNAUTHORIZED, "worker identity mismatch"); + } + match authenticate_sequence( + session.secret_digest, + &mut session.last_sequence, + &mut session.last_request_id, + &request, + ) { + Ok(_) => { + session.lease_expires_at_unix_ms = now_unix_ms().saturating_add(WORKER_LEASE_MS); + StatusCode::NO_CONTENT.into_response() + } + Err(response) => response, + } +} + +async fn public_proxy(state: State>, request: Request) -> Response { + let catalog = + request.method() == Method::GET && matches!(request.uri().path(), "/models" | "/v1/models"); + let mut response = public_proxy_inner(state, request).await; + if catalog { + // Shared catalog URLs are keyed by a private credential, not by their URI. Override + // provider cache policy for GET responses, without changing LLM streaming headers. + response.headers_mut().insert( + axum::http::header::CACHE_CONTROL, + HeaderValue::from_static("no-store"), + ); + } + response +} + +async fn public_proxy_inner( + State(state): State>, + mut request: Request, +) -> Response { + let Some(route) = PublicRoute::from_path(request.uri().path()) else { + return StatusCode::NOT_FOUND.into_response(); + }; + let credential = match public_credential(request.headers()) { + Ok(credential) => credential, + Err(response) => return response, + }; + strip_public_relay_headers(request.headers_mut(), route); + if responses_websocket_probe(&request) { + return StatusCode::UPGRADE_REQUIRED.into_response(); + } + if !public_method_allowed(request.method(), request.uri().path()) { + return StatusCode::METHOD_NOT_ALLOWED.into_response(); + } + let target = match state.registry.resolve_target(&credential.digest()) { + Ok(target) => target, + Err(ResolveError::UnknownToken) => { + return control_message(StatusCode::UNAUTHORIZED, "invalid route credential"); + } + Err(ResolveError::Unavailable(_)) => return unavailable_response(), + }; + match (target, route) { + (ResolvedTarget::PassThrough, PublicRoute::Hook(hook)) => { + let mut response = Response::new(Body::from(hook.pass_through_body())); + response + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + response + } + (ResolvedTarget::PassThrough, PublicRoute::Provider(provider)) => { + forward_to_provider(&state, request, provider).await + } + (ResolvedTarget::Worker(worker), _) => { + forward_to_worker(Arc::clone(&state), request, worker).await + } + } +} + +fn responses_websocket_probe(request: &Request) -> bool { + request.method() == http::Method::GET + && matches!( + request.uri().path(), + "/responses" | "/v1/responses" | "/backend-api/codex/responses" + ) + && request + .headers() + .get(axum::http::header::UPGRADE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.eq_ignore_ascii_case("websocket")) +} + +fn public_method_allowed(method: &Method, path: &str) -> bool { + if matches!(path, "/models" | "/v1/models") { + method == Method::GET + } else { + method == Method::POST + } +} + +fn strip_public_relay_headers(headers: &mut HeaderMap, route: PublicRoute) { + let keep_named_upstream = matches!(route, PublicRoute::Provider(_)); + let private_names = headers + .keys() + .filter(|name| { + name.as_str().starts_with("x-nemo-relay-") + && name.as_str() != CLIENT_TOKEN_HEADER + && !(keep_named_upstream + && name.as_str() == crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER) + }) + .cloned() + .collect::>(); + for name in private_names { + headers.remove(name); + } +} + +async fn forward_to_provider( + state: &DaemonState, + mut request: Request, + route: ProviderRoute, +) -> Response { + let allow_environment_provider_auth = + crate::gateway::daemon_allows_environment_provider_auth(request.headers()); + let path_and_query = request + .uri() + .path_and_query() + .map_or("/", |value| value.as_str()); + let destination = match crate::gateway::daemon_provider_upstream_url( + request.headers(), + path_and_query, + &state.config, + ) { + Ok(Some(destination)) => destination, + Ok(None) => route.upstream_url(&state.config, path_and_query), + Err(error) => return error.into_response(), + }; + if let Some(aligned) = crate::gateway::daemon_provider_forward_headers( + request.headers(), + request.uri().path(), + &state.config, + ) { + *request.headers_mut() = aligned; + } + request + .headers_mut() + .remove(crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER); + if allow_environment_provider_auth { + inject_provider_auth(request.headers_mut(), route, &state.config); + } + forward(&state.upstream, request, &destination, None, None) + .await + .response +} + +fn inject_provider_auth(headers: &mut HeaderMap, route: ProviderRoute, config: &GatewayConfig) { + if crate::provider_auth::has_provider_credential(headers) { + return; + } + let configured = match route { + ProviderRoute::OpenAi => config.openai_auth_header.as_deref(), + ProviderRoute::Anthropic => config.anthropic_auth_header.as_deref(), + }; + if let Some(configured) = configured.and_then(|value| HeaderValue::from_str(value).ok()) { + headers.insert(AUTHORIZATION, configured); + return; + } + let (name, value) = match route { + ProviderRoute::OpenAi => { + let Some(key) = nonempty_environment("OPENAI_API_KEY") else { + return; + }; + (AUTHORIZATION, format!("Bearer {key}")) + } + ProviderRoute::Anthropic => { + let Some(key) = nonempty_environment("ANTHROPIC_API_KEY") else { + return; + }; + (HeaderName::from_static("x-api-key"), key) + } + }; + if let Ok(value) = HeaderValue::from_str(&value) { + headers.insert(name, value); + } +} + +fn nonempty_environment(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +async fn forward_to_worker( + state: Arc, + request: Request, + worker: WorkerRequest, +) -> Response { + let fingerprint = worker.fingerprint(); + let worker_id = worker.target().worker_id().to_owned(); + let path_and_query = request + .uri() + .path_and_query() + .map_or("/", |value| value.as_str()); + let destination = format!( + "{}{}", + worker.target().endpoint().trim_end_matches('/'), + path_and_query + ); + let token = worker.session_token().to_owned(); + let client = worker.target().client().clone(); + let mut outcome = forward( + &client, + request, + &destination, + Some((HeaderName::from_static(WORKER_TOKEN_HEADER), token)), + Some(worker), + ) + .await; + let route_failure = take_worker_route_failure(&mut outcome.response); + if outcome.communication_failure || route_failure { + handle_worker_communication_failure(&state, fingerprint, &worker_id); + return outcome.response; + } + let (parts, body) = outcome.response.into_parts(); + let observed = ErrorObservedBody { + body, + on_error: Some(move || { + handle_worker_communication_failure(&state, fingerprint, &worker_id); + }), + }; + Response::from_parts(parts, Body::new(observed)) +} + +fn take_worker_route_failure(response: &mut Response) -> bool { + response + .headers_mut() + .remove(WORKER_ROUTE_FAILURE_HEADER) + .is_some() +} + +struct ErrorObservedBody { + body: B, + on_error: Option, +} + +impl HttpBody for ErrorObservedBody +where + B: HttpBody + Unpin, + F: FnOnce() + Unpin, +{ + type Data = Bytes; + type Error = B::Error; + + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let frame = Pin::new(&mut self.body).poll_frame(context); + if matches!(frame, Poll::Ready(Some(Err(_)))) + && let Some(on_error) = self.on_error.take() + { + on_error(); + } + frame + } + + fn is_end_stream(&self) -> bool { + self.body.is_end_stream() + } + + fn size_hint(&self) -> hyper::body::SizeHint { + self.body.size_hint() + } +} + +struct ForwardOutcome { + response: Response, + communication_failure: bool, +} + +impl ForwardOutcome { + fn response(response: Response) -> Self { + Self { + response, + communication_failure: false, + } + } + + fn communication_failure(response: Response) -> Self { + Self { + response, + communication_failure: true, + } + } +} + +async fn forward( + client: &PooledClient, + request: Request, + destination: &str, + authentication: Option<(HeaderName, String)>, + hold: Option, +) -> ForwardOutcome { + let destination = match destination.parse::() { + Ok(destination) => destination, + Err(_) => { + return ForwardOutcome::response(control_message( + StatusCode::BAD_GATEWAY, + "invalid upstream destination", + )); + } + }; + let strip = [ + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderName::from_static(WORKER_TOKEN_HEADER), + ]; + let mut request = match prepare_forward_request(request, destination, &strip) { + Ok(request) => request.map(box_body), + Err(error) => { + return ForwardOutcome::response(control_error(StatusCode::BAD_REQUEST, error)); + } + }; + if let Some((name, value)) = authentication { + let Ok(value) = HeaderValue::from_str(&value) else { + return ForwardOutcome::response(control_message( + StatusCode::INTERNAL_SERVER_ERROR, + "invalid worker credential", + )); + }; + request.headers_mut().insert(name, value); + } + let response = match tokio::time::timeout(RESPONSE_HEAD_TIMEOUT, client.request(request)).await + { + Ok(Ok(response)) => response, + Ok(Err(error)) => { + let response = control_error(StatusCode::BAD_GATEWAY, &error); + return if is_request_body_failure(&error) { + ForwardOutcome::response(response) + } else { + ForwardOutcome::communication_failure(response) + }; + } + Err(_) => { + return ForwardOutcome::communication_failure(control_message( + StatusCode::GATEWAY_TIMEOUT, + "response-head timeout", + )); + } + }; + let response = match prepare_forward_response(response, &strip) { + Ok(response) => response, + Err(error) => { + return ForwardOutcome::communication_failure(control_error( + StatusCode::BAD_GATEWAY, + error, + )); + } + }; + let (parts, body) = response.into_parts(); + let body: RelayBody = match hold { + Some(hold) => hold_body(body, hold), + None => box_body(body), + }; + ForwardOutcome::response(Response::from_parts(parts, Body::new(body))) +} + +fn is_request_body_failure(error: &(dyn StdError + 'static)) -> bool { + let mut source = Some(error); + while let Some(current) = source { + if current + .downcast_ref::() + .is_some_and(|error| error.is_user() || error.is_body_write_aborted()) + { + return true; + } + source = current.source(); + } + false +} + +fn handle_worker_communication_failure( + state: &Arc, + fingerprint: Fingerprint, + worker_id: &str, +) { + // Remove the session and copy its generation before scheduling durable work. No filesystem + // operation or publication mutex acquisition may happen on the body-polling thread. + let generation = lock(&state.worker_sessions) + .remove(worker_id) + .map(|session| (session.fingerprint, session.generation_grant.generation_id)); + if let Some((fingerprint, generation_id)) = generation { + let state = Arc::clone(state); + tokio::spawn(async move { + if let Err(error) = tokio::task::spawn_blocking(move || { + revoke_active_worker_generation(&state, fingerprint, &generation_id); + }) + .await + { + log::error!( + target: "nemo_relay.daemon", + event = "worker_generation_revocation_join_failed", + error_kind = if error.is_panic() { "panic" } else { "cancelled" }; + "Worker generation revocation task failed after communication loss" + ); + } + }); + } + let Ok(canceled_activation) = state + .registry + .mark_worker_communication_failed(fingerprint, worker_id) + else { + return; + }; + if let Some(activation_id) = canceled_activation { + revoke_activation(state, &activation_id); + } + let fingerprint = fingerprint.to_string(); + log::error!( + target: "nemo_relay.daemon", + event = "worker_communication_failed", + fingerprint = fingerprint.as_str(), + worker_id = worker_id; + "Worker communication failed; route changed to pass-through" + ); +} + +struct AuthenticatedMcp { + fingerprint: Fingerprint, + session_id: McpSessionId, + duplicate: bool, + cached_heartbeat: Option, + released: bool, +} + +#[allow(clippy::result_large_err)] +fn authenticate_mcp( + state: &DaemonState, + request: &SessionRequest, + renew_lease_expires_at_unix_ms: Option, +) -> Result> { + let mut sessions = lock(&state.mcp_sessions); + let Some(session) = sessions.get_mut(&request.session_id) else { + return Err(control_message( + StatusCode::UNAUTHORIZED, + "unknown MCP session", + )); + }; + let duplicate = authenticate_sequence( + session.secret_digest, + &mut session.last_sequence, + &mut session.last_request_id, + request, + )?; + if !session.released + && !duplicate + && let Some(lease_expires_at_unix_ms) = renew_lease_expires_at_unix_ms + { + session.lease_expires_at_unix_ms = lease_expires_at_unix_ms; + } + let cached_heartbeat = cached_heartbeat_response(session, request, duplicate); + let session_id = McpSessionId::new(request.session_id.clone()) + .map_err(|error| control_error(StatusCode::BAD_REQUEST, error))?; + Ok(AuthenticatedMcp { + fingerprint: session.fingerprint, + session_id, + duplicate, + cached_heartbeat, + released: session.released, + }) +} + +fn cached_heartbeat_response( + session: &McpControlSession, + request: &SessionRequest, + duplicate: bool, +) -> Option { + duplicate + .then_some(session.last_heartbeat.as_ref()) + .flatten() + .filter(|cached| { + cached.sequence == request.sequence && cached.request_id == request.request_id + }) + .map(|cached| cached.response.clone()) +} + +#[allow(clippy::result_large_err)] +fn authenticate_sequence( + expected_secret: TokenDigest, + last_sequence: &mut u64, + last_request_id: &mut String, + request: &SessionRequest, +) -> Result> { + if !expected_secret.matches(&TokenDigest::from_token( + request.session_token.expose().as_bytes(), + )) || !request.validate_payload_hash() + || request.request_id.is_empty() + || request.request_id.len() > 128 + { + return Err(control_message( + StatusCode::UNAUTHORIZED, + "invalid control message authentication", + )); + } + if request.sequence == *last_sequence && request.request_id == *last_request_id { + return Ok(true); + } + if request.sequence != last_sequence.saturating_add(1) { + return Err(control_message( + StatusCode::CONFLICT, + "control sequence is stale or out of order", + )); + } + *last_sequence = request.sequence; + *last_request_id = request.request_id.clone(); + Ok(false) +} + +#[allow(clippy::result_large_err)] +fn validate_registration( + state: &DaemonState, + proof: &crate::daemon::common::control::RegistrationProof, +) -> Result> { + let transcript = &proof.transcript; + if !has_required_transport_capabilities(&transcript.initiator) { + return Err(control_message( + StatusCode::UPGRADE_REQUIRED, + "component lacks required lossless streaming and trailer capabilities", + )); + } + let pending = lock(&state.challenges).remove(&transcript.challenge_id); + let Some(mut pending) = pending else { + return Err(control_message( + StatusCode::UNAUTHORIZED, + "unknown, expired, or replayed challenge", + )); + }; + if let Err(error) = pending + .record + .consume(&transcript.challenge_id, now_unix_ms()) + { + return Err(control_error(StatusCode::UNAUTHORIZED, error)); + } + let request = pending.request; + let selected = request + .initiator + .protocol + .negotiate(state.descriptor.protocol) + .map_err(|error| control_error(StatusCode::UNAUTHORIZED, error))?; + if transcript.daemon_target != state.public_origin + || transcript.initiator != request.initiator + || transcript.responder != state.descriptor + || transcript.initiator_instance_id != request.initiator_instance_id + || transcript.responder_instance_id != state.instance_id + || transcript.selected_protocol != selected + || transcript.initiator_public_identity != request.initiator_public_identity + || transcript.responder_public_identity != state.identity.public_identity() + || transcript.initiator_fingerprint != request.initiator_fingerprint + || transcript.responder_fingerprint != state.identity.fingerprint() + || transcript.initiator_nonce != request.initiator_nonce + || transcript.responder_nonce != pending.record.challenge().nonce + || proof.initiator_proof.signer != request.initiator.role + { + return Err(control_message( + StatusCode::UNAUTHORIZED, + "handshake transcript does not match the issued challenge", + )); + } + transcript + .verify(&proof.initiator_proof) + .map_err(|error| control_error(StatusCode::UNAUTHORIZED, error))?; + transcript + .sign(ComponentRole::Daemon, &state.identity) + .map_err(|error| control_error(StatusCode::INTERNAL_SERVER_ERROR, error)) +} + +fn has_required_transport_capabilities( + descriptor: &crate::daemon::common::protocol::ComponentDescriptor, +) -> bool { + descriptor + .capabilities + .includes(&Capabilities::streaming_transport()) +} + +fn fresh_launch(worker_network: WorkerNetworkHint) -> Result { + worker_network.validate()?; + let now = now_unix_ms(); + let loopback = worker_network.is_loopback(); + Ok(WorkerLaunch { + activation_id: random_secret(16)?, + activation_token: SensitiveString::new(random_secret(32)?) + .map_err(|error| CliError::Launch(error.to_string()))?, + deadline_unix_ms: now.saturating_add(ACTIVATION_LIFETIME_MS), + bind_ip: if loopback { + Ipv4Addr::LOCALHOST + } else { + Ipv4Addr::UNSPECIFIED + }, + port: worker_network.port.unwrap_or(0), + advertise_address: (!loopback).then_some(worker_network.advertised_host), + }) +} + +fn remember_activation(state: &DaemonState, fingerprint: Fingerprint, directive: &BrokerDirective) { + if let BrokerDirective::LaunchWorker { + activation_id, + activation_token, + deadline_unix_ms, + bind_ip, + port, + advertise_address, + .. + } = directive + { + lock(&state.activations) + .entry(activation_id.clone()) + .or_insert_with(|| Activation { + fingerprint, + secret_digest: TokenDigest::from_token(activation_token.expose().as_bytes()), + deadline_unix_ms: *deadline_unix_ms, + consumed: false, + bind_ip: *bind_ip, + port: *port, + advertise_address: advertise_address.clone(), + }); + } +} + +fn revoke_activation(state: &DaemonState, activation_id: &str) { + lock(&state.activations).remove(activation_id); + lock(&state.pending_directives).retain(|_, directive| { + !matches!( + directive, + BrokerDirective::LaunchWorker { + activation_id: pending, + .. + } if pending == activation_id + ) + }); +} + +fn expire_activation_routes(state: &DaemonState, now_unix_ms: u64) { + for ExpiredActivation { + fingerprint, + activation_id, + } in state.registry.expire_activations(now_unix_ms) + { + revoke_activation(state, &activation_id); + let fingerprint = fingerprint.to_string(); + log::error!( + target: "nemo_relay.daemon", + event = "worker_activation_expired", + fingerprint = fingerprint.as_str(); + "Worker activation expired; route changed to pass-through" + ); + } +} + +fn handle_release_action(state: Arc, fingerprint: Fingerprint, action: ReleaseAction) { + match action { + ReleaseAction::NoChange => {} + ReleaseAction::CancelActivation { activation_id } => { + revoke_activation(&state, &activation_id); + } + ReleaseAction::BeginDrain { + target, + deadline_unix_ms, + } => { + tokio::spawn(async move { + let revoke_state = Arc::clone(&state); + let worker_id = target.worker_id().to_owned(); + if let Err(error) = tokio::task::spawn_blocking(move || { + revoke_worker_generation(&revoke_state, &worker_id); + }) + .await + { + log::error!( + target: "nemo_relay.daemon", + event = "worker_generation_revocation_join_failed", + error_kind = if error.is_panic() { "panic" } else { "cancelled" }; + "Worker generation revocation task failed during drain; continuing cleanup" + ); + } + request_worker_drain(&state, &target, deadline_unix_ms).await; + loop { + let now = now_unix_ms(); + if target.in_flight() == 0 || now >= deadline_unix_ms { + let _ = state.registry.finish_draining(fingerprint, now); + lock(&state.worker_sessions).remove(target.worker_id()); + break; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }); + } + ReleaseAction::TransferActivation { + session_id, + directive, + } => { + lock(&state.pending_directives).insert(session_id.as_str().to_owned(), directive); + } + ReleaseAction::NominateMcp { session_id } => { + nominate_relaunch(&state, fingerprint, session_id); + } + } +} + +fn revoke_worker_generation(state: &DaemonState, worker_id: &str) { + let generation = worker_generation(state, worker_id); + if let Some((fingerprint, generation_id)) = generation { + revoke_active_worker_generation(state, fingerprint, &generation_id); + } +} + +fn worker_generation(state: &DaemonState, worker_id: &str) -> Option<(Fingerprint, String)> { + // End the session-lock scope before callers acquire the publication lock. + lock(&state.worker_sessions).get(worker_id).map(|session| { + ( + session.fingerprint, + session.generation_grant.generation_id.clone(), + ) + }) +} + +fn revoke_active_worker_generation( + state: &DaemonState, + fingerprint: Fingerprint, + generation_id: &str, +) -> bool { + let _generation_publication = lock(&state.worker_generation_publication); + match state + .active_worker_generations + .revoke_if_matches(fingerprint, generation_id) + { + Ok(revoked) => revoked, + Err(error) => { + let fingerprint = fingerprint.to_string(); + log::error!( + target: "nemo_relay.daemon", + event = "worker_generation_revocation_failed", + fingerprint = fingerprint.as_str(), + error_kind = error.log_kind(); + "Failed to durably revoke worker generation" + ); + true + } + } +} + +fn nominate_relaunch(state: &Arc, fingerprint: Fingerprint, session_id: McpSessionId) { + let worker_network = lock(&state.mcp_sessions) + .get(session_id.as_str()) + .filter(|session| !session.released) + .map(|session| session.worker_network.clone()); + let Some(worker_network) = worker_network else { + return; + }; + let Ok(launch) = fresh_launch(worker_network) else { + return; + }; + let Ok(directive) = state + .registry + .begin_relaunch(fingerprint, &session_id, launch) + else { + return; + }; + remember_activation(state, fingerprint, &directive); + lock(&state.pending_directives).insert(session_id.as_str().to_owned(), directive); +} + +async fn request_worker_drain( + state: &DaemonState, + target: &Arc, + deadline_unix_ms: u64, +) { + let request = { + let mut sessions = lock(&state.worker_sessions); + let Some(session) = sessions.get_mut(target.worker_id()) else { + return; + }; + let Some(sequence) = session.next_daemon_sequence.checked_add(1) else { + return; + }; + session.next_daemon_sequence = sequence; + SessionRequest::new( + target.worker_id().to_owned(), + session.secret.clone(), + sequence, + WorkerDrainRequest { + worker_id: target.worker_id().to_owned(), + deadline_unix_ms, + timeout_ms: Some( + deadline_unix_ms + .saturating_sub(now_unix_ms()) + .min(DRAIN_LIFETIME_MS), + ), + }, + ) + }; + let Ok(request) = request else { + return; + }; + let uri = format!( + "{}{}", + target.endpoint().trim_end_matches('/'), + WORKER_DRAIN_PATH + ); + let Ok(uri) = uri.parse::() else { + return; + }; + let payload = match serde_json::to_vec(&request) { + Ok(payload) => payload, + Err(_) => return, + }; + let payload = Bytes::from(payload); + for _ in 0..2 { + let request = match Request::post(uri.clone()) + .header(CONTENT_TYPE, "application/json") + .header(WORKER_TOKEN_HEADER, target.session_token()) + .body(box_body(http_body_util::Full::new(payload.clone()))) + { + Ok(request) => request, + Err(_) => return, + }; + match tokio::time::timeout(Duration::from_secs(2), target.client().request(request)).await { + Ok(Ok(response)) if response.status() == StatusCode::NO_CONTENT => return, + _ => {} + } + } +} + +fn spawn_maintenance(state: Arc) { + tokio::spawn(async move { + let mut interval = tokio::time::interval(MAINTENANCE_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + interval.tick().await; + let now = now_unix_ms(); + expire_activation_routes(&state, now); + lock(&state.activations).retain(|_, activation| activation.deadline_unix_ms > now); + let actions = state + .registry + .expire_mcp_leases(now, now.saturating_add(DRAIN_LIFETIME_MS)); + for (fingerprint, action) in actions { + handle_release_action(Arc::clone(&state), fingerprint, action); + } + prune_expired_mcp_control_state( + &mut lock(&state.mcp_sessions), + &mut lock(&state.pending_directives), + now, + ); + let expired_workers: Vec<_> = { + let mut sessions = lock(&state.worker_sessions); + let expired = sessions + .iter() + .filter(|(_, session)| session.lease_expires_at_unix_ms <= now) + .map(|(id, session)| { + ( + id.clone(), + session.fingerprint, + session.generation_grant.generation_id.clone(), + ) + }) + .collect::>(); + for (id, _, _) in &expired { + sessions.remove(id); + } + expired + }; + for (worker_id, fingerprint, generation_id) in expired_workers { + let revoke_state = Arc::clone(&state); + if let Err(error) = tokio::task::spawn_blocking(move || { + revoke_active_worker_generation(&revoke_state, fingerprint, &generation_id) + }) + .await + { + log::error!( + target: "nemo_relay.daemon", + event = "worker_generation_revocation_join_failed", + error_kind = if error.is_panic() { "panic" } else { "cancelled" }; + "Worker generation revocation task failed during expiry; continuing recovery" + ); + } + if let Ok(WorkerFailureAction::NominateMcp { session_id }) = + state.registry.worker_failed( + fingerprint, + &worker_id, + now.saturating_add(RECOVERY_LIFETIME_MS), + ) + { + nominate_relaunch(&state, fingerprint, session_id); + } + } + } + }); +} + +fn prune_expired_mcp_control_state( + sessions: &mut HashMap, + pending_directives: &mut HashMap, + now_unix_ms: u64, +) { + let expired = sessions + .iter() + .filter(|(_, session)| session.lease_expires_at_unix_ms <= now_unix_ms) + .map(|(session_id, _)| session_id.clone()) + .collect::>(); + for session_id in expired { + sessions.remove(&session_id); + pending_directives.remove(&session_id); + } +} + +#[allow(clippy::result_large_err)] +fn public_credential(headers: &HeaderMap) -> Result> { + let values = headers.get_all(CLIENT_TOKEN_HEADER); + if values.iter().count() != 1 { + return Err(control_message( + StatusCode::UNAUTHORIZED, + "exactly one route credential is required", + )); + } + let value = values + .iter() + .next() + .and_then(|value| value.to_str().ok()) + .ok_or_else(|| control_message(StatusCode::UNAUTHORIZED, "invalid route credential"))?; + RouteCredential::parse(value.to_owned()) + .map_err(|_| control_message(StatusCode::UNAUTHORIZED, "invalid route credential")) +} + +fn validate_worker_endpoint( + endpoint: &str, + tls_root_certificate: Option<&str>, +) -> Result<(), CliError> { + let explicit_port = endpoint + .parse::() + .ok() + .and_then(|uri| uri.authority().and_then(http::uri::Authority::port_u16)); + let url = reqwest::Url::parse(endpoint) + .map_err(|error| CliError::Config(format!("invalid worker endpoint: {error}")))?; + if !matches!(url.scheme(), "http" | "https") + || explicit_port.is_none() + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + || url.host_str() == Some("0.0.0.0") + { + return Err(CliError::Config("invalid worker endpoint origin".into())); + } + let host_is_loopback = url.host_str().is_some_and(|host| { + host.eq_ignore_ascii_case("localhost") + || host + .parse::() + .is_ok_and(|address| address.is_loopback()) + }); + match (url.scheme(), host_is_loopback, tls_root_certificate) { + ("http", true, None) | ("https", _, Some(_)) => Ok(()), + ("http", false, _) => Err(CliError::Config( + "non-loopback worker endpoints must use pinned TLS".into(), + )), + ("http", true, Some(_)) | ("https", _, None) => Err(CliError::Config( + "worker endpoint scheme and TLS trust anchor do not match".into(), + )), + _ => Err(CliError::Config("invalid worker endpoint origin".into())), + } +} + +fn activation_endpoint_matches(endpoint: &str, activation: &Activation) -> bool { + // URL normalization erases explicit default ports. Read the original authority instead. + let Ok(uri) = endpoint.parse::() else { + return false; + }; + let Ok(url) = reqwest::Url::parse(endpoint) else { + return false; + }; + let (Some(host), Some(port)) = (url.host_str(), uri.port_u16()) else { + return false; + }; + let expected_host = activation + .advertise_address + .as_deref() + .unwrap_or("127.0.0.1"); + let expected_scheme = if activation.bind_ip.is_unspecified() { + "https" + } else { + "http" + }; + url.scheme() == expected_scheme + && host == expected_host + && port != 0 + && (activation.port == 0 || activation.port == port) +} + +fn daemon_origin(options: &ServerOptions, local: SocketAddr) -> Result { + if let Some(advertised) = options.advertise_address.as_deref() { + let url = daemon_url(advertised)?; + let host_is_loopback = url.host_str().is_some_and(|host| { + let host = host.trim_matches(['[', ']']); + host.eq_ignore_ascii_case("localhost") + || host + .parse::() + .is_ok_and(|address| address.is_loopback()) + }); + if options.bind == Ipv4Addr::UNSPECIFIED && host_is_loopback { + return Err(CliError::Config( + "a daemon bound to 0.0.0.0 must advertise a concrete reachable host".into(), + )); + } + if options.tls_cert.is_some() && url.scheme() != "https" { + return Err(CliError::Config( + "the advertised daemon URL must use https when native TLS is configured".into(), + )); + } + return Ok(url.as_str().trim_end_matches('/').to_owned()); + } + if options.bind == Ipv4Addr::UNSPECIFIED { + return Err(CliError::Config( + "--advertise-address is required when the daemon binds to 0.0.0.0".into(), + )); + } + let origin = format!( + "{}://{local}", + if options.tls_cert.is_some() { + "https" + } else { + "http" + } + ); + // Signed transcript origins use URL canonicalization, unlike process argv which retains + // explicit ports. Match the client even when the listener uses HTTP :80 or HTTPS :443. + Ok(daemon_url(&origin)? + .as_str() + .trim_end_matches('/') + .to_owned()) +} + +fn load_tls_config( + certificate_path: &Path, + key_path: &Path, +) -> Result, CliError> { + let certificate_pem = crate::filesystem::bounded::read_bounded_regular_file( + certificate_path, + "daemon TLS certificate", + ) + .map_err(CliError::Config)?; + let key_pem = + crate::filesystem::bounded::read_bounded_regular_file(key_path, "daemon TLS private key") + .map_err(CliError::Config)?; + let certificates = decode_pem_blocks(&certificate_pem, "CERTIFICATE")? + .into_iter() + .map(CertificateDer::from) + .collect::>(); + if certificates.is_empty() { + return Err(CliError::Config(format!( + "daemon TLS certificate {} contains no CERTIFICATE blocks", + certificate_path.display() + ))); + } + let mut keys = decode_pem_blocks(&key_pem, "PRIVATE KEY")?; + if keys.len() != 1 { + return Err(CliError::Config(format!( + "daemon TLS key {} must contain exactly one unencrypted PKCS#8 PRIVATE KEY block", + key_path.display() + ))); + } + let key = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(keys.remove(0))); + let _ = rustls::crypto::ring::default_provider().install_default(); + let mut config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certificates, key) + .map_err(|error| CliError::Config(format!("invalid daemon TLS identity: {error}")))?; + config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + Ok(Arc::new(config)) +} + +fn decode_pem_blocks(bytes: &[u8], label: &str) -> Result>, CliError> { + let text = std::str::from_utf8(bytes) + .map_err(|error| CliError::Config(format!("TLS PEM is not UTF-8: {error}")))?; + let begin = format!("-----BEGIN {label}-----"); + let end = format!("-----END {label}-----"); + let mut blocks = Vec::new(); + let mut remainder = text; + while let Some((_, after_begin)) = remainder.split_once(&begin) { + let Some((encoded, after_end)) = after_begin.split_once(&end) else { + return Err(CliError::Config(format!( + "TLS PEM has an unterminated {label} block" + ))); + }; + let compact = encoded + .chars() + .filter(|character| !character.is_ascii_whitespace()) + .collect::(); + let decoded = base64::engine::general_purpose::STANDARD + .decode(compact) + .map_err(|_| CliError::Config(format!("TLS PEM contains invalid {label} base64")))?; + if decoded.is_empty() { + return Err(CliError::Config(format!( + "TLS PEM contains an empty {label} block" + ))); + } + blocks.push(decoded); + remainder = after_end; + } + Ok(blocks) +} + +async fn serve_tls( + listener: TcpListener, + app: Router, + config: Arc, +) -> Result<(), CliError> { + let acceptor = tokio_rustls::TlsAcceptor::from(config); + let handshake_permits = Arc::new(Semaphore::new(MAX_CONCURRENT_TLS_HANDSHAKES)); + let (shutdown_tx, shutdown_rx) = tokio::sync::watch::channel(false); + let mut connections = tokio::task::JoinSet::new(); + let mut shutdown = Box::pin(shutdown_signal()); + loop { + tokio::select! { + biased; + _ = &mut shutdown => break, + _ = connections.join_next(), if !connections.is_empty() => {} + accepted = listener.accept() => { + let (stream, peer) = match accepted { + Ok(accepted) => accepted, + Err(error) => { + log::warn!( + target: "nemo_relay.daemon", + event = "tls_listener_accept_failed"; + "TLS listener accept failed: {error}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + continue; + } + }; + let _ = stream.set_nodelay(true); + let Ok(handshake_permit) = Arc::clone(&handshake_permits).try_acquire_owned() else { + continue; + }; + let acceptor = acceptor.clone(); + let service = app.clone().layer(axum::Extension(ConnectInfo(peer))); + let mut shutdown_rx = shutdown_rx.clone(); + connections.spawn(async move { + let Ok(Ok(stream)) = tokio::time::timeout( + TLS_HANDSHAKE_TIMEOUT, + acceptor.accept(stream), + ).await else { + return; + }; + drop(handshake_permit); + let builder = ConnectionBuilder::new(TokioExecutor::new()); + let connection = builder.serve_connection_with_upgrades( + TokioIo::new(stream), + TowerToHyperService::new(service), + ); + tokio::pin!(connection); + tokio::select! { + _ = &mut connection => {} + changed = shutdown_rx.changed() => { + if changed.is_ok() { + connection.as_mut().graceful_shutdown(); + let _ = connection.await; + } + } + } + }); + } + } + } + let _ = shutdown_tx.send(true); + let drain = async { while connections.join_next().await.is_some() {} }; + let _ = tokio::time::timeout(Duration::from_millis(DRAIN_LIFETIME_MS), drain).await; + Ok(()) +} + +fn unavailable_response() -> Response { + let mut response = control_message(StatusCode::SERVICE_UNAVAILABLE, "route is not ready"); + response + .headers_mut() + .insert(RETRY_AFTER, HeaderValue::from_static("1")); + response +} + +fn registry_error(error: RegistryError) -> Response { + let status = match error { + RegistryError::TokenAlreadyBound | RegistryError::FingerprintTokenMismatch => { + StatusCode::UNAUTHORIZED + } + RegistryError::UnknownRoute | RegistryError::UnknownMcpSession => StatusCode::NOT_FOUND, + RegistryError::RouteCapacityReached | RegistryError::McpReferenceCapacityReached => { + StatusCode::TOO_MANY_REQUESTS + } + RegistryError::ActivationMismatch + | RegistryError::WorkerMismatch + | RegistryError::NoLiveMcpReferences + | RegistryError::NotLaunchOwner + | RegistryError::RecoveryNotAuthorized + | RegistryError::RecoveryGenerationChanged + | RegistryError::InvalidState { .. } + | RegistryError::DrainInProgress => StatusCode::CONFLICT, + #[cfg(test)] + RegistryError::RecoveryInProgress => StatusCode::CONFLICT, + }; + control_error(status, error) +} + +fn control_error(status: StatusCode, error: impl std::fmt::Display) -> Response { + control_message(status, &error.to_string()) +} + +fn control_message(status: StatusCode, message: &str) -> Response { + (status, Json(json!({ "error": { "message": message } }))).into_response() +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(|error| error.into_inner()) +} + +async fn shutdown_signal() { + #[cfg(unix)] + { + let mut terminate = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("installing SIGTERM handler should succeed"); + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = terminate.recv() => {} + } + } + #[cfg(windows)] + { + let mut shutdown = tokio::signal::windows::ctrl_shutdown() + .expect("installing shutdown handler should succeed"); + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = shutdown.recv() => {} + } + } + #[cfg(not(any(unix, windows)))] + let _ = tokio::signal::ctrl_c().await; +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/server_tests.rs"] +mod tests; + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/daemon_worker_e2e_tests.rs"] +mod daemon_worker_e2e_tests; diff --git a/crates/cli/src/daemon/common/address.rs b/crates/cli/src/daemon/common/address.rs new file mode 100644 index 000000000..7542f2dc1 --- /dev/null +++ b/crates/cli/src/daemon/common/address.rs @@ -0,0 +1,163 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + +use reqwest::Url; + +use crate::error::CliError; + +pub(crate) const DEFAULT_DAEMON_PORT: u16 = 47_632; +pub(crate) const DEFAULT_DAEMON_BIND: Ipv4Addr = Ipv4Addr::LOCALHOST; +pub(crate) const DEFAULT_WORKER_BIND: Ipv4Addr = Ipv4Addr::LOCALHOST; + +pub(crate) fn validate_bind_ip(ip: Ipv4Addr, component: &str) -> Result<(), CliError> { + if matches!(ip, Ipv4Addr::LOCALHOST | Ipv4Addr::UNSPECIFIED) { + return Ok(()); + } + Err(CliError::Config(format!( + "{component} bind address must be 127.0.0.1 or 0.0.0.0, got {ip}" + ))) +} + +pub(crate) fn daemon_url(raw: &str) -> Result { + let explicit_port = raw + .parse::() + .ok() + .and_then(|uri| uri.authority().and_then(http::uri::Authority::port_u16)); + let url = Url::parse(raw) + .map_err(|error| CliError::Config(format!("invalid daemon address {raw:?}: {error}")))?; + if !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + || !matches!(url.path(), "" | "/") + { + return Err(CliError::Config( + "daemon address must be an origin URL without credentials, path, query, or fragment" + .into(), + )); + } + if explicit_port.is_none() { + return Err(CliError::Config( + "daemon address must include an explicit port".into(), + )); + } + let host = url + .host_str() + .ok_or_else(|| CliError::Config("daemon address is missing a host".into()))?; + let normalized_host = host.trim_matches(['[', ']']); + if normalized_host + .parse::() + .is_ok_and(|address| address.is_unspecified()) + { + return Err(CliError::Config( + "0.0.0.0 is a bind address and cannot be used as a daemon target".into(), + )); + } + let loopback = normalized_host.eq_ignore_ascii_case("localhost") + || normalized_host + .parse::() + .is_ok_and(|address| address.is_loopback()); + match url.scheme() { + "https" => {} + "http" if loopback => {} + "http" => { + return Err(CliError::Config( + "non-loopback daemon addresses must use https".into(), + )); + } + scheme => { + return Err(CliError::Config(format!( + "daemon address scheme must be http or https, got {scheme}" + ))); + } + } + Ok(url) +} + +/// Canonical process/control address retaining the explicit port required by CLI validation. +/// URL serialization otherwise removes :80 and :443, breaking later validation and worker argv. +pub(crate) fn explicit_daemon_origin(raw: &str) -> Result { + let url = daemon_url(raw)?; + Ok(format!( + "{}://{}:{}", + url.scheme(), + url.host_str().expect("daemon_url validated the host"), + url.port_or_known_default() + .expect("daemon_url validated HTTP(S)") + )) +} + +pub(crate) fn worker_socket(bind: Ipv4Addr, port: Option) -> Result { + validate_bind_ip(bind, "worker")?; + if port == Some(0) { + return Err(CliError::Config( + "an explicitly supplied worker port must be between 1 and 65535; omit --port for automatic allocation" + .into(), + )); + } + Ok(SocketAddr::new(IpAddr::V4(bind), port.unwrap_or(0))) +} + +pub(crate) fn worker_advertised_address( + local: SocketAddr, + configured: Option<&str>, +) -> Result { + if local.ip().is_unspecified() { + let host = configured.ok_or_else(|| { + CliError::Config( + "--advertise-address is required when the worker binds to 0.0.0.0".into(), + ) + })?; + let host = host.trim(); + let normalized_host = host.trim_matches(['[', ']']); + let address = normalized_host.parse::().ok(); + let valid_hostname = !host.contains([':', '[', ']']) + && host.len() <= 253 + && host.split('.').all(|label| { + !label.is_empty() + && label.len() <= 63 + && !label.starts_with('-') + && !label.ends_with('-') + && label + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-') + }); + let valid_ip = address.is_some() + && (host == normalized_host || host == format!("[{normalized_host}]")); + if !valid_hostname && !valid_ip { + return Err(CliError::Config( + "worker advertised address must be a host or IP without a port".into(), + )); + } + if normalized_host.is_empty() + || normalized_host + .parse::() + .is_ok_and(|address| address.is_unspecified()) + { + return Err(CliError::Config( + "worker advertised address must be a concrete host or IP, not 0.0.0.0".into(), + )); + } + return Ok(format_host_port(host, local.port())); + } + if configured.is_some() { + return Err(CliError::Config( + "--advertise-address is only valid with --bind 0.0.0.0".into(), + )); + } + Ok(local.to_string()) +} + +fn format_host_port(host: &str, port: u16) -> String { + if host.starts_with('[') || !host.contains(':') { + format!("{host}:{port}") + } else { + format!("[{host}]:{port}") + } +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/address_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/client.rs b/crates/cli/src/daemon/common/client.rs new file mode 100644 index 000000000..39f5cb2cd --- /dev/null +++ b/crates/cli/src/daemon/common/client.rs @@ -0,0 +1,471 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared authenticated control-plane client used by MCP and worker processes. + +use std::future::Future; +use std::time::Duration; + +use bytes::{Bytes, BytesMut}; +use futures_util::StreamExt; +use reqwest::{Client, Response, StatusCode}; +use serde::Serialize; +use serde::de::DeserializeOwned; + +use super::address::daemon_url; +use super::control::{ + CHALLENGE_PATH, CLIENT_TOKEN_HEADER, ChallengeRequest, ChallengeResponse, RegistrationProof, + descriptor, fresh_nonce, +}; +use super::identity::{MachineIdentity, TokenDigest}; +use super::protocol::{ComponentRole, HandshakeTranscript}; +use super::state::verify_or_store_daemon_pin; +use crate::error::CliError; + +const CONTROL_CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +const CONTROL_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); +const MAX_CONTROL_RESPONSE_BYTES: usize = 256 * 1024; + +/// Limits retries for one idempotent, session-authenticated control request. +/// +/// The request is serialized once before the first attempt. Every retry therefore carries the +/// same session sequence, request ID, payload hash, and JSON bytes. +#[derive(Debug, Clone, Copy)] +pub(crate) struct ControlRetryPolicy { + attempt_timeout: Duration, + total_timeout: Duration, + retry_delay: Duration, +} + +impl ControlRetryPolicy { + pub(crate) const fn new( + attempt_timeout: Duration, + total_timeout: Duration, + retry_delay: Duration, + ) -> Self { + Self { + attempt_timeout, + total_timeout, + retry_delay, + } + } +} + +struct ControlAttemptError { + error: CliError, + transient: bool, + retry_after: Option, +} + +impl ControlAttemptError { + fn permanent(error: CliError) -> Self { + Self { + error, + transient: false, + retry_after: None, + } + } + + fn transient(error: CliError) -> Self { + Self { + error, + transient: true, + retry_after: None, + } + } + + fn with_retry_after(mut self, retry_after: Option) -> Self { + self.retry_after = retry_after; + self + } +} + +pub(crate) struct ClientHandshake { + pub(crate) proof: RegistrationProof, + daemon_origin: String, +} + +impl ClientHandshake { + /// Verifies the daemon's signature before TOFU-pinning its public identity. + pub(crate) fn authenticate_daemon( + &self, + proof: &super::protocol::HandshakeProof, + ) -> Result<(), CliError> { + if proof.signer != ComponentRole::Daemon { + return Err(CliError::Unauthorized( + "daemon registration proof used the wrong role".into(), + )); + } + self.proof + .transcript + .verify(proof) + .map_err(|error| CliError::Unauthorized(error.to_string()))?; + verify_or_store_daemon_pin( + &self.daemon_origin, + self.proof.transcript.responder_public_identity, + ) + } +} + +pub(crate) fn control_client() -> Result { + Client::builder() + .connect_timeout(CONTROL_CONNECT_TIMEOUT) + .timeout(CONTROL_REQUEST_TIMEOUT) + .redirect(reqwest::redirect::Policy::none()) + .http2_keep_alive_interval(Duration::from_secs(15)) + .build() + .map_err(CliError::Upstream) +} + +pub(crate) async fn begin_handshake( + client: &Client, + daemon_address: &str, + role: ComponentRole, + identity: &MachineIdentity, + instance_id: &str, + route_token_digest: Option, +) -> Result { + if role == ComponentRole::Daemon { + return Err(CliError::Config( + "a daemon cannot initiate a daemon client handshake".into(), + )); + } + let daemon = daemon_url(daemon_address)?; + let daemon_origin = daemon.as_str().trim_end_matches('/').to_owned(); + let initiator = descriptor(role); + let initiator_nonce = fresh_nonce()?; + let request = ChallengeRequest { + initiator: initiator.clone(), + initiator_instance_id: instance_id.to_owned(), + initiator_public_identity: identity.public_identity(), + initiator_fingerprint: identity.fingerprint(), + initiator_nonce, + }; + let challenge: ChallengeResponse = post_json( + client, + &format!("{daemon_origin}{CHALLENGE_PATH}"), + &request, + None, + ) + .await?; + challenge + .daemon + .validate() + .map_err(|error| CliError::Unauthorized(error.to_string()))?; + if challenge.daemon.role != ComponentRole::Daemon + || challenge.daemon_public_identity.fingerprint() != challenge.daemon_fingerprint + || challenge.daemon_instance_id.is_empty() + { + return Err(CliError::Unauthorized( + "daemon returned an invalid service identity".into(), + )); + } + challenge.verify_attestation(&request)?; + // Authenticate and TOFU-pin the daemon before a subsequent registration request can disclose + // the reusable route credential. First contact retains the normal limitations of TOFU. + verify_or_store_daemon_pin(&daemon_origin, challenge.daemon_public_identity)?; + let selected_protocol = initiator + .protocol + .negotiate(challenge.daemon.protocol) + .map_err(|error| CliError::Unauthorized(error.to_string()))?; + let transcript = HandshakeTranscript { + daemon_target: daemon_origin.clone(), + initiator, + responder: challenge.daemon, + initiator_instance_id: instance_id.to_owned(), + responder_instance_id: challenge.daemon_instance_id, + selected_protocol, + initiator_public_identity: identity.public_identity(), + responder_public_identity: challenge.daemon_public_identity, + initiator_fingerprint: identity.fingerprint(), + responder_fingerprint: challenge.daemon_fingerprint, + challenge_id: challenge.challenge.id, + initiator_nonce, + responder_nonce: challenge.challenge.nonce, + route_token_digest, + }; + let initiator_proof = transcript + .sign(role, identity) + .map_err(|error| CliError::Unauthorized(error.to_string()))?; + Ok(ClientHandshake { + proof: RegistrationProof { + transcript, + initiator_proof, + }, + daemon_origin, + }) +} + +pub(crate) async fn post_json( + client: &Client, + url: &str, + payload: &T, + route_token: Option<&str>, +) -> Result +where + T: Serialize + ?Sized, + R: DeserializeOwned, +{ + let body = encode_control_request(payload)?; + post_json_encoded(client, url, body, route_token) + .await + .map_err(|failure| failure.error) +} + +pub(crate) async fn post_json_idempotent( + client: &Client, + url: &str, + payload: &T, + route_token: Option<&str>, + policy: ControlRetryPolicy, +) -> Result +where + T: Serialize + ?Sized, + R: DeserializeOwned, +{ + let body = encode_control_request(payload)?; + retry_control(policy, || { + post_json_encoded(client, url, body.clone(), route_token) + }) + .await +} + +pub(crate) async fn post_empty_idempotent( + client: &Client, + url: &str, + payload: &T, + policy: ControlRetryPolicy, +) -> Result<(), CliError> { + let body = encode_control_request(payload)?; + retry_control(policy, || post_empty_encoded(client, url, body.clone())).await +} + +fn encode_control_request(payload: &T) -> Result { + serde_json::to_vec(payload) + .map(Bytes::from) + .map_err(|error| { + CliError::Launch(format!("failed to encode daemon control request: {error}")) + }) +} + +async fn post_json_encoded( + client: &Client, + url: &str, + body: Bytes, + route_token: Option<&str>, +) -> Result { + let response = send_control_request(client, url, body, route_token).await?; + let status = response.status(); + let retry_after = retry_after(&response); + let bytes = read_bounded_control_response(response).await?; + if !status.is_success() { + return Err(status_error(status, &bytes).with_retry_after(retry_after)); + } + serde_json::from_slice(&bytes).map_err(|error| { + ControlAttemptError::permanent(CliError::Launch(format!( + "invalid daemon control response: {error}" + ))) + }) +} + +async fn post_empty_encoded( + client: &Client, + url: &str, + body: Bytes, +) -> Result<(), ControlAttemptError> { + let response = send_control_request(client, url, body, None).await?; + let status = response.status(); + let retry_after = retry_after(&response); + if status.is_success() { + return Ok(()); + } + Err(if status == StatusCode::UNAUTHORIZED { + ControlAttemptError::permanent(CliError::Unauthorized( + "daemon rejected the control session credential".into(), + )) + } else { + let error = CliError::Launch(format!("daemon control request failed with HTTP {status}")); + if is_transient_status(status) { + ControlAttemptError::transient(error) + } else { + ControlAttemptError::permanent(error) + } + .with_retry_after(retry_after) + }) +} + +async fn send_control_request( + client: &Client, + url: &str, + body: Bytes, + route_token: Option<&str>, +) -> Result { + let mut request = client + .post(url) + .header(reqwest::header::CONTENT_TYPE, "application/json") + .body(body); + if let Some(token) = route_token { + request = request.header(CLIENT_TOKEN_HEADER, token); + } + request + .send() + .await + .map_err(|error| ControlAttemptError::transient(CliError::Upstream(error))) +} + +async fn read_bounded_control_response(response: Response) -> Result { + if response + .content_length() + .is_some_and(|length| length > MAX_CONTROL_RESPONSE_BYTES as u64) + { + return Err(response_too_large()); + } + let initial_capacity = response + .content_length() + .and_then(|length| usize::try_from(length).ok()) + .unwrap_or(0) + .min(MAX_CONTROL_RESPONSE_BYTES); + read_bounded_control_chunks(response.bytes_stream(), initial_capacity, |error| { + ControlAttemptError::transient(CliError::Upstream(error)) + }) + .await +} + +async fn read_bounded_control_chunks( + stream: S, + initial_capacity: usize, + map_error: F, +) -> Result +where + S: futures_util::Stream>, + F: Fn(E) -> ControlAttemptError, +{ + let mut bytes = BytesMut::with_capacity(initial_capacity.min(MAX_CONTROL_RESPONSE_BYTES)); + futures_util::pin_mut!(stream); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(&map_error)?; + if chunk.len() > MAX_CONTROL_RESPONSE_BYTES.saturating_sub(bytes.len()) { + return Err(response_too_large()); + } + bytes.extend_from_slice(&chunk); + } + Ok(bytes.freeze()) +} + +fn response_too_large() -> ControlAttemptError { + ControlAttemptError::permanent(CliError::Launch(format!( + "daemon control response exceeded {MAX_CONTROL_RESPONSE_BYTES} bytes" + ))) +} + +fn status_error(status: StatusCode, bytes: &[u8]) -> ControlAttemptError { + let message = serde_json::from_slice::(bytes) + .ok() + .and_then(|value| { + value + .pointer("/error/message") + .and_then(serde_json::Value::as_str) + .map(ToOwned::to_owned) + }) + .unwrap_or_else(|| "daemon rejected the control request".into()); + if status == StatusCode::UNAUTHORIZED { + return ControlAttemptError::permanent(CliError::Unauthorized(message)); + } + let error = CliError::Launch(format!( + "daemon control request failed with HTTP {status}: {message}" + )); + if is_transient_status(status) { + ControlAttemptError::transient(error) + } else { + ControlAttemptError::permanent(error) + } +} + +fn is_transient_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_MANY_REQUESTS + | StatusCode::INTERNAL_SERVER_ERROR + | StatusCode::BAD_GATEWAY + | StatusCode::SERVICE_UNAVAILABLE + | StatusCode::GATEWAY_TIMEOUT + ) || status.as_u16() == 425 +} + +fn retry_after(response: &Response) -> Option { + if response.status() != StatusCode::SERVICE_UNAVAILABLE { + return None; + } + response + .headers() + .get(reqwest::header::RETRY_AFTER)? + .to_str() + .ok()? + .trim() + .parse::() + .ok() + .map(Duration::from_secs) +} + +fn retry_backoff(base: Duration, attempt: u32) -> Duration { + const MAX_BACKOFF: Duration = Duration::from_secs(5); + if base.is_zero() { + return Duration::ZERO; + } + let scaled = base + .saturating_mul(1_u32 << attempt.min(6)) + .min(MAX_BACKOFF); + let jitter_bound = scaled / 4; + let jitter = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .ok() + .map_or(Duration::ZERO, |elapsed| { + let bound = jitter_bound.as_nanos(); + if bound == 0 { + Duration::ZERO + } else { + Duration::from_nanos((elapsed.as_nanos() % bound).min(u64::MAX as u128) as u64) + } + }); + scaled + .saturating_sub(jitter_bound / 2) + .saturating_add(jitter) +} + +async fn retry_control( + policy: ControlRetryPolicy, + mut operation: Operation, +) -> Result +where + Operation: FnMut() -> Attempt, + Attempt: Future>, +{ + let deadline = tokio::time::Instant::now() + policy.total_timeout; + let mut retry_attempt = 0; + loop { + let now = tokio::time::Instant::now(); + let attempt_deadline = deadline.min(now + policy.attempt_timeout); + let result = tokio::time::timeout_at(attempt_deadline, operation()).await; + let (error, retry_after) = match result { + Ok(Ok(value)) => return Ok(value), + Ok(Err(failure)) if !failure.transient => return Err(failure.error), + Ok(Err(failure)) => (failure.error, failure.retry_after), + Err(_) => ( + CliError::Launch("daemon control request attempt timed out".into()), + None, + ), + }; + let now = tokio::time::Instant::now(); + if now >= deadline { + return Err(error); + } + let delay = retry_after.unwrap_or_else(|| retry_backoff(policy.retry_delay, retry_attempt)); + retry_attempt = retry_attempt.saturating_add(1); + tokio::time::sleep_until(deadline.min(now + delay)).await; + } +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/client_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/control.rs b/crates/cli/src/daemon/common/control.rs new file mode 100644 index 000000000..cdbf10565 --- /dev/null +++ b/crates/cli/src/daemon/common/control.rs @@ -0,0 +1,562 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Versioned control-plane wire messages shared by daemon, MCP, and worker processes. + +use std::net::Ipv4Addr; +use std::time::{SystemTime, UNIX_EPOCH}; + +use base64::Engine; +use ring::rand::{SecureRandom, SystemRandom}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use super::identity::{ + Challenge, ChallengeId, ChallengeNonce, Ed25519Signature, Fingerprint, MachineIdentity, + PublicIdentity, encode_transcript, +}; +use super::protocol::{ + BrokerDirective, ComponentDescriptor, ComponentRole, HandshakeProof, HandshakeTranscript, + SensitiveString, +}; +use crate::error::CliError; + +pub(crate) const CHALLENGE_PATH: &str = "/_nemo-relay/control/v1/challenge"; +pub(crate) const MCP_REGISTER_PATH: &str = "/_nemo-relay/control/v1/mcp/register"; +pub(crate) const MCP_HEARTBEAT_PATH: &str = "/_nemo-relay/control/v1/mcp/heartbeat"; +pub(crate) const MCP_RELEASE_PATH: &str = "/_nemo-relay/control/v1/mcp/release"; +pub(crate) const MCP_ACTIVATION_FAILED_PATH: &str = "/_nemo-relay/control/v1/mcp/activation-failed"; +pub(crate) const WORKER_REGISTER_PATH: &str = "/_nemo-relay/control/v1/worker/register"; +pub(crate) const WORKER_RECOVER_PATH: &str = "/_nemo-relay/control/v1/worker/recover"; +pub(crate) const WORKER_READY_PATH: &str = "/_nemo-relay/control/v1/worker/ready"; +pub(crate) const WORKER_HEARTBEAT_PATH: &str = "/_nemo-relay/control/v1/worker/heartbeat"; +pub(crate) const WORKER_DRAIN_PATH: &str = "/_nemo-relay/control/v1/worker/drain"; +pub(crate) const WORKER_PROBE_PATH: &str = "/_nemo-relay/worker/v1/ready"; + +pub(crate) const CLIENT_TOKEN_HEADER: &str = "x-nemo-relay-client-token"; +pub(crate) const WORKER_TOKEN_HEADER: &str = "x-nemo-relay-worker-token"; +/// Private worker-to-daemon signal that a route-wide invariant failed after authentication. +/// The daemon consumes this field and never exposes it on the public response. +pub(crate) const WORKER_ROUTE_FAILURE_HEADER: &str = "x-nemo-relay-worker-route-failure"; +pub(crate) const MAX_CONTROL_BODY_BYTES: usize = 256 * 1024; +pub(crate) const CHALLENGE_LIFETIME_MS: u64 = 15_000; +pub(crate) const MCP_HEARTBEAT_INTERVAL_MS: u64 = 10_000; +pub(crate) const MCP_LEASE_MS: u64 = 30_000; +pub(crate) const WORKER_HEARTBEAT_INTERVAL_MS: u64 = 5_000; +pub(crate) const WORKER_LEASE_MS: u64 = 20_000; +pub(crate) const ACTIVATION_LIFETIME_MS: u64 = 15_000; +pub(crate) const DRAIN_LIFETIME_MS: u64 = 120_000; +pub(crate) const RECOVERY_LIFETIME_MS: u64 = 120_000; +const WORKER_NETWORK_HINT_DOMAIN: &[u8] = b"nemo-relay/worker-network-hint/v1"; +const WORKER_GENERATION_DOMAIN: &[u8] = b"nemo-relay/worker-generation/v1"; +const DAEMON_CHALLENGE_DOMAIN: &[u8] = b"nemo-relay/daemon-challenge/v1"; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct ChallengeRequest { + pub(crate) initiator: ComponentDescriptor, + pub(crate) initiator_instance_id: String, + pub(crate) initiator_public_identity: PublicIdentity, + pub(crate) initiator_fingerprint: Fingerprint, + pub(crate) initiator_nonce: ChallengeNonce, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct ChallengeResponse { + pub(crate) daemon: ComponentDescriptor, + pub(crate) daemon_instance_id: String, + pub(crate) daemon_public_identity: PublicIdentity, + pub(crate) daemon_fingerprint: Fingerprint, + pub(crate) challenge: Challenge, + /// Daemon signature over this challenge and the complete initiating request. + pub(crate) daemon_challenge_proof: Ed25519Signature, +} + +impl ChallengeResponse { + pub(crate) fn verify_attestation(&self, request: &ChallengeRequest) -> Result<(), CliError> { + let canonical = daemon_challenge_bytes(request, self)?; + self.daemon_public_identity + .verify(&canonical, &self.daemon_challenge_proof) + .map_err(|_| CliError::Unauthorized("daemon challenge signature did not verify".into())) + } +} + +pub(crate) fn daemon_challenge_bytes( + request: &ChallengeRequest, + response: &ChallengeResponse, +) -> Result, CliError> { + let initiator = serde_json::to_vec(&request.initiator).map_err(|error| { + CliError::Launch(format!("failed to encode initiator descriptor: {error}")) + })?; + let daemon = serde_json::to_vec(&response.daemon).map_err(|error| { + CliError::Launch(format!("failed to encode daemon descriptor: {error}")) + })?; + let issued_at = response.challenge.issued_at_unix_ms.to_be_bytes(); + let expires_at = response.challenge.expires_at_unix_ms.to_be_bytes(); + encode_transcript( + DAEMON_CHALLENGE_DOMAIN, + &[ + ("initiator", initiator.as_slice()), + ( + "initiator_instance_id", + request.initiator_instance_id.as_bytes(), + ), + ( + "initiator_public_identity", + request.initiator_public_identity.as_bytes().as_slice(), + ), + ( + "initiator_fingerprint", + request.initiator_fingerprint.as_bytes().as_slice(), + ), + ( + "initiator_nonce", + request.initiator_nonce.as_bytes().as_slice(), + ), + ("daemon", daemon.as_slice()), + ("daemon_instance_id", response.daemon_instance_id.as_bytes()), + ( + "daemon_public_identity", + response.daemon_public_identity.as_bytes().as_slice(), + ), + ( + "daemon_fingerprint", + response.daemon_fingerprint.as_bytes().as_slice(), + ), + ("challenge_id", response.challenge.id.as_bytes().as_slice()), + ( + "challenge_nonce", + response.challenge.nonce.as_bytes().as_slice(), + ), + ("issued_at_unix_ms", issued_at.as_slice()), + ("expires_at_unix_ms", expires_at.as_slice()), + ], + ) + .map_err(|error| CliError::Launch(format!("failed to encode daemon challenge: {error}"))) +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct RegistrationProof { + pub(crate) transcript: HandshakeTranscript, + pub(crate) initiator_proof: HandshakeProof, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct McpRegisterRequest { + pub(crate) proof: RegistrationProof, + pub(crate) worker_network: WorkerNetworkHintProof, +} + +/// The MCP machine's daemon-reachable IPv4 address and optional prescribed worker port. +/// +/// The daemon validates this signed hint and remains authoritative over the resulting bind and +/// advertise arguments in [`BrokerDirective::LaunchWorker`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct WorkerNetworkHint { + pub(crate) advertised_host: String, + pub(crate) port: Option, +} + +impl WorkerNetworkHint { + pub(crate) fn new( + advertised_host: impl Into, + port: Option, + ) -> Result { + let hint = Self { + advertised_host: advertised_host.into().to_ascii_lowercase(), + port, + }; + hint.validate()?; + Ok(hint) + } + + pub(crate) fn validate(&self) -> Result<(), CliError> { + let host = self.advertised_host.as_str(); + let ipv4 = host.parse::().ok(); + let valid_hostname = host.len() <= 253 + && !host.is_empty() + && host.is_ascii() + && !host.contains(['/', ':', '@', '[', ']', '?', '#']) + && host.split('.').all(|label| { + !label.is_empty() + && label.len() <= 63 + && !label.starts_with('-') + && !label.ends_with('-') + && label + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-') + }); + if ipv4.is_some_and(|address| address.is_unspecified()) + || (ipv4.is_none() && !valid_hostname) + || self.port == Some(0) + { + return Err(CliError::Config( + "worker network hint requires a concrete hostname or IPv4 address and a nonzero prescribed port" + .into(), + )); + } + Ok(()) + } + + pub(crate) fn is_loopback(&self) -> bool { + self.advertised_host.eq_ignore_ascii_case("localhost") + || self + .advertised_host + .parse::() + .is_ok_and(|address| address.is_loopback()) + } +} + +/// A worker network hint bound to the authenticated MCP registration challenge. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct WorkerNetworkHintProof { + pub(crate) hint: WorkerNetworkHint, + pub(crate) signature: Ed25519Signature, +} + +impl WorkerNetworkHintProof { + pub(crate) fn sign( + hint: WorkerNetworkHint, + daemon_target: &str, + mcp_instance_id: &str, + challenge_id: &ChallengeId, + fingerprint: &Fingerprint, + identity: &MachineIdentity, + ) -> Result { + hint.validate()?; + let canonical = worker_network_hint_bytes( + &hint, + daemon_target, + mcp_instance_id, + challenge_id, + fingerprint, + )?; + Ok(Self { + hint, + signature: identity.sign(&canonical), + }) + } + + pub(crate) fn verify( + &self, + daemon_target: &str, + mcp_instance_id: &str, + challenge_id: &ChallengeId, + fingerprint: &Fingerprint, + identity: &PublicIdentity, + ) -> Result<(), CliError> { + self.hint.validate()?; + let canonical = worker_network_hint_bytes( + &self.hint, + daemon_target, + mcp_instance_id, + challenge_id, + fingerprint, + )?; + identity + .verify(&canonical, &self.signature) + .map_err(|_| CliError::Unauthorized("invalid signed worker network hint".into())) + } +} + +fn worker_network_hint_bytes( + hint: &WorkerNetworkHint, + daemon_target: &str, + mcp_instance_id: &str, + challenge_id: &ChallengeId, + fingerprint: &Fingerprint, +) -> Result, CliError> { + let port_present = [u8::from(hint.port.is_some())]; + let port = hint.port.unwrap_or_default().to_be_bytes(); + encode_transcript( + WORKER_NETWORK_HINT_DOMAIN, + &[ + ("daemon_target", daemon_target.as_bytes()), + ("mcp_instance_id", mcp_instance_id.as_bytes()), + ("challenge_id", challenge_id.as_bytes().as_slice()), + ("fingerprint", fingerprint.as_bytes().as_slice()), + ("advertised_host", hint.advertised_host.as_bytes()), + ("port_present", port_present.as_slice()), + ("port", port.as_slice()), + ], + ) + .map_err(|error| CliError::Unauthorized(error.to_string())) +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct McpRegisterResponse { + pub(crate) daemon_proof: HandshakeProof, + pub(crate) session_token: SensitiveString, + pub(crate) heartbeat_interval_ms: u64, + pub(crate) directive: BrokerDirective, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct WorkerBootstrap { + pub(crate) activation_id: String, + pub(crate) activation_token: SensitiveString, + pub(crate) deadline_unix_ms: u64, + pub(crate) bind_ip: Ipv4Addr, + pub(crate) port: u16, + pub(crate) advertise_address: Option, +} + +impl WorkerBootstrap { + pub(crate) fn from_directive(directive: BrokerDirective) -> Option { + let BrokerDirective::LaunchWorker { + activation_id, + activation_token, + deadline_unix_ms, + bind_ip, + port, + advertise_address, + } = directive + else { + return None; + }; + Some(Self { + activation_id, + activation_token, + deadline_unix_ms, + bind_ip, + port, + advertise_address, + }) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct WorkerRegisterRequest { + pub(crate) proof: RegistrationProof, + pub(crate) worker_id: String, + pub(crate) endpoint: String, + pub(crate) activation_id: String, + pub(crate) activation_token: SensitiveString, + pub(crate) tls_root_certificate: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct WorkerRecoverRequest { + pub(crate) proof: RegistrationProof, + pub(crate) worker_id: String, + pub(crate) endpoint: String, + pub(crate) tls_root_certificate: Option, + pub(crate) generation_grant: WorkerGenerationGrant, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct WorkerRegisterResponse { + pub(crate) daemon_proof: HandshakeProof, + pub(crate) session_token: SensitiveString, + pub(crate) data_token: SensitiveString, + pub(crate) heartbeat_interval_ms: u64, + pub(crate) generation_grant: WorkerGenerationGrant, +} + +/// A daemon-signed proof binding one worker generation to its endpoint and TLS trust anchor. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct WorkerGenerationGrant { + pub(crate) generation_id: String, + pub(crate) worker_id: String, + pub(crate) fingerprint: Fingerprint, + pub(crate) endpoint: String, + pub(crate) tls_root_digest: Option<[u8; 32]>, + pub(crate) signature: Ed25519Signature, +} + +impl WorkerGenerationGrant { + pub(crate) fn issue( + worker_id: &str, + fingerprint: Fingerprint, + endpoint: &str, + tls_root_certificate: Option<&str>, + daemon_identity: &MachineIdentity, + ) -> Result { + let generation_id = random_secret(16)?; + let tls_root_digest = + tls_root_certificate.map(|root| Sha256::digest(root.as_bytes()).into()); + let canonical = worker_generation_bytes( + &generation_id, + worker_id, + &fingerprint, + endpoint, + tls_root_digest.as_ref(), + )?; + Ok(Self { + generation_id, + worker_id: worker_id.to_owned(), + fingerprint, + endpoint: endpoint.to_owned(), + tls_root_digest, + signature: daemon_identity.sign(&canonical), + }) + } + + pub(crate) fn verify( + &self, + worker_id: &str, + fingerprint: Fingerprint, + endpoint: &str, + tls_root_certificate: Option<&str>, + daemon_identity: &PublicIdentity, + ) -> Result<(), CliError> { + let expected_root = tls_root_certificate.map(|root| Sha256::digest(root.as_bytes()).into()); + if self.generation_id.is_empty() + || self.worker_id != worker_id + || self.fingerprint != fingerprint + || self.endpoint != endpoint + || self.tls_root_digest != expected_root + { + return Err(CliError::Unauthorized( + "worker generation grant does not match recovery".into(), + )); + } + let canonical = worker_generation_bytes( + &self.generation_id, + &self.worker_id, + &self.fingerprint, + &self.endpoint, + self.tls_root_digest.as_ref(), + )?; + daemon_identity + .verify(&canonical, &self.signature) + .map_err(|_| CliError::Unauthorized("invalid worker generation grant".into())) + } +} + +fn worker_generation_bytes( + generation_id: &str, + worker_id: &str, + fingerprint: &Fingerprint, + endpoint: &str, + tls_root_digest: Option<&[u8; 32]>, +) -> Result, CliError> { + let root_present = [u8::from(tls_root_digest.is_some())]; + encode_transcript( + WORKER_GENERATION_DOMAIN, + &[ + ("generation_id", generation_id.as_bytes()), + ("worker_id", worker_id.as_bytes()), + ("fingerprint", fingerprint.as_bytes().as_slice()), + ("endpoint", endpoint.as_bytes()), + ("tls_root_present", root_present.as_slice()), + ( + "tls_root_digest", + tls_root_digest.map_or(&[][..], |digest| digest.as_slice()), + ), + ], + ) + .map_err(|error| CliError::Unauthorized(error.to_string())) +} + +/// A session-authenticated message. Sequence numbers are strictly increasing per session. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct SessionRequest { + pub(crate) session_id: String, + pub(crate) session_token: SensitiveString, + pub(crate) request_id: String, + pub(crate) sequence: u64, + pub(crate) payload_sha256: [u8; 32], + pub(crate) payload: T, +} + +impl SessionRequest { + pub(crate) fn new( + session_id: String, + session_token: SensitiveString, + sequence: u64, + payload: T, + ) -> Result { + let encoded = serde_json::to_vec(&payload).map_err(|error| { + CliError::Launch(format!("failed to encode daemon control payload: {error}")) + })?; + Ok(Self { + session_id, + session_token, + request_id: random_secret(16)?, + sequence, + payload_sha256: Sha256::digest(encoded).into(), + payload, + }) + } + + pub(crate) fn validate_payload_hash(&self) -> bool { + serde_json::to_vec(&self.payload) + .map(|encoded| { + let actual: [u8; 32] = Sha256::digest(encoded).into(); + subtle::ConstantTimeEq::ct_eq(actual.as_slice(), self.payload_sha256.as_slice()) + .into() + }) + .unwrap_or(false) + } +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub(crate) struct EmptyPayload {} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct ActivationFailedPayload { + pub(crate) activation_id: String, + pub(crate) reason: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct McpHeartbeatResponse { + pub(crate) directive: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct WorkerHeartbeatPayload { + pub(crate) worker_id: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct WorkerReadyPayload { + pub(crate) worker_id: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub(crate) struct WorkerDrainRequest { + pub(crate) worker_id: String, + /// Daemon wall-clock deadline retained for protocol-v1 compatibility and audit logs. + pub(crate) deadline_unix_ms: u64, + /// Relative lifetime enforced against the worker's local monotonic clock. + #[serde(default)] + pub(crate) timeout_ms: Option, +} + +pub(crate) fn random_secret(bytes: usize) -> Result { + let mut value = vec![0_u8; bytes]; + SystemRandom::new() + .fill(&mut value) + .map_err(|_| CliError::Launch("failed to generate daemon session credential".into()))?; + Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(value)) +} + +pub(crate) fn fresh_nonce() -> Result { + let record = super::identity::ChallengeRecord::generate(now_unix_ms(), 1).map_err(|error| { + CliError::Launch(format!("failed to generate handshake nonce: {error}")) + })?; + Ok(record.challenge().nonce) +} + +pub(crate) fn now_unix_ms() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .try_into() + .unwrap_or(u64::MAX) +} + +pub(crate) fn descriptor(role: ComponentRole) -> ComponentDescriptor { + ComponentDescriptor::nemo_relay( + role, + super::protocol::ProtocolRange::default(), + super::protocol::Capabilities::streaming_transport(), + env!("CARGO_PKG_VERSION"), + ) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/control_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/identity.rs b/crates/cli/src/daemon/common/identity.rs new file mode 100644 index 000000000..739da1558 --- /dev/null +++ b/crates/cli/src/daemon/common/identity.rs @@ -0,0 +1,397 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use ring::rand::{SecureRandom, SystemRandom}; +use ring::signature::{ED25519, Ed25519KeyPair, KeyPair, UnparsedPublicKey}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use subtle::ConstantTimeEq; +use thiserror::Error; + +const ED25519_PUBLIC_KEY_BYTES: usize = 32; +const CHALLENGE_ID_BYTES: usize = 16; +const CHALLENGE_NONCE_BYTES: usize = 32; +const TRANSCRIPT_MAGIC: &[u8] = b"NEMO-RELAY-SIGNED-TRANSCRIPT\0"; + +/// An Ed25519 identity used by one daemon component. +#[derive(Clone)] +pub(crate) struct MachineIdentity { + key_pair: Arc, +} + +impl MachineIdentity { + /// Generates an identity and returns its PKCS#8 document for owner-private storage. + pub(crate) fn generate() -> Result { + let random = SystemRandom::new(); + let document = + Ed25519KeyPair::generate_pkcs8(&random).map_err(|_| IdentityError::KeyGeneration)?; + let identity = Self::from_pkcs8(document.as_ref())?; + Ok(GeneratedMachineIdentity { + identity, + pkcs8: document.as_ref().to_vec(), + }) + } + + /// Loads an identity from an unencrypted PKCS#8 Ed25519 document. + pub(crate) fn from_pkcs8(pkcs8: &[u8]) -> Result { + let key_pair = + Ed25519KeyPair::from_pkcs8(pkcs8).map_err(|_| IdentityError::InvalidPrivateKey)?; + Ok(Self { + key_pair: Arc::new(key_pair), + }) + } + + /// Returns the public half of this identity. + pub(crate) fn public_identity(&self) -> PublicIdentity { + let bytes = self + .key_pair + .public_key() + .as_ref() + .try_into() + .expect("ring Ed25519 public keys have a fixed length"); + PublicIdentity(bytes) + } + + /// Returns the stable SHA-256 fingerprint of the public identity. + pub(crate) fn fingerprint(&self) -> Fingerprint { + self.public_identity().fingerprint() + } + + /// Signs already-canonical transcript bytes. + pub(crate) fn sign(&self, transcript: &[u8]) -> Ed25519Signature { + Ed25519Signature(self.key_pair.sign(transcript).as_ref().to_vec()) + } +} + +impl fmt::Debug for MachineIdentity { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("MachineIdentity") + .field("fingerprint", &self.fingerprint()) + .finish_non_exhaustive() + } +} + +/// A newly generated identity and the private document that must be persisted securely. +pub(crate) struct GeneratedMachineIdentity { + pub(crate) identity: MachineIdentity, + pub(crate) pkcs8: Vec, +} + +/// An Ed25519 public identity suitable for control-protocol serialization. +#[derive(Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct PublicIdentity([u8; ED25519_PUBLIC_KEY_BYTES]); + +impl PublicIdentity { + /// Parses an Ed25519 public key. + #[cfg(test)] + pub(crate) fn from_bytes(bytes: &[u8]) -> Result { + let bytes = bytes + .try_into() + .map_err(|_| IdentityError::InvalidPublicKey)?; + Ok(Self(bytes)) + } + + /// Returns the raw Ed25519 public-key bytes. + pub(crate) const fn as_bytes(&self) -> &[u8; ED25519_PUBLIC_KEY_BYTES] { + &self.0 + } + + /// Returns the stable SHA-256 fingerprint of this public key. + pub(crate) fn fingerprint(&self) -> Fingerprint { + Fingerprint(sha256(&self.0)) + } + + /// Verifies a signature over already-canonical transcript bytes. + pub(crate) fn verify( + &self, + transcript: &[u8], + signature: &Ed25519Signature, + ) -> Result<(), IdentityError> { + UnparsedPublicKey::new(&ED25519, self.0) + .verify(transcript, signature.as_bytes()) + .map_err(|_| IdentityError::SignatureVerification) + } +} + +impl fmt::Debug for PublicIdentity { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_tuple("PublicIdentity") + .field(&self.fingerprint()) + .finish() + } +} + +/// A serialized Ed25519 signature. +#[derive(Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct Ed25519Signature(Vec); + +impl Ed25519Signature { + /// Returns the signature bytes. + pub(crate) fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl fmt::Debug for Ed25519Signature { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("Ed25519Signature") + .field("length", &self.0.len()) + .finish_non_exhaustive() + } +} + +/// A stable public-key fingerprint used as the broker route key. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct Fingerprint([u8; 32]); + +impl Fingerprint { + /// Returns the raw SHA-256 digest. + pub(crate) const fn as_bytes(&self) -> &[u8; 32] { + &self.0 + } +} + +impl fmt::Debug for Fingerprint { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "Fingerprint({self})") + } +} + +impl fmt::Display for Fingerprint { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write_hex(formatter, &self.0) + } +} + +/// The SHA-256 digest of the per-user-machine route token. +#[derive(Clone, Copy, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct TokenDigest([u8; 32]); + +impl TokenDigest { + /// Hashes the exact token bytes received from the environment or HTTP header. + pub(crate) fn from_token(token: &[u8]) -> Self { + Self(sha256(token)) + } + + /// Returns the raw SHA-256 digest. + pub(crate) const fn as_bytes(&self) -> &[u8; 32] { + &self.0 + } + + /// Compares two token digests in constant time. + pub(crate) fn matches(&self, other: &Self) -> bool { + bool::from(self.0.ct_eq(&other.0)) + } +} + +impl PartialEq for TokenDigest { + fn eq(&self, other: &Self) -> bool { + self.matches(other) + } +} + +impl Eq for TokenDigest {} + +impl Hash for TokenDigest { + fn hash(&self, state: &mut H) { + self.0.hash(state); + } +} + +impl fmt::Debug for TokenDigest { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(formatter, "TokenDigest({self})") + } +} + +impl fmt::Display for TokenDigest { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + write_hex(formatter, &self.0) + } +} + +/// A random identifier for one daemon challenge. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct ChallengeId([u8; CHALLENGE_ID_BYTES]); + +impl ChallengeId { + /// Returns the identifier bytes. + pub(crate) const fn as_bytes(&self) -> &[u8; CHALLENGE_ID_BYTES] { + &self.0 + } +} + +/// A random challenge nonce. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct ChallengeNonce([u8; CHALLENGE_NONCE_BYTES]); + +impl ChallengeNonce { + /// Returns the nonce bytes. + pub(crate) const fn as_bytes(&self) -> &[u8; CHALLENGE_NONCE_BYTES] { + &self.0 + } +} + +/// The wire-safe portion of a challenge record. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct Challenge { + pub(crate) id: ChallengeId, + pub(crate) nonce: ChallengeNonce, + pub(crate) issued_at_unix_ms: u64, + pub(crate) expires_at_unix_ms: u64, +} + +/// A one-use local challenge record with explicit expiry handling. +#[derive(Debug)] +pub(crate) struct ChallengeRecord { + challenge: Challenge, + consumed: bool, +} + +impl ChallengeRecord { + /// Creates a random challenge using caller-supplied wall-clock values. + pub(crate) fn generate( + issued_at_unix_ms: u64, + lifetime_ms: u64, + ) -> Result { + let expires_at_unix_ms = issued_at_unix_ms + .checked_add(lifetime_ms) + .ok_or(IdentityError::ChallengeLifetimeOverflow)?; + let random = SystemRandom::new(); + let mut id = [0_u8; CHALLENGE_ID_BYTES]; + let mut nonce = [0_u8; CHALLENGE_NONCE_BYTES]; + random + .fill(&mut id) + .and_then(|()| random.fill(&mut nonce)) + .map_err(|_| IdentityError::ChallengeGeneration)?; + Ok(Self::from_challenge(Challenge { + id: ChallengeId(id), + nonce: ChallengeNonce(nonce), + issued_at_unix_ms, + expires_at_unix_ms, + })) + } + + /// Wraps a challenge for tracking. Primarily useful when restoring an issued challenge. + pub(crate) fn from_challenge(challenge: Challenge) -> Self { + Self { + challenge, + consumed: false, + } + } + + /// Returns the challenge sent to the peer. + pub(crate) const fn challenge(&self) -> Challenge { + self.challenge + } + + /// Consumes this challenge exactly once before its expiry time. + pub(crate) fn consume( + &mut self, + presented_id: &ChallengeId, + now_unix_ms: u64, + ) -> Result { + if !bool::from(self.challenge.id.0.ct_eq(&presented_id.0)) { + return Err(ChallengeError::IdentifierMismatch); + } + if self.consumed { + return Err(ChallengeError::Replay); + } + if now_unix_ms >= self.challenge.expires_at_unix_ms { + self.consumed = true; + return Err(ChallengeError::Expired); + } + self.consumed = true; + Ok(self.challenge) + } +} + +/// Identity and transcript construction failures. +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum IdentityError { + #[error("failed to generate an Ed25519 identity")] + KeyGeneration, + #[error("the Ed25519 private key is invalid")] + InvalidPrivateKey, + #[cfg(test)] + #[error("the Ed25519 public key is invalid")] + InvalidPublicKey, + #[error("the Ed25519 signature did not verify")] + SignatureVerification, + #[error("failed to generate a handshake challenge")] + ChallengeGeneration, + #[error("the handshake challenge lifetime overflowed")] + ChallengeLifetimeOverflow, + #[error("a signed transcript field is too large")] + TranscriptFieldTooLarge, +} + +/// Challenge rejection reasons that callers can map to typed protocol errors. +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum ChallengeError { + #[error("the challenge identifier does not match")] + IdentifierMismatch, + #[error("the challenge has already been consumed")] + Replay, + #[error("the challenge has expired")] + Expired, +} + +/// Encodes a signed transcript without relying on JSON map order or host endianness. +pub(crate) fn encode_transcript( + domain: &[u8], + fields: &[(&str, &[u8])], +) -> Result, IdentityError> { + let mut encoded = Vec::with_capacity( + TRANSCRIPT_MAGIC.len() + + domain.len() + + fields + .iter() + .map(|(name, value)| name.len() + value.len() + 16) + .sum::(), + ); + encoded.extend_from_slice(TRANSCRIPT_MAGIC); + append_length_prefixed(&mut encoded, domain)?; + let field_count = + u32::try_from(fields.len()).map_err(|_| IdentityError::TranscriptFieldTooLarge)?; + encoded.extend_from_slice(&field_count.to_be_bytes()); + for (name, value) in fields { + append_length_prefixed(&mut encoded, name.as_bytes())?; + append_length_prefixed(&mut encoded, value)?; + } + Ok(encoded) +} + +fn append_length_prefixed(encoded: &mut Vec, value: &[u8]) -> Result<(), IdentityError> { + let length = u64::try_from(value.len()).map_err(|_| IdentityError::TranscriptFieldTooLarge)?; + encoded.extend_from_slice(&length.to_be_bytes()); + encoded.extend_from_slice(value); + Ok(()) +} + +fn sha256(value: &[u8]) -> [u8; 32] { + Sha256::digest(value).into() +} + +fn write_hex(formatter: &mut fmt::Formatter<'_>, bytes: &[u8]) -> fmt::Result { + for byte in bytes { + write!(formatter, "{byte:02x}")?; + } + Ok(()) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/identity_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/mod.rs b/crates/cli/src/daemon/common/mod.rs new file mode 100644 index 000000000..d5bbdc1f3 --- /dev/null +++ b/crates/cli/src/daemon/common/mod.rs @@ -0,0 +1,14 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared daemon control-plane and transport primitives. + +pub(crate) mod address; +pub(crate) mod client; +pub(crate) mod control; +pub(crate) mod identity; +pub(crate) mod protocol; +pub(crate) mod routes; +pub(crate) mod state; +pub(crate) mod transport; +pub(crate) mod worker_tls; diff --git a/crates/cli/src/daemon/common/protocol.rs b/crates/cli/src/daemon/common/protocol.rs new file mode 100644 index 000000000..466e876d0 --- /dev/null +++ b/crates/cli/src/daemon/common/protocol.rs @@ -0,0 +1,530 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeSet; +use std::fmt; +use std::net::Ipv4Addr; + +use serde::de::Error as _; +use serde::{Deserialize, Deserializer, Serialize}; +use subtle::ConstantTimeEq; +use thiserror::Error; + +use super::identity::{ + ChallengeId, ChallengeNonce, Ed25519Signature, Fingerprint, IdentityError, MachineIdentity, + PublicIdentity, TokenDigest, encode_transcript, +}; + +pub(crate) const SERVICE_NAME: &str = "nemo-relay"; +pub(crate) const PROTOCOL_V1: u16 = 1; +const HANDSHAKE_DOMAIN: &[u8] = b"nemo-relay/daemon-handshake/v1"; +const MAX_CAPABILITIES: usize = 64; +const MAX_CAPABILITY_BYTES: usize = 128; +const MAX_BINARY_VERSION_BYTES: usize = 256; +const MAX_INSTANCE_ID_BYTES: usize = 256; +const MAX_DAEMON_TARGET_BYTES: usize = 2_048; + +/// The authenticated role of a daemon-protocol participant. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum ComponentRole { + Daemon, + Mcp, + Worker, +} + +impl ComponentRole { + const fn transcript_name(self) -> &'static str { + match self { + Self::Daemon => "daemon", + Self::Mcp => "mcp", + Self::Worker => "worker", + } + } +} + +/// An inclusive range of daemon-protocol versions. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ProtocolRange { + pub(crate) minimum: u16, + pub(crate) maximum: u16, +} + +impl ProtocolRange { + /// Constructs a validated inclusive protocol range. + #[cfg(test)] + pub(crate) fn new(minimum: u16, maximum: u16) -> Result { + let range = Self { minimum, maximum }; + range.validate()?; + Ok(range) + } + + /// Returns the highest mutually supported protocol version. + pub(crate) fn negotiate(self, peer: Self) -> Result { + self.validate()?; + peer.validate()?; + let minimum = self.minimum.max(peer.minimum); + let maximum = self.maximum.min(peer.maximum); + (minimum <= maximum) + .then_some(maximum) + .ok_or(ProtocolError::NoProtocolOverlap) + } + + /// Reports whether this range contains one protocol version. + pub(crate) const fn contains(self, version: u16) -> bool { + version >= self.minimum && version <= self.maximum + } + + fn validate(self) -> Result<(), ProtocolError> { + if self.minimum == 0 || self.minimum > self.maximum { + return Err(ProtocolError::InvalidProtocolRange); + } + Ok(()) + } +} + +impl Default for ProtocolRange { + fn default() -> Self { + Self { + minimum: PROTOCOL_V1, + maximum: PROTOCOL_V1, + } + } +} + +/// A forward-compatible, deterministically ordered set of protocol capabilities. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(transparent)] +pub(crate) struct Capabilities(BTreeSet); + +impl Capabilities { + /// Constructs a capability set and validates every capability name. + pub(crate) fn new( + names: impl IntoIterator>, + ) -> Result { + let capabilities = Self(names.into_iter().map(Into::into).collect()); + capabilities.validate()?; + Ok(capabilities) + } + + /// Returns the baseline lossless HTTP transport capabilities. + pub(crate) fn streaming_transport() -> Self { + Self::new([ + "http1", + "http2", + "streaming_body_frames", + "sse_passthrough", + "trailers", + ]) + .expect("built-in capability names are valid") + } + + #[cfg(test)] + pub(crate) fn contains(&self, name: &str) -> bool { + self.0.contains(name) + } + + /// Reports whether this set includes every required capability. + pub(crate) fn includes(&self, required: &Self) -> bool { + required.0.is_subset(&self.0) + } + + fn validate(&self) -> Result<(), ProtocolError> { + if self.0.len() > MAX_CAPABILITIES + || self + .0 + .iter() + .any(|name| name.len() > MAX_CAPABILITY_BYTES || !valid_capability_name(name)) + { + return Err(ProtocolError::InvalidCapability); + } + Ok(()) + } + + fn canonical_bytes(&self) -> Result, ProtocolError> { + self.validate()?; + let mut encoded = Vec::new(); + let count = u32::try_from(self.0.len()).map_err(|_| ProtocolError::FieldTooLarge)?; + encoded.extend_from_slice(&count.to_be_bytes()); + for capability in &self.0 { + let length = + u32::try_from(capability.len()).map_err(|_| ProtocolError::FieldTooLarge)?; + encoded.extend_from_slice(&length.to_be_bytes()); + encoded.extend_from_slice(capability.as_bytes()); + } + Ok(encoded) + } +} + +/// Authenticated metadata describing one protocol component. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct ComponentDescriptor { + pub(crate) service: String, + pub(crate) role: ComponentRole, + pub(crate) protocol: ProtocolRange, + pub(crate) capabilities: Capabilities, + pub(crate) binary_version: String, +} + +impl ComponentDescriptor { + /// Constructs a descriptor for a real NeMo Relay component. + pub(crate) fn nemo_relay( + role: ComponentRole, + protocol: ProtocolRange, + capabilities: Capabilities, + binary_version: impl Into, + ) -> Self { + Self { + service: SERVICE_NAME.to_owned(), + role, + protocol, + capabilities, + binary_version: binary_version.into(), + } + } + + /// Validates invariants that must hold regardless of binary release version. + pub(crate) fn validate(&self) -> Result<(), ProtocolError> { + if self.service != SERVICE_NAME { + return Err(ProtocolError::WrongService); + } + self.protocol.validate()?; + self.capabilities.validate()?; + if self.binary_version.is_empty() { + return Err(ProtocolError::MissingBinaryVersion); + } + if self.binary_version.len() > MAX_BINARY_VERSION_BYTES { + return Err(ProtocolError::BinaryVersionTooLong); + } + Ok(()) + } +} + +/// A sensitive wire value whose debug output is always redacted. +#[derive(Clone, Serialize)] +#[serde(transparent)] +pub(crate) struct SensitiveString(String); + +impl SensitiveString { + /// Constructs a non-empty sensitive string. + pub(crate) fn new(value: impl Into) -> Result { + let value = value.into(); + if value.is_empty() { + return Err(ProtocolError::MissingSensitiveValue); + } + Ok(Self(value)) + } + + /// Exposes the value only at the protocol boundary that consumes it. + pub(crate) fn expose(&self) -> &str { + &self.0 + } +} + +impl PartialEq for SensitiveString { + fn eq(&self, other: &Self) -> bool { + bool::from(self.0.as_bytes().ct_eq(other.0.as_bytes())) + } +} + +impl Eq for SensitiveString {} + +impl<'de> Deserialize<'de> for SensitiveString { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + Self::new(value).map_err(D::Error::custom) + } +} + +impl fmt::Debug for SensitiveString { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("[REDACTED]") + } +} + +/// The canonical transcript signed by both sides of MCP or worker registration. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct HandshakeTranscript { + pub(crate) daemon_target: String, + pub(crate) initiator: ComponentDescriptor, + pub(crate) responder: ComponentDescriptor, + pub(crate) initiator_instance_id: String, + pub(crate) responder_instance_id: String, + pub(crate) selected_protocol: u16, + pub(crate) initiator_public_identity: PublicIdentity, + pub(crate) responder_public_identity: PublicIdentity, + pub(crate) initiator_fingerprint: Fingerprint, + pub(crate) responder_fingerprint: Fingerprint, + pub(crate) challenge_id: ChallengeId, + pub(crate) initiator_nonce: ChallengeNonce, + pub(crate) responder_nonce: ChallengeNonce, + pub(crate) route_token_digest: Option, +} + +impl HandshakeTranscript { + /// Validates identities, service names, roles, and negotiated protocol values. + pub(crate) fn validate(&self) -> Result<(), ProtocolError> { + self.initiator.validate()?; + self.responder.validate()?; + if self.initiator.role == ComponentRole::Daemon + || self.responder.role != ComponentRole::Daemon + { + return Err(ProtocolError::InvalidRolePair); + } + if self.daemon_target.is_empty() + || self.daemon_target.len() > MAX_DAEMON_TARGET_BYTES + || self.initiator_instance_id.is_empty() + || self.initiator_instance_id.len() > MAX_INSTANCE_ID_BYTES + || self.responder_instance_id.is_empty() + || self.responder_instance_id.len() > MAX_INSTANCE_ID_BYTES + { + return Err(ProtocolError::MissingTranscriptIdentity); + } + if !self.initiator.protocol.contains(self.selected_protocol) + || !self.responder.protocol.contains(self.selected_protocol) + { + return Err(ProtocolError::InvalidSelectedProtocol); + } + if self.initiator_public_identity.fingerprint() != self.initiator_fingerprint + || self.responder_public_identity.fingerprint() != self.responder_fingerprint + { + return Err(ProtocolError::FingerprintMismatch); + } + if self.initiator.role == ComponentRole::Mcp && self.route_token_digest.is_none() { + return Err(ProtocolError::MissingRouteTokenDigest); + } + Ok(()) + } + + /// Encodes all signed fields deterministically and independently of JSON serialization. + pub(crate) fn canonical_bytes(&self) -> Result, ProtocolError> { + self.validate()?; + let initiator_protocol_minimum = self.initiator.protocol.minimum.to_be_bytes(); + let initiator_protocol_maximum = self.initiator.protocol.maximum.to_be_bytes(); + let responder_protocol_minimum = self.responder.protocol.minimum.to_be_bytes(); + let responder_protocol_maximum = self.responder.protocol.maximum.to_be_bytes(); + let selected_protocol = self.selected_protocol.to_be_bytes(); + let initiator_capabilities = self.initiator.capabilities.canonical_bytes()?; + let responder_capabilities = self.responder.capabilities.canonical_bytes()?; + let route_token_present = [u8::from(self.route_token_digest.is_some())]; + let route_token_digest = self + .route_token_digest + .as_ref() + .map_or(&[][..], |digest| digest.as_bytes().as_slice()); + let fields = [ + ("daemon_target", self.daemon_target.as_bytes()), + ("initiator_service", self.initiator.service.as_bytes()), + ( + "initiator_role", + self.initiator.role.transcript_name().as_bytes(), + ), + ( + "initiator_protocol_minimum", + initiator_protocol_minimum.as_slice(), + ), + ( + "initiator_protocol_maximum", + initiator_protocol_maximum.as_slice(), + ), + ("initiator_capabilities", initiator_capabilities.as_slice()), + ( + "initiator_binary_version", + self.initiator.binary_version.as_bytes(), + ), + ("responder_service", self.responder.service.as_bytes()), + ( + "responder_role", + self.responder.role.transcript_name().as_bytes(), + ), + ( + "responder_protocol_minimum", + responder_protocol_minimum.as_slice(), + ), + ( + "responder_protocol_maximum", + responder_protocol_maximum.as_slice(), + ), + ("responder_capabilities", responder_capabilities.as_slice()), + ( + "responder_binary_version", + self.responder.binary_version.as_bytes(), + ), + ( + "initiator_instance_id", + self.initiator_instance_id.as_bytes(), + ), + ( + "responder_instance_id", + self.responder_instance_id.as_bytes(), + ), + ("selected_protocol", selected_protocol.as_slice()), + ( + "initiator_public_identity", + self.initiator_public_identity.as_bytes().as_slice(), + ), + ( + "responder_public_identity", + self.responder_public_identity.as_bytes().as_slice(), + ), + ( + "initiator_fingerprint", + self.initiator_fingerprint.as_bytes().as_slice(), + ), + ( + "responder_fingerprint", + self.responder_fingerprint.as_bytes().as_slice(), + ), + ("challenge_id", self.challenge_id.as_bytes().as_slice()), + ( + "initiator_nonce", + self.initiator_nonce.as_bytes().as_slice(), + ), + ( + "responder_nonce", + self.responder_nonce.as_bytes().as_slice(), + ), + ("route_token_present", route_token_present.as_slice()), + ("route_token_digest", route_token_digest), + ]; + encode_transcript(HANDSHAKE_DOMAIN, &fields).map_err(ProtocolError::Transcript) + } + + /// Signs the canonical transcript for one of its declared participants. + pub(crate) fn sign( + &self, + signer: ComponentRole, + identity: &MachineIdentity, + ) -> Result { + let expected_identity = self.identity_for_role(signer)?; + if identity.public_identity() != expected_identity { + return Err(ProtocolError::SignerIdentityMismatch); + } + Ok(HandshakeProof { + signer, + signature: identity.sign(&self.canonical_bytes()?), + }) + } + + /// Verifies that a proof signs this exact canonical transcript. + pub(crate) fn verify(&self, proof: &HandshakeProof) -> Result<(), ProtocolError> { + self.identity_for_role(proof.signer)? + .verify(&self.canonical_bytes()?, &proof.signature) + .map_err(ProtocolError::Identity) + } + + fn identity_for_role(&self, role: ComponentRole) -> Result { + if self.initiator.role == role { + return Ok(self.initiator_public_identity); + } + if self.responder.role == role { + return Ok(self.responder_public_identity); + } + Err(ProtocolError::UnknownSignerRole) + } +} + +/// A participant's signature over a complete handshake transcript. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct HandshakeProof { + pub(crate) signer: ComponentRole, + pub(crate) signature: Ed25519Signature, +} + +/// A daemon-issued plan for launching one worker. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub(crate) struct WorkerLaunch { + pub(crate) activation_id: String, + pub(crate) activation_token: SensitiveString, + pub(crate) deadline_unix_ms: u64, + pub(crate) bind_ip: Ipv4Addr, + pub(crate) port: u16, + pub(crate) advertise_address: Option, +} + +impl WorkerLaunch { + /// Converts the launch plan into its wire directive. + pub(crate) fn into_directive(self) -> BrokerDirective { + BrokerDirective::LaunchWorker { + activation_id: self.activation_id, + activation_token: self.activation_token, + deadline_unix_ms: self.deadline_unix_ms, + bind_ip: self.bind_ip, + port: self.port, + advertise_address: self.advertise_address, + } + } +} + +/// The daemon's authoritative instruction for an MCP session. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "directive", rename_all = "snake_case")] +pub(crate) enum BrokerDirective { + ReuseWorker { + endpoint: String, + }, + WaitForWorker { + retry_after_ms: u64, + }, + LaunchWorker { + activation_id: String, + activation_token: SensitiveString, + deadline_unix_ms: u64, + bind_ip: Ipv4Addr, + port: u16, + advertise_address: Option, + }, + UsePassThrough, +} + +/// Protocol construction or verification failures. +#[derive(Debug, Error, PartialEq, Eq)] +pub(crate) enum ProtocolError { + #[error("the peer did not identify itself as nemo-relay")] + WrongService, + #[error("the protocol version range is invalid")] + InvalidProtocolRange, + #[error("the peers have no overlapping daemon protocol version")] + NoProtocolOverlap, + #[error("the selected protocol version is not supported by both peers")] + InvalidSelectedProtocol, + #[error("the component binary version is missing")] + MissingBinaryVersion, + #[error("the component binary version exceeds the protocol limit")] + BinaryVersionTooLong, + #[error("the capability set contains an invalid name")] + InvalidCapability, + #[error("the handshake role pair must be MCP/daemon or worker/daemon")] + InvalidRolePair, + #[error("the handshake is missing a daemon target or component instance ID")] + MissingTranscriptIdentity, + #[error("a public identity does not match its advertised fingerprint")] + FingerprintMismatch, + #[error("an MCP handshake is missing its route-token digest")] + MissingRouteTokenDigest, + #[error("the signing key does not match the transcript participant")] + SignerIdentityMismatch, + #[error("the proof signer is not a participant in this handshake")] + UnknownSignerRole, + #[error("a required sensitive protocol value is empty")] + MissingSensitiveValue, + #[error("a protocol field is too large")] + FieldTooLarge, + #[error(transparent)] + Identity(IdentityError), + #[error("failed to encode the signed transcript: {0}")] + Transcript(IdentityError), +} + +fn valid_capability_name(name: &str) -> bool { + !name.is_empty() + && name.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"._-".contains(&byte) + }) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/protocol_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/routes.rs b/crates/cli/src/daemon/common/routes.rs new file mode 100644 index 000000000..40d0a969c --- /dev/null +++ b/crates/cli/src/daemon/common/routes.rs @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Stable public route classification shared by daemon pass-through and workers. + +use crate::configuration::GatewayConfig; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PublicRoute { + Hook(HookRoute), + Provider(ProviderRoute), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum HookRoute { + Codex, + Claude, + Pi, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum ProviderRoute { + OpenAi, + Anthropic, +} + +impl PublicRoute { + pub(crate) fn from_path(path: &str) -> Option { + match path { + "/hooks/codex" => Some(Self::Hook(HookRoute::Codex)), + "/hooks/claude-code" => Some(Self::Hook(HookRoute::Claude)), + "/hooks/pi" => Some(Self::Hook(HookRoute::Pi)), + "/responses" + | "/chat/completions" + | "/models" + | "/v1/responses" + | "/backend-api/codex/responses" + | "/v1/chat/completions" + | "/v1/images/generations" + | "/v1/models" => Some(Self::Provider(ProviderRoute::OpenAi)), + "/v1/messages" | "/v1/messages/count_tokens" => { + Some(Self::Provider(ProviderRoute::Anthropic)) + } + _ => None, + } + } +} + +impl HookRoute { + pub(crate) const fn pass_through_body(self) -> &'static [u8] { + match self { + Self::Codex | Self::Pi => b"{}", + Self::Claude => br#"{"continue":true}"#, + } + } +} + +impl ProviderRoute { + pub(crate) fn upstream_url(self, config: &GatewayConfig, path_and_query: &str) -> String { + let base = match self { + Self::OpenAi => config.openai_base_url.as_str(), + Self::Anthropic => config.anthropic_base_url.as_str(), + } + .trim_end_matches('/'); + let path = match self { + Self::OpenAi => canonical_openai_path(path_and_query), + Self::Anthropic => path_and_query.to_owned(), + }; + let path = normalize_v1_path(base, &path); + format!("{base}{path}") + } +} + +fn canonical_openai_path(path_and_query: &str) -> String { + path_and_query + .strip_prefix("/backend-api/codex/responses") + .map_or_else( + || path_and_query.to_owned(), + |suffix| format!("/responses{suffix}"), + ) +} + +fn normalize_v1_path(base: &str, path_and_query: &str) -> String { + match (base.ends_with("/v1"), path_and_query.starts_with("/v1/")) { + (true, true) => path_and_query + .strip_prefix("/v1") + .expect("prefix was checked") + .to_owned(), + (false, false) => format!("/v1{path_and_query}"), + _ => path_and_query.to_owned(), + } +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/routes_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/state.rs b/crates/cli/src/daemon/common/state.rs new file mode 100644 index 000000000..2039dfc04 --- /dev/null +++ b/crates/cli/src/daemon/common/state.rs @@ -0,0 +1,543 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Owner-private daemon identity, trust, and environment state. + +use std::collections::HashMap; +use std::fmt; +use std::fs::{self, OpenOptions}; +use std::io::Read; +use std::path::{Path, PathBuf}; +use std::thread; +use std::time::{Duration, Instant}; + +use base64::Engine; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +use super::identity::{Fingerprint, MachineIdentity, PublicIdentity, TokenDigest}; +use crate::error::CliError; +use crate::filesystem::{LockAttempt, atomic_write_private, try_lock_exclusive, unlock_file}; + +pub(crate) const ROUTE_TOKEN_ENV: &str = "NEMO_RELAY_CLIENT_TOKEN"; +const IDENTITY_LOCK_TIMEOUT: Duration = Duration::from_secs(5); +const MAX_IDENTITY_BYTES: u64 = 4 * 1024; +const ACTIVE_WORKER_GENERATIONS_FILENAME: &str = "active-worker-generations.json"; +const ACTIVE_WORKER_GENERATIONS_SCHEMA_VERSION: u32 = 1; +const MAX_ACTIVE_WORKER_GENERATIONS: usize = 4_096; +const MAX_ACTIVE_WORKER_GENERATIONS_BYTES: u64 = 2 * 1024 * 1024; +const MAX_GENERATION_ID_BYTES: usize = 128; + +#[derive(Clone)] +pub(crate) struct RouteCredential { + value: String, + digest: TokenDigest, +} + +impl fmt::Debug for RouteCredential { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("RouteCredential") + .field("digest", &self.digest) + .finish_non_exhaustive() + } +} + +impl RouteCredential { + pub(crate) fn from_environment() -> Result { + let value = std::env::var(ROUTE_TOKEN_ENV).map_err(|_| { + CliError::Config(format!( + "managed daemon integration requires {ROUTE_TOKEN_ENV}; contact the managed environment administrator" + )) + })?; + Self::parse(value) + } + + pub(crate) fn parse(value: String) -> Result { + if value.trim() != value || value.bytes().any(|byte| byte.is_ascii_whitespace()) { + return Err(CliError::Config(format!( + "{ROUTE_TOKEN_ENV} must be an unpadded base64url credential without whitespace" + ))); + } + let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(&value) + .map_err(|_| { + CliError::Config(format!( + "{ROUTE_TOKEN_ENV} must be an unpadded base64url credential" + )) + })?; + if decoded.len() != 32 { + return Err(CliError::Config(format!( + "{ROUTE_TOKEN_ENV} must decode to exactly 32 bytes" + ))); + } + let digest = TokenDigest::from_token(value.as_bytes()); + Ok(Self { value, digest }) + } + + pub(crate) fn expose(&self) -> &str { + &self.value + } + + pub(crate) const fn digest(&self) -> TokenDigest { + self.digest + } +} + +pub(crate) fn load_or_create_machine_identity() -> Result { + load_or_create_identity(&daemon_state_dir()?.join("machine-identity.pk8")) +} + +pub(crate) fn load_or_create_daemon_identity() -> Result { + load_or_create_identity(&daemon_state_dir()?.join("daemon-identity.pk8")) +} + +/// Owner-private durable record of the only worker generation allowed to recover per route. +#[derive(Debug)] +pub(crate) struct ActiveWorkerGenerations { + path: PathBuf, +} + +impl ActiveWorkerGenerations { + /// Loads and validates the durable generation record before the daemon accepts traffic. + pub(crate) fn load() -> Result { + Self::load_from_path(daemon_state_dir()?.join(ACTIVE_WORKER_GENERATIONS_FILENAME)) + } + + fn load_from_path(path: PathBuf) -> Result { + let state = Self { path }; + state.with_locked_generations(|_| Ok(((), false)))?; + Ok(state) + } + + /// Returns whether `generation_id` is the exact active generation for `fingerprint`. + pub(crate) fn matches( + &self, + fingerprint: Fingerprint, + generation_id: &str, + ) -> Result { + validate_generation_id(generation_id)?; + self.with_locked_generations(|generations| { + Ok(( + generations + .get(&fingerprint) + .is_some_and(|active| active == generation_id), + false, + )) + }) + } + + /// Publishes a ready generation, atomically replacing any prior generation for the route. + pub(crate) fn publish( + &self, + fingerprint: Fingerprint, + generation_id: &str, + ) -> Result, CliError> { + validate_generation_id(generation_id)?; + self.with_locked_generations(|generations| { + if !generations.contains_key(&fingerprint) + && generations.len() >= MAX_ACTIVE_WORKER_GENERATIONS + { + return Err(CliError::Config(format!( + "active worker generation state exceeds {MAX_ACTIVE_WORKER_GENERATIONS} routes" + ))); + } + let previous = generations.insert(fingerprint, generation_id.to_owned()); + let changed = previous.as_deref() != Some(generation_id); + Ok((previous, changed)) + }) + } + + /// Revokes a generation only if it is still active, protecting a newer replacement. + pub(crate) fn revoke_if_matches( + &self, + fingerprint: Fingerprint, + generation_id: &str, + ) -> Result { + validate_generation_id(generation_id)?; + self.with_locked_generations(|generations| { + let matches = generations + .get(&fingerprint) + .is_some_and(|active| active == generation_id); + if matches { + generations.remove(&fingerprint); + } + Ok((matches, matches)) + }) + } + + /// Restores the prior value if a broker publication loses a race after durable publication. + pub(crate) fn restore_if_matches( + &self, + fingerprint: Fingerprint, + expected_generation_id: &str, + previous_generation_id: Option<&str>, + ) -> Result { + validate_generation_id(expected_generation_id)?; + if let Some(previous) = previous_generation_id { + validate_generation_id(previous)?; + } + self.with_locked_generations(|generations| { + let matches = generations + .get(&fingerprint) + .is_some_and(|active| active == expected_generation_id); + if !matches { + return Ok((false, false)); + } + match previous_generation_id { + Some(previous) => { + generations.insert(fingerprint, previous.to_owned()); + } + None => { + generations.remove(&fingerprint); + } + } + Ok((true, true)) + }) + } + + fn with_locked_generations( + &self, + operation: impl FnOnce(&mut HashMap) -> Result<(T, bool), CliError>, + ) -> Result { + let parent = self + .path + .parent() + .ok_or_else(|| CliError::Config("worker generation state path has no parent".into()))?; + create_private_directory(parent)?; + let lock_path = self.path.with_extension("lock"); + let lock = open_private_lock(&lock_path)?; + acquire_private_lock(&lock, &lock_path)?; + let result = (|| { + let mut generations = read_active_worker_generations(&self.path)?; + let (output, changed) = operation(&mut generations)?; + if changed { + write_active_worker_generations(&self.path, &generations)?; + } + Ok(output) + })(); + let _ = unlock_file(&lock); + result + } + + #[cfg(test)] + pub(crate) fn load_for_test(path: PathBuf) -> Result { + Self::load_from_path(path) + } +} + +#[derive(Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct PersistedActiveWorkerGenerations { + schema_version: u32, + generations: Vec, +} + +#[derive(Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct PersistedActiveWorkerGeneration { + fingerprint: Fingerprint, + generation_id: String, +} + +fn read_active_worker_generations(path: &Path) -> Result, CliError> { + let Some(bytes) = read_bounded( + path, + MAX_ACTIVE_WORKER_GENERATIONS_BYTES, + "active worker generation state", + )? + else { + return Ok(HashMap::new()); + }; + let persisted: PersistedActiveWorkerGenerations = + serde_json::from_slice(&bytes).map_err(|_| { + CliError::Config(format!( + "active worker generation state {} is corrupt", + path.display() + )) + })?; + if persisted.schema_version != ACTIVE_WORKER_GENERATIONS_SCHEMA_VERSION { + return Err(CliError::Config(format!( + "active worker generation state {} has unsupported schema version {}", + path.display(), + persisted.schema_version + ))); + } + if persisted.generations.len() > MAX_ACTIVE_WORKER_GENERATIONS { + return Err(CliError::Config(format!( + "active worker generation state {} exceeds {MAX_ACTIVE_WORKER_GENERATIONS} routes", + path.display() + ))); + } + let mut generations = HashMap::with_capacity(persisted.generations.len()); + for entry in persisted.generations { + validate_generation_id(&entry.generation_id)?; + if generations + .insert(entry.fingerprint, entry.generation_id) + .is_some() + { + return Err(CliError::Config(format!( + "active worker generation state {} contains a duplicate fingerprint", + path.display() + ))); + } + } + Ok(generations) +} + +fn write_active_worker_generations( + path: &Path, + generations: &HashMap, +) -> Result<(), CliError> { + let mut entries = generations + .iter() + .map( + |(fingerprint, generation_id)| PersistedActiveWorkerGeneration { + fingerprint: *fingerprint, + generation_id: generation_id.clone(), + }, + ) + .collect::>(); + entries.sort_by_key(|entry| entry.fingerprint.to_string()); + let document = PersistedActiveWorkerGenerations { + schema_version: ACTIVE_WORKER_GENERATIONS_SCHEMA_VERSION, + generations: entries, + }; + let bytes = serde_json::to_vec(&document).map_err(|error| { + CliError::Config(format!( + "failed to serialize active worker generation state: {error}" + )) + })?; + if bytes.len() as u64 > MAX_ACTIVE_WORKER_GENERATIONS_BYTES { + return Err(CliError::Config(format!( + "active worker generation state exceeds {MAX_ACTIVE_WORKER_GENERATIONS_BYTES} bytes" + ))); + } + atomic_write_private(path, &bytes).map_err(CliError::Config)?; + sync_parent_directory(path) +} + +fn validate_generation_id(generation_id: &str) -> Result<(), CliError> { + if generation_id.is_empty() || generation_id.len() > MAX_GENERATION_ID_BYTES { + return Err(CliError::Config( + "active worker generation ID is invalid".into(), + )); + } + Ok(()) +} + +pub(crate) fn verify_or_store_daemon_pin( + daemon_origin: &str, + identity: PublicIdentity, +) -> Result<(), CliError> { + let name = hex_digest(daemon_origin.as_bytes()); + let path = daemon_state_dir()? + .join("pins") + .join(format!("{name}.ed25519")); + let parent = path + .parent() + .ok_or_else(|| CliError::Config("daemon pin path has no parent".into()))?; + create_private_directory(parent)?; + let lock_path = path.with_extension("lock"); + let lock = open_private_lock(&lock_path)?; + acquire_private_lock(&lock, &lock_path)?; + let result = match read_bounded(&path, MAX_IDENTITY_BYTES, "daemon trust pin")? { + Some(existing) if existing == identity.as_bytes() => Ok(()), + Some(_) => Err(CliError::Unauthorized(format!( + "daemon identity changed for {daemon_origin}; remove the owner-private trust pin only after verifying the daemon replacement" + ))), + None => atomic_write_private(&path, identity.as_bytes()).map_err(CliError::Config), + }; + let _ = unlock_file(&lock); + result +} + +fn load_or_create_identity(path: &Path) -> Result { + let parent = path + .parent() + .ok_or_else(|| CliError::Config("daemon identity path has no parent".into()))?; + create_private_directory(parent)?; + let lock_path = path.with_extension("lock"); + let lock = open_private_lock(&lock_path)?; + acquire_private_lock(&lock, &lock_path)?; + let result = match read_bounded(path, MAX_IDENTITY_BYTES, "daemon identity state")? { + Some(bytes) => MachineIdentity::from_pkcs8(&bytes) + .map_err(|error| CliError::Config(format!("invalid daemon identity: {error}"))), + None => { + let generated = MachineIdentity::generate().map_err(|error| { + CliError::Config(format!("failed to generate identity: {error}")) + })?; + atomic_write_private(path, &generated.pkcs8).map_err(CliError::Config)?; + Ok(generated.identity) + } + }; + let _ = unlock_file(&lock); + result +} + +fn open_private_lock(path: &Path) -> Result { + #[cfg(windows)] + let file = crate::filesystem::open_private_windows_file(path).map_err(|error| { + CliError::Config(format!( + "failed to open owner-private daemon lock {}: {error}", + path.display() + )) + })?; + #[cfg(not(windows))] + let file = { + let mut options = OpenOptions::new(); + options.create(true).read(true).write(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.mode(0o600).custom_flags(libc::O_NOFOLLOW); + } + options.open(path).map_err(|error| { + CliError::Config(format!( + "failed to open owner-private daemon lock {}: {error}", + path.display() + )) + })? + }; + validate_private_file(&file, path)?; + Ok(file) +} + +fn acquire_private_lock(file: &std::fs::File, path: &Path) -> Result<(), CliError> { + let deadline = Instant::now() + IDENTITY_LOCK_TIMEOUT; + loop { + match try_lock_exclusive(file) { + Ok(LockAttempt::Acquired) => return Ok(()), + Ok(LockAttempt::Contended) if Instant::now() < deadline => { + thread::sleep(Duration::from_millis(25)); + } + Ok(LockAttempt::Contended) => { + return Err(CliError::Config(format!( + "timed out waiting for daemon identity lock {}", + path.display() + ))); + } + Err(error) => return Err(CliError::Io(error)), + } + } +} + +fn read_bounded( + path: &Path, + max_bytes: u64, + description: &str, +) -> Result>, CliError> { + let mut options = OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.custom_flags(libc::O_NOFOLLOW); + } + let mut file = match options.open(path) { + Ok(file) => file, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(CliError::Io(error)), + }; + validate_private_file(&file, path)?; + let length = file.metadata()?.len(); + if length > max_bytes { + return Err(CliError::Config(format!( + "{description} {} exceeds {max_bytes} bytes", + path.display(), + ))); + } + let mut bytes = Vec::with_capacity(length as usize); + file.read_to_end(&mut bytes)?; + Ok(Some(bytes)) +} + +fn daemon_state_dir() -> Result { + crate::configuration::user_config_dir() + .map(|directory| directory.join("daemon")) + .ok_or_else(|| { + CliError::Config( + "cannot determine the per-user daemon state directory; set HOME or USERPROFILE" + .into(), + ) + }) +} + +fn create_private_directory(path: &Path) -> Result<(), CliError> { + match fs::symlink_metadata(path) { + Ok(metadata) if !metadata.file_type().is_dir() => { + return Err(CliError::Config(format!( + "daemon state directory {} must be a real directory", + path.display() + ))); + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => fs::create_dir_all(path)?, + Err(error) => return Err(CliError::Io(error)), + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + + use std::os::unix::fs::MetadataExt; + let metadata = fs::symlink_metadata(path)?; + if metadata.uid() != unsafe { libc::geteuid() } { + return Err(CliError::Config(format!( + "daemon state directory {} is not owned by the current user", + path.display() + ))); + } + fs::set_permissions(path, fs::Permissions::from_mode(0o700))?; + } + #[cfg(windows)] + crate::filesystem::protect_private_windows_path(path)?; + Ok(()) +} + +fn validate_private_file(file: &std::fs::File, path: &Path) -> Result<(), CliError> { + let metadata = file.metadata()?; + if !metadata.is_file() { + return Err(CliError::Config(format!( + "daemon state {} must be a regular file", + path.display() + ))); + } + #[cfg(unix)] + { + use std::os::unix::fs::{MetadataExt, PermissionsExt}; + if metadata.uid() != unsafe { libc::geteuid() } { + return Err(CliError::Config(format!( + "daemon state {} is not owned by the current user", + path.display() + ))); + } + if metadata.mode() & 0o077 != 0 { + file.set_permissions(fs::Permissions::from_mode(0o600))?; + } + } + #[cfg(windows)] + crate::filesystem::protect_private_windows_path(path)?; + Ok(()) +} + +fn sync_parent_directory(_path: &Path) -> Result<(), CliError> { + #[cfg(unix)] + { + let parent = _path + .parent() + .ok_or_else(|| CliError::Config("worker generation state path has no parent".into()))?; + std::fs::File::open(parent)?.sync_all()?; + } + Ok(()) +} + +fn hex_digest(bytes: &[u8]) -> String { + Sha256::digest(bytes) + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/state_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/transport.rs b/crates/cli/src/daemon/common/transport.rs new file mode 100644 index 000000000..d6e8268c1 --- /dev/null +++ b/crates/cli/src/daemon/common/transport.rs @@ -0,0 +1,308 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Lossless, pull-driven HTTP transport used between daemon data-plane hops. +//! +//! The adapters in this module deliberately operate on [`Frame`] values rather than on decoded +//! payloads. Data and trailer frames therefore stay under Hyper's normal demand-driven +//! backpressure, and dropping the downstream body cancels the upstream body without a forwarding +//! task or intermediate queue. + +use std::error::Error; +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::time::Duration; + +use bytes::Bytes; +use http::header::{CONNECTION, HOST, HeaderName, HeaderValue, TE, UPGRADE}; +use http::{HeaderMap, Method, Request, Response, StatusCode, Uri, Version}; +use http_body_util::BodyExt; +use http_body_util::combinators::UnsyncBoxBody; +use hyper::body::Body; +use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder}; +use hyper_util::client::legacy::Client; +use hyper_util::client::legacy::connect::HttpConnector; +use hyper_util::rt::{TokioExecutor, TokioTimer}; +use thiserror::Error; + +/// Error type shared by transport bodies after their concrete body implementation is erased. +pub(crate) type BoxError = Box; + +/// A pull-driven body that preserves both data and trailer frames. +pub(crate) type RelayBody = UnsyncBoxBody; + +/// A pooled client supporting cleartext HTTP and rustls-backed HTTPS with HTTP/1.1 and HTTP/2. +pub(crate) type PooledClient = Client, RelayBody>; + +/// A pooled cleartext client. The normal builder uses HTTP/1.1; the h2c builder uses HTTP/2 prior +/// knowledge. +#[cfg(test)] +pub(crate) type PooledHttpClient = Client; + +pub(crate) const CONNECT_TIMEOUT: Duration = Duration::from_secs(10); +pub(crate) const MAX_IDLE_CONNECTIONS_PER_HOST: usize = 256; + +#[derive(Debug, Error)] +pub(crate) enum TransportError { + #[error("CONNECT and HTTP Upgrade are not supported by the daemon data plane")] + UnsupportedTunnel, + #[error("forward destination must contain an HTTP or HTTPS scheme and an authority")] + InvalidDestination, + #[error("invalid Connection header value")] + InvalidConnectionHeader, + #[error("forward destination authority is not a valid Host header")] + InvalidHost(#[source] http::header::InvalidHeaderValue), + #[error("failed to load native TLS trust roots")] + NativeRoots(#[source] std::io::Error), +} + +/// Erases a body's implementation and error while retaining its pull-based [`Body::poll_frame`] +/// behavior. This function does not spawn a forwarding task, decode frames, or queue bytes. +pub(crate) fn box_body(body: B) -> RelayBody +where + B: Body + Send + 'static, + B::Error: Into, +{ + body.map_err(Into::into).boxed_unsync() +} + +/// Keeps request accounting or another lifetime guard alive until a body completes or is dropped. +/// The body is still polled directly; no forwarding task or queue is introduced. +pub(crate) fn hold_body(body: B, hold: H) -> RelayBody +where + B: Body + Send + Unpin + 'static, + B::Error: Into, + H: Send + Unpin + 'static, +{ + box_body(HeldBody { + body, + hold: Some(hold), + }) +} + +struct HeldBody { + body: B, + hold: Option, +} + +impl Body for HeldBody +where + B: Body + Unpin, + B::Error: Into, + H: Unpin, +{ + type Data = Bytes; + type Error = BoxError; + + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let frame = Pin::new(&mut self.body).poll_frame(context); + if matches!(frame, Poll::Ready(None)) { + self.hold.take(); + } + frame.map(|frame| frame.map(|result| result.map_err(Into::into))) + } + + fn is_end_stream(&self) -> bool { + self.body.is_end_stream() + } + + fn size_hint(&self) -> hyper::body::SizeHint { + self.body.size_hint() + } +} + +/// Removes fields scoped to one HTTP connection. +/// +/// `Trailer` is intentionally not in the fixed hop-by-hop list. It declares the fields carried by +/// a later trailer frame and remains valid across a framing-preserving intermediary. It is removed +/// only when an incoming `Connection` field explicitly nominates it. +pub(crate) fn strip_hop_by_hop_headers(headers: &mut HeaderMap) -> Result<(), TransportError> { + let nominated = connection_nominated_headers(headers)?; + + for name in nominated { + headers.remove(name); + } + for name in [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "proxy-connection", + "te", + "transfer-encoding", + "upgrade", + ] { + headers.remove(name); + } + + Ok(()) +} + +/// Rewrites a request head for one proxy hop while transferring ownership of its body unchanged. +/// +/// `destination` is the complete URI selected by the router, including its path and query. Names +/// in `additional_strip` are routing or authentication fields consumed by the current hop. +pub(crate) fn prepare_forward_request( + mut request: Request, + destination: Uri, + additional_strip: &[HeaderName], +) -> Result, TransportError> { + if request.method() == Method::CONNECT || request.headers().contains_key(UPGRADE) { + return Err(TransportError::UnsupportedTunnel); + } + + let scheme = destination + .scheme_str() + .filter(|scheme| matches!(*scheme, "http" | "https")); + let authority = destination.authority().cloned(); + if scheme.is_none() || authority.is_none() { + return Err(TransportError::InvalidDestination); + } + let authority = authority.expect("authority was checked above"); + if scheme == Some("http") + && !authority.host().eq_ignore_ascii_case("localhost") + && !authority + .host() + .trim_matches(['[', ']']) + .parse::() + .is_ok_and(|address| address.is_loopback()) + { + return Err(TransportError::InvalidDestination); + } + let host = HeaderValue::from_str(authority.as_str()).map_err(TransportError::InvalidHost)?; + + strip_hop_by_hop_headers(request.headers_mut())?; + for name in additional_strip { + request.headers_mut().remove(name); + } + // `TE` is scoped to one connection, but Relay accepts and relays trailer frames. Advertise + // that capability independently on every upstream hop after consuming the caller's value. + request + .headers_mut() + .insert(TE, HeaderValue::from_static("trailers")); + request.headers_mut().insert(HOST, host); + *request.uri_mut() = destination; + // The protocol version belongs to the connection on which this request arrived. It is not a + // requirement for the next proxy hop: leaving HTTP/2 here makes Hyper reject an H2 ingress + // request when the selected upstream only speaks HTTP/1.1. HTTP/1.1 is the neutral request + // value; ALPN or an H2-only client still selects HTTP/2 independently. + *request.version_mut() = Version::HTTP_11; + + Ok(request) +} + +/// Filters a response head for one proxy hop while transferring ownership of its body unchanged. +pub(crate) fn prepare_forward_response( + mut response: Response, + additional_strip: &[HeaderName], +) -> Result, TransportError> { + if response.status() == StatusCode::SWITCHING_PROTOCOLS + || response.headers().contains_key(UPGRADE) + { + return Err(TransportError::UnsupportedTunnel); + } + + strip_hop_by_hop_headers(response.headers_mut())?; + for name in additional_strip { + response.headers_mut().remove(name); + } + Ok(response) +} + +/// Builds a pooled HTTP(S) client. Callers should construct this once per process and clone its +/// lightweight handle rather than building one per request. +pub(crate) fn pooled_client() -> Result { + // The workspace enables more than one rustls backend through unrelated integrations. Select + // Relay's direct `ring` dependency before rustls tries to infer a process-wide provider. + let _ = rustls::crypto::ring::default_provider().install_default(); + let connector = pooled_connector()?; + let builder = pooled_builder(); + Ok(builder.build(connector)) +} + +/// Builds a pooled cleartext HTTP/1.1 client with persistent connections and `TCP_NODELAY`. +#[cfg(test)] +pub(crate) fn pooled_http_client() -> PooledHttpClient { + let mut connector = HttpConnector::new(); + connector.enforce_http(true); + connector.set_nodelay(true); + connector.set_connect_timeout(Some(CONNECT_TIMEOUT)); + let builder = pooled_builder(); + builder.build(connector) +} + +/// Builds a pooled cleartext HTTP/2 client using prior knowledge rather than an Upgrade exchange. +#[cfg(test)] +pub(crate) fn pooled_h2c_client() -> PooledHttpClient { + let mut connector = HttpConnector::new(); + connector.enforce_http(true); + connector.set_nodelay(true); + connector.set_connect_timeout(Some(CONNECT_TIMEOUT)); + + let mut builder = pooled_builder(); + builder.http2_only(true); + builder.build(connector) +} + +/// Builds the production cleartext worker pool using HTTP/2 prior knowledge. +/// +/// Loopback workers selected by this client must support HTTP/2 without an Upgrade exchange. +pub(crate) fn pooled_worker_h2c_client() -> Result { + let _ = rustls::crypto::ring::default_provider().install_default(); + let connector = pooled_connector()?; + let mut builder = pooled_builder(); + builder.http2_only(true); + Ok(builder.build(connector)) +} + +fn pooled_connector() -> Result, TransportError> { + let mut http = HttpConnector::new(); + http.enforce_http(false); + http.set_nodelay(true); + http.set_connect_timeout(Some(CONNECT_TIMEOUT)); + + Ok(HttpsConnectorBuilder::new() + .with_native_roots() + .map_err(TransportError::NativeRoots)? + .https_or_http() + .enable_http1() + .enable_http2() + .wrap_connector(http)) +} + +pub(crate) fn pooled_builder() -> hyper_util::client::legacy::Builder { + let mut builder = Client::builder(TokioExecutor::new()); + builder.timer(TokioTimer::new()); + builder.pool_idle_timeout(Duration::from_secs(120)); + builder.pool_max_idle_per_host(MAX_IDLE_CONNECTIONS_PER_HOST); + builder.http2_keep_alive_interval(Duration::from_secs(15)); + builder.http2_keep_alive_timeout(Duration::from_secs(5)); + builder.http2_keep_alive_while_idle(true); + builder +} + +fn connection_nominated_headers(headers: &HeaderMap) -> Result, TransportError> { + let mut nominated = Vec::new(); + for value in headers.get_all(CONNECTION) { + let value = value + .to_str() + .map_err(|_| TransportError::InvalidConnectionHeader)?; + for token in value + .split(',') + .map(str::trim) + .filter(|token| !token.is_empty()) + { + let name = HeaderName::from_bytes(token.as_bytes()) + .map_err(|_| TransportError::InvalidConnectionHeader)?; + nominated.push(name); + } + } + Ok(nominated) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/transport_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/common/worker_tls.rs b/crates/cli/src/daemon/common/worker_tls.rs new file mode 100644 index 000000000..4ba69287c --- /dev/null +++ b/crates/cli/src/daemon/common/worker_tls.rs @@ -0,0 +1,258 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Ephemeral, daemon-pinned TLS identities for network-reachable workers. + +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex, MutexGuard}; + +use base64::Engine; +use hyper_rustls::HttpsConnectorBuilder; +use hyper_util::client::legacy::connect::HttpConnector; +use rcgen::{ + BasicConstraints, CertificateParams, DnType, ExtendedKeyUsagePurpose, IsCa, KeyPair, + KeyUsagePurpose, +}; +use rustls::pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer}; +use sha2::{Digest, Sha256}; + +use super::transport::{CONNECT_TIMEOUT, PooledClient, pooled_builder, pooled_worker_h2c_client}; +use crate::error::CliError; + +const MAX_ROOT_CERTIFICATE_BYTES: usize = 64 * 1024; +const MAX_CACHED_TLS_WORKER_POOLS: usize = 256; + +/// The daemon's process-wide connection pools for authenticated worker data-plane traffic. +/// +/// Cleartext loopback workers share one HTTP/2 prior-knowledge pool. Network-reachable workers +/// share a pool only when they present the same pinned TLS root; the bounded LRU cache prevents +/// unbounded retention while active targets keep their selected pool alive through [`Arc`]. +pub(crate) struct WorkerClientPool { + cleartext_h2c: Arc, + tls: Mutex, +} + +impl WorkerClientPool { + pub(crate) fn new() -> Result { + Self::with_tls_capacity(MAX_CACHED_TLS_WORKER_POOLS) + } + + fn with_tls_capacity(capacity: usize) -> Result { + assert!(capacity > 0, "worker TLS pool capacity must be positive"); + Ok(Self { + cleartext_h2c: Arc::new( + pooled_worker_h2c_client().map_err(|error| CliError::Launch(error.to_string()))?, + ), + tls: Mutex::new(TlsClientCache::new(capacity)), + }) + } + + /// Selects a shared pool without weakening the per-root TLS trust boundary. + pub(crate) fn client( + &self, + tls_root_certificate: Option<&str>, + ) -> Result, CliError> { + let Some(root_certificate) = tls_root_certificate else { + return Ok(Arc::clone(&self.cleartext_h2c)); + }; + let root_der = decode_worker_tls_root(root_certificate)?; + let root_id = WorkerTlsRootId::from_der(&root_der); + let mut cache = lock(&self.tls); + if let Some(client) = cache.get(root_id) { + return Ok(client); + } + let client = Arc::new(pooled_worker_tls_client_from_der(root_der)?); + cache.insert(root_id, Arc::clone(&client)); + Ok(client) + } +} + +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +struct WorkerTlsRootId([u8; 32]); + +impl WorkerTlsRootId { + fn from_der(der: &[u8]) -> Self { + Self(Sha256::digest(der).into()) + } +} + +struct TlsClientCache { + capacity: usize, + entries: HashMap>, + recency: VecDeque, +} + +impl TlsClientCache { + fn new(capacity: usize) -> Self { + Self { + capacity, + entries: HashMap::with_capacity(capacity), + recency: VecDeque::with_capacity(capacity), + } + } + + fn get(&mut self, root_id: WorkerTlsRootId) -> Option> { + let client = self.entries.get(&root_id).cloned()?; + self.touch(root_id); + Some(client) + } + + fn insert(&mut self, root_id: WorkerTlsRootId, client: Arc) { + if self.entries.len() == self.capacity + && let Some(evicted) = self.recency.pop_front() + { + self.entries.remove(&evicted); + } + self.entries.insert(root_id, client); + self.recency.push_back(root_id); + } + + fn touch(&mut self, root_id: WorkerTlsRootId) { + if let Some(index) = self + .recency + .iter() + .position(|candidate| *candidate == root_id) + { + self.recency.remove(index); + } + self.recency.push_back(root_id); + } +} + +/// A worker-only server identity and the public trust anchor sent to the daemon. +pub(crate) struct WorkerTlsIdentity { + server_config: Arc, + root_certificate: String, +} + +impl WorkerTlsIdentity { + /// Generates a private root and a server leaf whose SAN matches the advertised worker host. + /// The root private key is discarded before this value is returned. + pub(crate) fn generate(advertised_host: &str) -> Result { + let advertised_host = advertised_host.trim_matches(['[', ']']); + let invalid_ip = advertised_host + .parse::() + .is_ok_and(|address| address.is_unspecified() || address.is_multicast()); + if advertised_host.is_empty() || invalid_ip { + return Err(CliError::Config( + "worker TLS requires a concrete advertised host or IP".into(), + )); + } + + let mut root_params = CertificateParams::new(Vec::::new()) + .map_err(|error| tls_error("create worker root parameters", error))?; + root_params + .distinguished_name + .push(DnType::CommonName, "NeMo Relay ephemeral worker root"); + root_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + root_params.key_usages = vec![ + KeyUsagePurpose::KeyCertSign, + KeyUsagePurpose::DigitalSignature, + ]; + let root_key = + KeyPair::generate().map_err(|error| tls_error("generate worker root key", error))?; + let root = root_params + .self_signed(&root_key) + .map_err(|error| tls_error("sign worker root certificate", error))?; + + let mut leaf_params = CertificateParams::new(vec![advertised_host.to_owned()]) + .map_err(|error| tls_error("create worker leaf parameters", error))?; + leaf_params + .distinguished_name + .push(DnType::CommonName, "nemo-relay worker"); + leaf_params.key_usages = vec![KeyUsagePurpose::DigitalSignature]; + leaf_params.extended_key_usages = vec![ExtendedKeyUsagePurpose::ServerAuth]; + let leaf_key = + KeyPair::generate().map_err(|error| tls_error("generate worker leaf key", error))?; + let leaf = leaf_params + .signed_by(&leaf_key, &root, &root_key) + .map_err(|error| tls_error("sign worker leaf certificate", error))?; + + let root_der = root.der().to_vec(); + let certificate_chain = vec![leaf.der().clone(), root.der().clone()]; + let private_key = PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(leaf_key.serialize_der())); + let _ = rustls::crypto::ring::default_provider().install_default(); + let mut server_config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certificate_chain, private_key) + .map_err(|error| tls_error("build worker TLS server", error))?; + server_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + + Ok(Self { + server_config: Arc::new(server_config), + root_certificate: base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(root_der), + }) + } + + pub(crate) fn server_config(&self) -> Arc { + Arc::clone(&self.server_config) + } + + pub(crate) fn root_certificate(&self) -> &str { + &self.root_certificate + } +} + +/// Builds one long-lived daemon-to-worker pool pinned to the worker's ephemeral root. +#[cfg(test)] +pub(crate) fn pooled_worker_tls_client(root_certificate: &str) -> Result { + pooled_worker_tls_client_from_der(decode_worker_tls_root(root_certificate)?) +} + +fn decode_worker_tls_root(root_certificate: &str) -> Result, CliError> { + let encoded = root_certificate.as_bytes(); + if encoded.len() > encoded_certificate_limit() { + return Err(CliError::Unauthorized( + "worker TLS root certificate exceeds its size limit".into(), + )); + } + let der = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(encoded) + .map_err(|_| CliError::Unauthorized("worker TLS root certificate is invalid".into()))?; + if der.is_empty() || der.len() > MAX_ROOT_CERTIFICATE_BYTES { + return Err(CliError::Unauthorized( + "worker TLS root certificate has an invalid size".into(), + )); + } + Ok(der) +} + +fn pooled_worker_tls_client_from_der(der: Vec) -> Result { + let _ = rustls::crypto::ring::default_provider().install_default(); + let mut roots = rustls::RootCertStore::empty(); + roots + .add(CertificateDer::from(der)) + .map_err(|_| CliError::Unauthorized("worker TLS root certificate is invalid".into()))?; + let client_config = rustls::ClientConfig::builder() + .with_root_certificates(roots) + .with_no_client_auth(); + + let mut http = HttpConnector::new(); + http.enforce_http(false); + http.set_nodelay(true); + http.set_connect_timeout(Some(CONNECT_TIMEOUT)); + let connector = HttpsConnectorBuilder::new() + .with_tls_config(client_config) + .https_only() + .enable_http1() + .enable_http2() + .wrap_connector(http); + let builder = pooled_builder(); + Ok(builder.build(connector)) +} + +const fn encoded_certificate_limit() -> usize { + MAX_ROOT_CERTIFICATE_BYTES.div_ceil(3) * 4 +} + +fn tls_error(context: &str, error: impl std::fmt::Display) -> CliError { + CliError::Launch(format!("failed to {context}: {error}")) +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(|error| error.into_inner()) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/worker_tls_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/hook/mod.rs b/crates/cli/src/daemon/hook/mod.rs new file mode 100644 index 000000000..911f45a80 --- /dev/null +++ b/crates/cli/src/daemon/hook/mod.rs @@ -0,0 +1,203 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Managed hook forwarding through an explicitly selected daemon. + +use std::io::{Read, Write}; +use std::time::Duration; + +use futures_util::StreamExt; +use reqwest::header::{CONTENT_TYPE, HeaderValue}; +use serde_json::Value; + +use crate::agents::CodingAgent; +use crate::daemon::common::state::{ROUTE_TOKEN_ENV, RouteCredential}; +use crate::error::CliError; +use crate::hooks::HookFailurePolicy; + +pub(crate) const CLIENT_TOKEN_ENV: &str = ROUTE_TOKEN_ENV; +const CLIENT_TOKEN_HEADER: &str = crate::configuration::BOOTSTRAP_CLIENT_TOKEN_HEADER; +const HOOK_FORWARD_TIMEOUT: Duration = Duration::from_secs(30); +const MAX_HOOK_RESPONSE_BYTES: usize = 1024 * 1024; + +#[derive(Debug, Clone)] +pub(crate) struct Options { + pub(crate) agent: CodingAgent, + pub(crate) daemon_address: String, + pub(crate) failure_policy: HookFailurePolicy, +} + +/// Reads one native hook payload, sends it to the daemon, and relays the response to stdout. +pub(crate) async fn run(options: Options) -> Result<(), CliError> { + let payload = read_hook_payload(std::io::stdin()); + let fail_closed = effective_fail_closed(options.failure_policy, payload.as_deref().ok()); + let result: Result<(), CliError> = async { + let token = route_token_from_environment()?; + let payload = payload?; + let body = forward(&options, payload, token).await?; + if !body.is_empty() { + std::io::stdout().write_all(&body)?; + } + Ok(()) + } + .await; + + match result { + Ok(()) => Ok(()), + Err(error) if error.guardrail_rejection_reason().is_some() => Err(error), + Err(error) => handle_delivery_failure(error, fail_closed), + } +} + +fn effective_fail_closed(policy: HookFailurePolicy, payload: Option<&[u8]>) -> bool { + match policy { + HookFailurePolicy::FailOpen => false, + HookFailurePolicy::FailClosed => true, + HookFailurePolicy::Default => { + if policy.fail_closed() { + return true; + } + payload + .and_then(|payload| serde_json::from_slice::(payload).ok()) + .and_then(|payload| { + ["hook_event_name", "event_name", "event", "type"] + .into_iter() + .find_map(|name| payload.get(name).and_then(Value::as_str)) + .map(crate::hooks::event_requires_fail_closed) + }) + .unwrap_or(false) + } + } +} + +fn read_hook_payload(mut reader: impl Read) -> Result, CliError> { + let limit = crate::configuration::DEFAULT_MAX_HOOK_PAYLOAD_BYTES; + let mut payload = Vec::new(); + reader + .by_ref() + .take(limit.saturating_add(1) as u64) + .read_to_end(&mut payload)?; + if payload.len() > limit { + return Err(CliError::PayloadTooLarge(format!( + "hook payload exceeds the {limit}-byte limit" + ))); + } + std::str::from_utf8(&payload) + .map_err(|error| CliError::InvalidPayload(format!("hook payload is not UTF-8: {error}")))?; + if payload.iter().all(u8::is_ascii_whitespace) { + Ok(b"{}".to_vec()) + } else { + Ok(payload) + } +} + +fn route_token_from_environment() -> Result { + let credential = RouteCredential::from_environment()?; + HeaderValue::from_str(credential.expose()) + .map_err(|_| CliError::Config(format!("{CLIENT_TOKEN_ENV} is not valid HTTP header text"))) +} + +#[cfg(test)] +fn route_token(value: &str) -> Result { + let credential = RouteCredential::parse(value.to_owned())?; + HeaderValue::from_str(credential.expose()) + .map_err(|_| CliError::Config(format!("{CLIENT_TOKEN_ENV} is not valid HTTP header text"))) +} + +async fn forward( + options: &Options, + payload: Vec, + token: HeaderValue, +) -> Result, CliError> { + let endpoint = hook_endpoint(&options.daemon_address, options.agent)?; + let response = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .timeout(HOOK_FORWARD_TIMEOUT) + .build()? + .post(endpoint) + .header(CONTENT_TYPE, "application/json") + .header(CLIENT_TOKEN_HEADER, token) + .body(payload) + .send() + .await?; + let status = response.status(); + let body = read_response(response).await?; + if status.is_success() { + return Ok(body); + } + if let Some(reason) = guardrail_rejection_reason(&body) { + return Err(CliError::GuardrailRejected(reason)); + } + Err(CliError::Install(format!( + "daemon hook forward failed with HTTP {status}" + ))) +} + +fn hook_endpoint(daemon_address: &str, agent: CodingAgent) -> Result { + let mut url = crate::daemon::common::address::daemon_url(daemon_address)?; + url.set_path(agent.hook_path()); + Ok(url) +} + +async fn read_response(response: reqwest::Response) -> Result, CliError> { + let mut stream = response.bytes_stream(); + let mut body = Vec::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk?; + if body.len().saturating_add(chunk.len()) > MAX_HOOK_RESPONSE_BYTES { + return Err(CliError::PayloadTooLarge(format!( + "daemon hook response exceeds the {MAX_HOOK_RESPONSE_BYTES}-byte limit" + ))); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + +fn guardrail_rejection_reason(body: &[u8]) -> Option { + let value: Value = serde_json::from_slice(body).ok()?; + let error = value.get("error")?; + (error.get("type").and_then(Value::as_str) == Some("nemo_relay_guardrail_rejected")) + .then(|| { + error + .get("reason") + .and_then(Value::as_str) + .or_else(|| error.get("message").and_then(Value::as_str)) + .map(ToOwned::to_owned) + }) + .flatten() +} + +fn handle_delivery_failure(error: CliError, fail_closed: bool) -> Result<(), CliError> { + let mode = if fail_closed { + "fail_closed" + } else { + "fail_open" + }; + if fail_closed { + log::error!( + target: "nemo_relay.hook", + event = "daemon_hook_delivery_failed", + mode, + error_kind = error.log_kind(); + "Managed daemon hook delivery failed" + ); + Err(CliError::HookDelivery { + source: Box::new(error), + }) + } else { + log::warn!( + target: "nemo_relay.hook", + event = "daemon_hook_delivery_failed", + mode, + error_kind = error.log_kind(); + "Managed daemon hook delivery failed open" + ); + eprintln!("nemo-relay daemon hook failed: {error}"); + Ok(()) + } +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/hook_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/managed/mod.rs b/crates/cli/src/daemon/managed/mod.rs new file mode 100644 index 000000000..1ca5ae4cc --- /dev/null +++ b/crates/cli/src/daemon/managed/mod.rs @@ -0,0 +1,956 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Immutable deployment artifacts for administrator-managed daemon integrations. +//! +//! Personal `nemo-relay install` deliberately remains separate. A managed bundle is rendered +//! once for one deployment and then distributed by the administrator. Refresh and doctor only +//! validate it: changing a deployed v1 artifact in place is an error, and an incompatible contract +//! after v1 publication needs a separately named v2 bundle. + +use std::collections::BTreeSet; +use std::fs::{self, OpenOptions}; +use std::io::{Read, Write}; +use std::path::{Component, Path}; +use std::str::FromStr; + +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +use super::common::state::{ROUTE_TOKEN_ENV, RouteCredential}; +use crate::error::CliError; + +pub(crate) const BUNDLE_FAMILY: &str = "nemo-relay-managed-v1"; +pub(crate) const MANIFEST_FILE: &str = "nemo-relay-managed-v1.manifest.json"; +const SCHEMA_VERSION: u32 = 1; +const MAX_MANIFEST_BYTES: u64 = 1024 * 1024; +const MAX_ARTIFACTS: usize = 256; +const CLAUDE_CUSTOM_HEADERS_ENV: &str = "ANTHROPIC_CUSTOM_HEADERS"; +const ROUTE_TOKEN_HEADER: &str = "x-nemo-relay-client-token"; +const PI_DAEMON_ADDRESS_PLACEHOLDER: &str = "__NEMO_RELAY_DAEMON_ADDRESS__"; +const PI_DISPATCHER_PLACEHOLDER: &str = "__NEMO_RELAY_DISPATCHER_COMMAND__"; +const CODEX_ARTIFACT_PATHS: &[&str] = &[ + "codex/plugin-v1/.codex-plugin/plugin.json", + "codex/plugin-v1/.mcp.json", + "codex/plugin-v1/hooks/hooks.json", + "codex/settings-v1/config.toml", +]; +const CLAUDE_ARTIFACT_PATHS: &[&str] = &[ + "claude-code/plugin-v1/.claude-plugin/plugin.json", + "claude-code/plugin-v1/.mcp.json", + "claude-code/plugin-v1/hooks/hooks.json", + "claude-code/settings-v1/managed-settings.json", +]; +const PI_ARTIFACT_PATHS: &[&str] = &[ + "pi/extension-v1/README.md", + "pi/extension-v1/index.ts", + "pi/extension-v1/managed-config.json", + "pi/extension-v1/package.json", + "pi/extension-v1/tsconfig.json", +]; + +// Part of the v1 artifact family. Once v1 is published, new host events belong in a v2 family: +// silently changing these lists would make an upgrade rewrite enterprise-managed plugin bytes. +const CODEX_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "PreToolUse", + "PostToolUse", + "PermissionRequest", + "SubagentStart", + "SubagentStop", + "Stop", + "PreCompact", + "PostCompact", +]; +const CLAUDE_HOOK_EVENTS: &[&str] = &[ + "SessionStart", + "UserPromptSubmit", + "UserPromptExpansion", + "PreToolUse", + "PostToolUse", + "PostToolUseFailure", + "PermissionRequest", + "SubagentStart", + "SubagentStop", + "Notification", + "Stop", + "PreCompact", + "PostCompact", + "SessionEnd", +]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub(crate) enum ManagedAgent { + Codex, + ClaudeCode, + Pi, +} + +impl ManagedAgent { + const fn hook_argument(self) -> &'static str { + match self { + Self::Codex => "codex", + Self::ClaudeCode => "claude", + Self::Pi => "pi", + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub(crate) enum ManagedPlatform { + Linux, + Macos, + Windows, +} + +impl ManagedPlatform { + pub(crate) const fn as_str(self) -> &'static str { + match self { + Self::Linux => "linux", + Self::Macos => "macos", + Self::Windows => "windows", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ManagedBundleSpec { + daemon_address: String, + dispatcher_command: String, + platform: ManagedPlatform, + agents: BTreeSet, +} + +impl ManagedBundleSpec { + pub(crate) fn new( + daemon_address: impl Into, + dispatcher_command: impl Into, + platform: ManagedPlatform, + agents: impl IntoIterator, + ) -> Result { + let daemon_address = normalize_daemon_address(&daemon_address.into())?; + let dispatcher_command = dispatcher_command.into(); + validate_dispatcher(&dispatcher_command, platform)?; + let agents = agents.into_iter().collect::>(); + if agents.is_empty() { + return Err(CliError::Config( + "a managed daemon bundle must target at least one agent".into(), + )); + } + Ok(Self { + daemon_address, + dispatcher_command, + platform, + agents, + }) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ManagedBundleValidation { + pub(crate) artifact_count: usize, + pub(crate) daemon_address: String, + pub(crate) platform: ManagedPlatform, + pub(crate) sha256: ManagedBundleDigest, +} + +/// SHA-256 over the canonical, length-prefixed manifest and artifact byte stream. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ManagedBundleDigest(String); + +impl std::fmt::Display for ManagedBundleDigest { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(&self.0) + } +} + +impl FromStr for ManagedBundleDigest { + type Err = String; + + fn from_str(value: &str) -> Result { + if value.len() != 64 + || !value + .bytes() + .all(|byte| byte.is_ascii_digit() || matches!(byte, b'a'..=b'f')) + { + return Err( + "managed bundle SHA-256 must be exactly 64 lowercase hexadecimal characters".into(), + ); + } + Ok(Self(value.to_owned())) + } +} + +#[derive(Debug, Clone)] +struct RenderedArtifact { + agent: ManagedAgent, + path: &'static str, + bytes: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct ManagedManifest { + schema_version: u32, + family: String, + daemon_address: String, + dispatcher_command: String, + platform: ManagedPlatform, + agents: Vec, + artifacts: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct ManifestArtifact { + agent: ManagedAgent, + path: String, + byte_length: u64, + sha256: String, +} + +struct RenderedBundle { + manifest: Vec, + artifacts: Vec, +} + +/// Creates a bundle only when the destination does not exist. +/// +/// If the destination already exists, it is validated byte-for-byte and left untouched. The +/// returned digest is suitable for separately provisioning `doctor --managed-bundle-sha256`. +pub(crate) fn write_new_bundle( + root: &Path, + spec: &ManagedBundleSpec, +) -> Result { + let expected = render_bundle(spec)?; + let expected_digest = rendered_bundle_digest(&expected); + if root.exists() { + validate_bundle_files(root, false, None)?; + let actual_manifest = read_bounded(&root.join(MANIFEST_FILE), MAX_MANIFEST_BYTES)?; + if actual_manifest != expected.manifest { + return Err(CliError::Config(format!( + "refused to replace existing managed bundle {} with different deployment bytes; use a separately named artifact", + root.display() + ))); + } + return Ok(expected_digest); + } + let parent = root.parent().ok_or_else(|| { + CliError::Config(format!( + "managed bundle destination {} has no parent", + root.display() + )) + })?; + fs::create_dir_all(parent)?; + let name = root + .file_name() + .and_then(|value| value.to_str()) + .ok_or_else(|| { + CliError::Config("managed bundle destination is not valid Unicode".into()) + })?; + let stage = parent.join(format!(".{name}.{}.tmp", Uuid::now_v7())); + fs::create_dir(&stage)?; + let result = write_rendered_bundle(&stage, expected).and_then(|()| { + fs::rename(&stage, root).map_err(|error| { + CliError::Config(format!( + "failed to publish managed bundle {}: {error}", + root.display() + )) + }) + }); + if result.is_err() { + let _ = fs::remove_dir_all(&stage); + } + result.map(|()| expected_digest) +} + +/// Managed refresh is validation-only; it never regenerates or overwrites v1 artifacts. +pub(crate) fn refresh_bundle( + root: &Path, + expected_sha256: &ManagedBundleDigest, +) -> Result { + validate_bundle(root, expected_sha256) +} + +/// Validates the administrator-provided digest, every exact artifact byte, and managed env. +pub(crate) fn validate_bundle( + root: &Path, + expected_sha256: &ManagedBundleDigest, +) -> Result { + validate_bundle_files(root, true, Some(expected_sha256)) +} + +fn validate_bundle_files( + root: &Path, + validate_environment: bool, + expected_sha256: Option<&ManagedBundleDigest>, +) -> Result { + reject_non_directory_or_symlink(root)?; + let manifest_path = root.join(MANIFEST_FILE); + reject_symlink(&manifest_path)?; + let manifest_bytes = read_bounded(&manifest_path, MAX_MANIFEST_BYTES)?; + let manifest: ManagedManifest = serde_json::from_slice(&manifest_bytes).map_err(|error| { + CliError::Config(format!( + "managed bundle manifest {} is invalid: {error}", + manifest_path.display() + )) + })?; + if manifest.schema_version != SCHEMA_VERSION || manifest.family != BUNDLE_FAMILY { + return Err(CliError::Config(format!( + "managed bundle {} is not the supported {BUNDLE_FAMILY} artifact family", + root.display() + ))); + } + if manifest.artifacts.len() > MAX_ARTIFACTS { + return Err(CliError::Config(format!( + "managed bundle manifest has more than {MAX_ARTIFACTS} artifacts" + ))); + } + let spec = ManagedBundleSpec::new( + manifest.daemon_address, + manifest.dispatcher_command, + manifest.platform, + manifest.agents, + )?; + let expected = render_bundle(&spec)?; + if manifest_bytes != expected.manifest { + return Err(CliError::Config(format!( + "managed bundle manifest {} differs from the canonical {BUNDLE_FAMILY} bytes", + manifest_path.display() + ))); + } + let expected_paths = expected + .artifacts + .iter() + .map(|artifact| artifact.path) + .chain(std::iter::once(MANIFEST_FILE)) + .collect::>(); + let actual_paths = bundle_files(root)?; + if actual_paths != expected_paths { + let missing = expected_paths + .difference(&actual_paths) + .copied() + .collect::>(); + let unexpected = actual_paths + .difference(&expected_paths) + .copied() + .collect::>(); + return Err(CliError::Config(format!( + "managed bundle file set differs from its immutable manifest (missing: {}; unexpected: {})", + display_paths(&missing), + display_paths(&unexpected) + ))); + } + for artifact in &expected.artifacts { + let path = root.join(artifact.path); + reject_symlink(&path)?; + let actual = read_bounded(&path, artifact.bytes.len() as u64)?; + if actual != artifact.bytes { + return Err(CliError::Config(format!( + "managed artifact {} differs from its exact canonical bytes", + path.display() + ))); + } + } + let sha256 = rendered_bundle_digest(&expected); + if let Some(expected_sha256) = expected_sha256 + && expected_sha256 != &sha256 + { + return Err(CliError::Config(format!( + "managed bundle SHA-256 mismatch: expected {expected_sha256}, calculated {sha256}" + ))); + } + if validate_environment { + validate_managed_environment(&spec)?; + } + Ok(ManagedBundleValidation { + artifact_count: expected.artifacts.len(), + daemon_address: spec.daemon_address, + platform: spec.platform, + sha256, + }) +} + +fn validate_managed_environment(spec: &ManagedBundleSpec) -> Result<(), CliError> { + let credential = RouteCredential::from_environment()?; + if !spec.agents.contains(&ManagedAgent::ClaudeCode) { + return Ok(()); + } + let custom_headers = std::env::var(CLAUDE_CUSTOM_HEADERS_ENV).map_err(|_| { + CliError::Config(format!( + "managed Claude Code integration requires {CLAUDE_CUSTOM_HEADERS_ENV}; enterprise bootstrap must derive it from {ROUTE_TOKEN_ENV}" + )) + })?; + let matches = custom_headers + .lines() + .filter_map(|line| line.split_once(':')) + .filter(|(name, _)| name.trim().eq_ignore_ascii_case(ROUTE_TOKEN_HEADER)) + .map(|(_, value)| value.trim()) + .collect::>(); + if matches.as_slice() != [credential.expose()] { + return Err(CliError::Config(format!( + "{CLAUDE_CUSTOM_HEADERS_ENV} must contain exactly one {ROUTE_TOKEN_HEADER} header whose value matches {ROUTE_TOKEN_ENV}" + ))); + } + Ok(()) +} + +fn render_bundle(spec: &ManagedBundleSpec) -> Result { + let mut artifacts = Vec::new(); + for agent in &spec.agents { + artifacts.extend(render_agent(*agent, spec)?); + } + artifacts.sort_by_key(|artifact| artifact.path); + let manifest = ManagedManifest { + schema_version: SCHEMA_VERSION, + family: BUNDLE_FAMILY.into(), + daemon_address: spec.daemon_address.clone(), + dispatcher_command: spec.dispatcher_command.clone(), + platform: spec.platform, + agents: spec.agents.iter().copied().collect(), + artifacts: artifacts + .iter() + .map(|artifact| ManifestArtifact { + agent: artifact.agent, + path: artifact.path.into(), + byte_length: artifact.bytes.len() as u64, + sha256: sha256_hex(&artifact.bytes), + }) + .collect(), + }; + Ok(RenderedBundle { + manifest: json_bytes(&manifest)?, + artifacts, + }) +} + +fn render_agent( + agent: ManagedAgent, + spec: &ManagedBundleSpec, +) -> Result, CliError> { + match agent { + ManagedAgent::Codex => render_codex(spec), + ManagedAgent::ClaudeCode => render_claude(spec), + ManagedAgent::Pi => render_pi(spec), + } +} + +fn render_codex(spec: &ManagedBundleSpec) -> Result, CliError> { + let mcp = json!({ + "nemo-relay": { + "command": spec.dispatcher_command, + "args": ["daemon", "mcp", "--daemon-address", spec.daemon_address], + "env_vars": [ROUTE_TOKEN_ENV], + "required": true, + "startup_timeout_sec": 20 + } + }); + let plugin = plugin_manifest("codex"); + let settings = format!( + "model_provider = \"nemo-relay-managed-v1\"\n\n[model_providers.nemo-relay-managed-v1]\nname = \"NeMo Relay Managed\"\nbase_url = {}\nwire_api = \"responses\"\nrequires_openai_auth = true\nsupports_websockets = false\nenv_http_headers = {{ {} = {} }}\n", + toml_string(&format!("{}/v1", spec.daemon_address)), + toml_string(ROUTE_TOKEN_HEADER), + toml_string(ROUTE_TOKEN_ENV), + ); + Ok(vec![ + artifact( + ManagedAgent::Codex, + CODEX_ARTIFACT_PATHS[0], + json_bytes(&plugin)?, + ), + artifact( + ManagedAgent::Codex, + CODEX_ARTIFACT_PATHS[1], + json_bytes(&mcp)?, + ), + artifact( + ManagedAgent::Codex, + CODEX_ARTIFACT_PATHS[2], + hook_bytes(ManagedAgent::Codex, CODEX_HOOK_EVENTS, spec)?, + ), + artifact( + ManagedAgent::Codex, + CODEX_ARTIFACT_PATHS[3], + settings.into_bytes(), + ), + ]) +} + +fn render_claude(spec: &ManagedBundleSpec) -> Result, CliError> { + let mcp = json!({ + "mcpServers": { + "nemo-relay": { + "command": spec.dispatcher_command, + "args": ["daemon", "mcp", "--daemon-address", spec.daemon_address], + "env": { (ROUTE_TOKEN_ENV): format!("${{{ROUTE_TOKEN_ENV}}}") }, + "alwaysLoad": true + } + } + }); + let settings = json!({ + "$schema": "https://json.schemastore.org/claude-code-settings.json", + "env": { "ANTHROPIC_BASE_URL": spec.daemon_address } + }); + Ok(vec![ + artifact( + ManagedAgent::ClaudeCode, + CLAUDE_ARTIFACT_PATHS[0], + json_bytes(&plugin_manifest("claude-code"))?, + ), + artifact( + ManagedAgent::ClaudeCode, + CLAUDE_ARTIFACT_PATHS[1], + json_bytes(&mcp)?, + ), + artifact( + ManagedAgent::ClaudeCode, + CLAUDE_ARTIFACT_PATHS[2], + hook_bytes(ManagedAgent::ClaudeCode, CLAUDE_HOOK_EVENTS, spec)?, + ), + artifact( + ManagedAgent::ClaudeCode, + CLAUDE_ARTIFACT_PATHS[3], + json_bytes(&settings)?, + ), + ]) +} + +fn render_pi(spec: &ManagedBundleSpec) -> Result, CliError> { + let config = render_pi_config(spec)?; + Ok(vec![ + artifact( + ManagedAgent::Pi, + PI_ARTIFACT_PATHS[0], + canonical_embedded_text(include_str!("pi_extension/README.md")).into_bytes(), + ), + artifact( + ManagedAgent::Pi, + PI_ARTIFACT_PATHS[1], + canonical_embedded_text(include_str!("pi_extension/index.ts")).into_bytes(), + ), + artifact(ManagedAgent::Pi, PI_ARTIFACT_PATHS[2], config), + artifact( + ManagedAgent::Pi, + PI_ARTIFACT_PATHS[3], + canonical_embedded_text(include_str!("pi_extension/package.json")).into_bytes(), + ), + artifact( + ManagedAgent::Pi, + PI_ARTIFACT_PATHS[4], + canonical_embedded_text(include_str!("pi_extension/tsconfig.json")).into_bytes(), + ), + ]) +} + +fn render_pi_config(spec: &ManagedBundleSpec) -> Result, CliError> { + let template = canonical_embedded_text(include_str!("pi_extension/managed-config.json")); + let rendered = replace_json_string_value( + &template, + PI_DAEMON_ADDRESS_PLACEHOLDER, + &spec.daemon_address, + )?; + let rendered = replace_json_string_value( + &rendered, + PI_DISPATCHER_PLACEHOLDER, + &spec.dispatcher_command, + )?; + if rendered.contains("__NEMO_RELAY_") { + return Err(CliError::Config( + "managed Pi configuration contains an unrendered deployment placeholder".into(), + )); + } + serde_json::from_str::(&rendered).map_err(|error| { + CliError::Config(format!( + "rendered managed Pi configuration is invalid JSON: {error}" + )) + })?; + Ok(rendered.into_bytes()) +} + +fn canonical_embedded_text(template: &str) -> String { + template.replace("\r\n", "\n") +} + +fn replace_json_string_value( + template: &str, + placeholder: &str, + value: &str, +) -> Result { + let placeholder = serde_json::to_string(placeholder).map_err(|error| { + CliError::Config(format!("failed to encode managed Pi placeholder: {error}")) + })?; + if template.matches(&placeholder).count() != 1 { + return Err(CliError::Config( + "managed Pi configuration must contain each deployment placeholder exactly once".into(), + )); + } + let value = serde_json::to_string(value).map_err(|error| { + CliError::Config(format!( + "failed to encode managed Pi deployment value: {error}" + )) + })?; + Ok(template.replacen(&placeholder, &value, 1)) +} + +fn plugin_manifest(agent: &str) -> Value { + json!({ + "name": "nemo-relay-managed-v1", + "version": "1.0.0", + "description": format!("Immutable NeMo Relay managed integration for {agent}."), + "author": { "name": "NVIDIA Corporation and Affiliates" }, + "license": "Apache-2.0", + "mcpServers": "./.mcp.json" + }) +} + +fn hook_bytes( + agent: ManagedAgent, + events: &[&str], + spec: &ManagedBundleSpec, +) -> Result, CliError> { + let fail_open = hook_command(agent, spec, false); + let fail_closed = hook_command(agent, spec, true); + let hooks = events + .iter() + .map(|event| { + let mut group = serde_json::Map::new(); + if matches!( + *event, + "PreToolUse" | "PostToolUse" | "PostToolUseFailure" | "PermissionRequest" + ) { + group.insert("matcher".into(), json!("*")); + } + let command = if matches!(*event, "PreToolUse" | "PermissionRequest") { + &fail_closed + } else { + &fail_open + }; + group.insert( + "hooks".into(), + json!([{ "type": "command", "command": command, "timeout": 30 }]), + ); + ( + (*event).to_string(), + Value::Array(vec![Value::Object(group)]), + ) + }) + .collect::>(); + json_bytes(&json!({ "hooks": hooks })) +} + +fn hook_command(agent: ManagedAgent, spec: &ManagedBundleSpec, fail_closed: bool) -> String { + format!( + "{} daemon hook {} --daemon-address {} {}", + spec.dispatcher_command, + agent.hook_argument(), + spec.daemon_address, + if fail_closed { + "--fail-closed" + } else { + "--fail-open" + } + ) +} + +fn artifact(agent: ManagedAgent, path: &'static str, bytes: Vec) -> RenderedArtifact { + RenderedArtifact { agent, path, bytes } +} + +fn write_rendered_bundle(root: &Path, bundle: RenderedBundle) -> Result<(), CliError> { + for artifact in bundle.artifacts { + write_new_file(&root.join(artifact.path), &artifact.bytes)?; + } + write_new_file(&root.join(MANIFEST_FILE), &bundle.manifest) +} + +fn write_new_file(path: &Path, bytes: &[u8]) -> Result<(), CliError> { + let parent = path.parent().ok_or_else(|| { + CliError::Config(format!("managed artifact {} has no parent", path.display())) + })?; + fs::create_dir_all(parent)?; + let mut file = OpenOptions::new() + .write(true) + .create_new(true) + .open(path) + .map_err(|error| { + CliError::Config(format!( + "refused to overwrite managed artifact {}: {error}", + path.display() + )) + })?; + file.write_all(bytes)?; + file.sync_all()?; + drop(file); + let mut permissions = fs::metadata(path)?.permissions(); + permissions.set_readonly(true); + fs::set_permissions(path, permissions)?; + Ok(()) +} + +fn normalize_daemon_address(raw: &str) -> Result { + let raw = raw.trim_end_matches('/'); + super::common::address::daemon_url(raw)?; + Ok(raw.to_string()) +} + +fn validate_dispatcher(command: &str, platform: ManagedPlatform) -> Result<(), CliError> { + let normalized = command.replace('\\', "/").to_ascii_lowercase(); + let windows_drive_absolute = normalized.as_bytes().get(1) == Some(&b':') + && normalized + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphabetic) + && normalized.as_bytes().get(2) == Some(&b'/'); + let absolute = match platform { + ManagedPlatform::Linux | ManagedPlatform::Macos => command.starts_with('/'), + ManagedPlatform::Windows => windows_drive_absolute || normalized.starts_with("//"), + }; + let platform_separators_are_valid = + matches!(platform, ManagedPlatform::Windows) || !command.contains('\\'); + let parts = normalized.split('/').filter(|part| !part.is_empty()); + let has_relative_component = parts + .clone() + .any(|component| matches!(component, "." | "..")); + let forbidden_root = match platform { + ManagedPlatform::Linux | ManagedPlatform::Macos => [ + "/tmp", + "/var/tmp", + "/private/tmp", + "/home", + "/users", + "/root", + "/run/user", + "/var/folders", + ] + .iter() + .any(|root| normalized == *root || normalized.starts_with(&format!("{root}/"))), + ManagedPlatform::Windows => { + let path = if windows_drive_absolute { + &normalized[2..] + } else { + normalized.as_str() + }; + ["/users", "/temp", "/tmp", "/windows/temp"] + .iter() + .any(|root| path == *root || path.starts_with(&format!("{root}/"))) + } + }; + let points_to_directory = normalized.ends_with('/'); + if command.is_empty() + || !absolute + || !platform_separators_are_valid + || command.chars().any(char::is_whitespace) + || has_relative_component + || forbidden_root + || points_to_directory + || command.chars().any(|character| { + !character.is_ascii_alphanumeric() + && !matches!(character, '/' | '\\' | ':' | '.' | '_' | '-') + }) + { + return Err(CliError::Config( + "managed dispatcher command must be an absolute, platform-appropriate, shell-safe administrator path outside user and temporary directories" + .into(), + )); + } + Ok(()) +} + +fn json_bytes(value: &impl Serialize) -> Result, CliError> { + let mut bytes = serde_json::to_vec_pretty(value) + .map_err(|error| CliError::Config(format!("failed to render managed artifact: {error}")))?; + bytes.push(b'\n'); + Ok(bytes) +} + +fn toml_string(value: &str) -> String { + assert!( + value + .chars() + .all(|character| character.is_ascii_graphic() || character == ' '), + "managed TOML values must be validated ASCII" + ); + format!("{value:?}") +} + +fn sha256_hex(bytes: &[u8]) -> String { + lowercase_hex(&Sha256::digest(bytes)) +} + +fn lowercase_hex(bytes: &[u8]) -> String { + bytes.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn rendered_bundle_digest(bundle: &RenderedBundle) -> ManagedBundleDigest { + let mut hasher = Sha256::new(); + update_bundle_digest(&mut hasher, MANIFEST_FILE.as_bytes(), &bundle.manifest); + for artifact in &bundle.artifacts { + update_bundle_digest(&mut hasher, artifact.path.as_bytes(), &artifact.bytes); + } + ManagedBundleDigest(lowercase_hex(&hasher.finalize())) +} + +fn update_bundle_digest(hasher: &mut Sha256, name: &[u8], bytes: &[u8]) { + hasher.update( + u64::try_from(name.len()) + .expect("managed artifact names fit in u64") + .to_be_bytes(), + ); + hasher.update(name); + hasher.update( + u64::try_from(bytes.len()) + .expect("managed artifacts fit in u64") + .to_be_bytes(), + ); + hasher.update(bytes); +} + +fn read_bounded(path: &Path, maximum: u64) -> Result, CliError> { + let mut file = OpenOptions::new().read(true).open(path).map_err(|error| { + CliError::Config(format!( + "failed to read managed artifact {}: {error}", + path.display() + )) + })?; + let length = file.metadata()?.len(); + if length > maximum { + return Err(CliError::Config(format!( + "managed artifact {} exceeds its expected size limit", + path.display() + ))); + } + let mut bytes = Vec::with_capacity(length as usize); + file.read_to_end(&mut bytes)?; + Ok(bytes) +} + +fn reject_symlink(path: &Path) -> Result<(), CliError> { + let metadata = fs::symlink_metadata(path).map_err(|error| { + CliError::Config(format!( + "failed to inspect managed artifact {}: {error}", + path.display() + )) + })?; + if metadata.file_type().is_symlink() || !metadata.is_file() { + return Err(CliError::Config(format!( + "managed artifact {} must be a regular file, not a symlink", + path.display() + ))); + } + Ok(()) +} + +fn reject_non_directory_or_symlink(path: &Path) -> Result<(), CliError> { + let metadata = fs::symlink_metadata(path).map_err(|error| { + CliError::Config(format!( + "failed to inspect managed bundle {}: {error}", + path.display() + )) + })?; + if metadata.file_type().is_symlink() || !metadata.is_dir() { + return Err(CliError::Config(format!( + "managed bundle {} must be a directory, not a symlink", + path.display() + ))); + } + Ok(()) +} + +fn bundle_files(root: &Path) -> Result, CliError> { + let expected = all_artifact_paths(); + let expected_lookup = expected.iter().copied().collect::>(); + let mut actual = BTreeSet::new(); + let mut pending = vec![root.to_path_buf()]; + let mut visited = 0_usize; + while let Some(directory) = pending.pop() { + for entry in fs::read_dir(&directory)? { + visited += 1; + if visited > MAX_ARTIFACTS * 4 { + return Err(CliError::Config( + "managed bundle contains too many filesystem entries".into(), + )); + } + let entry = entry?; + let metadata = entry.file_type()?; + if metadata.is_symlink() { + return Err(CliError::Config(format!( + "managed bundle entry {} must not be a symlink", + entry.path().display() + ))); + } + if metadata.is_dir() { + pending.push(entry.path()); + continue; + } + if !metadata.is_file() { + return Err(CliError::Config(format!( + "managed bundle entry {} must be a regular file", + entry.path().display() + ))); + } + let path = entry.path(); + let relative = path + .strip_prefix(root) + .map_err(|_| CliError::Config("managed bundle entry escaped its root".into()))?; + if relative + .components() + .any(|component| !matches!(component, Component::Normal(_))) + { + return Err(CliError::Config(format!( + "managed bundle entry {} has an invalid path", + relative.display() + ))); + } + let relative = relative.to_string_lossy().replace('\\', "/"); + let canonical = expected_lookup + .get(relative.as_str()) + .copied() + .ok_or_else(|| { + CliError::Config(format!( + "managed bundle contains unexpected artifact {relative}" + )) + })?; + actual.insert(canonical); + } + } + Ok(actual) +} + +fn all_artifact_paths() -> BTreeSet<&'static str> { + std::iter::once(MANIFEST_FILE) + .chain( + [ + ManagedAgent::Codex, + ManagedAgent::ClaudeCode, + ManagedAgent::Pi, + ] + .into_iter() + .flat_map(|agent| artifact_paths_for(agent).iter().copied()), + ) + .collect() +} + +fn artifact_paths_for(agent: ManagedAgent) -> &'static [&'static str] { + match agent { + ManagedAgent::Codex => CODEX_ARTIFACT_PATHS, + ManagedAgent::ClaudeCode => CLAUDE_ARTIFACT_PATHS, + ManagedAgent::Pi => PI_ARTIFACT_PATHS, + } +} + +fn display_paths(paths: &[&str]) -> String { + if paths.is_empty() { + "none".into() + } else { + paths.join(", ") + } +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/managed_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/managed/pi_extension/README.md b/crates/cli/src/daemon/managed/pi_extension/README.md new file mode 100644 index 000000000..5ebe24d0b --- /dev/null +++ b/crates/cli/src/daemon/managed/pi_extension/README.md @@ -0,0 +1,41 @@ + + +# NeMo Relay Managed Pi Extension v1 + +This directory is the immutable version 1 managed Pi package. Pi loads +`index.ts` directly; the package has no runtime dependencies beyond Node.js. + +The administrator-owned launcher must load only this managed extension: + +```bash +pi --no-extensions -e /srv/nemo-relay/nemo-relay-managed-v1/pi/extension-v1/index.ts +``` + +`--no-extensions` is required. It preserves the explicitly loaded `-e` +extension while preventing user, project, and discovered extensions from +running alongside the managed policy boundary. + +The managed-bundle renderer replaces the two complete JSON string values in +`managed-config.json`: + +- `__NEMO_RELAY_DAEMON_ADDRESS__` becomes the fixed root daemon URL. +- `__NEMO_RELAY_DISPATCHER_COMMAND__` becomes the fixed administrator-owned + dispatcher path. + +The renderer must JSON-encode replacement values. It must not perform raw text +substitution. An extension installed with either placeholder still present +fails closed. + +At runtime, the extension reads `NEMO_RELAY_CLIENT_TOKEN` from the process +environment. No credential, user identity, machine fingerprint, generation +identifier, or Relay binary version is written into this package. + +The extension redirects the selected Pi provider when every known sibling model +uses an OpenAI Completions, OpenAI Responses, or Anthropic Messages API. This +includes custom provider names, but the administrator-managed worker upstream +remains authoritative: the extension never sends a per-user upstream URL or +routing header. Providers containing an unsupported API remain on their +original endpoint. diff --git a/crates/cli/src/daemon/managed/pi_extension/index.ts b/crates/cli/src/daemon/managed/pi_extension/index.ts new file mode 100644 index 000000000..ffdf0f8c8 --- /dev/null +++ b/crates/cli/src/daemon/managed/pi_extension/index.ts @@ -0,0 +1,1124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +/** + * Immutable managed NeMo Relay integration for Pi. + * + * Deployment-specific values live in managed-config.json. Per-user state is + * read only at runtime. The process-global MCP lease deliberately outlives Pi + * extension reloads and new/resume/fork session replacement. + */ +import { spawn, type ChildProcess } from 'node:child_process'; +import { readFileSync } from 'node:fs'; +import { isAbsolute } from 'node:path'; +import { createInterface } from 'node:readline'; + +const CLIENT_TOKEN_ENV = 'NEMO_RELAY_CLIENT_TOKEN'; +const CLIENT_TOKEN_HEADER = 'x-nemo-relay-client-token'; +const UPSTREAM_BASE_URL_HEADER = 'x-nemo-relay-upstream-base-url'; +const CONFIG_SCHEMA = 'nemo-relay-managed-pi-v1'; +const CONFIG_PLACEHOLDER_PREFIX = '__NEMO_RELAY_'; +const MCP_SINGLETON = Symbol.for('nemo-relay.managed-pi.mcp.v1'); +const MCP_INITIALIZE_ID = 'nemo-relay-managed-pi-ready-v1'; +const MCP_PROTOCOL_VERSION = '2025-11-25'; +// The broker may legally keep a new MCP waiting through a 120-second worker drain before it can +// establish the next route generation. Leave reconciliation margin rather than killing a healthy +// lifecycle client during that window. +const MCP_READY_TIMEOUT_MS = 180_000; +const MCP_RELEASE_TIMEOUT_MS = 5_000; +const HOOK_TIMEOUT_MS = 30_000; +const MAX_HOOK_PAYLOAD_BYTES = 20 * 1024 * 1024; +const MAX_HOOK_RESPONSE_BYTES = 1024 * 1024; +const MAX_RESULT_CHARS = 2_000; +const SHARED_LEASE_KIND = 'nemo-relay-managed-pi-mcp-lease-v1'; +const SERVICEABLE_APIS = new Set(['openai-completions', 'openai-responses', 'anthropic-messages']); + +type DeploymentConfig = { + schema: typeof CONFIG_SCHEMA; + daemonAddress: string; + dispatcherCommand: string; +}; + +type McpLease = { + kind: typeof SHARED_LEASE_KIND; + daemonAddress: string; + dispatcherCommand: string; + credential: string; + ensureReady(): Promise; + release(): Promise; +}; + +type HookOutcome = + | { kind: 'allow'; body: Record } + | { kind: 'block'; reason: string } + | { kind: 'fault'; reason: string }; + +type PiModel = { + id: string; + api: string; + provider: string; + baseUrl: string; +}; + +type ProviderRedirectDecision = + | { kind: 'redirect'; reason: string; upstream: string } + | { kind: 'skip'; code: string; reason: string }; + +type ExtensionContext = { + cwd: string; + model?: PiModel; + modelRegistry?: { getAll?(): PiModel[] }; + sessionManager?: { getSessionId?(): string }; +}; + +type SessionStartEvent = { + type: 'session_start'; + reason: 'startup' | 'reload' | 'new' | 'resume' | 'fork'; + previousSessionFile?: string; +}; + +type SessionShutdownEvent = { + type: 'session_shutdown'; + reason: 'quit' | 'reload' | 'new' | 'resume' | 'fork'; + targetSessionFile?: string; +}; + +type AgentEndEvent = { type: 'agent_end'; messages?: unknown[] }; +type TurnEvent = { type: 'turn_start' | 'turn_end'; turnIndex: number }; +type CompactEvent = { + type: 'session_before_compact' | 'session_compact'; + reason: string; + willRetry: boolean; + fromExtension?: boolean; + preparation?: { tokensBefore?: number; isSplitTurn?: boolean }; + compactionEntry?: { tokensBefore?: number }; +}; +type ToolExecutionStartEvent = { + type: 'tool_execution_start'; + toolCallId: string; + toolName: string; +}; +type ToolExecutionEndEvent = { + type: 'tool_execution_end'; + toolCallId: string; + toolName: string; + result: unknown; + isError: boolean; +}; +type ToolCallEvent = { + type: 'tool_call'; + toolCallId: string; + toolName: string; + input: Record; +}; +type ToolCallResult = { block: true; reason: string }; +type UserBashEvent = { + type: 'user_bash'; + command: string; + cwd: string; + excludeFromContext: boolean; +}; +type UserBashResult = { + result: { + output: string; + exitCode: 126; + cancelled: boolean; + truncated: boolean; + }; +}; + +type ExtensionHandler = ( + event: TEvent, + context: ExtensionContext, +) => TResult | undefined | Promise; + +type ExtensionAPI = { + on(event: 'session_start', handler: ExtensionHandler): void; + on(event: 'session_shutdown', handler: ExtensionHandler): void; + on(event: 'agent_start', handler: ExtensionHandler<{ type: 'agent_start' }>): void; + on(event: 'agent_end', handler: ExtensionHandler): void; + on(event: 'agent_settled', handler: ExtensionHandler<{ type: 'agent_settled' }>): void; + on(event: 'turn_start', handler: ExtensionHandler): void; + on(event: 'turn_end', handler: ExtensionHandler): void; + on(event: 'session_before_compact', handler: ExtensionHandler): void; + on(event: 'session_compact', handler: ExtensionHandler): void; + on(event: 'tool_execution_start', handler: ExtensionHandler): void; + on(event: 'tool_execution_end', handler: ExtensionHandler): void; + on(event: 'tool_call', handler: ExtensionHandler): void; + on(event: 'user_bash', handler: ExtensionHandler): void; + on(event: 'model_select', handler: ExtensionHandler<{ type: 'model_select'; model: PiModel }>): void; + registerProvider(name: string, config: { baseUrl: string; headers: Record }): void; +}; + +type Runtime = { + config: DeploymentConfig; + credential: string; + lease: McpLease; +}; + +export default function managedNemoRelayPi(pi: ExtensionAPI): void { + let runtimePromise: Promise | undefined; + const redirectedProviders = new Set(); + let hookQueue: Promise = Promise.resolve(); + let attemptIndex = 0; + let turnSequence = 0; + let userBashSequence = 0; + const toolNames = new Map(); + + async function runtime(): Promise { + runtimePromise ??= initializeRuntime(); + const active = await runtimePromise; + // Re-establish the route if a previously ready child exited. All callers + // share the same in-flight restart promise inside the global lease. + await active.lease.ensureReady(); + return active; + } + + function enqueue(operation: () => Promise): Promise { + const result = hookQueue.then(operation, operation); + hookQueue = result.then( + () => undefined, + () => undefined, + ); + return result; + } + + function attribution(): { attempt_index: number; turn_seq: number } { + return { + attempt_index: Math.max(0, attemptIndex - 1), + turn_seq: Math.max(0, turnSequence - 1), + }; + } + + function payload( + context: ExtensionContext, + hookEventName: string, + fields: Record = {}, + ): Record { + return { + session_id: sessionId(context), + hook_event_name: hookEventName, + ...fields, + }; + } + + async function sendObservation(body: Record): Promise { + try { + const active = await runtime(); + const outcome = await postHook(active, body); + if (outcome.kind !== 'allow') { + console.error(`NeMo Relay managed hook was not accepted: ${outcome.reason}`); + } + } catch (error) { + console.error(`NeMo Relay managed hook failed: ${safeError(error)}`); + } + } + + function observe(body: Record): void { + void enqueue(() => sendObservation(body)); + } + + async function observeOrdered(body: Record): Promise { + await enqueue(() => sendObservation(body)); + } + + function registerManagedProvider( + active: Runtime, + model: PiModel | undefined, + context: ExtensionContext, + source: 'session_start' | 'model_select', + ): void { + if (!model || redirectedProviders.has(model.provider)) return; + const catalog = context.modelRegistry?.getAll?.(); + const decision = decideManagedProviderRedirect(model, catalog); + if (decision.kind === 'skip') { + observe( + payload(context, 'model_redirect', { + source, + outcome: 'skip', + code: decision.code, + reason: decision.reason, + provider: model.provider, + model_api: model.api, + model_id: model.id, + ...attribution(), + }), + ); + return; + } + const providerConfig = { + baseUrl: active.config.daemonAddress, + headers: { + [CLIENT_TOKEN_HEADER]: active.credential, + [UPSTREAM_BASE_URL_HEADER]: decision.upstream, + }, + }; + // Pi resolves the concrete provider endpoint before Relay redirects the provider. Preserve + // that endpoint on the provider-wide registration so the authenticated daemon route reaches + // the same destination Pi would have called directly, including custom providers. + pi.registerProvider(model.provider, providerConfig); + redirectedProviders.add(model.provider); + observe( + payload(context, 'model_redirect', { + source, + outcome: 'redirect', + reason: decision.reason, + provider: model.provider, + model_api: model.api, + model_id: model.id, + ...attribution(), + }), + ); + } + + function refuseUserBash(context: ExtensionContext, callId: string, reason: string): UserBashResult { + observe( + payload(context, 'user_bash_end', { + tool_call_id: callId, + tool_name: 'user_bash', + status: 'error', + result: { content: reason }, + ...attribution(), + }), + ); + return refusedBash(reason); + } + + pi.on('session_start', async (event, context) => { + const active = await runtime(); + await observeOrdered( + payload(context, 'session_start', { + reason: event.reason, + cwd: context.cwd, + ...(event.previousSessionFile ? { previous_session_file: event.previousSessionFile } : {}), + }), + ); + // Registration happens after the MCP initialize response. A provider call + // therefore cannot reach the daemon before this process owns a broker route. + registerManagedProvider(active, context.model, context, 'session_start'); + }); + + pi.on('model_select', async (event, context) => { + const active = await runtime(); + registerManagedProvider(active, event.model, context, 'model_select'); + }); + + pi.on('session_shutdown', async (event, context) => { + if (event.reason === 'reload') { + await hookQueue; + return; + } + if (event.reason === 'quit') { + await hookQueue; + // Use only the existing lease: shutdown must not initialize or restart MCP. + const active = runtimePromise ? await runtimePromise : undefined; + if (active) { + try { + const outcome = await postHook(active, payload(context, 'session_shutdown', { reason: event.reason })); + if (outcome.kind !== 'allow') { + console.error(`NeMo Relay managed hook was not accepted: ${outcome.reason}`); + } + } finally { + await active.lease.release(); + } + } + return; + } + await observeOrdered( + payload(context, 'session_shutdown', { + reason: event.reason, + ...(event.targetSessionFile ? { target_session_file: event.targetSessionFile } : {}), + }), + ); + await hookQueue; + // New/resume/fork replace the Pi session inside the same process. Keeping + // the global lease avoids a zero-reference drain and worker restart. + }); + + pi.on('agent_start', async (_event, context) => { + observe(payload(context, 'agent_start', { attempt_index: attemptIndex })); + attemptIndex += 1; + }); + + pi.on('agent_end', async (event, context) => { + observe( + payload(context, 'agent_end', { + attempt_index: Math.max(0, attemptIndex - 1), + message_count: event.messages?.length ?? 0, + }), + ); + }); + + pi.on('agent_settled', async (_event, context) => { + observe( + payload(context, 'agent_settled', { + attempts: attemptIndex, + ...attribution(), + }), + ); + attemptIndex = 0; + }); + + pi.on('turn_start', async (event, context) => { + const sequence = turnSequence; + turnSequence += 1; + await observeOrdered( + payload(context, 'turn_start', { + turn_index: event.turnIndex, + turn_seq: sequence, + attempt_index: Math.max(0, attemptIndex - 1), + }), + ); + }); + + pi.on('turn_end', async (event, context) => { + await observeOrdered( + payload(context, 'turn_end', { + turn_index: event.turnIndex, + ...attribution(), + }), + ); + }); + + pi.on('session_before_compact', async (event, context) => { + observe( + payload(context, 'session_before_compact', { + reason: event.reason, + will_retry: event.willRetry, + tokens_before: event.preparation?.tokensBefore, + is_split_turn: event.preparation?.isSplitTurn, + ...attribution(), + }), + ); + }); + + pi.on('session_compact', async (event, context) => { + observe( + payload(context, 'session_compact', { + reason: event.reason, + will_retry: event.willRetry, + from_extension: event.fromExtension, + tokens_before: event.compactionEntry?.tokensBefore, + ...attribution(), + }), + ); + }); + + pi.on('tool_execution_start', async (event) => { + toolNames.set(event.toolCallId, event.toolName); + }); + + pi.on('tool_call', async (event, context) => { + let outcome: HookOutcome; + try { + outcome = await enqueue(async () => { + const active = await runtime(); + return postHook( + active, + payload(context, 'tool_call', { + tool_call_id: event.toolCallId, + tool_name: event.toolName, + input: event.input, + ...attribution(), + }), + ); + }); + } catch (error) { + return blockedInfrastructure(event.toolName, safeError(error)); + } + + if (outcome.kind === 'fault') { + return blockedInfrastructure(event.toolName, outcome.reason); + } + if (outcome.kind === 'block') { + return { block: true, reason: outcome.reason }; + } + const transformed = decideManagedToolTransform(outcome.body, event.toolCallId, event.input); + if (transformed.kind === 'invalid') { + return { + block: true, + reason: `NeMo Relay returned an invalid argument rewrite: ${transformed.reason}`, + }; + } + if (transformed.kind === 'replace') { + // Pi executes the same object after this hook and does not revalidate it. The recursive + // shape check below proves every assignment preserves the already-validated structure. + Object.assign(event.input, transformed.input); + observe( + payload(context, 'tool_arguments_transformed', { + tool_call_id: event.toolCallId, + tool_name: event.toolName, + ...attribution(), + }), + ); + } + return undefined; + }); + + pi.on('tool_execution_end', async (event, context) => { + const toolName = event.toolName || toolNames.get(event.toolCallId) || 'unknown'; + toolNames.delete(event.toolCallId); + observe( + payload(context, 'tool_execution_end', { + tool_call_id: event.toolCallId, + tool_name: toolName, + result: summarizeManagedToolResult(event.result, event.isError), + status: event.isError ? 'error' : 'ok', + ...attribution(), + }), + ); + }); + + pi.on('user_bash', async (event, context) => { + const callId = `user-bash-${userBashSequence++}`; + let outcome: HookOutcome; + try { + outcome = await enqueue(async () => { + const active = await runtime(); + return postHook( + active, + payload(context, 'user_bash', { + tool_call_id: callId, + tool_name: 'user_bash', + input: { + command: event.command, + cwd: event.cwd, + exclude_from_context: event.excludeFromContext, + }, + ...attribution(), + }), + ); + }); + } catch (error) { + return refuseUserBash(context, callId, blockedInfrastructure('user_bash', safeError(error)).reason); + } + if (outcome.kind === 'fault') { + return refuseUserBash(context, callId, blockedInfrastructure('user_bash', outcome.reason).reason); + } + if (outcome.kind === 'block') { + return refuseUserBash(context, callId, outcome.reason); + } + const original = { + command: event.command, + cwd: event.cwd, + exclude_from_context: event.excludeFromContext, + }; + const transformed = decideManagedToolTransform(outcome.body, callId, original); + if (transformed.kind === 'invalid') { + return refuseUserBash(context, callId, `NeMo Relay returned an invalid argument rewrite: ${transformed.reason}`); + } + if (transformed.kind === 'replace') { + return refuseUserBash( + context, + callId, + 'NeMo Relay rewrote this inline command, but Pi cannot safely apply inline-shell rewrites.', + ); + } + observe( + payload(context, 'user_bash_end', { + tool_call_id: callId, + tool_name: 'user_bash', + status: 'policy-allowed', + result: { content: 'Allowed by policy; Pi does not expose the command outcome.' }, + ...attribution(), + }), + ); + return undefined; + }); +} + +async function initializeRuntime(): Promise { + const config = readDeploymentConfig(); + const credential = readCredential(); + const lease = sharedLease(config, credential); + return { config, credential, lease }; +} + +function readDeploymentConfig(): DeploymentConfig { + let parsed: unknown; + try { + parsed = JSON.parse(readFileSync(new URL('./managed-config.json', import.meta.url), 'utf8')); + } catch (error) { + throw new Error(`managed Pi configuration is unreadable: ${safeError(error)}`); + } + if (!isRecord(parsed) || parsed.schema !== CONFIG_SCHEMA) { + throw new Error(`managed Pi configuration must use schema ${CONFIG_SCHEMA}`); + } + const daemonAddress = requiredRenderedString(parsed.daemonAddress, 'daemonAddress'); + const dispatcherCommand = requiredRenderedString(parsed.dispatcherCommand, 'dispatcherCommand'); + if (!isAbsolute(dispatcherCommand)) { + throw new Error('managed Pi dispatcherCommand must be an absolute administrator-owned path'); + } + let daemon: URL; + try { + daemon = new URL(daemonAddress); + } catch { + throw new Error('managed Pi daemonAddress must be an absolute HTTP(S) URL'); + } + if ( + !['http:', 'https:'].includes(daemon.protocol) || + daemon.username !== '' || + daemon.password !== '' || + daemon.search !== '' || + daemon.hash !== '' || + !['', '/'].includes(daemon.pathname) + ) { + throw new Error('managed Pi daemonAddress must be a root HTTP(S) URL without credentials, query, or fragment'); + } + const host = daemon.hostname.replace(/^\[|\]$/g, ''); + const loopback = host === 'localhost' || host === '::1' || /^127(?:\.\d{1,3}){3}$/.test(host); + if (daemon.protocol === 'http:' && !loopback) { + throw new Error('managed Pi daemonAddress must use HTTPS unless it names a loopback host'); + } + return { + schema: CONFIG_SCHEMA, + daemonAddress: daemon.href.replace(/\/$/, ''), + dispatcherCommand, + }; +} + +function requiredRenderedString(value: unknown, field: string): string { + if ( + typeof value !== 'string' || + value.length === 0 || + value.includes('\0') || + value.startsWith(CONFIG_PLACEHOLDER_PREFIX) + ) { + throw new Error(`managed Pi ${field} was not rendered by the administrator`); + } + return value; +} + +function readCredential(): string { + const value = process.env[CLIENT_TOKEN_ENV]; + if ( + value === undefined || + !/^[A-Za-z0-9_-]{43}$/.test(value) || + Buffer.from(value, 'base64url').length !== 32 || + Buffer.from(value, 'base64url').toString('base64url') !== value + ) { + throw new Error(`${CLIENT_TOKEN_ENV} must be an unpadded base64url credential encoding exactly 32 bytes`); + } + return value; +} + +function sharedLease(config: DeploymentConfig, credential: string): McpLease { + const registry = globalThis as unknown as Record; + const existing = registry[MCP_SINGLETON]; + if (existing !== undefined) { + if (!isSharedLease(existing)) { + throw new Error('the managed Pi process-global MCP slot is already occupied'); + } + if ( + existing.daemonAddress !== config.daemonAddress || + existing.dispatcherCommand !== config.dispatcherCommand || + existing.credential !== credential + ) { + throw new Error('managed Pi configuration changed while the process-global MCP lease was active'); + } + return existing; + } + + const created = createSharedLease(config, credential, () => { + if (registry[MCP_SINGLETON] === created) delete registry[MCP_SINGLETON]; + }); + registry[MCP_SINGLETON] = created; + return created; +} + +function isSharedLease(value: unknown): value is McpLease { + return ( + isRecord(value) && + value.kind === SHARED_LEASE_KIND && + typeof value.daemonAddress === 'string' && + typeof value.dispatcherCommand === 'string' && + typeof value.credential === 'string' && + typeof value.ensureReady === 'function' && + typeof value.release === 'function' + ); +} + +function createSharedLease(config: DeploymentConfig, credential: string, removeFromRegistry: () => void): McpLease { + let child: ChildProcess | undefined; + let initialized = false; + let starting: Promise | undefined; + let releasing: Promise | undefined; + let released = false; + + const lease: McpLease = { + kind: SHARED_LEASE_KIND, + daemonAddress: config.daemonAddress, + dispatcherCommand: config.dispatcherCommand, + credential, + ensureReady(): Promise { + if (released) return Promise.reject(new Error('managed Pi MCP lease was released')); + if (initialized && child && child.exitCode === null && child.signalCode === null) { + return Promise.resolve(); + } + if (starting) return starting; + starting = launch().then( + () => { + starting = undefined; + }, + (error: unknown) => { + starting = undefined; + throw error; + }, + ); + return starting; + }, + release(): Promise { + releasing ??= releaseActive(); + return releasing; + }, + }; + + async function launch(): Promise { + const launched = spawn(config.dispatcherCommand, ['daemon', 'mcp', '--daemon-address', config.daemonAddress], { + shell: false, + windowsHide: true, + stdio: ['pipe', 'pipe', 'inherit'], + env: { ...process.env, [CLIENT_TOKEN_ENV]: credential }, + }); + child = launched; + launched.stdin?.on('error', () => undefined); + launched.once('exit', () => { + if (child === launched) { + child = undefined; + initialized = false; + starting = undefined; + } + }); + try { + await initializeMcp(launched); + if (launched.exitCode !== null || launched.signalCode !== null) { + throw new Error('managed Pi MCP exited during initialization'); + } + initialized = true; + } catch (error) { + if (child === launched) child = undefined; + initialized = false; + launched.kill(); + throw error; + } + } + + async function releaseActive(): Promise { + released = true; + if (starting) await starting.catch(() => undefined); + const active = child; + initialized = false; + child = undefined; + if (active) { + active.stdin?.end(); + if (!(await waitForExit(active, MCP_RELEASE_TIMEOUT_MS))) { + active.kill(); + await waitForExit(active, 1_000); + } + } + removeFromRegistry(); + } + + return lease; +} + +function initializeMcp(child: ChildProcess): Promise { + const stdin = child.stdin; + const stdout = child.stdout; + if (!stdin || !stdout) { + return Promise.reject(new Error('managed Pi MCP did not expose stdio pipes')); + } + return new Promise((resolve, reject) => { + const lines = createInterface({ input: stdout, crlfDelay: Infinity }); + const timer = setTimeout( + () => finish(new Error('managed Pi MCP initialize response timed out')), + MCP_READY_TIMEOUT_MS, + ); + let settled = false; + + const onError = (error: Error): void => finish(error); + const onExit = (code: number | null, signal: NodeJS.Signals | null): void => { + finish(new Error(`managed Pi MCP exited before initialization (${code ?? signal ?? 'unknown'})`)); + }; + const onLine = (line: string): void => { + let message: unknown; + try { + message = JSON.parse(line); + } catch { + return; + } + if (!isRecord(message) || message.id !== MCP_INITIALIZE_ID) return; + if (isRecord(message.error)) { + finish(new Error('managed Pi MCP rejected the initialize request')); + return; + } + const result = message.result; + if ( + message.jsonrpc !== '2.0' || + !isRecord(result) || + typeof result.protocolVersion !== 'string' || + !isRecord(result.serverInfo) || + result.serverInfo.name !== 'nemo-relay' + ) { + finish(new Error('managed Pi MCP returned an invalid initialize response')); + return; + } + stdin.write(`${JSON.stringify({ jsonrpc: '2.0', method: 'notifications/initialized' })}\n`, (error) => + finish(error ?? undefined), + ); + }; + + function finish(error?: Error): void { + if (settled) return; + settled = true; + clearTimeout(timer); + child.off('error', onError); + child.off('exit', onExit); + lines.off('line', onLine); + lines.close(); + if (error) reject(error); + else resolve(); + } + + child.once('error', onError); + child.once('exit', onExit); + lines.on('line', onLine); + stdin.write( + `${JSON.stringify({ + jsonrpc: '2.0', + id: MCP_INITIALIZE_ID, + method: 'initialize', + params: { + protocolVersion: MCP_PROTOCOL_VERSION, + capabilities: {}, + clientInfo: { name: 'nemo-relay-managed-pi', version: '1.0.0' }, + }, + })}\n`, + (error) => { + if (error) finish(error); + }, + ); + }); +} + +function waitForExit(child: ChildProcess, timeoutMs: number): Promise { + if (child.exitCode !== null || child.signalCode !== null) return Promise.resolve(true); + return new Promise((resolve) => { + const timer = setTimeout(() => finish(false), timeoutMs); + const onExit = (): void => finish(true); + const finish = (exited: boolean): void => { + clearTimeout(timer); + child.off('exit', onExit); + resolve(exited); + }; + child.once('exit', onExit); + }); +} + +async function postHook(runtime: Runtime, payload: Record): Promise { + let encoded: string; + try { + encoded = JSON.stringify(payload); + } catch (error) { + return { kind: 'fault', reason: `hook payload is not JSON-safe: ${safeError(error)}` }; + } + if (Buffer.byteLength(encoded) > MAX_HOOK_PAYLOAD_BYTES) { + return { kind: 'fault', reason: 'hook payload exceeds the managed payload limit' }; + } + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), HOOK_TIMEOUT_MS); + try { + const response = await fetch(`${runtime.config.daemonAddress}/hooks/pi`, { + method: 'POST', + redirect: 'error', + headers: { + 'content-type': 'application/json', + [CLIENT_TOKEN_HEADER]: runtime.credential, + }, + body: encoded, + signal: controller.signal, + }); + const decoded = await boundedJson(response); + if (response.ok) { + if (!isRecord(decoded)) { + return { kind: 'fault', reason: 'daemon returned a non-object success body' }; + } + return { kind: 'allow', body: decoded }; + } + if (response.status === 403 && isRecord(decoded)) { + const detail = decoded.error; + if (isRecord(detail) && detail.type === 'nemo_relay_guardrail_rejected' && typeof detail.reason === 'string') { + return { kind: 'block', reason: detail.reason }; + } + } + return { kind: 'fault', reason: `daemon returned HTTP ${response.status}` }; + } catch (error) { + const reason = + error instanceof Error && error.name === 'AbortError' + ? `daemon did not answer within ${HOOK_TIMEOUT_MS}ms` + : `daemon hook request failed: ${safeError(error)}`; + return { kind: 'fault', reason }; + } finally { + clearTimeout(timer); + } +} + +async function boundedJson(response: Response): Promise { + if (!response.body) return null; + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let length = 0; + while (true) { + const part = await reader.read(); + if (part.done) break; + length += part.value.byteLength; + if (length > MAX_HOOK_RESPONSE_BYTES) { + await reader.cancel(); + throw new Error('daemon hook response exceeds the managed response limit'); + } + chunks.push(part.value); + } + const bytes = new Uint8Array(length); + let offset = 0; + for (const chunk of chunks) { + bytes.set(chunk, offset); + offset += chunk.byteLength; + } + try { + return JSON.parse(new TextDecoder('utf-8', { fatal: true }).decode(bytes)); + } catch { + return null; + } +} + +export function decideManagedProviderRedirect( + model: PiModel | undefined, + catalog: readonly PiModel[] | undefined, +): ProviderRedirectDecision { + if (!model) return { kind: 'skip', code: 'no-model', reason: 'no model is selected' }; + if (!SERVICEABLE_APIS.has(model.api)) { + return { + kind: 'skip', + code: 'unserviceable-api', + reason: `the managed daemon serves no route for the ${model.api} API`, + }; + } + if (!catalog) { + return { + kind: 'skip', + code: 'model-registry-unavailable', + reason: 'Pi did not expose the provider catalog required for a safe provider-wide redirect', + }; + } + const siblings = catalog.filter((candidate) => candidate.provider === model.provider); + if (siblings.length === 0) { + return { + kind: 'skip', + code: 'provider-catalog-missing', + reason: `Pi's model registry contains no models for provider ${model.provider}`, + }; + } + const unsupported = siblings.find((candidate) => !SERVICEABLE_APIS.has(candidate.api)); + if (unsupported) { + return { + kind: 'skip', + code: 'provider-mixed-apis', + reason: + `redirecting ${model.provider} would also move its unsupported ` + `${unsupported.api} model ${unsupported.id}`, + }; + } + const upstream = normalizeBaseUrl(model.baseUrl); + const mismatched = siblings.find((candidate) => normalizeBaseUrl(candidate.baseUrl) !== upstream); + if (mismatched) { + return { + kind: 'skip', + code: 'provider-mixed-endpoints', + reason: + `redirecting ${model.provider} would also move ${mismatched.id}, which targets ` + + `${mismatched.baseUrl} rather than ${model.baseUrl}`, + }; + } + return { + kind: 'redirect', + upstream: model.baseUrl, + reason: 'provider uses only daemon-supported APIs and every model shares its endpoint', + }; +} + +function normalizeBaseUrl(value: string): string { + const trimmed = value.trim().replace(/\/+$/, ''); + try { + const url = new URL(trimmed); + const path = url.pathname.replace(/\/+$/, ''); + return `${url.protocol}//${url.host.toLowerCase()}${path}`; + } catch { + return trimmed.toLowerCase(); + } +} + +export function decideManagedToolTransform( + body: Record, + callId: string, + current: Record, +): { kind: 'none' } | { kind: 'replace'; input: Record } | { kind: 'invalid'; reason: string } { + const toolCall = body.tool_call; + if (toolCall === undefined) return { kind: 'none' }; + if (!isRecord(toolCall)) return { kind: 'invalid', reason: 'tool_call is not an object' }; + if (toolCall.input === undefined) return { kind: 'none' }; + if (typeof toolCall.tool_call_id !== 'string' || toolCall.tool_call_id !== callId) { + return { kind: 'invalid', reason: 'tool_call_id does not match the active call' }; + } + if (!isRecord(toolCall.input)) { + return { kind: 'invalid', reason: 'tool_call.input is not an object' }; + } + const violation = shapeViolation(current, toolCall.input); + if (violation) return { kind: 'invalid', reason: violation }; + return { kind: 'replace', input: toolCall.input }; +} + +function jsonType(value: unknown): string { + if (value === null) return 'null'; + if (Array.isArray(value)) return 'array'; + return typeof value; +} + +function shapeViolation(current: unknown, next: unknown, path = 'input'): string | null { + const currentType = jsonType(current); + const nextType = jsonType(next); + if (currentType !== nextType) { + return `${path} changed type from ${currentType} to ${nextType}`; + } + if (currentType === 'object') { + const currentRecord = current as Record; + const nextRecord = next as Record; + const currentKeys = Object.keys(currentRecord).sort((left, right) => left.localeCompare(right)); + const nextKeys = Object.keys(nextRecord).sort((left, right) => left.localeCompare(right)); + const added = nextKeys.filter((key) => !currentKeys.includes(key)); + const removed = currentKeys.filter((key) => !nextKeys.includes(key)); + if (added.length > 0) return `${path} added ${added.join(', ')}`; + if (removed.length > 0) return `${path} removed ${removed.join(', ')}`; + for (const key of currentKeys) { + const violation = shapeViolation(currentRecord[key], nextRecord[key], `${path}.${key}`); + if (violation) return violation; + } + } + if (currentType === 'array') { + const currentItems = current as unknown[]; + const nextItems = next as unknown[]; + if (currentItems.length !== nextItems.length) { + return `${path} changed length from ${currentItems.length} to ${nextItems.length}`; + } + for (const [index, item] of currentItems.entries()) { + const violation = shapeViolation(item, nextItems[index], `${path}[${index}]`); + if (violation) return violation; + } + } + return null; +} + +function blockedInfrastructure(toolName: string, detail: string): ToolCallResult { + return { + block: true, + reason: + `The managed NeMo Relay service could not authorize this ${toolName} call, so it was ` + + `blocked rather than allowed through unchecked. Details: ${detail}`, + }; +} + +function refusedBash(reason: string): UserBashResult { + return { + result: { + output: `NeMo Relay blocked this inline shell command: ${reason}`, + exitCode: 126, + cancelled: false, + truncated: false, + }, + }; +} + +function sessionId(context: ExtensionContext): string { + try { + return context.sessionManager?.getSessionId?.() ?? 'unknown-session'; + } catch { + return 'unknown-session'; + } +} + +export function summarizeManagedToolResult(result: unknown, isError: boolean): Record { + if (result === null || result === undefined) { + return { content: isError ? 'Tool failed with no result.' : 'Tool completed with no result.' }; + } + if (typeof result === 'string') return { content: truncate(result) }; + if (isRecord(result)) { + const content = result.content ?? result.output ?? result.text; + const text = toolResultText(content); + return { + content: text === null ? `Tool ${isError ? 'failed' : 'completed'}.` : text, + result_keys: Object.keys(result).slice(0, 20), + }; + } + return { content: primitiveSummary(result, isError) }; +} + +function primitiveSummary(result: unknown, isError: boolean): string { + switch (typeof result) { + case 'boolean': + case 'number': + case 'bigint': + case 'symbol': + return truncate(String(result)); + default: + return `Tool ${isError ? 'failed' : 'completed'} with an unsupported result type.`; + } +} + +function toolResultText(content: unknown): string | null { + if (typeof content === 'string') return truncate(content); + if (!Array.isArray(content)) return null; + + let text = ''; + let omittedChars = 0; + let foundText = false; + const append = (value: string): void => { + const kept = sliceAtCodePointBoundary(value, Math.max(0, MAX_RESULT_CHARS - text.length)); + text += kept; + omittedChars += value.length - kept.length; + }; + for (const part of content) { + if (!isRecord(part) || part.type !== 'text' || typeof part.text !== 'string') continue; + if (foundText) append('\n'); + append(part.text); + foundText = true; + } + if (!foundText) return null; + return omittedChars === 0 ? text : `${text}... [truncated ${omittedChars} chars]`; +} + +function sliceAtCodePointBoundary(value: string, limit: number): string { + let end = Math.min(value.length, limit); + if ( + end > 0 && + end < value.length && + value.charCodeAt(end - 1) >= 0xd800 && + value.charCodeAt(end - 1) <= 0xdbff && + value.charCodeAt(end) >= 0xdc00 && + value.charCodeAt(end) <= 0xdfff + ) { + end -= 1; + } + return value.slice(0, end); +} + +function truncate(value: string): string { + if (value.length <= MAX_RESULT_CHARS) return value; + const kept = sliceAtCodePointBoundary(value, MAX_RESULT_CHARS); + return `${kept}... [truncated ${value.length - kept.length} chars]`; +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function safeError(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} diff --git a/crates/cli/src/daemon/managed/pi_extension/managed-config.json b/crates/cli/src/daemon/managed/pi_extension/managed-config.json new file mode 100644 index 000000000..f8955b9c6 --- /dev/null +++ b/crates/cli/src/daemon/managed/pi_extension/managed-config.json @@ -0,0 +1,5 @@ +{ + "schema": "nemo-relay-managed-pi-v1", + "daemonAddress": "__NEMO_RELAY_DAEMON_ADDRESS__", + "dispatcherCommand": "__NEMO_RELAY_DISPATCHER_COMMAND__" +} diff --git a/crates/cli/src/daemon/managed/pi_extension/package.json b/crates/cli/src/daemon/managed/pi_extension/package.json new file mode 100644 index 000000000..63f056eb3 --- /dev/null +++ b/crates/cli/src/daemon/managed/pi_extension/package.json @@ -0,0 +1,20 @@ +{ + "name": "nemo-relay-managed-pi-v1", + "version": "1.0.0", + "private": true, + "description": "Immutable NeMo Relay managed integration for pi.", + "type": "module", + "main": "./index.ts", + "pi": { + "extensions": [ + "./index.ts" + ] + }, + "scripts": { + "typecheck": "tsc -p tsconfig.json" + }, + "devDependencies": { + "@types/node": "24.12.4", + "typescript": "5.9.3" + } +} diff --git a/crates/cli/src/daemon/managed/pi_extension/tsconfig.json b/crates/cli/src/daemon/managed/pi_extension/tsconfig.json new file mode 100644 index 000000000..6ac36f647 --- /dev/null +++ b/crates/cli/src/daemon/managed/pi_extension/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "allowImportingTsExtensions": true, + "exactOptionalPropertyTypes": true, + "lib": ["ES2023", "DOM", "DOM.Iterable"], + "module": "NodeNext", + "moduleResolution": "NodeNext", + "noEmit": true, + "noImplicitOverride": true, + "noUncheckedIndexedAccess": true, + "strict": true, + "types": ["node"], + "verbatimModuleSyntax": true + }, + "include": ["index.ts"] +} diff --git a/crates/cli/src/daemon/mcp/mod.rs b/crates/cli/src/daemon/mcp/mod.rs new file mode 100644 index 000000000..b9d2e3790 --- /dev/null +++ b/crates/cli/src/daemon/mcp/mod.rs @@ -0,0 +1,701 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Broker-attached MCP stdio process. It advertises no MCP tools. + +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::process::Stdio; +use std::time::Duration; + +use reqwest::Client; +use tokio::io::AsyncWriteExt; +use tokio::net::UdpSocket; +use tokio::process::{Child, Command}; + +use super::common::address::{daemon_url, explicit_daemon_origin}; +use super::common::client::{ + ControlRetryPolicy, begin_handshake, control_client, post_empty_idempotent, post_json, + post_json_idempotent, +}; +use super::common::control::{ + ACTIVATION_LIFETIME_MS, ActivationFailedPayload, EmptyPayload, MCP_ACTIVATION_FAILED_PATH, + MCP_HEARTBEAT_INTERVAL_MS, MCP_HEARTBEAT_PATH, MCP_LEASE_MS, MCP_REGISTER_PATH, + McpHeartbeatResponse, McpRegisterRequest, McpRegisterResponse, SessionRequest, WorkerBootstrap, + WorkerNetworkHint, WorkerNetworkHintProof, +}; +use super::common::identity::MachineIdentity; +use super::common::protocol::{BrokerDirective, ComponentRole, SensitiveString}; +use super::common::state::{ROUTE_TOKEN_ENV, RouteCredential, load_or_create_machine_identity}; +use crate::error::CliError; + +// Includes the full two-minute legal drain plus reconciliation margin before a replacement +// activation is issued. +const ACTIVATION_POLL_MAX: Duration = Duration::from_secs(150); +const REGISTRATION_RETRY_MAX: Duration = Duration::from_secs(30); +const REGISTRATION_RETRY_DELAY: Duration = Duration::from_millis(100); +const MIN_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(1); +const MAX_HEARTBEAT_INTERVAL: Duration = Duration::from_millis(MCP_LEASE_MS / 3); +const HEARTBEAT_RETRY_WINDOW_MS: u64 = MCP_LEASE_MS - MCP_HEARTBEAT_INTERVAL_MS - 5_000; +const HEARTBEAT_RETRY_POLICY: ControlRetryPolicy = ControlRetryPolicy::new( + Duration::from_secs(2), + Duration::from_millis(HEARTBEAT_RETRY_WINDOW_MS), + Duration::from_millis(250), +); +const RELEASE_RETRY_POLICY: ControlRetryPolicy = ControlRetryPolicy::new( + Duration::from_millis(500), + Duration::from_secs(2), + Duration::from_millis(100), +); +const WORKER_ADVERTISE_ENV: &str = "NEMO_RELAY_WORKER_ADVERTISE_ADDRESS"; +const WORKER_PORT_ENV: &str = "NEMO_RELAY_WORKER_PORT"; + +#[derive(Debug, Clone)] +pub(crate) struct Options { + pub(crate) daemon_address: String, +} + +struct McpLease { + client: Client, + daemon_origin: String, + route_credential: RouteCredential, + identity: MachineIdentity, + session_id: String, + session_token: SensitiveString, + heartbeat_interval: Duration, + sequence: u64, + pending_heartbeat: Option>, +} + +struct Registration { + directive: BrokerDirective, + session_token: SensitiveString, + heartbeat_interval: Duration, +} + +pub(crate) async fn run(options: Options) -> Result<(), CliError> { + let daemon_origin = explicit_daemon_origin(&options.daemon_address)?; + let client = control_client()?; + let route_credential = RouteCredential::from_environment()?; + let identity = load_or_create_machine_identity()?; + let session_id = uuid::Uuid::now_v7().to_string(); + let registration = register( + &client, + &daemon_origin, + &route_credential, + &identity, + &session_id, + ) + .await?; + let mut lease = McpLease { + client, + daemon_origin, + route_credential, + identity, + session_id, + session_token: registration.session_token, + heartbeat_interval: registration.heartbeat_interval, + sequence: 0, + pending_heartbeat: None, + }; + make_route_ready(&mut lease, registration.directive).await?; + + log::info!( + target: "nemo_relay.daemon.mcp", + event = "daemon_mcp_ready"; + "Broker reference acquired; MCP protocol is ready" + ); + let result = { + let protocol = crate::mcp::serve_daemon_stdio(); + let control = maintain_lease(&mut lease); + tokio::pin!(protocol); + tokio::pin!(control); + tokio::select! { + result = &mut protocol => result, + result = &mut control => result, + } + }; + release(&mut lease).await; + result +} + +async fn register( + client: &Client, + daemon_origin: &str, + credential: &RouteCredential, + identity: &MachineIdentity, + session_id: &str, +) -> Result { + let deadline = tokio::time::Instant::now() + REGISTRATION_RETRY_MAX; + loop { + match register_once(client, daemon_origin, credential, identity, session_id).await { + Ok(registration) => return Ok(registration), + Err(error @ CliError::Upstream(_)) => { + let now = tokio::time::Instant::now(); + if now >= deadline { + return Err(error); + } + tokio::time::sleep_until(deadline.min(now + REGISTRATION_RETRY_DELAY)).await; + } + Err(error) => return Err(error), + } + } +} + +async fn register_once( + client: &Client, + daemon_origin: &str, + credential: &RouteCredential, + identity: &MachineIdentity, + session_id: &str, +) -> Result { + let handshake = begin_handshake( + client, + daemon_origin, + ComponentRole::Mcp, + identity, + session_id, + Some(credential.digest()), + ) + .await?; + let worker_network = worker_network_hint(daemon_origin).await?; + let worker_network = WorkerNetworkHintProof::sign( + worker_network, + &handshake.proof.transcript.daemon_target, + session_id, + &handshake.proof.transcript.challenge_id, + &identity.fingerprint(), + identity, + )?; + let response: McpRegisterResponse = post_json( + client, + &format!("{daemon_origin}{MCP_REGISTER_PATH}"), + &McpRegisterRequest { + proof: handshake.proof.clone(), + worker_network, + }, + Some(credential.expose()), + ) + .await?; + handshake.authenticate_daemon(&response.daemon_proof)?; + Ok(Registration { + directive: response.directive, + session_token: response.session_token, + heartbeat_interval: validate_heartbeat_interval(response.heartbeat_interval_ms)?, + }) +} + +async fn worker_network_hint(daemon_origin: &str) -> Result { + let advertised_override = optional_environment(WORKER_ADVERTISE_ENV)?; + let port_override = optional_environment(WORKER_PORT_ENV)?; + let (advertised_override, port) = + parse_worker_network_overrides(advertised_override.as_deref(), port_override.as_deref())?; + let daemon = daemon_url(daemon_origin)?; + let daemon_addresses = tokio::net::lookup_host(( + daemon + .host_str() + .ok_or_else(|| CliError::Config("daemon address is missing a host".into()))?, + daemon + .port_or_known_default() + .ok_or_else(|| CliError::Config("daemon address is missing a port".into()))?, + )) + .await + .map_err(|error| CliError::Launch(format!("failed to resolve daemon IPv4 route: {error}")))? + .filter_map(|address| match address { + SocketAddr::V4(address) => Some(address), + SocketAddr::V6(_) => None, + }) + .collect::>(); + let daemon_address = daemon_addresses + .iter() + .copied() + .find(|address| !address.ip().is_loopback()) + .or_else(|| daemon_addresses.first().copied()) + .ok_or_else(|| { + CliError::Config( + "daemon target has no IPv4 route; daemon workers support IPv4 networking only" + .into(), + ) + })?; + let advertised_host = match advertised_override { + Some(address) => address, + None if daemon_address.ip().is_loopback() => Ipv4Addr::LOCALHOST.to_string(), + None => { + let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0)).await?; + socket.connect(daemon_address).await?; + match socket.local_addr()?.ip() { + IpAddr::V4(address) if !address.is_unspecified() => address.to_string(), + _ => { + return Err(CliError::Launch( + "failed to determine a concrete local IPv4 route to the daemon".into(), + )); + } + } + } + }; + let advertised_is_loopback = advertised_host.eq_ignore_ascii_case("localhost") + || advertised_host + .parse::() + .is_ok_and(|address| address.is_loopback()); + if !daemon_address.ip().is_loopback() && advertised_is_loopback { + return Err(CliError::Config(format!( + "{WORKER_ADVERTISE_ENV} cannot be loopback for a remote daemon" + ))); + } + WorkerNetworkHint::new(advertised_host, port) +} + +fn parse_worker_network_overrides( + advertised: Option<&str>, + port: Option<&str>, +) -> Result<(Option, Option), CliError> { + let advertised = advertised + .map(str::trim) + .map(|value| { + WorkerNetworkHint::new(value, None) + .map(|hint| hint.advertised_host) + .map_err(|_| { + CliError::Config(format!( + "{WORKER_ADVERTISE_ENV} must be a concrete hostname or IPv4 address" + )) + }) + }) + .transpose()?; + let port = port + .map(str::trim) + .map(|value| { + value + .parse::() + .ok() + .filter(|port| *port != 0) + .ok_or_else(|| { + CliError::Config(format!( + "{WORKER_PORT_ENV} must be an integer between 1 and 65535" + )) + }) + }) + .transpose()?; + Ok((advertised, port)) +} + +fn optional_environment(name: &str) -> Result, CliError> { + std::env::var_os(name) + .map(|value| { + value + .into_string() + .map_err(|_| CliError::Config(format!("{name} must contain valid Unicode text"))) + }) + .transpose() +} + +fn validate_heartbeat_interval(milliseconds: u64) -> Result { + let interval = Duration::from_millis(milliseconds); + if !(MIN_HEARTBEAT_INTERVAL..=MAX_HEARTBEAT_INTERVAL).contains(&interval) { + return Err(CliError::Unauthorized( + "daemon returned an invalid MCP heartbeat interval".into(), + )); + } + Ok(interval) +} + +/// A pending worker must not survive failed activation or cancellation. Readiness transfers +/// ownership to the broker; only that success path disarms this guard. +struct ActivationChild { + child: Child, + published: bool, +} + +impl Drop for ActivationChild { + fn drop(&mut self) { + if !self.published { + let _ = self.child.start_kill(); + } + } +} + +type PendingLaunch = Option<(String, ActivationChild, tokio::time::Instant)>; + +async fn stop_pending_launch(launched: &mut PendingLaunch) -> Result<(), CliError> { + if let Some((_, mut child, _)) = launched.take() { + child.child.kill().await.map_err(CliError::Io)?; + } + Ok(()) +} + +async fn make_route_ready( + lease: &mut McpLease, + mut directive: BrokerDirective, +) -> Result<(), CliError> { + let started = tokio::time::Instant::now(); + let mut launched: PendingLaunch = None; + let result = async { + loop { + if started.elapsed() > ACTIVATION_POLL_MAX + && !matches!( + directive, + BrokerDirective::ReuseWorker { .. } | BrokerDirective::UsePassThrough + ) + { + return Err(CliError::Launch( + "timed out waiting for the broker route to become ready".into(), + )); + } + match directive { + BrokerDirective::ReuseWorker { .. } => { + if let Some((_, child, _)) = launched.as_mut() { + child.published = true; + } + launched.take(); + return Ok(()); + } + BrokerDirective::UsePassThrough => return Ok(()), + BrokerDirective::LaunchWorker { .. } => { + let bootstrap = WorkerBootstrap::from_directive(directive.clone()) + .expect("launch directive was matched"); + let already_launched = launched + .as_ref() + .is_some_and(|(id, _, _)| id == &bootstrap.activation_id); + if !already_launched { + stop_pending_launch(&mut launched).await?; + match launch_worker(&lease.daemon_origin, &bootstrap).await { + Ok(child) => { + launched = Some(( + bootstrap.activation_id.clone(), + child, + tokio::time::Instant::now(), + )); + } + Err(error) => { + report_activation_failed(lease, &bootstrap.activation_id, &error) + .await?; + directive = refresh_registration(lease).await?.directive; + continue; + } + } + } + if let Some((activation_id, child, _)) = launched.as_mut() + && activation_id == &bootstrap.activation_id + && let Some(status) = child.child.try_wait().map_err(CliError::Io)? + { + let error = CliError::Launch(format!( + "activated worker exited before readiness with {status}" + )); + report_activation_failed(lease, &bootstrap.activation_id, &error).await?; + directive = refresh_registration(lease).await?.directive; + continue; + } + if launched + .as_ref() + .is_some_and(|(activation_id, _, started)| { + activation_timed_out( + &bootstrap.activation_id, + activation_id, + *started, + tokio::time::Instant::now(), + ) + }) + { + let error = CliError::Launch( + "activated worker did not register within 15 seconds".into(), + ); + stop_pending_launch(&mut launched).await?; + report_activation_failed(lease, &bootstrap.activation_id, &error).await?; + directive = refresh_registration(lease).await?.directive; + continue; + } + // Poll readiness through the authenticated lease rather than repeating the full + // signed registration handshake while the worker starts. + if let Some(updated) = poll_worker_activation(lease).await? { + directive = updated; + continue; + } + } + BrokerDirective::WaitForWorker { retry_after_ms } => { + tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 1_000))) + .await; + } + } + if !matches!(directive, BrokerDirective::LaunchWorker { .. }) { + directive = refresh_registration(lease).await?.directive; + } + } + } + .await; + let cleanup = stop_pending_launch(&mut launched).await; + result.and(cleanup) +} + +async fn poll_worker_activation(lease: &mut McpLease) -> Result, CliError> { + tokio::time::sleep(Duration::from_secs(1)).await; + Ok(renew_lease_with(lease, HEARTBEAT_RETRY_POLICY) + .await? + .directive) +} + +fn activation_timed_out( + current_activation_id: &str, + launched_activation_id: &str, + launched_at: tokio::time::Instant, + now: tokio::time::Instant, +) -> bool { + current_activation_id == launched_activation_id + && now.saturating_duration_since(launched_at) + >= Duration::from_millis(ACTIVATION_LIFETIME_MS) +} + +async fn refresh_registration(lease: &mut McpLease) -> Result { + let registration = register( + &lease.client, + &lease.daemon_origin, + &lease.route_credential, + &lease.identity, + &lease.session_id, + ) + .await?; + apply_registration(lease, ®istration); + Ok(registration) +} + +fn apply_registration(lease: &mut McpLease, registration: &Registration) { + let session_rotated = lease.session_token != registration.session_token; + lease.session_token = registration.session_token.clone(); + lease.heartbeat_interval = registration.heartbeat_interval; + if session_rotated { + lease.sequence = 0; + lease.pending_heartbeat = None; + } +} + +async fn launch_worker( + daemon_origin: &str, + bootstrap: &WorkerBootstrap, +) -> Result { + let executable = std::env::current_exe().map_err(|error| { + CliError::Launch(format!( + "failed to resolve the nemo-relay executable: {error}" + )) + })?; + let mut command = worker_command(&executable, daemon_origin, bootstrap); + let child = command + .spawn() + .map_err(|error| CliError::Launch(format!("failed to launch daemon worker: {error}")))?; + let mut child = ActivationChild { + child, + published: false, + }; + let transfer = async { + let mut stdin = child.child.stdin.take().ok_or_else(|| { + CliError::Launch("failed to create the protected worker activation pipe".into()) + })?; + let payload = serde_json::to_vec(bootstrap).map_err(|error| { + CliError::Launch(format!("failed to encode worker activation grant: {error}")) + })?; + stdin.write_all(&payload).await.map_err(|error| { + CliError::Launch(format!( + "failed to transfer worker activation grant: {error}" + )) + })?; + stdin.shutdown().await.map_err(|error| { + CliError::Launch(format!("failed to close worker activation pipe: {error}")) + })?; + Ok::<(), CliError>(()) + } + .await; + if let Err(error) = transfer { + child.child.kill().await.map_err(CliError::Io)?; + return Err(error); + } + Ok(child) +} + +fn worker_command( + executable: &std::path::Path, + daemon_origin: &str, + bootstrap: &WorkerBootstrap, +) -> Command { + let mut command = Command::new(executable); + command + .arg("daemon") + .arg("worker") + .arg("--daemon-address") + .arg(daemon_origin) + .arg("--bind") + .arg(bootstrap.bind_ip.to_string()) + .stdin(Stdio::piped()) + .stdout(Stdio::null()) + .stderr(Stdio::inherit()) + .env_remove(ROUTE_TOKEN_ENV) + .kill_on_drop(false); + if bootstrap.port != 0 { + command.arg("--port").arg(bootstrap.port.to_string()); + } + if let Some(advertise_address) = bootstrap.advertise_address.as_deref() { + command.arg("--advertise-address").arg(advertise_address); + } + command +} + +async fn report_activation_failed( + lease: &mut McpLease, + activation_id: &str, + error: &CliError, +) -> Result<(), CliError> { + log::error!( + target: "nemo_relay.daemon.mcp", + event = "worker_launch_failed", + error_kind = error.log_kind(); + "MCP could not activate the broker-selected worker" + ); + lease.sequence = lease.sequence.saturating_add(1); + let request = SessionRequest::new( + lease.session_id.clone(), + lease.session_token.clone(), + lease.sequence, + ActivationFailedPayload { + activation_id: activation_id.to_owned(), + reason: error.to_string(), + }, + )?; + post_empty_idempotent( + &lease.client, + &format!("{}{}", lease.daemon_origin, MCP_ACTIVATION_FAILED_PATH), + &request, + RELEASE_RETRY_POLICY, + ) + .await +} + +async fn maintain_lease(lease: &mut McpLease) -> Result<(), CliError> { + let mut interval = heartbeat_interval(lease.heartbeat_interval); + interval.tick().await; + loop { + interval.tick().await; + let response = match renew_lease_with(lease, HEARTBEAT_RETRY_POLICY).await { + Ok(response) => response, + Err(CliError::Unauthorized(_)) => { + // A daemon restart invalidates its in-memory session token. Re-authenticate using + // the pinned daemon identity and the same user-machine identity instead of + // tearing down an otherwise healthy MCP stdio session. + let registration = refresh_registration(lease).await?; + make_route_ready(lease, registration.directive).await?; + interval = heartbeat_interval(lease.heartbeat_interval); + continue; + } + Err(error) => return Err(error), + }; + if let Some(directive) = response.directive { + make_route_ready(lease, directive).await?; + interval = heartbeat_interval(lease.heartbeat_interval); + } + } +} + +async fn renew_lease_with( + lease: &mut McpLease, + retry_policy: ControlRetryPolicy, +) -> Result { + if lease.pending_heartbeat.is_none() { + lease.sequence = lease + .sequence + .checked_add(1) + .ok_or_else(|| CliError::Launch("daemon MCP control sequence was exhausted".into()))?; + lease.pending_heartbeat = Some(SessionRequest::new( + lease.session_id.clone(), + lease.session_token.clone(), + lease.sequence, + EmptyPayload::default(), + )?); + } + let request = lease + .pending_heartbeat + .as_ref() + .expect("pending MCP heartbeat was initialized"); + let response = post_json_idempotent( + &lease.client, + &format!("{}{}", lease.daemon_origin, MCP_HEARTBEAT_PATH), + request, + None, + retry_policy, + ) + .await?; + lease.pending_heartbeat = None; + Ok(response) +} + +fn heartbeat_interval(duration: Duration) -> tokio::time::Interval { + let mut interval = tokio::time::interval(duration); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + interval +} + +async fn release(lease: &mut McpLease) { + if let Some(request) = lease.pending_heartbeat.as_ref() + && let Err(error) = post_json_idempotent::<_, McpHeartbeatResponse>( + &lease.client, + &format!("{}{}", lease.daemon_origin, MCP_HEARTBEAT_PATH), + request, + None, + RELEASE_RETRY_POLICY, + ) + .await + { + log::warn!( + target: "nemo_relay.daemon.mcp", + event = "mcp_release_failed", + error_kind = error.log_kind(); + "Failed to settle the pending MCP heartbeat before release" + ); + return; + } + lease.pending_heartbeat = None; + lease.sequence = match lease.sequence.checked_add(1) { + Some(sequence) => sequence, + None => { + log::warn!( + target: "nemo_relay.daemon.mcp", + event = "mcp_release_failed"; + "Daemon MCP control sequence was exhausted before release" + ); + return; + } + }; + let request = match SessionRequest::new( + lease.session_id.clone(), + lease.session_token.clone(), + lease.sequence, + EmptyPayload::default(), + ) { + Ok(request) => request, + Err(error) => { + log::warn!( + target: "nemo_relay.daemon.mcp", + event = "mcp_release_failed", + error_kind = error.log_kind(); + "Failed to construct the MCP release message" + ); + return; + } + }; + if let Err(error) = post_empty_idempotent( + &lease.client, + &format!( + "{}{}", + lease.daemon_origin, + super::common::control::MCP_RELEASE_PATH + ), + &request, + RELEASE_RETRY_POLICY, + ) + .await + { + log::warn!( + target: "nemo_relay.daemon.mcp", + event = "mcp_release_failed", + error_kind = error.log_kind(); + "Failed to release the daemon MCP reference" + ); + } +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/mcp_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/mod.rs b/crates/cli/src/daemon/mod.rs new file mode 100644 index 000000000..642eb1b9b --- /dev/null +++ b/crates/cli/src/daemon/mod.rs @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Brokered daemon runtime for managed coding-agent integrations. + +use std::net::Ipv4Addr; +use std::path::PathBuf; + +use crate::error::CliError; + +pub(crate) mod broker; +pub(crate) mod common; +pub(crate) mod hook; +pub(crate) mod managed; +pub(crate) mod mcp; +pub(crate) mod worker; + +/// Runtime options for the public daemon listener. +#[derive(Debug, Clone)] +pub(crate) struct ServerOptions { + pub(crate) bind: Ipv4Addr, + pub(crate) port: u16, + pub(crate) advertise_address: Option, + pub(crate) pass_through: bool, + pub(crate) gateway: crate::server::GatewayOverrides, + pub(crate) tls_cert: Option, + pub(crate) tls_key: Option, +} + +pub(crate) async fn serve(options: ServerOptions) -> Result<(), CliError> { + broker::server::serve(options).await +} diff --git a/crates/cli/src/daemon/worker/control.rs b/crates/cli/src/daemon/worker/control.rs new file mode 100644 index 000000000..c265106ab --- /dev/null +++ b/crates/cli/src/daemon/worker/control.rs @@ -0,0 +1,330 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Worker-side authenticated daemon control session. + +use std::time::Duration; + +use reqwest::Client; + +use super::super::common::client::{ + ControlRetryPolicy, begin_handshake, control_client, post_empty_idempotent, post_json, +}; +use super::super::common::control::{ + SessionRequest, WORKER_HEARTBEAT_PATH, WORKER_READY_PATH, WORKER_RECOVER_PATH, + WORKER_REGISTER_PATH, WorkerBootstrap, WorkerGenerationGrant, WorkerHeartbeatPayload, + WorkerReadyPayload, WorkerRecoverRequest, WorkerRegisterRequest, WorkerRegisterResponse, +}; +use super::super::common::identity::{MachineIdentity, TokenDigest}; +use super::super::common::protocol::{ComponentRole, SensitiveString}; +use crate::error::CliError; + +const MIN_HEARTBEAT_INTERVAL: Duration = Duration::from_millis(100); +const MAX_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(20); +const CONTROL_RETRY_POLICY: ControlRetryPolicy = ControlRetryPolicy::new( + Duration::from_secs(1), + Duration::from_secs(14), + Duration::from_millis(100), +); +const REGISTRATION_RETRY_MAX: Duration = Duration::from_secs(30); +const REGISTRATION_RETRY_DELAY: Duration = Duration::from_millis(100); + +pub(super) struct Registration { + client: Client, + session_token: SensitiveString, + data_token: SensitiveString, + heartbeat_interval: Duration, + next_sequence: u64, + pending_ready: Option>, + pending_heartbeat: Option>, + generation_grant: WorkerGenerationGrant, +} + +impl Registration { + pub(super) fn data_token_digest(&self) -> TokenDigest { + TokenDigest::from_token(self.data_token.expose().as_bytes()) + } + + pub(super) fn session_token_digest(&self) -> TokenDigest { + TokenDigest::from_token(self.session_token.expose().as_bytes()) + } + + pub(super) const fn heartbeat_interval(&self) -> Duration { + self.heartbeat_interval + } + + pub(super) const fn generation_grant(&self) -> &WorkerGenerationGrant { + &self.generation_grant + } + + pub(super) async fn ready( + &mut self, + daemon_origin: &str, + worker_id: &str, + ) -> Result<(), CliError> { + if self.pending_ready.is_none() { + self.pending_ready = Some(SessionRequest::new( + worker_id.to_owned(), + self.session_token.clone(), + self.next_sequence, + WorkerReadyPayload { + worker_id: worker_id.to_owned(), + }, + )?); + } + let request = self + .pending_ready + .as_ref() + .expect("pending readiness message was initialized"); + post_empty_idempotent( + &self.client, + &format!("{daemon_origin}{WORKER_READY_PATH}"), + request, + CONTROL_RETRY_POLICY, + ) + .await?; + self.pending_ready = None; + self.advance_sequence() + } + + pub(super) async fn heartbeat( + &mut self, + daemon_origin: &str, + worker_id: &str, + ) -> Result<(), CliError> { + if self.pending_heartbeat.is_none() { + self.pending_heartbeat = Some(SessionRequest::new( + worker_id.to_owned(), + self.session_token.clone(), + self.next_sequence, + WorkerHeartbeatPayload { + worker_id: worker_id.to_owned(), + }, + )?); + } + let request = self + .pending_heartbeat + .as_ref() + .expect("pending heartbeat was initialized"); + post_empty_idempotent( + &self.client, + &format!("{daemon_origin}{WORKER_HEARTBEAT_PATH}"), + request, + CONTROL_RETRY_POLICY, + ) + .await?; + self.pending_heartbeat = None; + self.advance_sequence() + } + + fn advance_sequence(&mut self) -> Result<(), CliError> { + self.next_sequence = self.next_sequence.checked_add(1).ok_or_else(|| { + CliError::Launch("daemon worker control sequence was exhausted".into()) + })?; + Ok(()) + } +} + +pub(super) async fn register( + daemon_origin: &str, + identity: &MachineIdentity, + worker_id: &str, + endpoint: &str, + bootstrap: WorkerBootstrap, + tls_root_certificate: Option, +) -> Result { + let deadline = tokio::time::Instant::now() + REGISTRATION_RETRY_MAX; + loop { + match register_once( + daemon_origin, + identity, + worker_id, + endpoint, + bootstrap.clone(), + tls_root_certificate.clone(), + ) + .await + { + Ok(registration) => return Ok(registration), + Err(error @ CliError::Upstream(_)) => { + let now = tokio::time::Instant::now(); + if now >= deadline { + return Err(error); + } + tokio::time::sleep_until(deadline.min(now + REGISTRATION_RETRY_DELAY)).await; + } + Err(error) => return Err(error), + } + } +} + +async fn register_once( + daemon_origin: &str, + identity: &MachineIdentity, + worker_id: &str, + endpoint: &str, + bootstrap: WorkerBootstrap, + tls_root_certificate: Option, +) -> Result { + let client = control_client()?; + let handshake = begin_handshake( + &client, + daemon_origin, + ComponentRole::Worker, + identity, + worker_id, + None, + ) + .await?; + let request = WorkerRegisterRequest { + proof: handshake.proof.clone(), + worker_id: worker_id.to_owned(), + endpoint: endpoint.to_owned(), + activation_id: bootstrap.activation_id, + activation_token: bootstrap.activation_token, + tls_root_certificate, + }; + let response: WorkerRegisterResponse = post_json( + &client, + &format!("{daemon_origin}{WORKER_REGISTER_PATH}"), + &request, + None, + ) + .await?; + handshake.authenticate_daemon(&response.daemon_proof)?; + registration(client, response) +} + +pub(super) async fn recover( + daemon_origin: &str, + identity: &MachineIdentity, + worker_id: &str, + endpoint: &str, + tls_root_certificate: Option<&str>, + generation_grant: WorkerGenerationGrant, +) -> Result { + let deadline = tokio::time::Instant::now() + REGISTRATION_RETRY_MAX; + loop { + match recover_once( + daemon_origin, + identity, + worker_id, + endpoint, + tls_root_certificate, + generation_grant.clone(), + ) + .await + { + Ok(registration) => return Ok(registration), + Err(error @ CliError::Upstream(_)) => { + let now = tokio::time::Instant::now(); + if now >= deadline { + return Err(error); + } + tokio::time::sleep_until(deadline.min(now + REGISTRATION_RETRY_DELAY)).await; + } + Err(error) => return Err(error), + } + } +} + +async fn recover_once( + daemon_origin: &str, + identity: &MachineIdentity, + worker_id: &str, + endpoint: &str, + tls_root_certificate: Option<&str>, + generation_grant: WorkerGenerationGrant, +) -> Result { + let client = control_client()?; + let handshake = begin_handshake( + &client, + daemon_origin, + ComponentRole::Worker, + identity, + worker_id, + None, + ) + .await?; + let request = WorkerRecoverRequest { + proof: handshake.proof.clone(), + worker_id: worker_id.to_owned(), + endpoint: endpoint.to_owned(), + tls_root_certificate: tls_root_certificate.map(ToOwned::to_owned), + generation_grant, + }; + let response: WorkerRegisterResponse = post_json( + &client, + &format!("{daemon_origin}{WORKER_RECOVER_PATH}"), + &request, + None, + ) + .await?; + handshake.authenticate_daemon(&response.daemon_proof)?; + registration(client, response) +} + +fn registration( + client: Client, + response: WorkerRegisterResponse, +) -> Result { + let heartbeat_interval = validate_heartbeat_interval(response.heartbeat_interval_ms)?; + Ok(Registration { + client, + session_token: response.session_token, + data_token: response.data_token, + heartbeat_interval, + next_sequence: 1, + pending_ready: None, + pending_heartbeat: None, + generation_grant: response.generation_grant, + }) +} + +fn validate_heartbeat_interval(milliseconds: u64) -> Result { + let interval = Duration::from_millis(milliseconds); + if !(MIN_HEARTBEAT_INTERVAL..=MAX_HEARTBEAT_INTERVAL).contains(&interval) { + return Err(CliError::Unauthorized( + "daemon returned an invalid worker heartbeat interval".into(), + )); + } + Ok(interval) +} + +#[cfg(test)] +pub(super) fn test_registration(data_token: &str, session_token: &str) -> Registration { + let identity = MachineIdentity::generate().expect("test identity").identity; + let generation_grant = WorkerGenerationGrant::issue( + "worker-one", + identity.fingerprint(), + "http://127.0.0.1:1", + None, + &identity, + ) + .expect("test generation grant"); + Registration { + client: control_client().expect("test control client"), + session_token: SensitiveString::new(session_token).expect("test session token"), + data_token: SensitiveString::new(data_token).expect("test data token"), + heartbeat_interval: Duration::from_secs(5), + next_sequence: 1, + pending_ready: None, + pending_heartbeat: None, + generation_grant, + } +} + +#[cfg(test)] +pub(super) fn test_registration_values( + registration: Registration, +) -> (String, String, WorkerGenerationGrant) { + ( + registration.data_token.expose().to_owned(), + registration.session_token.expose().to_owned(), + registration.generation_grant, + ) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/worker_control_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/worker/managed.rs b/crates/cli/src/daemon/worker/managed.rs new file mode 100644 index 000000000..8ca31c888 --- /dev/null +++ b/crates/cli/src/daemon/worker/managed.rs @@ -0,0 +1,1405 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Managed worker execution that keeps response delivery on the raw frame path. + +use std::collections::BTreeSet; +use std::error::Error as _; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use axum::Json; +use axum::body::{Body, Bytes}; +use axum::http::header::{ACCEPT, AUTHORIZATION, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE}; +use axum::http::{HeaderMap, HeaderName, HeaderValue, Method, Request, Response, StatusCode, Uri}; +use axum::response::IntoResponse; +use http_body_util::LengthLimitError; +use hyper::body::{Body as HttpBody, Frame, SizeHint}; +use nemo_relay::api::llm::{ + LlmCallEndParams, LlmCallParams, LlmRequest, llm_call, llm_call_end, llm_conditional_execution, + llm_request_intercepts, +}; +use nemo_relay::api::registry::{ + RuntimeRegistrationIdentity, RuntimeRegistrationKind, list_runtime_registrations, +}; +use nemo_relay::api::runtime::TASK_SCOPE_STACK; +use nemo_relay::codec::resolve::{ProviderSurface, request_codec, response_codec, streaming_codec}; +use nemo_relay::codec::streaming::SseEventDecoder; +use nemo_relay::error::FlowError; +use serde_json::{Value, json}; +use tokio::sync::{Notify, mpsc}; + +use super::super::common::control::{ + CLIENT_TOKEN_HEADER, WORKER_ROUTE_FAILURE_HEADER, WORKER_TOKEN_HEADER, +}; +use super::super::common::routes::{HookRoute, ProviderRoute}; +use super::super::common::transport::{ + BoxError, PooledClient, RelayBody, box_body, prepare_forward_request, prepare_forward_response, +}; +use crate::agents::shared::adapters::{claude_code, codex, pi}; +use crate::configuration::GatewayConfig; +use crate::error::CliError; +use crate::plugins::lifecycle::ActiveDynamicPluginComponent; +use crate::server::ServerPluginActivation; +use crate::sessions::{GatewayCallPrep, GatewaySessionFinish, SessionManager}; + +/// Owns session accounting across middleware, dispatch, and off-path observation. Cleanup is +/// spawned before it is awaited, so cancellation of either delivery or cleanup cannot lose it. +struct GatewayCallCleanup { + sessions: Option, + session_id: String, + finish: GatewaySessionFinish, + runtime: tokio::runtime::Handle, +} + +impl GatewayCallCleanup { + fn new(sessions: SessionManager, prep: &GatewayCallPrep) -> Self { + Self { + sessions: Some(sessions), + session_id: prep.session_id.clone(), + finish: prep.session_finish, + runtime: tokio::runtime::Handle::current(), + } + } + + fn start_cleanup(&mut self) -> Option> { + let sessions = self.sessions.take()?; + let session_id = self.session_id.clone(); + let finish = self.finish; + Some(self.runtime.spawn(async move { + sessions.finish_gateway_call(&session_id, finish).await; + })) + } + + async fn finish(mut self) { + if let Some(cleanup) = self.start_cleanup() { + let _ = cleanup.await; + } + } +} + +impl Drop for GatewayCallCleanup { + fn drop(&mut self) { + // Dropping a JoinHandle detaches the cleanup rather than cancelling it. + drop(self.start_cleanup()); + } +} + +const RESPONSE_HEAD_TIMEOUT: Duration = Duration::from_secs(60); +/// Observation is off-path and has its own budget for long-running completions. +const OBSERVATION_COMPLETION_TIMEOUT: Duration = Duration::from_secs(900); +const OBSERVATION_QUEUE_FRAMES: usize = 32; +const DEFAULT_OBSERVATION_CAPTURE_BYTES: usize = 4 * 1024 * 1024; +const OBSERVATION_CAPTURE_BYTES_ENV: &str = "NEMO_RELAY_DAEMON_OBSERVATION_CAPTURE_BYTES"; +const OBSERVATION_ACTIVE: u8 = 0; +const OBSERVATION_COMPLETE: u8 = 1; +const OBSERVATION_BODY_ERROR: u8 = 2; +const OBSERVATION_CANCELLED: u8 = 3; +const STREAM_MODE_MUTATION_ERROR: &str = + "daemon worker request middleware cannot change stream mode"; +const INTERNAL_HEADER_PREFIX: &str = "x-nemo-relay-"; +const INTERNAL_DISPATCH_URL_HEADER: &str = "x-nemo-relay-internal-dispatch-url"; +const INTERNAL_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-route"; +const INTERNAL_DISPATCH_BACKEND_HEADER: &str = "x-nemo-relay-internal-dispatch-backend"; +const INTERNAL_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware"; + +/// Runtime-owned plugin activation, hook sessions, and response observation. +pub(super) struct ManagedRuntime { + config: GatewayConfig, + sessions: SessionManager, + owner: String, + observation_capture_bytes: usize, + activation: Mutex>, +} + +impl ManagedRuntime { + pub(super) async fn initialize( + config: GatewayConfig, + dynamic_plugins: Vec, + owner: String, + ) -> Result { + let observation_capture_bytes = observation_capture_limit_from_environment()?; + let activation = + crate::server::initialize_plugin_host(config.plugin_config.clone(), dynamic_plugins) + .await?; + if let Err(error) = reject_incompatible_execution_middleware() { + if let Some(activation) = activation { + let _ = activation.clear(); + } + return Err(error); + } + let sessions = SessionManager::new(config.clone()); + sessions.start_idle_sweeper(); + Ok(Self { + config, + sessions, + owner, + observation_capture_bytes, + activation: Mutex::new(activation), + }) + } + + /// Rechecks the transport contract before a provider body is polled. Plugin activation is + /// normally static, but this also fails closed if a component installs middleware later. + pub(super) fn ensure_streaming_transport_compatible(&self) -> Result<(), CliError> { + reject_incompatible_execution_middleware() + } + + pub(super) async fn close(&self) -> Result<(), CliError> { + let sessions = self.sessions.close_all("daemon_worker_shutdown").await; + let subscribers = nemo_relay::api::runtime::flush_subscribers().map_err(CliError::from); + let activation = lock(&self.activation) + .take() + .map(ServerPluginActivation::clear); + sessions?; + subscribers?; + activation.transpose()?; + Ok(()) + } + + pub(super) async fn handle_hook( + &self, + route: HookRoute, + request: Request, + ) -> Response { + match self.handle_hook_inner(route, request).await { + Ok(response) => Json(response).into_response(), + Err(error) => error.into_response(), + } + } + + async fn handle_hook_inner( + &self, + route: HookRoute, + request: Request, + ) -> Result { + let (mut parts, body) = request.into_parts(); + strip_worker_headers(&mut parts.headers); + let bytes = axum::body::to_bytes(body, self.config.max_hook_payload_bytes) + .await + .map_err(body_read_error)?; + let payload = serde_json::from_slice::(&bytes) + .map_err(|error| CliError::InvalidPayload(error.to_string()))?; + match route { + HookRoute::Codex => { + let outcome = codex::adapt(payload, &parts.headers); + self.sessions + .apply_authenticated_events(&parts.headers, outcome.events, &self.owner) + .await?; + if let Some(permission) = outcome.permission + && let Err(error) = self.authorize_permission(permission).await + { + return Ok(json!({ + "decision": "deny", + "reason": permission_denial_reason(error), + })); + } + Ok(outcome.response) + } + HookRoute::Claude => { + let outcome = claude_code::adapt(payload, &parts.headers); + self.sessions + .apply_authenticated_events(&parts.headers, outcome.events, &self.owner) + .await?; + if let Some(permission) = outcome.permission { + let result = self.authorize_permission(permission).await; + return Ok(match result { + Ok(()) => json!({ + "continue": true, + "hookSpecificOutput": { + "hookEventName": "PermissionRequest", + "decision": { "behavior": "allow" }, + }, + }), + Err(error) => json!({ + "continue": true, + "hookSpecificOutput": { + "hookEventName": "PermissionRequest", + "decision": { + "behavior": "deny", + "message": permission_denial_reason(error), + }, + }, + }), + }); + } + Ok(outcome.response) + } + HookRoute::Pi => { + let outcome = pi::adapt(payload, &parts.headers); + // A daemon worker is already isolated to one authenticated machine owner. Keep + // pi's response-transform behavior while using that isolation as its ownership + // boundary, just as the personal gateway does for its local extension. + let effects = self + .sessions + .apply_events(&parts.headers, outcome.events) + .await?; + Ok(pi::response_with_effects(outcome.response, &effects)) + } + } + } + + async fn authorize_permission( + &self, + permission: Result, + ) -> Result<(), CliError> { + match permission { + Ok(permission) => { + self.sessions + .authorize_tool_permission(&permission, &self.owner) + .await + } + Err(reason) => Err(CliError::InvalidPayload(reason)), + } + } + + pub(super) async fn proxy_provider( + &self, + upstream: PooledClient, + mut request: Request, + route: ProviderRoute, + ) -> Result, CliError> { + let Some(surface) = provider_surface(request.uri().path()) else { + return dispatch_unmanaged(upstream, request, route, &self.config).await; + }; + if !request_body_decode_required()? { + strip_worker_headers(request.headers_mut()); + strip_untrusted_dispatch_headers(request.headers_mut()); + let streaming_hint = request_streaming_hint(request.headers()); + let start = crate::gateway::daemon_gateway_start( + request.headers(), + request.uri().path(), + Value::Null, + streaming_hint, + ) + .ok_or_else(|| CliError::InvalidPayload("unsupported provider path".into()))?; + let prep = self + .sessions + .prepare_gateway_call(request.headers(), start) + .await?; + return self + .proxy_unbuffered(upstream, request, route, surface, prep, streaming_hint) + .await; + } + let prepared = PreparedProviderRequest::read(request, &self.config).await?; + let start = crate::gateway::daemon_gateway_start( + &prepared.headers, + &prepared.path, + prepared.request_json.clone(), + prepared.streaming, + ) + .ok_or_else(|| CliError::InvalidPayload("unsupported provider path".into()))?; + let prep = self + .sessions + .prepare_gateway_call(&prepared.headers, start) + .await?; + if prep.bypass_managed_pipeline { + GatewayCallCleanup::new(self.sessions.clone(), &prep) + .finish() + .await; + return dispatch_observed( + upstream, + prepared, + route, + None, + &self.config, + self.observation_capture_bytes, + ) + .await + .map(|(response, _)| response); + } + self.proxy_managed(upstream, prepared, route, surface, prep) + .await + } + + async fn proxy_unbuffered( + &self, + upstream: PooledClient, + request: Request, + route: ProviderRoute, + surface: ProviderSurface, + prep: GatewayCallPrep, + streaming_hint: bool, + ) -> Result, CliError> { + let cleanup = GatewayCallCleanup::new(self.sessions.clone(), &prep); + let GatewayCallPrep { + scope_stack, + session_id, + provider_name, + request: request_for_event, + parent, + attributes, + metadata, + model_name, + owner_subagent_id, + bypass_managed_pipeline, + session_finish: _, + } = prep; + if bypass_managed_pipeline { + cleanup.finish().await; + return dispatch_unmanaged(upstream, request, route, &self.config).await; + } + + let handle = TASK_SCOPE_STACK + .scope(scope_stack, async { + llm_call( + LlmCallParams::builder() + .name(&provider_name) + .request(&request_for_event) + .parent_opt(parent.as_ref()) + .attributes(attributes) + .metadata(metadata.clone()) + .model_name_opt(model_name) + .build(), + ) + }) + .await; + let handle = match handle { + Ok(handle) => handle, + Err(error) => { + cleanup.finish().await; + return Err(error.into()); + } + }; + let response = dispatch_unbuffered_observed( + upstream, + request, + route, + &self.config, + self.observation_capture_bytes, + ) + .await; + let (response, observation, response_streaming) = match response { + Ok(result) => result, + Err(error) => { + finish_llm_after_dispatch_failure(&handle, metadata, &error); + cleanup.finish().await; + return Err(error); + } + }; + let sessions = self.sessions.clone(); + tokio::spawn(async move { + let observed = observation + .finish(surface, response_streaming || streaming_hint) + .await; + let response_value = observed.value.clone().unwrap_or(Value::Null); + let mut end_metadata = merge_object(metadata, observed.metadata()); + insert_metadata( + &mut end_metadata, + "daemon_worker_request_capture", + json!("head_only"), + ); + if observed.failure.is_none() { + insert_metadata(&mut end_metadata, "otel.status_code", json!("OK")); + } else if let Some(failure) = observed.failure.as_ref() { + insert_metadata(&mut end_metadata, "otel.status_code", json!("ERROR")); + insert_metadata(&mut end_metadata, "otel.status_description", json!(failure)); + } + if let Err(error) = llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response_value) + .metadata(end_metadata) + .response_codec_opt(observed.value.as_ref().map(|_| response_codec(surface))) + .build(), + ) { + log::warn!( + target: "nemo_relay.daemon.worker", + event = "worker_llm_observation_end_failed", + error_kind = "runtime"; + "Daemon worker failed to close an observed LLM lifecycle: {error}" + ); + } + if let Some(value) = observed.value { + sessions + .record_gateway_response_hints(&session_id, owner_subagent_id, value) + .await; + } + cleanup.finish().await; + }); + Ok(response) + } + + async fn proxy_managed( + &self, + upstream: PooledClient, + prepared: PreparedProviderRequest, + route: ProviderRoute, + surface: ProviderSurface, + prep: GatewayCallPrep, + ) -> Result, CliError> { + let cleanup = GatewayCallCleanup::new(self.sessions.clone(), &prep); + let GatewayCallPrep { + scope_stack, + session_id, + provider_name, + request, + parent, + attributes, + metadata, + model_name, + owner_subagent_id, + bypass_managed_pipeline: _, + session_finish: _, + } = prep; + let codec = request_codec(surface); + let middleware = TASK_SCOPE_STACK + .scope(scope_stack.clone(), async { + llm_conditional_execution(&request).await?; + let mut outcome = llm_request_intercepts(&provider_name, request).await?; + if let Some(annotated) = outcome.annotated_request.as_ref() { + outcome.request = codec.encode(annotated, &outcome.request)?; + } + Ok::<_, FlowError>(outcome) + }) + .await; + let outcome = match middleware { + Ok(outcome) => outcome, + Err(error) => { + cleanup.finish().await; + return Err(error.into()); + } + }; + if stream_mode(&outcome.request) != prepared.streaming { + cleanup.finish().await; + return Err(CliError::Flow(FlowError::InvalidArgument( + STREAM_MODE_MUTATION_ERROR.into(), + ))); + } + let annotated_request = outcome.annotated_request.clone().map(Arc::new); + let request_for_event = outcome.request.clone(); + let handle = TASK_SCOPE_STACK + .scope(scope_stack, async { + llm_call( + LlmCallParams::builder() + .name(&provider_name) + .request(&request_for_event) + .parent_opt(parent.as_ref()) + .attributes(attributes) + .metadata(metadata.clone()) + .model_name_opt(model_name) + .annotated_request_opt(annotated_request) + .build(), + ) + }) + .await; + let handle = match handle { + Ok(handle) => handle, + Err(error) => { + cleanup.finish().await; + return Err(CliError::Flow(error)); + } + }; + for contribution in outcome.optimization_contributions { + let _ = handle.optimization_recorder.record(contribution); + } + if !outcome.pending_marks.is_empty() { + log::warn!( + target: "nemo_relay.daemon.worker", + event = "worker_request_marks_unsupported", + pending_mark_count = outcome.pending_marks.len(); + "Daemon raw-stream observation cannot attach request-interceptor marks to the LLM handle" + ); + } + + let response = dispatch_observed( + upstream, + prepared, + route, + Some(&outcome.request), + &self.config, + self.observation_capture_bytes, + ) + .await; + let (response, observation) = match response { + Ok(result) => result, + Err(error) => { + finish_llm_after_dispatch_failure(&handle, metadata, &error); + cleanup.finish().await; + return Err(error); + } + }; + let sessions = self.sessions.clone(); + tokio::spawn(async move { + let observed = observation + .finish(surface, prepared_streaming(&request_for_event)) + .await; + let response_value = observed.value.clone().unwrap_or(Value::Null); + let mut end_metadata = merge_object(metadata, observed.metadata()); + if observed.failure.is_none() { + insert_metadata(&mut end_metadata, "otel.status_code", json!("OK")); + } else if let Some(failure) = observed.failure.as_ref() { + insert_metadata(&mut end_metadata, "otel.status_code", json!("ERROR")); + insert_metadata(&mut end_metadata, "otel.status_description", json!(failure)); + } + if let Err(error) = llm_call_end( + LlmCallEndParams::builder() + .handle(&handle) + .response(response_value) + .metadata(end_metadata) + .response_codec_opt(observed.value.as_ref().map(|_| response_codec(surface))) + .build(), + ) { + log::warn!( + target: "nemo_relay.daemon.worker", + event = "worker_llm_observation_end_failed", + error_kind = "runtime"; + "Daemon worker failed to close an observed LLM lifecycle: {error}" + ); + } + if let Some(value) = observed.value { + sessions + .record_gateway_response_hints(&session_id, owner_subagent_id, value) + .await; + } + cleanup.finish().await; + }); + Ok(response) + } +} + +pub(super) fn requires_route_pass_through(error: &CliError) -> bool { + matches!( + error, + CliError::Flow(FlowError::InvalidArgument(message)) + if message == STREAM_MODE_MUTATION_ERROR + ) +} + +async fn dispatch_unmanaged( + upstream: PooledClient, + mut request: Request, + route: ProviderRoute, + config: &GatewayConfig, +) -> Result, CliError> { + let allow_environment_provider_auth = + crate::gateway::daemon_allows_environment_provider_auth(request.headers()); + strip_worker_headers(request.headers_mut()); + strip_untrusted_dispatch_headers(request.headers_mut()); + let path_and_query = request + .uri() + .path_and_query() + .map_or_else(|| request.uri().path().to_owned(), ToString::to_string); + let destination = + crate::gateway::daemon_provider_upstream_url(request.headers(), &path_and_query, config)? + .unwrap_or_else(|| route.upstream_url(config, &path_and_query)) + .parse::() + .map_err(|_| CliError::InvalidPayload("invalid provider destination".into()))?; + strip_internal_headers(request.headers_mut()); + if allow_environment_provider_auth { + inject_provider_auth(request.headers_mut(), route, config); + } + let strip = [ + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderName::from_static(WORKER_ROUTE_FAILURE_HEADER), + ]; + let request = prepare_forward_request(request, destination, &strip) + .map_err(|error| CliError::InvalidPayload(error.to_string()))? + .map(box_body); + let response = tokio::time::timeout(RESPONSE_HEAD_TIMEOUT, upstream.request(request)) + .await + .map_err(|_| CliError::Launch("provider response-head timeout".into()))? + .map_err(|error| CliError::Launch(error.to_string()))?; + let response = prepare_forward_response(response, &strip) + .map_err(|error| CliError::Launch(error.to_string()))?; + let (parts, body) = response.into_parts(); + Ok(Response::from_parts(parts, box_body(body))) +} + +async fn dispatch_unbuffered_observed( + upstream: PooledClient, + mut request: Request, + route: ProviderRoute, + config: &GatewayConfig, + capture_limit: usize, +) -> Result<(Response, ObservationReceiver, bool), CliError> { + let allow_environment_provider_auth = + crate::gateway::daemon_allows_environment_provider_auth(request.headers()); + strip_worker_headers(request.headers_mut()); + strip_untrusted_dispatch_headers(request.headers_mut()); + let path_and_query = request + .uri() + .path_and_query() + .map_or_else(|| request.uri().path().to_owned(), ToString::to_string); + let destination = + crate::gateway::daemon_provider_upstream_url(request.headers(), &path_and_query, config)? + .unwrap_or_else(|| route.upstream_url(config, &path_and_query)) + .parse::() + .map_err(|_| CliError::InvalidPayload("invalid provider destination".into()))?; + if let Some(aligned) = crate::gateway::daemon_provider_forward_headers( + request.headers(), + request.uri().path(), + config, + ) { + *request.headers_mut() = aligned; + } + strip_internal_headers(request.headers_mut()); + if allow_environment_provider_auth { + inject_provider_auth(request.headers_mut(), route, config); + } + let strip = [ + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderName::from_static(WORKER_ROUTE_FAILURE_HEADER), + ]; + let request = prepare_forward_request(request, destination, &strip) + .map_err(|error| CliError::InvalidPayload(error.to_string()))? + .map(box_body); + let response = tokio::time::timeout(RESPONSE_HEAD_TIMEOUT, upstream.request(request)) + .await + .map_err(|_| CliError::Launch("provider response-head timeout".into()))? + .map_err(|error| CliError::Launch(error.to_string()))?; + let response = prepare_forward_response(response, &strip) + .map_err(|error| CliError::Launch(error.to_string()))?; + let status = response.status(); + let streaming = response_streaming(response.headers()); + let (parts, body) = response.into_parts(); + let (body, observation) = observe_body(body, status, capture_limit); + Ok((Response::from_parts(parts, body), observation, streaming)) +} + +impl Drop for ManagedRuntime { + fn drop(&mut self) { + if let Some(activation) = self + .activation + .get_mut() + .unwrap_or_else(|error| error.into_inner()) + .take() + { + let _ = activation.clear(); + } + } +} + +struct PreparedProviderRequest { + method: Method, + version: http::Version, + headers: HeaderMap, + path: String, + path_and_query: String, + body: Bytes, + request_json: Value, + streaming: bool, +} + +impl PreparedProviderRequest { + async fn read(request: Request, config: &GatewayConfig) -> Result { + let (mut parts, body) = request.into_parts(); + strip_worker_headers(&mut parts.headers); + // Dispatch controls are created only by worker-local middleware. Never allow a value that + // arrived on the authenticated daemon hop to become an upstream override. Correlation + // headers remain available to session normalization. + strip_untrusted_dispatch_headers(&mut parts.headers); + let bytes = axum::body::to_bytes(body, config.max_passthrough_body_bytes) + .await + .map_err(body_read_error)?; + let request_json = serde_json::from_slice::(&bytes).unwrap_or(Value::Null); + let streaming = request_json + .get("stream") + .and_then(Value::as_bool) + .unwrap_or(false); + let path = parts.uri.path().to_owned(); + let path_and_query = parts + .uri + .path_and_query() + .map_or_else(|| path.clone(), ToString::to_string); + Ok(Self { + method: parts.method, + version: parts.version, + headers: parts.headers, + path, + path_and_query, + body: bytes, + request_json, + streaming, + }) + } +} + +async fn dispatch_observed( + upstream: PooledClient, + prepared: PreparedProviderRequest, + route: ProviderRoute, + effective: Option<&LlmRequest>, + config: &GatewayConfig, + capture_limit: usize, +) -> Result<(Response, ObservationReceiver), CliError> { + let allow_environment_provider_auth = + crate::gateway::daemon_allows_environment_provider_auth(&prepared.headers); + let destination = effective_destination(&prepared, route, effective, config)?; + let (mut headers, body, explicit_target) = effective_request(&prepared, effective)?; + if !explicit_target + && let Some(aligned) = + crate::gateway::daemon_provider_forward_headers(&headers, &prepared.path, config) + { + headers = aligned; + } + let mut request = Request::builder() + .method(prepared.method.clone()) + .version(prepared.version) + .uri(destination.clone()) + .body(Body::from(body))?; + *request.headers_mut() = headers; + if !explicit_target && allow_environment_provider_auth { + inject_provider_auth(request.headers_mut(), route, config); + } + let strip = [ + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderName::from_static(WORKER_ROUTE_FAILURE_HEADER), + ]; + let request = prepare_forward_request(request, destination, &strip) + .map_err(|error| CliError::InvalidPayload(error.to_string()))? + .map(box_body); + let response = tokio::time::timeout(RESPONSE_HEAD_TIMEOUT, upstream.request(request)) + .await + .map_err(|_| CliError::Launch("provider response-head timeout".into()))? + .map_err(|error| CliError::Launch(error.to_string()))?; + let response = prepare_forward_response(response, &strip) + .map_err(|error| CliError::Launch(error.to_string()))?; + let status = response.status(); + let (parts, body) = response.into_parts(); + let (body, observation) = observe_body(body, status, capture_limit); + Ok((Response::from_parts(parts, body), observation)) +} + +fn effective_destination( + prepared: &PreparedProviderRequest, + route: ProviderRoute, + effective: Option<&LlmRequest>, + config: &GatewayConfig, +) -> Result { + let override_url = effective.and_then(|request| { + request + .headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case(INTERNAL_DISPATCH_URL_HEADER)) + .and_then(|(_, value)| value.as_str()) + .map(str::trim) + .filter(|value| !value.is_empty()) + }); + let destination = match override_url { + Some(destination) => destination.to_owned(), + None => crate::gateway::daemon_provider_upstream_url( + &prepared.headers, + &prepared.path_and_query, + config, + )? + .unwrap_or_else(|| route.upstream_url(config, &prepared.path_and_query)), + }; + destination + .parse::() + .map_err(|_| CliError::InvalidPayload("invalid provider destination".into())) +} + +fn effective_request( + prepared: &PreparedProviderRequest, + effective: Option<&LlmRequest>, +) -> Result<(HeaderMap, Bytes, bool), CliError> { + let mut headers = prepared.headers.clone(); + strip_internal_headers(&mut headers); + let Some(effective) = effective else { + headers.remove(CONTENT_LENGTH); + return Ok((headers, prepared.body.clone(), false)); + }; + let explicit_target = has_explicit_target(effective); + if explicit_target { + crate::provider_auth::remove_provider_credentials(&mut headers); + headers.remove(http::header::COOKIE); + } + let baseline = crate::gateway::daemon_observable_headers(&prepared.headers); + for name in baseline.keys() { + if !effective + .headers + .keys() + .any(|effective_name| effective_name.eq_ignore_ascii_case(name)) + && let Ok(name) = HeaderName::from_bytes(name.as_bytes()) + { + headers.remove(name); + } + } + for (name, value) in &effective.headers { + if name + .to_ascii_lowercase() + .starts_with(INTERNAL_HEADER_PREFIX) + { + continue; + } + let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else { + continue; + }; + if baseline.get(name.as_str()) == Some(value) { + continue; + } + let Some(value) = json_header_value(value) else { + continue; + }; + headers.insert(name, value); + } + strip_internal_headers(&mut headers); + headers.remove(CONTENT_LENGTH); + if effective.content == prepared.request_json || effective.content.is_null() { + return Ok((headers, prepared.body.clone(), explicit_target)); + } + let body = serde_json::to_vec(&effective.content) + .map(Bytes::from) + .map_err(|error| CliError::InvalidPayload(error.to_string()))?; + headers.remove(CONTENT_ENCODING); + Ok((headers, body, explicit_target)) +} + +fn has_explicit_target(request: &LlmRequest) -> bool { + request.headers.iter().any(|(name, value)| { + (name.eq_ignore_ascii_case(INTERNAL_DISPATCH_URL_HEADER) + || name.eq_ignore_ascii_case(INTERNAL_DISPATCH_ROUTE_HEADER)) + && value + .as_str() + .map(str::trim) + .is_some_and(|value| !value.is_empty()) + }) +} + +fn strip_worker_headers(headers: &mut HeaderMap) { + headers.remove(CLIENT_TOKEN_HEADER); + headers.remove(WORKER_TOKEN_HEADER); +} + +fn strip_internal_headers(headers: &mut HeaderMap) { + let names = headers + .keys() + .filter(|name| name.as_str().starts_with(INTERNAL_HEADER_PREFIX)) + .cloned() + .collect::>(); + for name in names { + headers.remove(name); + } +} + +fn strip_untrusted_dispatch_headers(headers: &mut HeaderMap) { + headers.remove(INTERNAL_DISPATCH_URL_HEADER); + headers.remove(INTERNAL_DISPATCH_ROUTE_HEADER); + headers.remove(INTERNAL_DISPATCH_BACKEND_HEADER); + headers.remove(INTERNAL_RETRY_AWARE_HEADER); +} + +fn body_read_error(error: axum::Error) -> CliError { + if error.source().is_some_and(|source| { + source.is::() + || source + .source() + .is_some_and(|source| source.is::()) + }) { + CliError::PayloadTooLarge(error.to_string()) + } else { + CliError::InvalidPayload(error.to_string()) + } +} + +fn inject_provider_auth(headers: &mut HeaderMap, route: ProviderRoute, config: &GatewayConfig) { + if crate::provider_auth::has_provider_credential(headers) { + return; + } + if let Some(configured) = match route { + ProviderRoute::OpenAi => config.openai_auth_header.as_deref(), + ProviderRoute::Anthropic => config.anthropic_auth_header.as_deref(), + } + .and_then(|value| HeaderValue::from_str(value).ok()) + { + headers.insert(AUTHORIZATION, configured); + return; + } + let (name, value) = match route { + ProviderRoute::OpenAi => { + let Some(key) = environment_value("OPENAI_API_KEY") else { + return; + }; + (AUTHORIZATION, format!("Bearer {key}")) + } + ProviderRoute::Anthropic => { + let Some(key) = environment_value("ANTHROPIC_API_KEY") else { + return; + }; + (HeaderName::from_static("x-api-key"), key) + } + }; + if let Ok(value) = HeaderValue::from_str(&value) { + headers.insert(name, value); + } +} + +fn environment_value(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn observation_capture_limit_from_environment() -> Result { + let Some(raw) = std::env::var_os(OBSERVATION_CAPTURE_BYTES_ENV) else { + return Ok(DEFAULT_OBSERVATION_CAPTURE_BYTES); + }; + let raw = raw.to_str().ok_or_else(|| { + CliError::Config(format!( + "{OBSERVATION_CAPTURE_BYTES_ENV} must be a positive integer" + )) + })?; + let value = raw.trim().parse::().map_err(|_| { + CliError::Config(format!( + "{OBSERVATION_CAPTURE_BYTES_ENV} must be a positive integer" + )) + })?; + if value == 0 { + return Err(CliError::Config(format!( + "{OBSERVATION_CAPTURE_BYTES_ENV} must be a positive integer" + ))); + } + Ok(value) +} + +fn json_header_value(value: &Value) -> Option { + let value = value + .as_str() + .map(ToOwned::to_owned) + .or_else(|| serde_json::to_string(value).ok())?; + HeaderValue::from_str(&value).ok() +} + +fn provider_surface(path: &str) -> Option { + match path { + "/responses" | "/v1/responses" | "/backend-api/codex/responses" => { + Some(ProviderSurface::OpenAIResponses) + } + "/chat/completions" | "/v1/chat/completions" => Some(ProviderSurface::OpenAIChat), + "/v1/messages" => Some(ProviderSurface::AnthropicMessages), + _ => None, + } +} + +fn stream_mode(request: &LlmRequest) -> bool { + request + .content + .get("stream") + .and_then(Value::as_bool) + .unwrap_or(false) +} + +fn request_streaming_hint(headers: &HeaderMap) -> bool { + headers.get_all(ACCEPT).iter().any(|value| { + value.to_str().ok().is_some_and(|value| { + value.split(',').any(|media_type| { + media_type + .split(';') + .next() + .is_some_and(|value| value.trim().eq_ignore_ascii_case("text/event-stream")) + }) + }) + }) +} + +fn response_streaming(headers: &HeaderMap) -> bool { + headers.get_all(CONTENT_TYPE).iter().any(|value| { + value.to_str().ok().is_some_and(|value| { + value + .split(';') + .next() + .is_some_and(|value| value.trim().eq_ignore_ascii_case("text/event-stream")) + }) + }) +} + +fn prepared_streaming(request: &LlmRequest) -> bool { + stream_mode(request) +} + +fn request_body_decode_required() -> Result { + let kinds = BTreeSet::from([ + RuntimeRegistrationKind::LlmSanitizeRequestGuardrail, + RuntimeRegistrationKind::LlmConditionalExecutionGuardrail, + RuntimeRegistrationKind::LlmRequestIntercept, + ]); + let registrations = list_runtime_registrations(Some(&kinds)).map_err(CliError::from)?; + Ok(registrations.iter().any(registration_reads_request_body)) +} + +fn registration_reads_request_body(registration: &RuntimeRegistrationIdentity) -> bool { + matches!( + registration.kind, + RuntimeRegistrationKind::LlmSanitizeRequestGuardrail + | RuntimeRegistrationKind::LlmConditionalExecutionGuardrail + | RuntimeRegistrationKind::LlmRequestIntercept + ) +} + +fn reject_incompatible_execution_middleware() -> Result<(), CliError> { + // Execution intercepts own the provider callback and may replace, suppress, retry, or mutate + // its result. The raw worker transport cannot safely invoke that contract while also returning + // the provider's response head and frames unchanged. Request intercepts and conditional + // execution guardrails remain supported above the transport boundary. + let kinds = BTreeSet::from([ + RuntimeRegistrationKind::LlmExecutionIntercept, + RuntimeRegistrationKind::LlmStreamExecutionIntercept, + ]); + let registrations = list_runtime_registrations(Some(&kinds)).map_err(CliError::from)?; + let incompatible = incompatible_registration_names(®istrations); + if incompatible.is_empty() { + return Ok(()); + } + Err(CliError::Config(format!( + "daemon worker raw delivery is incompatible with LLM execution middleware: {}", + incompatible.join(", ") + ))) +} + +fn incompatible_registration_names(registrations: &[RuntimeRegistrationIdentity]) -> Vec { + registrations + .iter() + .filter(|registration| { + matches!( + registration.kind, + RuntimeRegistrationKind::LlmExecutionIntercept + | RuntimeRegistrationKind::LlmStreamExecutionIntercept + ) + }) + .map(|registration| registration.effective_name.clone()) + .collect() +} + +fn finish_llm_after_dispatch_failure( + handle: &nemo_relay::api::llm::LlmHandle, + metadata: Value, + error: &CliError, +) { + let mut metadata = metadata; + insert_metadata(&mut metadata, "otel.status_code", json!("ERROR")); + insert_metadata( + &mut metadata, + "otel.status_description", + json!(error.to_string()), + ); + let _ = llm_call_end( + LlmCallEndParams::builder() + .handle(handle) + .response(Value::Null) + .metadata(metadata) + .build(), + ); +} + +fn permission_denial_reason(error: CliError) -> String { + error + .guardrail_rejection_reason() + .map(ToOwned::to_owned) + .unwrap_or_else(|| error.to_string()) +} + +fn merge_object(mut base: Value, extra: Value) -> Value { + if !base.is_object() { + base = json!({}); + } + if let (Some(base), Some(extra)) = (base.as_object_mut(), extra.as_object()) { + base.extend(extra.clone()); + } + base +} + +fn insert_metadata(metadata: &mut Value, name: &str, value: Value) { + if !metadata.is_object() { + *metadata = json!({}); + } + if let Some(metadata) = metadata.as_object_mut() { + metadata.insert(name.to_owned(), value); + } +} + +struct ObservationSignal { + terminal: AtomicU8, + truncated: AtomicBool, + notify: Notify, +} + +impl ObservationSignal { + fn new() -> Self { + Self { + terminal: AtomicU8::new(OBSERVATION_ACTIVE), + truncated: AtomicBool::new(false), + notify: Notify::new(), + } + } + + fn truncate(&self) { + self.truncated.store(true, Ordering::Release); + } + + fn finish(&self, terminal: u8) { + if self + .terminal + .compare_exchange( + OBSERVATION_ACTIVE, + terminal, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + self.notify.notify_waiters(); + } + } + + async fn wait(&self) -> u8 { + loop { + let notified = self.notify.notified(); + let terminal = self.terminal.load(Ordering::Acquire); + if terminal != OBSERVATION_ACTIVE { + return terminal; + } + notified.await; + } + } +} + +struct ObservedBody { + body: B, + sender: Option>, + signal: Arc, + scheduled_bytes: usize, + capture_limit: usize, +} + +impl HttpBody for ObservedBody +where + B: HttpBody + Unpin, + B::Error: Into, +{ + type Data = Bytes; + type Error = BoxError; + + fn poll_frame( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + match Pin::new(&mut self.body).poll_frame(context) { + Poll::Ready(Some(Ok(frame))) => { + if let Some(data) = frame.data_ref() + && let Some(next) = self.scheduled_bytes.checked_add(data.len()) + { + if next <= self.capture_limit { + let sent = self + .sender + .as_ref() + .is_some_and(|sender| sender.try_send(data.clone()).is_ok()); + if sent { + self.scheduled_bytes = next; + } else if self.sender.take().is_some() { + self.signal.truncate(); + } + } else if self.sender.take().is_some() { + self.signal.truncate(); + } + } + Poll::Ready(Some(Ok(frame))) + } + Poll::Ready(Some(Err(error))) => { + self.sender.take(); + self.signal.finish(OBSERVATION_BODY_ERROR); + Poll::Ready(Some(Err(error.into()))) + } + Poll::Ready(None) => { + self.sender.take(); + self.signal.finish(OBSERVATION_COMPLETE); + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } + + fn is_end_stream(&self) -> bool { + self.body.is_end_stream() + } + + fn size_hint(&self) -> SizeHint { + self.body.size_hint() + } +} + +impl Drop for ObservedBody { + fn drop(&mut self) { + self.sender.take(); + self.signal.finish(OBSERVATION_CANCELLED); + } +} + +struct ObservationReceiver { + receiver: mpsc::Receiver, + signal: Arc, + status: StatusCode, +} + +struct ObservedResponse { + value: Option, + truncated: bool, + terminal: u8, + status: StatusCode, + failure: Option, +} + +impl ObservedResponse { + fn metadata(&self) -> Value { + json!({ + "daemon_worker_observation": { + "truncated": self.truncated, + "http_status": self.status.as_u16(), + "terminal": match self.terminal { + OBSERVATION_COMPLETE => "complete", + OBSERVATION_BODY_ERROR => "body_error", + OBSERVATION_CANCELLED => "cancelled", + _ => "unknown", + }, + }, + }) + } +} + +impl ObservationReceiver { + async fn finish(self, surface: ProviderSurface, streaming: bool) -> ObservedResponse { + let status = self.status; + match tokio::time::timeout( + OBSERVATION_COMPLETION_TIMEOUT, + self.finish_inner(surface, streaming), + ) + .await + { + Ok(observed) => observed, + Err(_) => ObservedResponse { + value: None, + truncated: true, + terminal: OBSERVATION_ACTIVE, + status, + failure: Some("provider response observation timed out".to_owned()), + }, + } + } + + async fn finish_inner(mut self, surface: ProviderSurface, streaming: bool) -> ObservedResponse { + let value = if streaming { + self.finish_stream(surface).await + } else { + self.finish_buffered().await + }; + let terminal = self.signal.wait().await; + let truncated = self.signal.truncated.load(Ordering::Acquire); + let mut failure = match terminal { + OBSERVATION_COMPLETE => None, + OBSERVATION_BODY_ERROR => Some("provider response body failed".to_owned()), + OBSERVATION_CANCELLED => Some("downstream cancelled provider response".to_owned()), + _ => Some("provider response observation ended unexpectedly".to_owned()), + }; + if truncated { + failure = Some("provider response observation was truncated".to_owned()); + } else if !self.status.is_success() { + failure = Some(format!("provider returned HTTP {}", self.status.as_u16())); + } + ObservedResponse { + value: if truncated { None } else { value }, + truncated, + terminal, + status: self.status, + failure, + } + } + + async fn finish_buffered(&mut self) -> Option { + let mut bytes = Vec::new(); + while let Some(chunk) = self.receiver.recv().await { + bytes.extend_from_slice(&chunk); + } + if self.signal.truncated.load(Ordering::Acquire) { + return None; + } + serde_json::from_slice(&bytes).ok() + } + + async fn finish_stream(&mut self, surface: ProviderSurface) -> Option { + let mut decoder = SseEventDecoder::new(); + let codec = streaming_codec(surface); + let mut collector = codec.collector(); + let finalizer = codec.finalizer(); + let mut valid = true; + while let Some(chunk) = self.receiver.recv().await { + if !valid || self.signal.truncated.load(Ordering::Acquire) { + continue; + } + for event in decoder.push_bytes_results(&chunk) { + match event { + Ok(event) => { + if collector(event.data).is_ok() { + continue; + } + valid = false; + self.signal.truncate(); + break; + } + Err(_) => { + valid = false; + self.signal.truncate(); + break; + } + } + } + } + if valid + && !self.signal.truncated.load(Ordering::Acquire) + && let Ok(Some(event)) = decoder.finish() + && collector(event.data).is_err() + { + self.signal.truncate(); + valid = false; + } + (valid && !self.signal.truncated.load(Ordering::Acquire)).then(finalizer) + } +} + +fn observe_body( + body: B, + status: StatusCode, + capture_limit: usize, +) -> (RelayBody, ObservationReceiver) +where + B: HttpBody + Send + Unpin + 'static, + B::Error: Into, +{ + let signal = Arc::new(ObservationSignal::new()); + let (sender, receiver) = mpsc::channel(OBSERVATION_QUEUE_FRAMES); + let sender = if body.is_end_stream() { + signal.finish(OBSERVATION_COMPLETE); + None + } else { + Some(sender) + }; + let body = box_body(ObservedBody { + body, + sender, + signal: Arc::clone(&signal), + scheduled_bytes: 0, + capture_limit, + }); + ( + body, + ObservationReceiver { + receiver, + signal, + status, + }, + ) +} + +fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(|error| error.into_inner()) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/worker_managed_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/worker/mod.rs b/crates/cli/src/daemon/worker/mod.rs new file mode 100644 index 000000000..47b9f19e8 --- /dev/null +++ b/crates/cli/src/daemon/worker/mod.rs @@ -0,0 +1,173 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Daemon-attached per-user worker runtime. + +mod control; +mod managed; +mod runtime; + +#[cfg(test)] +pub(crate) use runtime::{TestWorkerHandle, test_router, test_router_with_control_tokens}; + +#[cfg(test)] +pub(crate) async fn test_recover_control_session( + daemon_origin: &str, + identity: &super::common::identity::MachineIdentity, + worker_id: &str, + endpoint: &str, + generation_grant: super::common::control::WorkerGenerationGrant, +) -> Result< + ( + String, + String, + super::common::control::WorkerGenerationGrant, + ), + CliError, +> { + let registration = control::recover( + daemon_origin, + identity, + worker_id, + endpoint, + None, + generation_grant, + ) + .await?; + Ok(control::test_registration_values(registration)) +} + +use std::net::{Ipv4Addr, SocketAddr}; + +use tokio::io::AsyncReadExt; +use tokio::net::TcpListener; + +use super::common::address::{explicit_daemon_origin, worker_advertised_address, worker_socket}; +use super::common::control::{MAX_CONTROL_BODY_BYTES, WorkerBootstrap}; +use super::common::state::load_or_create_machine_identity; +use super::common::worker_tls::WorkerTlsIdentity; +use crate::error::CliError; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct Options { + pub(crate) daemon_address: String, + pub(crate) bind: Ipv4Addr, + pub(crate) port: Option, + pub(crate) advertise_address: Option, +} + +pub(crate) async fn run(options: Options) -> Result<(), CliError> { + let daemon_origin = explicit_daemon_origin(&options.daemon_address)?; + let bootstrap = read_bootstrap().await?; + let bind = validate_bootstrap(&options, &bootstrap)?; + let identity = load_or_create_machine_identity()?; + let listener = TcpListener::bind(bind).await.map_err(|error| { + CliError::Launch(format!( + "failed to bind daemon worker listener {bind}: {error}" + )) + })?; + let local = listener.local_addr()?; + let advertised = worker_advertised_address(local, options.advertise_address.as_deref())?; + let worker_tls = if local.ip().is_unspecified() { + let advertised_host = options.advertise_address.as_deref().ok_or_else(|| { + CliError::Unauthorized( + "an unspecified worker bind requires an advertised address".into(), + ) + })?; + Some(WorkerTlsIdentity::generate(advertised_host)?) + } else { + None + }; + let endpoint = format!( + "{}://{advertised}", + if worker_tls.is_some() { + "https" + } else { + "http" + } + ); + let worker_tls_root = worker_tls + .as_ref() + .map(|identity| identity.root_certificate().to_owned()); + let tls_config = worker_tls.as_ref().map(WorkerTlsIdentity::server_config); + let worker_id = uuid::Uuid::now_v7().to_string(); + + let managed = crate::configuration::resolve_managed_worker_config()?; + let dynamic_plugins = crate::plugins::lifecycle::active_dynamic_plugin_components( + Some(&managed.plugin_config_path), + &managed.resolved, + )?; + let registration = control::register( + &daemon_origin, + &identity, + &worker_id, + &endpoint, + bootstrap, + worker_tls_root.clone(), + ) + .await?; + + runtime::serve( + listener, + runtime::RuntimeOptions { + daemon_origin, + identity, + worker_id, + endpoint, + worker_tls_root, + tls_config, + config: managed.resolved.gateway, + dynamic_plugins, + registration, + }, + ) + .await +} + +async fn read_bootstrap() -> Result { + let limit = u64::try_from(MAX_CONTROL_BODY_BYTES) + .unwrap_or(u64::MAX) + .saturating_add(1); + let mut encoded = Vec::new(); + tokio::io::stdin() + .take(limit) + .read_to_end(&mut encoded) + .await + .map_err(|error| { + CliError::Launch(format!( + "failed to read daemon worker activation channel: {error}" + )) + })?; + if encoded.len() > MAX_CONTROL_BODY_BYTES { + return Err(CliError::Launch(format!( + "daemon worker activation exceeded {MAX_CONTROL_BODY_BYTES} bytes" + ))); + } + if encoded.is_empty() { + return Err(CliError::Unauthorized( + "daemon worker requires a protected activation grant on standard input".into(), + )); + } + serde_json::from_slice(&encoded) + .map_err(|_| CliError::Unauthorized("daemon worker activation grant was invalid".into())) +} + +fn validate_bootstrap( + options: &Options, + bootstrap: &WorkerBootstrap, +) -> Result { + let requested_port = options.port.unwrap_or(0); + if options.bind != bootstrap.bind_ip + || requested_port != bootstrap.port + || options.advertise_address != bootstrap.advertise_address + { + return Err(CliError::Unauthorized( + "daemon worker network options do not match the activation grant".into(), + )); + } + worker_socket(options.bind, options.port) +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/worker_tests.rs"] +mod tests; diff --git a/crates/cli/src/daemon/worker/runtime.rs b/crates/cli/src/daemon/worker/runtime.rs new file mode 100644 index 000000000..212660298 --- /dev/null +++ b/crates/cli/src/daemon/worker/runtime.rs @@ -0,0 +1,889 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Authenticated worker listener and lossless provider data plane. + +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard}; +use std::time::Duration; + +use axum::body::Body; +use axum::extract::{DefaultBodyLimit, State}; +use axum::http::header::{AUTHORIZATION, CONTENT_TYPE}; +use axum::http::{HeaderMap, HeaderName, HeaderValue, Request, Response, StatusCode, Uri}; +use axum::middleware::{self, Next}; +use axum::response::IntoResponse; +use axum::routing::{get, post}; +use axum::serve::ListenerExt; +use axum::{Json, Router}; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use hyper_util::server::conn::auto::Builder as ConnectionBuilder; +use hyper_util::service::TowerToHyperService; +use serde_json::json; +use tokio::net::TcpListener; +use tokio::sync::{Notify, Semaphore}; + +use super::super::common::control::{ + CLIENT_TOKEN_HEADER, DRAIN_LIFETIME_MS, MAX_CONTROL_BODY_BYTES, RECOVERY_LIFETIME_MS, + SessionRequest, WORKER_DRAIN_PATH, WORKER_PROBE_PATH, WORKER_ROUTE_FAILURE_HEADER, + WORKER_TOKEN_HEADER, WorkerDrainRequest, now_unix_ms, +}; +use super::super::common::identity::{MachineIdentity, TokenDigest}; +use super::super::common::routes::{ProviderRoute, PublicRoute}; +use super::super::common::transport::{ + PooledClient, RelayBody, box_body, hold_body, pooled_client, prepare_forward_request, + prepare_forward_response, +}; +use super::control::{self, Registration}; +use crate::configuration::GatewayConfig; +use crate::error::CliError; +use crate::plugins::lifecycle::ActiveDynamicPluginComponent; + +use super::managed::ManagedRuntime; + +const RESPONSE_HEAD_TIMEOUT: Duration = Duration::from_secs(60); +const CONTROL_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(5); +const INITIAL_READY_TIMEOUT: Duration = Duration::from_secs(15); +const RECOVERY_RETRY_INTERVAL: Duration = Duration::from_secs(1); +const MAX_CONCURRENT_TLS_HANDSHAKES: usize = 256; +const TLS_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); + +pub(super) struct RuntimeOptions { + pub(super) daemon_origin: String, + pub(super) identity: MachineIdentity, + pub(super) worker_id: String, + pub(super) endpoint: String, + pub(super) worker_tls_root: Option, + pub(super) tls_config: Option>, + pub(super) config: GatewayConfig, + pub(super) dynamic_plugins: Vec, + pub(super) registration: Registration, +} + +struct AuthTokens { + data: TokenDigest, + pending_data: Option, + readiness_data: Option, + control: TokenDigest, + last_control_sequence: u64, + last_control_request_id: String, +} + +struct WorkerState { + worker_id: String, + config: GatewayConfig, + upstream: PooledClient, + managed: Option>, + auth: RwLock, + accepting: AtomicBool, + draining: AtomicBool, + exiting: AtomicBool, + in_flight: AtomicUsize, + drain_deadline: RwLock>, + lifecycle: Notify, +} + +impl WorkerState { + fn new( + worker_id: String, + config: GatewayConfig, + managed: Option>, + registration: &Registration, + ) -> Result { + let data_token = registration.data_token_digest(); + Ok(Self { + worker_id, + config, + upstream: pooled_client().map_err(|error| CliError::Launch(error.to_string()))?, + managed, + auth: RwLock::new(AuthTokens { + data: data_token, + pending_data: None, + readiness_data: Some(data_token), + control: registration.session_token_digest(), + last_control_sequence: 0, + last_control_request_id: String::new(), + }), + accepting: AtomicBool::new(false), + draining: AtomicBool::new(false), + exiting: AtomicBool::new(false), + in_flight: AtomicUsize::new(0), + drain_deadline: RwLock::new(None), + lifecycle: Notify::new(), + }) + } + + fn admit(self: &Arc) -> Option { + if !self.accepting.load(Ordering::Acquire) { + return None; + } + self.in_flight.fetch_add(1, Ordering::AcqRel); + if !self.accepting.load(Ordering::Acquire) { + self.release_in_flight(); + return None; + } + Some(InFlight { + state: Arc::clone(self), + }) + } + + fn release_in_flight(&self) { + let previous = self.in_flight.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "worker in-flight count underflowed"); + if previous == 1 && self.draining.load(Ordering::Acquire) { + self.lifecycle.notify_waiters(); + } + } + + fn authenticate_data(&self, headers: &HeaderMap) -> bool { + let mut values = headers.get_all(WORKER_TOKEN_HEADER).iter(); + let Some(value) = values.next() else { + return false; + }; + if values.next().is_some() { + return false; + } + let actual = TokenDigest::from_token(value.as_bytes()); + let auth = read_lock(&self.auth); + auth.data.matches(&actual) + || auth + .pending_data + .as_ref() + .is_some_and(|pending| pending.matches(&actual)) + } + + fn activate_pending_readiness(&self, headers: &HeaderMap) { + let mut values = headers.get_all(WORKER_TOKEN_HEADER).iter(); + let Some(value) = values.next() else { + return; + }; + if values.next().is_some() { + return; + } + let actual = TokenDigest::from_token(value.as_bytes()); + let activated = { + let mut auth = write_lock(&self.auth); + if auth + .readiness_data + .as_ref() + .is_some_and(|expected| expected.matches(&actual)) + { + auth.readiness_data = None; + true + } else { + false + } + }; + if activated + && !self.draining.load(Ordering::Acquire) + && !self.exiting.load(Ordering::Acquire) + { + self.accepting.store(true, Ordering::Release); + } + } + + fn authenticate_control(&self, request: &SessionRequest) -> bool { + if request.payload.worker_id != self.worker_id + || request.session_id != self.worker_id + || request.request_id.is_empty() + || request.request_id.len() > 128 + || !request.validate_payload_hash() + { + return false; + } + let actual = TokenDigest::from_token(request.session_token.expose().as_bytes()); + let mut auth = write_lock(&self.auth); + if !auth.control.matches(&actual) { + return false; + } + if request.sequence == auth.last_control_sequence + && request.request_id == auth.last_control_request_id + { + return true; + } + if request.sequence != auth.last_control_sequence.saturating_add(1) { + return false; + } + auth.last_control_sequence = request.sequence; + auth.last_control_request_id = request.request_id.clone(); + true + } + + fn begin_drain(&self, requested_timeout_ms: u64) { + let timeout = Duration::from_millis(requested_timeout_ms.min(DRAIN_LIFETIME_MS)); + *write_lock(&self.drain_deadline) = Some(tokio::time::Instant::now() + timeout); + self.draining.store(true, Ordering::Release); + self.accepting.store(false, Ordering::Release); + self.lifecycle.notify_waiters(); + } + + fn control_lost(&self) { + self.accepting.store(false, Ordering::Release); + } + + fn control_restored(&self, registration: &Registration) { + { + let mut auth = write_lock(&self.auth); + auth.data = registration.data_token_digest(); + auth.pending_data = None; + auth.readiness_data = None; + auth.control = registration.session_token_digest(); + auth.last_control_sequence = 0; + auth.last_control_request_id.clear(); + } + if !self.draining.load(Ordering::Acquire) && !self.exiting.load(Ordering::Acquire) { + self.accepting.store(true, Ordering::Release); + } + } + + fn stage_recovery_data_token(&self, registration: &Registration) { + let token = registration.data_token_digest(); + let mut auth = write_lock(&self.auth); + auth.pending_data = Some(token); + auth.readiness_data = Some(token); + } + + fn discard_recovery_data_token(&self) { + let mut auth = write_lock(&self.auth); + auth.pending_data = None; + auth.readiness_data = None; + self.accepting.store(false, Ordering::Release); + } + + fn request_exit(&self) { + self.exiting.store(true, Ordering::Release); + self.accepting.store(false, Ordering::Release); + self.lifecycle.notify_waiters(); + } + + async fn wait_until_stopped(&self) { + loop { + let notified = self.lifecycle.notified(); + if self.exiting.load(Ordering::Acquire) { + return; + } + if self.draining.load(Ordering::Acquire) { + if self.in_flight.load(Ordering::Acquire) == 0 { + return; + } + let deadline = read_lock(&self.drain_deadline) + .as_ref() + .copied() + .unwrap_or_else(tokio::time::Instant::now); + if deadline <= tokio::time::Instant::now() { + return; + } + tokio::select! { + _ = notified => {} + _ = tokio::time::sleep_until(deadline) => return, + } + } else { + notified.await; + } + } + } +} + +struct InFlight { + state: Arc, +} + +impl Drop for InFlight { + fn drop(&mut self) { + self.state.release_in_flight(); + } +} + +pub(super) async fn serve(listener: TcpListener, options: RuntimeOptions) -> Result<(), CliError> { + let RuntimeOptions { + daemon_origin, + identity, + worker_id, + endpoint, + worker_tls_root, + tls_config, + config, + dynamic_plugins, + mut registration, + } = options; + let managed = Arc::new( + ManagedRuntime::initialize( + config.clone(), + dynamic_plugins, + identity.fingerprint().to_string(), + ) + .await?, + ); + let state = Arc::new(WorkerState::new( + worker_id.clone(), + config, + Some(managed), + ®istration, + )?); + let app = router(Arc::clone(&state)); + let server = async move { + match tls_config { + Some(config) => serve_tls(listener, app, config).await, + None => axum::serve( + listener.tap_io(|stream| { + let _ = stream.set_nodelay(true); + }), + app, + ) + .await + .map_err(CliError::Io), + } + }; + tokio::pin!(server); + let readiness = tokio::time::timeout( + INITIAL_READY_TIMEOUT, + registration.ready(&daemon_origin, &worker_id), + ); + let readiness = tokio::select! { + result = &mut server => { + return result.and_then(|()| Err(CliError::Launch("worker listener stopped before readiness".into()))); + } + result = readiness => result, + }; + match readiness { + Ok(Ok(())) => state.control_restored(®istration), + Ok(Err(error)) => return Err(error), + Err(_) => { + return Err(CliError::Launch( + "daemon worker readiness acknowledgement timed out".into(), + )); + } + } + log::info!( + target: "nemo_relay.daemon.worker", + event = "worker_ready", + worker_id = worker_id.as_str(), + endpoint = endpoint.as_str(); + "Daemon worker is ready" + ); + let heartbeat = tokio::spawn(monitor_control( + Arc::clone(&state), + daemon_origin, + identity, + worker_id.clone(), + endpoint, + worker_tls_root, + registration, + )); + let signal_state = Arc::clone(&state); + let signal = tokio::spawn(async move { + shutdown_signal().await; + signal_state.request_exit(); + }); + let result = tokio::select! { + result = &mut server => result, + _ = state.wait_until_stopped() => Ok(()), + }; + state.request_exit(); + heartbeat.abort(); + signal.abort(); + if let Some(managed) = state.managed.as_ref() { + managed.close().await?; + } + result +} + +async fn serve_tls( + listener: TcpListener, + app: Router, + config: Arc, +) -> Result<(), CliError> { + let acceptor = tokio_rustls::TlsAcceptor::from(config); + let handshake_permits = Arc::new(Semaphore::new(MAX_CONCURRENT_TLS_HANDSHAKES)); + let mut connections = tokio::task::JoinSet::new(); + loop { + let (stream, _) = match listener.accept().await { + Ok(accepted) => accepted, + Err(error) => { + log::warn!( + target: "nemo_relay.daemon.worker", + event = "tls_listener_accept_failed"; + "TLS listener accept failed: {error}" + ); + tokio::time::sleep(Duration::from_millis(100)).await; + continue; + } + }; + let _ = stream.set_nodelay(true); + let Ok(handshake_permit) = Arc::clone(&handshake_permits).try_acquire_owned() else { + continue; + }; + let acceptor = acceptor.clone(); + let service = app.clone(); + connections.spawn(async move { + let Ok(Ok(stream)) = + tokio::time::timeout(TLS_HANDSHAKE_TIMEOUT, acceptor.accept(stream)).await + else { + return; + }; + drop(handshake_permit); + let builder = ConnectionBuilder::new(TokioExecutor::new()); + let _ = builder + .serve_connection_with_upgrades( + TokioIo::new(stream), + TowerToHyperService::new(service), + ) + .await; + }); + while connections.try_join_next().is_some() {} + } +} + +fn router(state: Arc) -> Router { + let control = Router::new() + .route(WORKER_DRAIN_PATH, post(drain)) + .route(WORKER_PROBE_PATH, get(readiness_probe)) + .layer(DefaultBodyLimit::max(MAX_CONTROL_BODY_BYTES)); + Router::new() + .merge(control) + .fallback(proxy) + .layer(middleware::from_fn_with_state( + Arc::clone(&state), + authenticate_daemon_request, + )) + .with_state(state) +} + +#[cfg(test)] +pub(crate) struct TestWorkerHandle { + state: Arc, +} + +#[cfg(test)] +impl TestWorkerHandle { + pub(crate) fn control_lost(&self) { + self.state.control_lost(); + } + + pub(crate) fn begin_drain(&self, deadline_unix_ms: u64) { + self.state + .begin_drain(deadline_unix_ms.saturating_sub(now_unix_ms())); + } + + pub(crate) fn in_flight(&self) -> usize { + self.state.in_flight.load(Ordering::Acquire) + } + + pub(crate) fn is_draining(&self) -> bool { + self.state.draining.load(Ordering::Acquire) + } + + pub(crate) fn stage_recovery_tokens(&self, data_token: &[u8], control_token: &[u8]) { + let mut auth = write_lock(&self.state.auth); + let data = TokenDigest::from_token(data_token); + auth.pending_data = Some(data); + auth.readiness_data = Some(data); + auth.control = TokenDigest::from_token(control_token); + auth.last_control_sequence = 0; + auth.last_control_request_id.clear(); + } +} + +/// Constructs the real authenticated worker router with an injected process-wide pool. This is a +/// narrow test seam for multi-hop network tests; request admission, authentication, routing, and +/// provider forwarding remain the production handlers above and below it. +#[cfg(test)] +pub(crate) fn test_router( + config: GatewayConfig, + upstream: PooledClient, + data_token: &[u8], +) -> (Router, TestWorkerHandle) { + test_router_with_control_tokens(config, upstream, data_token, b"unused-test-control-token") +} + +/// Constructs the authenticated worker router with both credentials issued by a broker control +/// registration. This keeps end-to-end control-plane tests on the production authentication and +/// drain handlers without weakening the normal runtime constructor. +#[cfg(test)] +pub(crate) fn test_router_with_control_tokens( + config: GatewayConfig, + upstream: PooledClient, + data_token: &[u8], + control_token: &[u8], +) -> (Router, TestWorkerHandle) { + let state = Arc::new(WorkerState { + worker_id: "test-worker".into(), + config, + upstream, + managed: None, + auth: RwLock::new(AuthTokens { + data: TokenDigest::from_token(data_token), + pending_data: None, + readiness_data: None, + control: TokenDigest::from_token(control_token), + last_control_sequence: 0, + last_control_request_id: String::new(), + }), + accepting: AtomicBool::new(true), + draining: AtomicBool::new(false), + exiting: AtomicBool::new(false), + in_flight: AtomicUsize::new(0), + drain_deadline: RwLock::new(None), + lifecycle: Notify::new(), + }); + (router(Arc::clone(&state)), TestWorkerHandle { state }) +} + +async fn readiness_probe( + State(state): State>, + headers: HeaderMap, +) -> Response { + if state.draining.load(Ordering::Acquire) || state.exiting.load(Ordering::Acquire) { + return message(StatusCode::SERVICE_UNAVAILABLE, "worker is stopping"); + } + // Only the exact token staged for this registration may open admission. A later health probe + // using an old, still-authenticated data token must not resurrect a worker after control loss. + // Opening before the response is returned keeps broker publication from racing local state. + state.activate_pending_readiness(&headers); + StatusCode::NO_CONTENT.into_response() +} + +async fn authenticate_daemon_request( + State(state): State>, + request: Request, + next: Next, +) -> Response { + if !state.authenticate_data(request.headers()) { + return message(StatusCode::UNAUTHORIZED, "invalid daemon worker credential"); + } + next.run(request).await +} + +async fn drain( + State(state): State>, + Json(request): Json>, +) -> Response { + if !state.authenticate_control(&request) { + return message( + StatusCode::UNAUTHORIZED, + "invalid daemon control credential", + ); + } + state.begin_drain(drain_timeout_ms(&request.payload)); + StatusCode::NO_CONTENT.into_response() +} + +fn drain_timeout_ms(request: &WorkerDrainRequest) -> u64 { + request + .timeout_ms + .unwrap_or_else(|| request.deadline_unix_ms.saturating_sub(now_unix_ms())) + .min(DRAIN_LIFETIME_MS) +} + +async fn proxy(State(state): State>, request: Request) -> Response { + let Some(route) = PublicRoute::from_path(request.uri().path()) else { + return StatusCode::NOT_FOUND.into_response(); + }; + if matches!(route, PublicRoute::Provider(_)) + && let Some(managed) = state.managed.as_ref() + && let Err(error) = managed.ensure_streaming_transport_compatible() + { + return route_failure_response(error); + } + let Some(in_flight) = state.admit() else { + let mut response = message( + StatusCode::SERVICE_UNAVAILABLE, + "worker is not accepting requests", + ); + if !state.draining.load(Ordering::Acquire) && !state.exiting.load(Ordering::Acquire) { + response.headers_mut().insert( + WORKER_ROUTE_FAILURE_HEADER, + HeaderValue::from_static("pass-through"), + ); + } + return response; + }; + match route { + PublicRoute::Hook(hook) => { + if let Some(managed) = state.managed.as_ref() { + let response = managed.handle_hook(hook, request).await; + drop(in_flight); + return response; + } + drop(in_flight); + let mut response = Response::new(Body::from(hook.pass_through_body())); + response + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); + response + } + PublicRoute::Provider(provider) => { + if let Some(managed) = state.managed.as_ref() { + let response = managed + .proxy_provider(state.upstream.clone(), request, provider) + .await; + return match response { + Ok(response) => { + let (parts, body) = response.into_parts(); + let body: RelayBody = hold_body(body, in_flight); + Response::from_parts(parts, Body::new(body)) + } + Err(error) if super::managed::requires_route_pass_through(&error) => { + route_failure_response(error) + } + Err(error) => error.into_response(), + }; + } + forward_to_provider(Arc::clone(&state), request, provider, in_flight).await + } + } +} + +fn route_failure_response(error: CliError) -> Response { + let mut response = error.into_response(); + response.headers_mut().insert( + WORKER_ROUTE_FAILURE_HEADER, + HeaderValue::from_static("pass-through"), + ); + response +} + +async fn forward_to_provider( + state: Arc, + mut request: Request, + route: ProviderRoute, + in_flight: InFlight, +) -> Response { + let allow_environment_provider_auth = + crate::gateway::daemon_allows_environment_provider_auth(request.headers()); + let path_and_query = request + .uri() + .path_and_query() + .map_or("/", |value| value.as_str()); + let destination = match crate::gateway::daemon_provider_upstream_url( + request.headers(), + path_and_query, + &state.config, + ) { + Ok(Some(destination)) => destination, + Ok(None) => route.upstream_url(&state.config, path_and_query), + Err(error) => return error.into_response(), + }; + if allow_environment_provider_auth { + inject_provider_auth(request.headers_mut(), route, &state.config); + } + let destination = match destination.parse::() { + Ok(destination) => destination, + Err(_) => return message(StatusCode::BAD_GATEWAY, "invalid provider destination"), + }; + let strip = [ + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderName::from_static(WORKER_ROUTE_FAILURE_HEADER), + HeaderName::from_static(crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER), + ]; + let request = match prepare_forward_request(request, destination, &strip) { + Ok(request) => request.map(box_body), + Err(error) => return message(StatusCode::BAD_REQUEST, &error.to_string()), + }; + let response = + match tokio::time::timeout(RESPONSE_HEAD_TIMEOUT, state.upstream.request(request)).await { + Ok(Ok(response)) => response, + Ok(Err(error)) => return message(StatusCode::BAD_GATEWAY, &error.to_string()), + Err(_) => { + return message( + StatusCode::GATEWAY_TIMEOUT, + "provider response-head timeout", + ); + } + }; + let response = match prepare_forward_response(response, &strip) { + Ok(response) => response, + Err(error) => return message(StatusCode::BAD_GATEWAY, &error.to_string()), + }; + let (parts, body) = response.into_parts(); + let body: RelayBody = hold_body(body, in_flight); + Response::from_parts(parts, Body::new(body)) +} + +fn inject_provider_auth(headers: &mut HeaderMap, route: ProviderRoute, config: &GatewayConfig) { + if crate::provider_auth::has_provider_credential(headers) { + return; + } + let configured = match route { + ProviderRoute::OpenAi => config.openai_auth_header.as_deref(), + ProviderRoute::Anthropic => config.anthropic_auth_header.as_deref(), + }; + if let Some(configured) = configured.and_then(header_value) { + headers.insert(AUTHORIZATION, configured); + return; + } + match route { + ProviderRoute::OpenAi => { + let Some(key) = nonempty_environment("OPENAI_API_KEY") else { + return; + }; + if let Some(value) = header_value(&format!("Bearer {key}")) { + headers.insert(AUTHORIZATION, value); + } + } + ProviderRoute::Anthropic => { + let Some(key) = nonempty_environment("ANTHROPIC_API_KEY") else { + return; + }; + if let Some(value) = header_value(&key) { + headers.insert(HeaderName::from_static("x-api-key"), value); + } + } + } +} + +fn nonempty_environment(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn header_value(value: &str) -> Option { + HeaderValue::from_str(value).ok() +} + +async fn monitor_control( + state: Arc, + daemon_origin: String, + identity: MachineIdentity, + worker_id: String, + endpoint: String, + worker_tls_root: Option, + mut registration: Registration, +) { + loop { + tokio::time::sleep(registration.heartbeat_interval()).await; + if state.draining.load(Ordering::Acquire) || state.exiting.load(Ordering::Acquire) { + return; + } + if heartbeat_attempt(&mut registration, &daemon_origin, &worker_id).await { + continue; + } + state.control_lost(); + log::error!( + target: "nemo_relay.daemon.worker", + event = "worker_control_lost", + worker_id = worker_id.as_str(); + "Worker lost its authenticated daemon control relationship" + ); + let recovery_deadline = + tokio::time::Instant::now() + Duration::from_millis(RECOVERY_LIFETIME_MS); + loop { + if state.draining.load(Ordering::Acquire) || state.exiting.load(Ordering::Acquire) { + return; + } + if heartbeat_attempt(&mut registration, &daemon_origin, &worker_id).await { + state.control_restored(®istration); + log::info!( + target: "nemo_relay.daemon.worker", + event = "worker_control_restored", + worker_id = worker_id.as_str(); + "Worker restored its daemon control relationship" + ); + break; + } + let recovered = tokio::time::timeout( + CONTROL_ATTEMPT_TIMEOUT, + control::recover( + &daemon_origin, + &identity, + &worker_id, + &endpoint, + worker_tls_root.as_deref(), + registration.generation_grant().clone(), + ), + ) + .await; + if let Ok(Ok(mut new_registration)) = recovered { + state.stage_recovery_data_token(&new_registration); + let ready = tokio::time::timeout( + CONTROL_ATTEMPT_TIMEOUT, + new_registration.ready(&daemon_origin, &worker_id), + ) + .await; + if matches!(ready, Ok(Ok(()))) { + registration = new_registration; + state.control_restored(®istration); + log::info!( + target: "nemo_relay.daemon.worker", + event = "worker_reregistered", + worker_id = worker_id.as_str(); + "Worker re-registered with its daemon" + ); + break; + } + state.discard_recovery_data_token(); + } + if tokio::time::Instant::now() >= recovery_deadline { + log::error!( + target: "nemo_relay.daemon.worker", + event = "worker_recovery_expired", + worker_id = worker_id.as_str(); + "Worker could not restore daemon control before the recovery deadline" + ); + state.request_exit(); + return; + } + tokio::time::sleep(RECOVERY_RETRY_INTERVAL).await; + } + } +} + +async fn heartbeat_attempt( + registration: &mut Registration, + daemon_origin: &str, + worker_id: &str, +) -> bool { + matches!( + tokio::time::timeout( + CONTROL_ATTEMPT_TIMEOUT, + registration.heartbeat(daemon_origin, worker_id), + ) + .await, + Ok(Ok(())) + ) +} + +fn message(status: StatusCode, text: &str) -> Response { + (status, Json(json!({ "error": { "message": text } }))).into_response() +} + +fn read_lock(lock: &RwLock) -> RwLockReadGuard<'_, T> { + lock.read().unwrap_or_else(|error| error.into_inner()) +} + +fn write_lock(lock: &RwLock) -> RwLockWriteGuard<'_, T> { + lock.write().unwrap_or_else(|error| error.into_inner()) +} + +async fn shutdown_signal() { + #[cfg(unix)] + { + let mut terminate = + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("installing SIGTERM handler should succeed"); + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = terminate.recv() => {} + } + } + #[cfg(windows)] + { + let mut shutdown = tokio::signal::windows::ctrl_shutdown() + .expect("installing shutdown handler should succeed"); + tokio::select! { + _ = tokio::signal::ctrl_c() => {} + _ = shutdown.recv() => {} + } + } + #[cfg(not(any(unix, windows)))] + { + let _ = tokio::signal::ctrl_c().await; + } +} + +#[cfg(test)] +#[path = "../../../tests/coverage/daemon/worker_runtime_tests.rs"] +mod tests; diff --git a/crates/cli/src/diagnostics/mod.rs b/crates/cli/src/diagnostics/mod.rs index 1e7e61ead..5603180ba 100644 --- a/crates/cli/src/diagnostics/mod.rs +++ b/crates/cli/src/diagnostics/mod.rs @@ -717,29 +717,43 @@ async fn pi_gateway_reachability_check( }, Some(crate::gateway::client::RelayHealth::Unavailable) => unreachable("is not answering"), // Non-loopback, or the blocking probe did not finish in budget. - None => match reqwest::Client::builder() - .timeout(NETWORK_TIMEOUT) - .build() - .ok() - { - Some(client) => match client.get(format!("{url}/healthz")).send().await { - Ok(response) if response.status().is_success() => Check { - name: NAME, - status: Status::Pass, - details: format!("{url} answered /healthz"), - }, - Ok(response) => Check { - name: NAME, - status: Status::Warn, - details: format!( - "{url} answered /healthz with HTTP {}; the pi extension posts hooks here", - response.status().as_u16() - ), - }, - Err(_) => unreachable("did not answer /healthz"), + None => remote_pi_gateway_reachability_check(&url).await, + } +} + +async fn remote_pi_gateway_reachability_check(url: &str) -> Check { + const NAME: &str = "pi gateway reachability"; + let unreachable = |detail: &str| Check { + name: NAME, + status: Status::Warn, + details: format!( + "{url} {detail}; start the gateway before pi, or every hook will fault -- and \ + because the extension fails open by default, pi keeps running with no policy applied" + ), + }; + + match reqwest::Client::builder() + .timeout(NETWORK_TIMEOUT) + .build() + .ok() + { + Some(client) => match client.get(format!("{url}/healthz")).send().await { + Ok(response) if response.status().is_success() => Check { + name: NAME, + status: Status::Pass, + details: format!("{url} answered /healthz"), + }, + Ok(response) => Check { + name: NAME, + status: Status::Warn, + details: format!( + "{url} answered /healthz with HTTP {}; the pi extension posts hooks here", + response.status().as_u16() + ), }, - None => unreachable("could not be probed"), + Err(_) => unreachable("did not answer /healthz"), }, + None => unreachable("could not be probed"), } } @@ -1952,6 +1966,62 @@ pub(crate) async fn run_doctor( } } +/// Runs the managed deployment doctor without loading or probing any personal runtime state. +pub(crate) fn run_managed_bundle_doctor( + path: &Path, + expected_sha256: &crate::daemon::managed::ManagedBundleDigest, + json: bool, +) -> Result { + let report = collect_managed_bundle_report(path, expected_sha256); + let failed = matches!(report.managed_bundle.status, Status::Fail); + if json { + print!("{}", format_managed_bundle_json(&report)?); + } else { + crate::banner::print_doctor_header(); + print!("{}", format_managed_bundle_human(&report)); + } + Ok(if failed { + std::process::ExitCode::FAILURE + } else { + std::process::ExitCode::SUCCESS + }) +} + +pub(crate) fn collect_managed_bundle_report( + path: &Path, + expected_sha256: &crate::daemon::managed::ManagedBundleDigest, +) -> ManagedBundleDoctorReport { + let expected_sha256_text = expected_sha256.to_string(); + let managed_bundle = match crate::daemon::managed::refresh_bundle(path, expected_sha256) { + Ok(validation) => ManagedBundleDoctorInfo { + status: Status::Pass, + path: path.display().to_string(), + expected_sha256: expected_sha256_text, + artifact_count: Some(validation.artifact_count), + daemon_address: Some(validation.daemon_address), + platform: Some(validation.platform.as_str().into()), + details: format!( + "{} immutable artifacts match the trusted bundle digest {}", + validation.artifact_count, validation.sha256 + ), + }, + Err(error) => ManagedBundleDoctorInfo { + status: Status::Fail, + path: path.display().to_string(), + expected_sha256: expected_sha256_text, + artifact_count: None, + daemon_address: None, + platform: None, + details: error.to_string(), + }, + }; + ManagedBundleDoctorReport { + schema_version: 1, + binary_version: env!("CARGO_PKG_VERSION"), + managed_bundle, + } +} + /// Top-level entry point invoked by `nemo-relay agents`. Always exits 0; the data drives caller /// decisions (e.g., CI gating on JSON output). pub(crate) async fn run_agents(json: bool) -> Result { diff --git a/crates/cli/src/diagnostics/model.rs b/crates/cli/src/diagnostics/model.rs index 397f15a00..54b54957e 100644 --- a/crates/cli/src/diagnostics/model.rs +++ b/crates/cli/src/diagnostics/model.rs @@ -41,6 +41,28 @@ pub(crate) struct DoctorReport { pub completions: Vec, } +/// Managed-only doctor output, intentionally separate from personal runtime diagnostics. +#[derive(Debug, Clone, Serialize)] +pub(crate) struct ManagedBundleDoctorReport { + pub schema_version: u32, + pub binary_version: &'static str, + pub managed_bundle: ManagedBundleDoctorInfo, +} + +#[derive(Debug, Clone, Serialize)] +pub(crate) struct ManagedBundleDoctorInfo { + pub status: Status, + pub path: String, + pub expected_sha256: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub artifact_count: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub daemon_address: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub platform: Option, + pub details: String, +} + #[derive(Debug, Clone, Serialize)] pub(crate) struct EnvironmentInfo { pub os: String, diff --git a/crates/cli/src/diagnostics/render.rs b/crates/cli/src/diagnostics/render.rs index b33f1f6a1..47d434d2c 100644 --- a/crates/cli/src/diagnostics/render.rs +++ b/crates/cli/src/diagnostics/render.rs @@ -76,6 +76,34 @@ pub(crate) fn format_human(report: &DoctorReport) -> String { out } +pub(crate) fn format_managed_bundle_human(report: &ManagedBundleDoctorReport) -> String { + let bundle = &report.managed_bundle; + let mut output = format!( + "\n NeMo Relay {}\n ─────────────────────────────────────────────\n\n Managed daemon bundle\n {} Managed bundle {}\n path {}\n sha256 {}\n\n", + report.binary_version, + format_status(bundle.status), + bundle.details, + bundle.path, + bundle.expected_sha256, + ); + if matches!(bundle.status, Status::Fail) { + output.push_str(" Managed bundle validation FAILED; see details above.\n"); + } else { + output.push_str(" Managed bundle validation passed.\n"); + } + output +} + +pub(crate) fn format_managed_bundle_json( + report: &ManagedBundleDoctorReport, +) -> Result { + serde_json::to_string_pretty(report).map_err(|error| { + CliError::Config(format!( + "could not serialize managed bundle doctor report: {error}" + )) + }) +} + pub(super) fn format_human_header(out: &mut String, report: &DoctorReport) { out.push_str(&format!("\n NeMo Relay {}\n", report.binary_version)); out.push_str(" ─────────────────────────────────────────────\n"); diff --git a/crates/cli/src/filesystem/snapshots.rs b/crates/cli/src/filesystem/snapshots.rs index e82fe9daf..1742ad2ec 100644 --- a/crates/cli/src/filesystem/snapshots.rs +++ b/crates/cli/src/filesystem/snapshots.rs @@ -234,3 +234,7 @@ fn create_symlink(_target: &Path, path: &Path) -> Result<(), String> { path.display() )) } + +#[cfg(test)] +#[path = "../../tests/coverage/shared/filesystem_snapshot_tests.rs"] +mod tests; diff --git a/crates/cli/src/gateway/mod.rs b/crates/cli/src/gateway/mod.rs index e22078e3e..e1268db41 100644 --- a/crates/cli/src/gateway/mod.rs +++ b/crates/cli/src/gateway/mod.rs @@ -44,6 +44,93 @@ use crate::error::CliError; use crate::server::AppState; use crate::sessions::{GatewayCallPrep, GatewaySessionFinish, SessionManager}; +/// Builds the same correlation input used by the foreground gateway without selecting its +/// response-reencoding delivery path. Daemon workers use this before handing the provider body to +/// the raw frame transport. +pub(crate) fn daemon_gateway_start( + headers: &HeaderMap, + path: &str, + request_json: Value, + streaming: bool, +) -> Option { + let provider = ProviderRoute::from_path(path)?; + Some(request::build_llm_gateway_start_from_parts( + headers, + path, + provider, + request_json, + streaming, + )) +} + +/// Resolves the provider destination for an authenticated daemon delivery. +/// +/// The explicit ChatGPT-shaped Responses path retains the personal gateway's Codex alignment +/// behavior for compatibility. Other paths honor Pi's validated named-upstream header; a bearer +/// token alone never selects a different authority. +pub(crate) fn daemon_provider_upstream_url( + headers: &HeaderMap, + path_and_query: &str, + config: &crate::configuration::GatewayConfig, +) -> Result, CliError> { + let path = path_and_query + .split_once('?') + .map_or(path_and_query, |(path, _)| path); + let Some(provider) = ProviderRoute::from_path(path) else { + return Ok(None); + }; + if path == "/backend-api/codex/responses" + && let Some(destination) = + gateway_upstream_url_override(provider, headers, path_and_query, true, config) + { + return Ok(Some(destination)); + } + match client_named_upstream_url(provider, headers, path_and_query, true) { + crate::agents::pi::alignment::NamedUpstream::Named(destination) => { + if !crate::provider_auth::has_provider_credential(headers) { + return Err(CliError::InvalidPayload( + "a client-named upstream requires a caller-supplied provider credential".into(), + )); + } + Ok(Some(destination)) + } + crate::agents::pi::alignment::NamedUpstream::Rejected(reason) => { + Err(CliError::InvalidPayload(reason.to_owned())) + } + crate::agents::pi::alignment::NamedUpstream::Absent => { + Ok(Some(provider.upstream_url(config, path_and_query))) + } + } +} + +/// Applies the personal gateway's agent-auth replacement rules to a daemon provider request. +pub(crate) fn daemon_provider_forward_headers( + headers: &HeaderMap, + path: &str, + config: &crate::configuration::GatewayConfig, +) -> Option { + let provider = ProviderRoute::from_path(path)?; + Some(strip_replaceable_agent_auth_headers( + headers, + provider, + daemon_allows_environment_provider_auth(headers), + provider.configured_auth_header(config), + )) +} + +pub(crate) fn daemon_allows_environment_provider_auth(headers: &HeaderMap) -> bool { + !headers.contains_key(crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER) +} + +/// Returns the exact credential-safe request-header view exposed to daemon-worker middleware. +/// +/// The worker uses this same view when applying middleware diffs to the original `HeaderMap`, so +/// headers hidden from middleware remain opaque and unchanged while unmodified multi-value +/// headers retain their original representation. +pub(crate) fn daemon_observable_headers(headers: &HeaderMap) -> serde_json::Map { + response::observable_headers(headers) +} + #[cfg(test)] #[path = "../../tests/coverage/shared/gateway_tests.rs"] mod tests; diff --git a/crates/cli/src/gateway/request.rs b/crates/cli/src/gateway/request.rs index 2c9d0f36d..08dfb7c25 100644 --- a/crates/cli/src/gateway/request.rs +++ b/crates/cli/src/gateway/request.rs @@ -212,18 +212,33 @@ fn passthrough_body_error(error: axum::Error) -> CliError { } pub(super) fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGatewayStart { + build_llm_gateway_start_from_parts( + &request.headers, + &request.path, + request.provider, + request.request_json.clone(), + request.streaming, + ) +} + +pub(super) fn build_llm_gateway_start_from_parts( + headers: &HeaderMap, + path: &str, + provider: ProviderRoute, + request_json: Value, + streaming: bool, +) -> LlmGatewayStart { LlmGatewayStart { - session_id: gateway_session_id(&request.headers, &request.request_json, request.provider), - provider: request.provider.name().to_string(), - model_name: request - .request_json + session_id: gateway_session_id(headers, &request_json, provider), + provider: provider.name().to_string(), + model_name: request_json .get("model") .and_then(Value::as_str) .map(ToOwned::to_owned), - subagent_id: gateway_subagent_id(&request.headers, &request.request_json, request.provider), + subagent_id: gateway_subagent_id(headers, &request_json, provider), conversation_id: gateway_identifier( - &request.headers, - &request.request_json, + headers, + &request_json, "x-nemo-relay-conversation-id", &[ &["conversation_id"], @@ -232,14 +247,14 @@ pub(super) fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGa ], ), generation_id: gateway_identifier( - &request.headers, - &request.request_json, + headers, + &request_json, "x-nemo-relay-generation-id", &[&["generation_id"], &["generationId"], &["generation", "id"]], ), request_id: gateway_identifier( - &request.headers, - &request.request_json, + headers, + &request_json, "x-nemo-relay-request-id", &[ &["request_id"], @@ -248,12 +263,12 @@ pub(super) fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGa &["metadata", "request_id"], ], ) - .or_else(|| crate::configuration::header_string(&request.headers, "x-request-id")), + .or_else(|| crate::configuration::header_string(headers, "x-request-id")), request: LlmRequest { - headers: observable_headers(&request.headers), - content: request.request_json.clone(), + headers: observable_headers(headers), + content: request_json, }, - streaming: request.streaming, - metadata: json!({ "gateway_path": request.path }), + streaming, + metadata: json!({ "gateway_path": path }), } } diff --git a/crates/cli/src/hooks/encoding.rs b/crates/cli/src/hooks/encoding.rs index a704f2677..8d6f51d1b 100644 --- a/crates/cli/src/hooks/encoding.rs +++ b/crates/cli/src/hooks/encoding.rs @@ -410,5 +410,14 @@ pub(crate) fn event_matches_tools(event: &str) -> bool { } pub(crate) fn event_requires_fail_closed(event: &str) -> bool { - matches!(event, "PreToolUse" | "PermissionRequest" | "pre_tool_call") + matches!( + event, + "PreToolUse" + | "PermissionRequest" + | "pre_tool_call" + | "tool_call" + | "toolCall" + | "user_bash" + | "userBash" + ) } diff --git a/crates/cli/src/hooks/mod.rs b/crates/cli/src/hooks/mod.rs index 5aceb355f..a2aa07b84 100644 --- a/crates/cli/src/hooks/mod.rs +++ b/crates/cli/src/hooks/mod.rs @@ -25,12 +25,12 @@ pub(crate) use encoding::decode_windows_hook_command; #[cfg(all(test, windows))] pub(crate) use encoding::windows_powershell_path; pub(crate) use encoding::{ - GeneratedHookCommands, generated_policy_hooks, persistent_hook_forward_commands, - transparent_hook_forward_commands, + GeneratedHookCommands, event_requires_fail_closed, generated_policy_hooks, + persistent_hook_forward_commands, transparent_hook_forward_commands, }; #[cfg(test)] pub(crate) use encoding::{ - encoded_windows_hook_command, event_matches_tools, event_requires_fail_closed, generated_hooks, + encoded_windows_hook_command, event_matches_tools, generated_hooks, persistent_hook_forward_commands_for_platform, transparent_hook_forward_commands_for_platform, }; #[cfg(test)] diff --git a/crates/cli/src/installation/marketplace/mod.rs b/crates/cli/src/installation/marketplace/mod.rs index cf8bab041..45f8471a7 100644 --- a/crates/cli/src/installation/marketplace/mod.rs +++ b/crates/cli/src/installation/marketplace/mod.rs @@ -42,6 +42,8 @@ use setup::{ }; #[cfg(test)] use setup::{run_plugin_doctor, run_plugin_setup}; +#[cfg(test)] +pub(crate) use state::register_managed_integration; use state::{ CanonicalizeOrSelf, HostRegistrationProgress, PluginInstallOptions, PluginLayout, PluginState, default_install_dir, mark_plugin_setup_installed, read_state, remove_path, state_path, diff --git a/crates/cli/src/installation/marketplace/state.rs b/crates/cli/src/installation/marketplace/state.rs index 8e019cab7..0795dba69 100644 --- a/crates/cli/src/installation/marketplace/state.rs +++ b/crates/cli/src/installation/marketplace/state.rs @@ -450,3 +450,7 @@ pub(super) fn remove_path(path: &Path, options: &PluginInstallOptions) -> Result } }) } + +#[cfg(test)] +#[path = "../../../tests/coverage/shared/marketplace_state_tests.rs"] +mod tests; diff --git a/crates/cli/src/lib.rs b/crates/cli/src/lib.rs index 1ce0fe721..792fd3f8c 100644 --- a/crates/cli/src/lib.rs +++ b/crates/cli/src/lib.rs @@ -8,6 +8,7 @@ mod banner; mod bootstrap; mod commands; mod configuration; +mod daemon; mod diagnostics; mod error; mod events; diff --git a/crates/cli/src/mcp/mod.rs b/crates/cli/src/mcp/mod.rs index 3b512349e..bbe6d6989 100644 --- a/crates/cli/src/mcp/mod.rs +++ b/crates/cli/src/mcp/mod.rs @@ -103,6 +103,12 @@ pub(crate) async fn run(server_args: &GatewayOverrides) -> Result Result<(), CliError> { + let frames = transport::spawn_stdin_reader()?; + session::run_without_gateway(frames, tokio::io::stdout()).await +} + /// Builds the host-independent persistent MCP launch contract. /// /// Host adapters add only schema-specific activation and environment-forwarding fields. Keeping diff --git a/crates/cli/src/mcp/session.rs b/crates/cli/src/mcp/session.rs index f35127c0e..b7b66ec82 100644 --- a/crates/cli/src/mcp/session.rs +++ b/crates/cli/src/mcp/session.rs @@ -12,24 +12,30 @@ use crate::error::CliError; pub(super) async fn run( mut lease: GatewayLease, - mut frames: FrameReceiver, + frames: FrameReceiver, + writer: W, +) -> Result<(), CliError> +where + W: AsyncWrite + Unpin, +{ + tokio::select! { + result = run_without_gateway(frames, writer) => result, + result = lease.wait() => result, + } +} + +/// Serves the no-tools MCP protocol when lifecycle is owned by the brokered daemon client. +pub(super) async fn run_without_gateway( + mut frames: super::transport::FrameReceiver, mut writer: W, ) -> Result<(), CliError> where W: AsyncWrite + Unpin, { - loop { - let received = tokio::select! { - frame = frames.recv() => frame, - result = lease.wait() => return result, - }; - let Some(frame) = received else { - return Ok(()); - }; - let frame = frame?; - let action = evaluate_frame(&frame); - write_response(action, &mut writer).await?; + while let Some(frame) = frames.recv().await { + write_response(evaluate_frame(&frame?), &mut writer).await?; } + Ok(()) } async fn write_response(action: FrameAction, writer: &mut W) -> Result<(), CliError> diff --git a/crates/cli/src/plugins/prompt.rs b/crates/cli/src/plugins/prompt.rs index f19eca699..cff0a47c4 100644 --- a/crates/cli/src/plugins/prompt.rs +++ b/crates/cli/src/plugins/prompt.rs @@ -1500,3 +1500,7 @@ pub(super) fn editor_error(err: dialoguer::Error) -> CliError { other => CliError::Config(format!("plugin edit error: {other}")), } } + +#[cfg(test)] +#[path = "../../tests/coverage/shared/plugin_prompt_tests.rs"] +mod tests; diff --git a/crates/cli/src/server/mod.rs b/crates/cli/src/server/mod.rs index 922856064..988676bca 100644 --- a/crates/cli/src/server/mod.rs +++ b/crates/cli/src/server/mod.rs @@ -926,7 +926,7 @@ where }) } -struct ServerPluginActivation { +pub(crate) struct ServerPluginActivation { host: PluginHostActivation, // The CLI attests and snapshots managed Python environments. The core host // owns plugin code and registration lifetimes; retaining snapshots here @@ -937,7 +937,7 @@ struct ServerPluginActivation { const REMOVED_SWITCHYARD_MESSAGE: &str = "the built-in Switchyard service integration was removed in NeMo Relay >=0.8.0; remove this `[[components]]` entry and refer to the NeMo Relay migration guides for current Switchyard migration information: https://docs.nvidia.com/nemo/relay/reference/migration-guides"; impl ServerPluginActivation { - fn clear(mut self) -> Result<(), CliError> { + pub(crate) fn clear(mut self) -> Result<(), CliError> { self.host .close() .map_err(|error| CliError::Config(format!("plugin teardown failed: {error}"))) @@ -1079,6 +1079,13 @@ async fn activate_server_plugins( })) } +pub(crate) async fn initialize_plugin_host( + config: Option, + dynamic_plugins: Vec, +) -> Result, CliError> { + activate_server_plugins(config, dynamic_plugins).await +} + // Normalizes a Codex hook payload, applies all resulting events before responding, and returns the // adapter's pass-through response body so hook delivery stays causally ordered with observability. async fn codex_hook( diff --git a/crates/cli/tests/architecture_tests.rs b/crates/cli/tests/architecture_tests.rs index edba893df..c2349e410 100644 --- a/crates/cli/tests/architecture_tests.rs +++ b/crates/cli/tests/architecture_tests.rs @@ -44,6 +44,153 @@ fn syntax_paths(source: &str) -> Vec { struct PathVisitor { paths: Vec, command_attributes: Vec, + test_attributes: Vec, +} + +#[derive(Default)] +struct StreamingSourceVisitor { + file: String, + function: Option, + await_depth: usize, + violations: Vec, +} + +impl StreamingSourceVisitor { + fn new(file: &str) -> Self { + Self { + file: file.to_owned(), + ..Self::default() + } + } + + fn function_name(&self) -> &str { + self.function.as_deref().unwrap_or("") + } + + fn allows_request_body_decode(&self) -> bool { + self.file == "daemon/worker/managed.rs" + && matches!(self.function_name(), "handle_hook_inner" | "read") + } + + fn allows_sse_observation(&self) -> bool { + self.file == "daemon/worker/managed.rs" && self.function_name() == "finish_stream" + } + + fn is_delivery_function(&self) -> bool { + matches!( + self.function_name(), + "public_proxy" + | "forward_to_provider" + | "forward_to_worker" + | "forward" + | "proxy" + | "proxy_provider" + | "proxy_managed" + | "dispatch_unmanaged" + | "dispatch_observed" + | "poll_frame" + ) + } + + fn reject(&mut self, operation: &str) { + self.violations.push(format!( + "{} uses {operation} in {}", + self.file, + self.function_name() + )); + } + + fn with_function(&mut self, name: String, visit: impl FnOnce(&mut Self)) { + let previous = self.function.replace(name); + visit(self); + self.function = previous; + } +} + +impl<'ast> Visit<'ast> for StreamingSourceVisitor { + fn visit_item_fn(&mut self, function: &'ast syn::ItemFn) { + self.with_function(function.sig.ident.to_string(), |visitor| { + syn::visit::visit_item_fn(visitor, function); + }); + } + + fn visit_impl_item_fn(&mut self, function: &'ast syn::ImplItemFn) { + self.with_function(function.sig.ident.to_string(), |visitor| { + syn::visit::visit_impl_item_fn(visitor, function); + }); + } + + fn visit_expr_method_call(&mut self, call: &'ast syn::ExprMethodCall) { + let method = call.method.to_string(); + match method.as_str() { + "bytes" | "text" | "json" | "bytes_stream" => self.reject(&format!(".{method}()")), + "collect" if self.await_depth > 0 => self.reject(".collect()"), + "push_bytes_results" if !self.allows_sse_observation() => { + self.reject("SSE decoding on the delivery path") + } + "extend_from_slice" | "extend" | "push_str" if self.is_delivery_function() => { + self.reject(&format!("response accumulation via .{method}()")) + } + _ => {} + } + syn::visit::visit_expr_method_call(self, call); + } + + fn visit_expr_await(&mut self, expression: &'ast syn::ExprAwait) { + self.await_depth += 1; + syn::visit::visit_expr_await(self, expression); + self.await_depth -= 1; + } + + fn visit_expr_call(&mut self, call: &'ast syn::ExprCall) { + if let syn::Expr::Path(path) = call.func.as_ref() { + let segments = path + .path + .segments + .iter() + .map(|segment| segment.ident.to_string()) + .collect::>(); + let last = segments.last().map(String::as_str).unwrap_or_default(); + if last == "to_bytes" && !self.allows_request_body_decode() { + self.reject("to_bytes() response aggregation"); + } + if last == "from_stream" { + self.reject("Body::from_stream()"); + } + if segments.iter().any(|segment| segment.contains("Sse")) + && !self.allows_sse_observation() + { + self.reject("SSE construction or decoding on the delivery path"); + } + let aggregate_constructor = segments + .iter() + .rev() + .take(2) + .map(String::as_str) + .collect::>(); + if self.is_delivery_function() + && matches!( + aggregate_constructor.as_slice(), + [constructor, container] + if matches!(*container, "Vec" | "String") + && matches!(*constructor, "new" | "with_capacity") + ) + { + self.reject("response-wide Vec/String construction"); + } + } + syn::visit::visit_expr_call(self, call); + } + + fn visit_use_rename(&mut self, rename: &'ast syn::UseRename) { + if matches!( + rename.ident.to_string().as_str(), + "to_bytes" | "from_stream" + ) { + self.reject("an alias for a forbidden aggregation API"); + } + syn::visit::visit_use_rename(self, rename); + } } impl<'ast> Visit<'ast> for PathVisitor { @@ -59,14 +206,19 @@ impl<'ast> Visit<'ast> for PathVisitor { } fn visit_attribute(&mut self, attribute: &'ast syn::Attribute) { - let name = attribute + let path = attribute .path() .segments - .last() + .iter() .map(|segment| segment.ident.to_string()) - .unwrap_or_default(); - if matches!(name.as_str(), "arg" | "command" | "value") { - self.command_attributes.push(name); + .collect::>() + .join("::"); + let name = path.rsplit("::").next().unwrap_or_default(); + if matches!(name, "arg" | "command" | "value") { + self.command_attributes.push(name.to_owned()); + } + if matches!(path.as_str(), "test" | "tokio::test") { + self.test_attributes.push(path); } syn::visit::visit_attribute(self, attribute); } @@ -171,10 +323,14 @@ fn tests_are_not_embedded_in_the_source_tree() { let src = source_root(); for path in rust_files(&src) { let source = fs::read_to_string(&path).unwrap(); + let file = syn::parse_file(&source).unwrap(); + let mut visitor = PathVisitor::default(); + visitor.visit_file(&file); assert!( !source.contains("#[cfg(test)]\nmod tests {") - && !source.contains("#[cfg(test)]\r\nmod tests {"), - "inline test module found under src: {}", + && !source.contains("#[cfg(test)]\r\nmod tests {") + && visitor.test_attributes.is_empty(), + "test body found under src instead of crates/cli/tests: {}", path.display() ); } @@ -268,6 +424,9 @@ const OPERATIONAL_LOG_TARGETS: &[&str] = &[ "nemo_relay.hook", "nemo_relay.installation", "nemo_relay.diagnostics", + "nemo_relay.daemon", + "nemo_relay.daemon.mcp", + "nemo_relay.daemon.worker", ]; #[derive(Default)] @@ -371,6 +530,7 @@ fn operational_direct_stderr_is_limited_to_emergency_and_ui_boundaries() { "src/hooks/delivery.rs", "src/hooks/response.rs", "src/plugins/lifecycle/render.rs", + "src/daemon/hook/mod.rs", ]; for path in rust_files(&crate_root.join("src")) { let source = fs::read_to_string(&path).unwrap(); @@ -416,6 +576,80 @@ fn operational_direct_stderr_is_limited_to_emergency_and_ui_boundaries() { } } +#[test] +fn daemon_streaming_modules_do_not_use_response_aggregation_apis() { + let src = source_root(); + let daemon = src.join("daemon"); + let streaming_modules = rust_files(&daemon) + .into_iter() + .filter(|path| { + let source = fs::read_to_string(path).unwrap(); + ["RelayBody", "Frame", "Response"] + .iter() + .any(|streaming_type| source.contains(streaming_type)) + }) + .collect::>(); + assert!( + !streaming_modules.is_empty(), + "no daemon streaming modules discovered" + ); + for path in streaming_modules { + let relative = path + .strip_prefix(&src) + .unwrap() + .components() + .map(|component| component.as_os_str().to_string_lossy()) + .collect::>() + .join("/"); + let source = fs::read_to_string(&path).unwrap(); + let file = syn::parse_file(&source).unwrap(); + let mut visitor = StreamingSourceVisitor::new(&relative); + visitor.visit_file(&file); + assert!( + visitor.violations.is_empty(), + "daemon streaming architecture violations:\n{}", + visitor.violations.join("\n") + ); + } +} + +#[test] +fn streaming_source_analysis_detects_formatted_aliased_and_manual_aggregation() { + let fixture = syn::parse_file( + r#" + use axum::body::to_bytes as aggregate; + async fn forward(body: Body) { + let _ = body.collect() + .await; + let mut response = Vec::new(); + response.extend_from_slice(b"data"); + let _ = Body::from_stream(response); + let _ = SseEventDecoder::new(); + } + "#, + ) + .unwrap(); + let mut visitor = StreamingSourceVisitor::new("daemon/broker/server.rs"); + visitor.visit_file(&fixture); + for expected in [ + "alias for a forbidden aggregation API", + ".collect()", + "response accumulation via .extend_from_slice()", + "response-wide Vec/String construction", + "Body::from_stream()", + "SSE construction or decoding", + ] { + assert!( + visitor + .violations + .iter() + .any(|violation| violation.contains(expected)), + "{expected} escaped streaming architecture analysis: {:?}", + visitor.violations + ); + } +} + #[test] fn shared_runtime_subsystems_do_not_dispatch_host_variants() { let src = source_root(); diff --git a/crates/cli/tests/cli_tests.rs b/crates/cli/tests/cli_tests.rs index fd3dbfd08..acb0cfff4 100644 --- a/crates/cli/tests/cli_tests.rs +++ b/crates/cli/tests/cli_tests.rs @@ -1487,6 +1487,14 @@ impl ChildGuard { let mut child = self.0.take().unwrap(); wait_child(&mut child) } + + fn child_mut(&mut self) -> &mut Child { + self.0.as_mut().unwrap() + } + + fn release(mut self) -> Child { + self.0.take().unwrap() + } } impl Drop for ChildGuard { @@ -1498,7 +1506,11 @@ impl Drop for ChildGuard { } } -fn wait_child_with_output(mut child: Child) -> Output { +fn wait_child_with_output(child: Child) -> Output { + wait_child_with_output_timeout(child, Duration::from_secs(CHILD_PROCESS_TIMEOUT_SECONDS)) +} + +fn wait_child_with_output_timeout(mut child: Child, timeout: Duration) -> Output { fn read_pipe( pipe: Option, ) -> mpsc::Receiver>> { @@ -1518,7 +1530,7 @@ fn wait_child_with_output(mut child: Child) -> Output { let stdout = read_pipe(child.stdout.take()); let stderr = read_pipe(child.stderr.take()); - let deadline = Instant::now() + Duration::from_secs(CHILD_PROCESS_TIMEOUT_SECONDS); + let deadline = Instant::now() + timeout; let status = loop { if let Some(status) = child.try_wait().unwrap() { break status; @@ -1526,7 +1538,7 @@ fn wait_child_with_output(mut child: Child) -> Output { if Instant::now() >= deadline { let _ = child.kill(); let _ = child.wait(); - panic!("child process did not exit within {CHILD_PROCESS_TIMEOUT_SECONDS} seconds"); + panic!("child process did not exit within {timeout:?}"); } thread::sleep(Duration::from_millis(20)); }; @@ -1616,7 +1628,10 @@ fn wait_for_port_closed(address: SocketAddr) { } fn wait_for_port_open(address: SocketAddr) { - let deadline = Instant::now() + Duration::from_secs(5); + // Coverage-instrumented Windows and macOS binaries can spend several seconds flushing or + // merging profiles while sibling process tests start. Keep the readiness bound deterministic + // without treating that CI-only startup cost as a daemon failure. + let deadline = Instant::now() + Duration::from_secs(15); loop { if TcpStream::connect_timeout(&address, Duration::from_millis(100)).is_ok() { return; @@ -5387,3 +5402,341 @@ fn cli_install_pi_refuses_to_add_a_copy_beside_a_project_scoped_one() { "nothing should have been written at user scope" ); } + +/// Exercises the deployed daemon topology through the real CLI processes. The MCP must complete +/// authenticated registration, launch its same-machine worker, wait for broker publication, and +/// expose the no-tools protocol only after the route is usable. A Pi hook then traverses the +/// daemon and worker using the same immutable managed command contract. +#[test] +fn cli_daemon_mcp_launches_worker_and_forwards_pi_hook() { + let temp = tempfile::tempdir().unwrap(); + let probe = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = probe.local_addr().unwrap(); + drop(probe); + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x6c_u8; 32]); + let config_home = temp.path().join("xdg"); + + let daemon = ChildGuard::new( + Command::new(gateway_bin()) + .current_dir(temp.path()) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", &config_home) + .env("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", "1") + .env_remove("NEMO_RELAY_CLIENT_TOKEN") + .args(["daemon", "--port", &address.port().to_string()]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(), + ); + wait_for_port_open(address); + + let daemon_origin = format!("http://{address}"); + let mcp_stderr_path = temp.path().join("daemon-mcp.stderr"); + let mut mcp = ChildGuard::new( + Command::new(gateway_bin()) + .current_dir(temp.path()) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", &config_home) + .env("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", "1") + .env("NEMO_RELAY_CLIENT_TOKEN", &token) + .args(["daemon", "mcp", "--daemon-address", &daemon_origin]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::from( + std::fs::File::create(&mcp_stderr_path).unwrap(), + )) + .spawn() + .unwrap(), + ); + mcp.child_mut().stdin + .as_mut() + .unwrap() + .write_all( + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-06-18\"}}\n", + ) + .unwrap(); + let stdout = mcp.child_mut().stdout.take().unwrap(); + let (response_sender, response_receiver) = mpsc::channel(); + thread::spawn(move || { + let mut response = String::new(); + let result = BufReader::new(stdout) + .read_line(&mut response) + .map(|_| response); + let _ = response_sender.send(result); + }); + let response = match response_receiver.recv_timeout(Duration::from_secs(20)) { + Ok(Ok(response)) => response, + Ok(Err(error)) => panic!("failed reading daemon MCP initialization response: {error}"), + Err(_) => { + let output = mcp.finish(); + panic!( + "daemon MCP did not initialize after worker activation:\n{}", + std::fs::read_to_string(&mcp_stderr_path) + .unwrap_or_else(|_| { String::from_utf8_lossy(&output.stderr).into_owned() }) + ); + } + }; + let response: serde_json::Value = serde_json::from_str(&response).unwrap(); + assert_eq!(response["result"]["serverInfo"]["name"], "nemo-relay"); + + let (provider_origin, provider_request) = spawn_single_request_server( + 200, + r#"{"id":"pi-managed","choices":[{"message":{"role":"assistant","content":"ok"}}]}"#, + ); + let llm_body = r#"{"model":"test-model","stream":true,"messages":[]}"#; + let mut stream = TcpStream::connect_timeout(&address, Duration::from_secs(2)).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + stream + .write_all( + format!( + "POST /v1/chat/completions HTTP/1.1\r\nHost: {address}\r\nContent-Type: application/json\r\nAuthorization: Bearer provider-token\r\nx-nemo-relay-client-token: {token}\r\nx-nemo-relay-upstream-base-url: {provider_origin}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{llm_body}", + llm_body.len() + ) + .as_bytes(), + ) + .unwrap(); + let mut llm_response = String::new(); + stream.read_to_string(&mut llm_response).unwrap(); + assert!( + llm_response.starts_with("HTTP/1.1 200"), + "Pi-selected LLM request did not traverse the daemon worker: {llm_response}" + ); + let provider_request = provider_request + .recv_timeout(Duration::from_secs(5)) + .expect("provider received Pi-selected request"); + assert!(provider_request.starts_with("POST /v1/chat/completions ")); + assert!( + !provider_request + .to_ascii_lowercase() + .contains("x-nemo-relay-client-token") + ); + assert!( + !provider_request + .to_ascii_lowercase() + .contains("x-nemo-relay-upstream-base-url") + ); + + let body = "{}"; + for (agent, path) in [ + ("codex", "/hooks/codex"), + ("claude", "/hooks/claude-code"), + ("pi", "/hooks/pi"), + ] { + let mut stream = TcpStream::connect_timeout(&address, Duration::from_secs(2)).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + stream + .write_all( + format!( + "POST {path} HTTP/1.1\r\nHost: {address}\r\nContent-Type: application/json\r\nx-nemo-relay-client-token: {token}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .unwrap(); + let mut hook_response = String::new(); + stream.read_to_string(&mut hook_response).unwrap(); + assert!( + hook_response.starts_with("HTTP/1.1 200"), + "{agent} hook did not traverse the daemon worker: {hook_response}" + ); + } + + drop(daemon); + wait_for_port_closed(address); + let daemon = ChildGuard::new( + Command::new(gateway_bin()) + .current_dir(temp.path()) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", &config_home) + .env("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", "1") + .env_remove("NEMO_RELAY_CLIENT_TOKEN") + .args(["daemon", "--port", &address.port().to_string()]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(), + ); + wait_for_port_open(address); + let recovery_deadline = Instant::now() + Duration::from_secs(20); + loop { + let recovered = TcpStream::connect_timeout(&address, Duration::from_secs(1)) + .and_then(|mut stream| { + stream.set_read_timeout(Some(Duration::from_secs(2)))?; + stream.write_all( + format!( + "POST /hooks/pi HTTP/1.1\r\nHost: {address}\r\nContent-Type: application/json\r\nx-nemo-relay-client-token: {token}\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{{}}" + ) + .as_bytes(), + )?; + let mut response = String::new(); + stream.read_to_string(&mut response)?; + Ok(response.starts_with("HTTP/1.1 200")) + }) + .unwrap_or(false); + if recovered { + break; + } + assert!( + Instant::now() < recovery_deadline, + "MCP and worker did not reattach after daemon restart" + ); + thread::sleep(Duration::from_millis(100)); + } + + drop(mcp.child_mut().stdin.take()); + let output = wait_child_with_output_timeout(mcp.release(), Duration::from_secs(30)); + assert!( + output.status.success(), + "daemon MCP shutdown failed:\n{}", + std::fs::read_to_string(&mcp_stderr_path) + .unwrap_or_else(|_| { String::from_utf8_lossy(&output.stderr).into_owned() }) + ); + drop(daemon); +} + +#[test] +fn cli_pass_through_daemon_serves_managed_hooks_and_pi_provider_routing() { + let temp = tempfile::tempdir().unwrap(); + let probe = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = probe.local_addr().unwrap(); + drop(probe); + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x3d_u8; 32]); + let daemon_origin = format!("http://{address}"); + let daemon = ChildGuard::new( + Command::new(gateway_bin()) + .current_dir(temp.path()) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", temp.path().join("xdg")) + .env("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", "1") + .env_remove("NEMO_RELAY_CLIENT_TOKEN") + .args([ + "daemon", + "--port", + &address.port().to_string(), + "--pass-through", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn() + .unwrap(), + ); + wait_for_port_open(address); + + let mut mcp = ChildGuard::new( + Command::new(gateway_bin()) + .current_dir(temp.path()) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", temp.path().join("xdg")) + .env("NEMO_RELAY_CLIENT_TOKEN", &token) + .args(["daemon", "mcp", "--daemon-address", &daemon_origin]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .spawn() + .unwrap(), + ); + mcp.child_mut().stdin.as_mut().unwrap().write_all( + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":\"2025-06-18\"}}\n", + ).unwrap(); + let stdout = mcp.child_mut().stdout.take().unwrap(); + let (sender, receiver) = mpsc::channel(); + thread::spawn(move || { + let mut response = String::new(); + let result = BufReader::new(stdout) + .read_line(&mut response) + .map(|_| response); + let _ = sender.send(result); + }); + let response = receiver + .recv_timeout(Duration::from_secs(20)) + .unwrap() + .unwrap(); + let response: serde_json::Value = serde_json::from_str(&response).unwrap(); + assert_eq!(response["result"]["serverInfo"]["name"], "nemo-relay"); + + for (agent, expected) in [ + ("codex", "{}"), + ("claude", r#"{"continue":true}"#), + ("pi", "{}"), + ] { + let mut hook = Command::new(gateway_bin()) + .current_dir(temp.path()) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", temp.path().join("xdg")) + .env("NEMO_RELAY_CLIENT_TOKEN", &token) + .args([ + "daemon", + "hook", + agent, + "--daemon-address", + &daemon_origin, + "--fail-closed", + ]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + hook.stdin.as_mut().unwrap().write_all(b"{}").unwrap(); + drop(hook.stdin.take()); + let output = wait_child_with_output(hook); + assert!( + output.status.success(), + "managed {agent} hook failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!(String::from_utf8(output.stdout).unwrap().trim(), expected); + } + + let unattached_worker = Command::new(gateway_bin()) + .current_dir(temp.path()) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", temp.path().join("xdg")) + .args(["daemon", "worker", "--daemon-address", &daemon_origin]) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + let unattached_worker = wait_child_with_output(unattached_worker); + assert!(!unattached_worker.status.success()); + assert!( + String::from_utf8_lossy(&unattached_worker.stderr) + .contains("requires a protected activation grant") + ); + + let (provider_origin, provider_request) = + spawn_single_request_server(200, r#"{"id":"pass-through"}"#); + let body = r#"{"model":"test-model","stream":true,"messages":[]}"#; + let mut stream = TcpStream::connect_timeout(&address, Duration::from_secs(2)).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + stream + .write_all( + format!( + "POST /v1/chat/completions HTTP/1.1\r\nHost: {address}\r\nContent-Type: application/json\r\nAuthorization: Bearer provider-token\r\nx-nemo-relay-client-token: {token}\r\nx-nemo-relay-upstream-base-url: {provider_origin}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .unwrap(); + let mut response = String::new(); + stream.read_to_string(&mut response).unwrap(); + assert!(response.starts_with("HTTP/1.1 200"), "{response}"); + let provider_request = provider_request + .recv_timeout(Duration::from_secs(5)) + .expect("pass-through provider request"); + assert!(provider_request.starts_with("POST /v1/chat/completions ")); + assert!( + !provider_request + .to_ascii_lowercase() + .contains("x-nemo-relay-client-token") + ); + drop(daemon); +} diff --git a/crates/cli/tests/coverage/agents/launcher_tests.rs b/crates/cli/tests/coverage/agents/launcher_tests.rs index f292b9cb9..9727be5e8 100644 --- a/crates/cli/tests/coverage/agents/launcher_tests.rs +++ b/crates/cli/tests/coverage/agents/launcher_tests.rs @@ -1575,8 +1575,10 @@ fn exit_code_preserves_normal_and_shell_wrapped_codes() { #[tokio::test] #[allow(clippy::await_holding_lock)] async fn run_starts_gateway_injects_env_and_returns_agent_exit_code() { + let _guard = crate::test_support::PLUGIN_CONFIG_TEST_LOCK.lock().await; let temp = tempfile::tempdir().unwrap(); let _cwd = crate::test_support::CwdTestScope::locked(); + let _env = EnvScope::without_managed_bootstrap(); let config = temp.path().join("config.toml"); std::fs::write(&config, "[upstream]\n").unwrap(); let output = temp.path().join("env.txt"); diff --git a/crates/cli/tests/coverage/agents/pi_alignment_tests.rs b/crates/cli/tests/coverage/agents/pi_alignment_tests.rs index 41dbac7eb..a2ff737e2 100644 --- a/crates/cli/tests/coverage/agents/pi_alignment_tests.rs +++ b/crates/cli/tests/coverage/agents/pi_alignment_tests.rs @@ -68,6 +68,16 @@ fn a_blank_header_names_nothing() { ))); } +#[test] +fn a_non_text_header_is_rejected() { + let mut headers = HeaderMap::new(); + headers.insert( + UPSTREAM_BASE_URL_HEADER, + HeaderValue::from_bytes(b"https://example.com/\xff").unwrap(), + ); + assert!(is_rejected(client_named_upstream_base(&headers, true))); +} + /// Each of these would reach somewhere the caller should not be able to send credentialed traffic. #[test] fn only_absolute_http_urls_with_a_bare_host_are_accepted() { diff --git a/crates/cli/tests/coverage/agents/plugin_install_tests.rs b/crates/cli/tests/coverage/agents/plugin_install_tests.rs index fe235a25c..45b0aa5d0 100644 --- a/crates/cli/tests/coverage/agents/plugin_install_tests.rs +++ b/crates/cli/tests/coverage/agents/plugin_install_tests.rs @@ -1031,6 +1031,36 @@ fn refresh_preflight_retires_multiple_directories_for_one_host() { } } +#[test] +fn refresh_preflight_debug_commit_and_unknown_target_operations_are_stable() { + let home = tempdir().unwrap(); + let _home = HomeScope::enter(home.path()); + let install = tempdir().unwrap(); + write_installed_state(CodingAgent::Codex, install.path()); + let layout = PluginLayout::new(CodingAgent::Codex, install.path()); + let mut preflight = + retire_integrations_for_refresh(&[(CodingAgent::Codex, install.path().to_path_buf())]) + .unwrap(); + + assert_eq!( + format!("{preflight:?}"), + "RefreshIntegrationsPreflight { retirement_count: 1 }" + ); + let absent = install.path().join("absent"); + preflight.commit_target(CodingAgent::ClaudeCode, &absent); + preflight + .restore_failed_target(CodingAgent::ClaudeCode, &absent) + .unwrap(); + preflight.commit_target(CodingAgent::Codex, install.path()); + drop(preflight); + + assert!( + std::fs::read_to_string(layout.generation_fence) + .unwrap() + .starts_with("retired:") + ); +} + #[test] fn refresh_preflight_restores_earlier_generations_when_a_target_is_invalid() { let home = tempdir().unwrap(); diff --git a/crates/cli/tests/coverage/commands/configure_command_tests.rs b/crates/cli/tests/coverage/commands/configure_command_tests.rs new file mode 100644 index 000000000..d31ef1983 --- /dev/null +++ b/crates/cli/tests/coverage/commands/configure_command_tests.rs @@ -0,0 +1,51 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::test_support::EnvScope; + +#[tokio::test] +async fn config_execute_resets_the_whole_file_or_one_agent_block() { + let directory = tempfile::tempdir().expect("config home"); + let _environment = EnvScope::set(&[ + ("HOME", Some(directory.path().as_os_str())), + ("XDG_CONFIG_HOME", Some(directory.path().as_os_str())), + ]); + let config_dir = directory.path().join("nemo-relay"); + std::fs::create_dir_all(&config_dir).unwrap(); + let config = config_dir.join("config.toml"); + std::fs::write( + &config, + "[agents.codex]\ncommand = \"codex\"\n\n[agents.claude]\ncommand = \"claude\"\n", + ) + .unwrap(); + let server = ServerArgs::default(); + + let status = execute( + ConfigCommand { + command: None, + agent: Some(AgentArg::Codex), + reset: true, + }, + &server, + ) + .await + .expect("agent reset"); + assert_eq!(status, ExitCode::SUCCESS); + let remaining = std::fs::read_to_string(&config).unwrap(); + assert!(!remaining.contains("agents.codex")); + assert!(remaining.contains("agents.claude")); + + let status = execute( + ConfigCommand { + command: None, + agent: None, + reset: true, + }, + &server, + ) + .await + .expect("whole reset"); + assert_eq!(status, ExitCode::SUCCESS); + assert!(!config.exists()); +} diff --git a/crates/cli/tests/coverage/commands/daemon_tests.rs b/crates/cli/tests/coverage/commands/daemon_tests.rs new file mode 100644 index 000000000..b19314b49 --- /dev/null +++ b/crates/cli/tests/coverage/commands/daemon_tests.rs @@ -0,0 +1,117 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn daemon_address_checks_bracketed_ipv6() { + assert!(parse_daemon_address("http://[::1]:47632").is_ok()); + assert!( + parse_daemon_address("https://[::]:47632") + .unwrap_err() + .contains("bind address") + ); + assert!( + parse_daemon_address("http://[2001:db8::1]:47632") + .unwrap_err() + .contains("must use https") + ); +} + +fn command(subcommand: Option) -> DaemonCommand { + DaemonCommand { + bind: Ipv4Addr::LOCALHOST, + port: 47632, + advertise_address: None, + tls_cert: None, + tls_key: None, + pass_through: false, + command: subcommand, + } +} + +#[tokio::test] +async fn daemon_execute_rejects_unspecified_listeners_without_advertisement() { + let server = crate::commands::serve::ServerArgs::default(); + let mut daemon = command(None); + daemon.bind = Ipv4Addr::UNSPECIFIED; + assert!(matches!( + execute(daemon, &server).await, + Err(CliError::Config(message)) if message.contains("requires --advertise-address") + )); + + let worker = command(Some(DaemonSubcommand::Worker(DaemonWorkerCommand { + daemon_address: "http://127.0.0.1:47632".into(), + bind: Ipv4Addr::UNSPECIFIED, + port: None, + advertise_address: None, + }))); + assert!(matches!( + execute(worker, &server).await, + Err(CliError::Config(message)) if message == "a worker bound to 0.0.0.0 requires --advertise-address" + )); +} + +#[tokio::test] +async fn managed_bundle_command_maps_every_agent_and_platform() { + let server = crate::commands::serve::ServerArgs::default(); + for (index, platform) in [ + ManagedPlatformArg::Linux, + ManagedPlatformArg::Macos, + ManagedPlatformArg::Windows, + ] + .into_iter() + .enumerate() + { + let directory = tempfile::tempdir().expect("bundle parent"); + let output = directory.path().join(format!("bundle-{index}")); + let dispatcher = match platform { + ManagedPlatformArg::Windows => r"C:\ProgramData\NVIDIA\nemo-relay.exe", + _ => "/opt/nvidia/bin/nemo-relay", + }; + let result = execute( + command(Some(DaemonSubcommand::ManagedBundle( + DaemonManagedBundleCommand { + output: output.clone(), + daemon_address: "https://relay.example.com:443".into(), + dispatcher_command: dispatcher.into(), + platform, + agents: vec![AgentArg::Codex, AgentArg::Claude, AgentArg::Pi], + }, + ))), + &server, + ) + .await; + assert_eq!(result.expect("managed bundle"), ExitCode::SUCCESS); + assert!(output.is_dir()); + assert!(std::fs::read_dir(&output).unwrap().next().is_some()); + } +} + +#[test] +fn daemon_value_parsers_cover_valid_and_invalid_address_shapes() { + assert_eq!( + parse_bind_address("127.0.0.1").unwrap(), + Ipv4Addr::LOCALHOST + ); + assert_eq!( + parse_bind_address("0.0.0.0").unwrap(), + Ipv4Addr::UNSPECIFIED + ); + assert!(parse_bind_address("localhost").is_err()); + assert!(parse_bind_address("192.0.2.1").is_err()); + assert_eq!(parse_nonzero_port("1").unwrap(), 1); + assert!(parse_nonzero_port("0").is_err()); + assert!(parse_nonzero_port("65536").is_err()); + + assert_eq!( + parse_daemon_address("http://localhost:47632/").unwrap(), + "http://localhost:47632" + ); + assert!(parse_daemon_address("ftp://relay.example:21").is_err()); + assert!(parse_daemon_address("https://relay.example").is_err()); + assert!(parse_daemon_address("https://user@relay.example:443").is_err()); + assert!(parse_daemon_address("https://relay.example:443/path").is_err()); + assert!(parse_daemon_address("https://0.0.0.0:443").is_err()); + assert!(parse_daemon_address("http://relay.example:80").is_err()); +} diff --git a/crates/cli/tests/coverage/commands/integrations_tests.rs b/crates/cli/tests/coverage/commands/integrations_tests.rs new file mode 100644 index 000000000..c0cd2bd03 --- /dev/null +++ b/crates/cli/tests/coverage/commands/integrations_tests.rs @@ -0,0 +1,191 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn explicit_refresh_directory_targets_only_marketplace_hosts() { + let directory = tempfile::tempdir().expect("temporary install directory"); + let targets = refresh_targets(Some(directory.path())).expect("refresh targets"); + assert_eq!(targets.len(), CodingAgent::MARKETPLACE_HOSTS.len()); + for (agent, install_dir) in targets { + assert!(CodingAgent::MARKETPLACE_HOSTS.contains(&agent)); + assert_ne!(agent, CodingAgent::Pi); + assert_eq!(install_dir, directory.path().canonicalize().unwrap()); + } +} + +#[test] +fn missing_explicit_refresh_directory_is_preserved_for_diagnostics() { + let directory = tempfile::tempdir().expect("temporary parent"); + let missing = directory.path().join("not-created"); + let targets = refresh_targets(Some(&missing)).expect("refresh targets"); + assert!( + targets + .iter() + .all(|(_, install_dir)| install_dir == &missing) + ); +} + +#[test] +fn dry_run_with_no_managed_install_is_a_successful_no_op() { + let directory = tempfile::tempdir().expect("temporary install directory"); + let command = IntegrationsCommand { + command: IntegrationsSubcommand::Refresh(RefreshCommand { + install_dir: Some(directory.path().to_path_buf()), + dry_run: true, + }), + }; + assert_eq!(execute(command).expect("refresh no-op"), ExitCode::SUCCESS); +} + +#[test] +fn dry_run_attempts_every_persisted_target_without_mutating_invalid_state() { + let directory = tempfile::tempdir().expect("temporary install directory"); + let state = crate::installation::marketplace::marketplace_state_path( + CodingAgent::Codex, + directory.path(), + ); + std::fs::write(&state, b"not valid state").unwrap(); + let result = execute(IntegrationsCommand { + command: IntegrationsSubcommand::Refresh(RefreshCommand { + install_dir: Some(directory.path().to_path_buf()), + dry_run: true, + }), + }) + .expect("dry-run refresh"); + assert_eq!(result, ExitCode::SUCCESS); + assert_eq!(std::fs::read(state).unwrap(), b"not valid state"); +} + +#[test] +fn unmanaged_local_marketplace_artifacts_remain_untouched() { + let directory = tempfile::tempdir().expect("temporary install directory"); + let (marketplace, _) = crate::installation::marketplace::marketplace_install_roots( + CodingAgent::ClaudeCode, + directory.path(), + ); + std::fs::create_dir_all(marketplace).unwrap(); + assert_eq!( + execute(IntegrationsCommand { + command: IntegrationsSubcommand::Refresh(RefreshCommand { + install_dir: Some(directory.path().to_path_buf()), + dry_run: true, + }), + }) + .unwrap(), + ExitCode::SUCCESS + ); +} + +#[test] +fn implicit_refresh_targets_include_each_default_marketplace_host_once() { + let directory = tempfile::tempdir().unwrap(); + let _environment = crate::test_support::EnvScope::set(&[ + ("XDG_CONFIG_HOME", Some(directory.path().as_os_str())), + ("HOME", Some(directory.path().as_os_str())), + ("USERPROFILE", Some(directory.path().as_os_str())), + ]); + let targets = refresh_targets(None).expect("implicit refresh targets"); + for agent in CodingAgent::MARKETPLACE_HOSTS { + assert_eq!( + targets.iter().filter(|(found, _)| *found == agent).count(), + 1 + ); + } +} + +#[test] +fn implicit_refresh_targets_include_registered_nondefault_directories_once() { + let config = tempfile::tempdir().unwrap(); + let install = tempfile::tempdir().unwrap(); + let _environment = crate::test_support::EnvScope::set(&[ + ("XDG_CONFIG_HOME", Some(config.path().as_os_str())), + ("HOME", Some(config.path().as_os_str())), + ("USERPROFILE", Some(config.path().as_os_str())), + ]); + crate::installation::marketplace::register_managed_integration( + CodingAgent::Codex, + install.path(), + ) + .unwrap(); + + let canonical = install.path().canonicalize().unwrap(); + let targets = refresh_targets(None).unwrap(); + assert_eq!( + targets + .iter() + .filter(|(agent, directory)| { + *agent == CodingAgent::Codex && directory == &canonical + }) + .count(), + 1 + ); +} + +#[test] +fn refresh_processes_a_managed_marketplace_install_without_host_cli() { + let directory = tempfile::tempdir().unwrap(); + let state = crate::installation::marketplace::marketplace_state_path( + CodingAgent::ClaudeCode, + directory.path(), + ); + std::fs::write(state, b"managed-state-marker").unwrap(); + + let mut installs = Vec::new(); + let result = refresh_with_installer( + RefreshCommand { + install_dir: Some(directory.path().to_path_buf()), + dry_run: true, + }, + |agent, request| { + installs.push((agent, request)); + Ok(ExitCode::SUCCESS) + }, + ) + .expect("dry-run refresh managed install"); + assert_eq!(result, ExitCode::SUCCESS); + assert_eq!(installs.len(), 1); + assert_eq!(installs[0].0, CodingAgent::ClaudeCode); + assert_eq!( + installs[0].1.install_dir.as_deref(), + Some(directory.path().canonicalize().unwrap().as_path()) + ); + assert!(installs[0].1.force); + assert!(installs[0].1.dry_run); + assert!(!installs[0].1.skip_doctor); +} + +#[test] +fn refresh_attempts_every_managed_install_and_aggregates_installer_failures() { + let directory = tempfile::tempdir().unwrap(); + for agent in CodingAgent::MARKETPLACE_HOSTS { + let state = + crate::installation::marketplace::marketplace_state_path(agent, directory.path()); + std::fs::write(state, b"managed-state-marker").unwrap(); + } + + let mut attempted = Vec::new(); + let error = refresh_with_installer( + RefreshCommand { + install_dir: Some(directory.path().to_path_buf()), + dry_run: true, + }, + |agent, _| { + attempted.push(agent); + match agent { + CodingAgent::Codex => Ok(ExitCode::FAILURE), + CodingAgent::ClaudeCode => Err(CliError::Install("host rejected refresh".into())), + CodingAgent::Pi => unreachable!("Pi is not a marketplace host"), + } + }, + ) + .unwrap_err(); + + assert_eq!(attempted, CodingAgent::MARKETPLACE_HOSTS); + let message = error.to_string(); + assert!(message.contains("Codex")); + assert!(message.contains("returned a nonzero status")); + assert!(message.contains("Claude Code")); + assert!(message.contains("host rejected refresh")); +} diff --git a/crates/cli/tests/coverage/commands/main_tests.rs b/crates/cli/tests/coverage/commands/main_tests.rs index 6c0d35ce2..c453214f4 100644 --- a/crates/cli/tests/coverage/commands/main_tests.rs +++ b/crates/cli/tests/coverage/commands/main_tests.rs @@ -14,11 +14,11 @@ use std::time::{Duration, Instant}; use listeners::{Listener, Process, Protocol}; -use super::completions::CompletionsCommand; use super::serve::ServerArgs; use super::*; use crate::commands::configure::ConfigSubcommand; -use crate::commands::model_pricing::{PricingSubcommand, PricingValidateCommand}; +use crate::commands::hook_forward::GatewayModeArg; +use crate::commands::model_pricing::{PricingScopeArgs, PricingSubcommand, PricingValidateCommand}; use crate::commands::plugins::{ PluginsCommand, PluginsEditCommand, PluginsInspectCommand, PluginsListCommand, PluginsScopeArgs, PluginsSubcommand, PluginsValidateCommand, @@ -90,17 +90,59 @@ fn easy_path_setup_inherits_explicit_plugin_target() { #[test] fn operational_command_names_cover_logging_exempt_commands() { for (args, expected) in [ + (vec!["nemo-relay", "daemon"], "daemon"), (vec!["nemo-relay", "codex"], "codex"), (vec!["nemo-relay", "pi"], "pi"), (vec!["nemo-relay", "config"], "config"), (vec!["nemo-relay", "gateway", "start"], "gateway"), (vec!["nemo-relay", "gateway", "stop"], "gateway"), + (vec!["nemo-relay", "uninstall", "pi"], "uninstall"), ] { let cli = Cli::try_parse_from(args).unwrap(); assert_eq!(cli.command.unwrap().log_name(), expected); } } +#[test] +fn command_scope_and_gateway_mode_conversions_cover_every_variant() { + assert_eq!( + crate::plugins::ConfigurationScope::from(PluginsScopeArgs { + user: true, + global: true, + }), + crate::plugins::ConfigurationScope::Invalid + ); + assert_eq!( + crate::plugins::ConfigurationScope::from(PricingScopeArgs { + user: true, + global: false, + }), + crate::plugins::ConfigurationScope::User + ); + assert_eq!( + crate::plugins::ConfigurationScope::from(PricingScopeArgs { + user: false, + global: true, + }), + crate::plugins::ConfigurationScope::Global + ); + assert_eq!( + crate::plugins::ConfigurationScope::from(PricingScopeArgs { + user: true, + global: true, + }), + crate::plugins::ConfigurationScope::Invalid + ); + assert_eq!( + crate::hooks::GatewayMode::from(GatewayModeArg::HookOnly), + crate::hooks::GatewayMode::HookOnly + ); + assert_eq!( + crate::hooks::GatewayMode::from(GatewayModeArg::Passthrough), + crate::hooks::GatewayMode::Passthrough + ); +} + #[test] fn gateway_stop_uses_the_daemon_bind_and_parses_force() { let default = Cli::try_parse_from(["nemo-relay", "gateway", "stop"]).unwrap(); @@ -365,6 +407,361 @@ fn cli_parses_native_mcp_subcommand_and_bind_override() { assert!(Cli::try_parse_from(["nemo-relay", "mcp", "--agent", "codex"]).is_err()); } +#[test] +fn cli_parses_daemon_server_defaults_and_pass_through() { + let cli = Cli::try_parse_from(["nemo-relay", "daemon"]).unwrap(); + let Some(Command::Daemon(command)) = cli.command else { + panic!("expected daemon command"); + }; + assert_eq!(command.bind, std::net::Ipv4Addr::LOCALHOST); + assert_eq!(command.port, 47_632); + assert!(!command.pass_through); + assert!(command.command.is_none()); + + let cli = Cli::try_parse_from(["nemo-relay", "daemon", "--pass-through"]).unwrap(); + let Some(Command::Daemon(command)) = cli.command else { + panic!("expected daemon command"); + }; + assert!(command.pass_through); + + assert!( + Cli::try_parse_from(["nemo-relay", "daemon", "--client-token-file", "/tmp/tokens"]) + .is_err() + ); + + let cli = Cli::try_parse_from([ + "nemo-relay", + "daemon", + "--advertise-address", + "https://relay.example.com:443", + "--tls-cert", + "/etc/nemo-relay/tls.crt", + "--tls-key", + "/etc/nemo-relay/tls.key", + ]) + .unwrap(); + let Some(Command::Daemon(command)) = cli.command else { + panic!("expected daemon command"); + }; + assert_eq!( + command.tls_cert.as_deref(), + Some(std::path::Path::new("/etc/nemo-relay/tls.crt")) + ); + assert_eq!( + command.tls_key.as_deref(), + Some(std::path::Path::new("/etc/nemo-relay/tls.key")) + ); + assert!( + Cli::try_parse_from([ + "nemo-relay", + "daemon", + "--tls-cert", + "/etc/nemo-relay/tls.crt", + ]) + .is_err() + ); +} + +#[test] +fn cli_requires_explicit_valid_daemon_targets_for_clients() { + for arguments in [ + vec!["nemo-relay", "daemon", "mcp"], + vec!["nemo-relay", "daemon", "hook", "codex"], + vec!["nemo-relay", "daemon", "worker"], + vec![ + "nemo-relay", + "daemon", + "mcp", + "--daemon-address", + "http://relay.example.com:47632", + ], + vec![ + "nemo-relay", + "daemon", + "mcp", + "--daemon-address", + "https://relay.example.com", + ], + vec![ + "nemo-relay", + "daemon", + "mcp", + "--daemon-address", + "https://0.0.0.0:47632", + ], + ] { + assert!(Cli::try_parse_from(arguments).is_err()); + } + + for address in [ + "http://127.0.0.1:47632", + "http://localhost:47632", + "https://relay.example.com:443", + "https://relay.example.com:8443", + ] { + assert!( + Cli::try_parse_from(["nemo-relay", "daemon", "mcp", "--daemon-address", address,]) + .is_ok(), + "address should be accepted: {address}" + ); + } +} + +#[test] +fn cli_rejects_daemon_listener_flags_for_daemon_clients() { + for arguments in [ + vec![ + "nemo-relay", + "daemon", + "--bind", + "127.0.0.1", + "mcp", + "--daemon-address", + "http://127.0.0.1:47632", + ], + vec![ + "nemo-relay", + "daemon", + "--port", + "47633", + "worker", + "--daemon-address", + "http://127.0.0.1:47632", + ], + vec![ + "nemo-relay", + "daemon", + "--advertise-address", + "https://relay.example.com:443", + "hook", + "codex", + "--daemon-address", + "http://127.0.0.1:47632", + ], + vec![ + "nemo-relay", + "daemon", + "--pass-through", + "mcp", + "--daemon-address", + "http://127.0.0.1:47632", + ], + ] { + assert!( + Cli::try_parse_from(arguments).is_err(), + "daemon listener flags must not be accepted by a daemon client subcommand" + ); + } +} + +#[test] +fn cli_parses_managed_hook_agent_and_failure_policy() { + let cli = Cli::try_parse_from([ + "nemo-relay", + "daemon", + "hook", + "claude", + "--daemon-address", + "https://relay.example.com:8443", + "--fail-closed", + ]) + .unwrap(); + let Some(Command::Daemon(command)) = cli.command else { + panic!("expected daemon command"); + }; + let Some(daemon::DaemonSubcommand::Hook(hook)) = command.command else { + panic!("expected daemon hook command"); + }; + assert_eq!(hook.agent, AgentArg::Claude); + assert!(hook.fail_closed); + assert!(!hook.fail_open); + + assert!( + Cli::try_parse_from([ + "nemo-relay", + "daemon", + "hook", + "pi", + "--daemon-address", + "http://127.0.0.1:47632", + "--fail-open", + "--fail-closed", + ]) + .is_err() + ); +} + +#[test] +fn cli_worker_defaults_to_loopback_with_an_implicit_ephemeral_port() { + let cli = Cli::try_parse_from([ + "nemo-relay", + "daemon", + "worker", + "--daemon-address", + "http://127.0.0.1:47632", + ]) + .unwrap(); + let Some(Command::Daemon(command)) = cli.command else { + panic!("expected daemon command"); + }; + let Some(daemon::DaemonSubcommand::Worker(worker)) = command.command else { + panic!("expected daemon worker command"); + }; + assert_eq!(worker.bind, std::net::Ipv4Addr::LOCALHOST); + assert_eq!(worker.port, None); + + for arguments in [ + vec![ + "nemo-relay", + "daemon", + "worker", + "--daemon-address", + "http://127.0.0.1:47632", + "--bind", + "192.0.2.1", + ], + vec![ + "nemo-relay", + "daemon", + "worker", + "--daemon-address", + "http://127.0.0.1:47632", + "--port", + "0", + ], + ] { + assert!(Cli::try_parse_from(arguments).is_err()); + } +} + +#[test] +fn cli_parses_managed_bundle_creation_as_an_ambient_config_free_command() { + let cli = Cli::try_parse_from([ + "nemo-relay", + "daemon", + "managed-bundle", + "--output", + "/srv/nemo-relay/bundle-v1", + "--daemon-address", + "https://relay.example.com:443", + "--dispatcher-command", + "/opt/nvidia/bin/nemo-relay-dispatch", + "--platform", + "linux", + "--agent", + "codex", + "--agent", + "claude", + ]) + .unwrap(); + let command = cli.command.unwrap(); + assert!(command.skips_logging()); + let Command::Daemon(command) = command else { + panic!("expected daemon command"); + }; + let Some(daemon::DaemonSubcommand::ManagedBundle(bundle)) = command.command else { + panic!("expected managed-bundle command"); + }; + assert_eq!(bundle.output, PathBuf::from("/srv/nemo-relay/bundle-v1")); + assert_eq!( + bundle.dispatcher_command, + "/opt/nvidia/bin/nemo-relay-dispatch" + ); + assert_eq!(bundle.platform, daemon::ManagedPlatformArg::Linux); + assert_eq!(bundle.agents, [AgentArg::Codex, AgentArg::Claude]); + + assert!( + Cli::try_parse_from([ + "nemo-relay", + "daemon", + "managed-bundle", + "--output", + "/srv/nemo-relay/bundle-v1", + "--daemon-address", + "https://relay.example.com:443", + "--dispatcher-command", + "/opt/nvidia/bin/nemo-relay-dispatch", + "--platform", + "linux", + ]) + .is_err(), + "at least one explicit managed agent is required" + ); +} + +#[tokio::test] +async fn daemon_execute_validates_server_and_worker_bind_contracts() { + let cli = Cli::try_parse_from(["nemo-relay", "daemon", "--bind", "0.0.0.0"]).unwrap(); + let Command::Daemon(command) = cli.command.unwrap() else { + panic!("daemon command"); + }; + assert!(matches!(daemon::execute(command, &cli.server).await, + Err(crate::error::CliError::Config(message)) if message.contains("requires --advertise-address"))); + + let cli = Cli::try_parse_from([ + "nemo-relay", + "daemon", + "worker", + "--daemon-address", + "http://127.0.0.1:47632", + "--bind", + "0.0.0.0", + ]) + .unwrap(); + let Command::Daemon(command) = cli.command.unwrap() else { + panic!("daemon worker command"); + }; + assert!(matches!(daemon::execute(command, &cli.server).await, + Err(crate::error::CliError::Config(message)) if message == "a worker bound to 0.0.0.0 requires --advertise-address")); +} + +#[tokio::test] +async fn daemon_execute_creates_an_immutable_managed_bundle() { + let directory = tempfile::tempdir().unwrap(); + let output = directory.path().join("bundle-v1"); + let cli = Cli::try_parse_from([ + "nemo-relay", + "daemon", + "managed-bundle", + "--output", + output.to_str().unwrap(), + "--daemon-address", + "https://relay.example.com:443", + "--dispatcher-command", + "/opt/nvidia/bin/nemo-relay-dispatch", + "--platform", + "linux", + "--agent", + "codex", + "--agent", + "pi", + ]) + .unwrap(); + let Command::Daemon(command) = cli.command.unwrap() else { + panic!("managed bundle command"); + }; + assert_eq!( + daemon::execute(command, &cli.server).await.unwrap(), + ExitCode::SUCCESS + ); + for relative in [ + "nemo-relay-managed-v1.manifest.json", + "codex/plugin-v1/.codex-plugin/plugin.json", + "codex/plugin-v1/.mcp.json", + "codex/plugin-v1/hooks/hooks.json", + "codex/settings-v1/config.toml", + "pi/extension-v1/README.md", + "pi/extension-v1/index.ts", + "pi/extension-v1/managed-config.json", + "pi/extension-v1/package.json", + "pi/extension-v1/tsconfig.json", + ] { + let metadata = std::fs::metadata(output.join(relative)).unwrap(); + assert!(metadata.is_file(), "missing {relative}"); + assert!(metadata.permissions().readonly(), "{relative} is writable"); + } +} + #[test] fn cli_logging_options_override_environment_source() { let _environment = crate::test_support::EnvScope::set(&[ @@ -630,7 +1027,7 @@ fn multi_agent_operations_attempt_every_target_before_reporting_errors() { } #[test] -fn safe_dispatch_helpers_cover_completions_and_plugins_paths() { +fn safe_dispatch_helpers_cover_plugins_paths() { let temp = tempfile::tempdir().unwrap(); let _env = EnvScope::hermetic(&temp); let config_path = temp.path().join("config.toml"); @@ -640,15 +1037,6 @@ fn safe_dispatch_helpers_cover_completions_and_plugins_paths() { ..ServerArgs::default() }; - assert_eq!( - run_completions(CompletionsCommand { - shell: Some(clap_complete::Shell::Bash), - install: false, - }) - .unwrap(), - ExitCode::SUCCESS - ); - assert_eq!( run_plugins( PluginsCommand { @@ -917,3 +1305,62 @@ fn install_dir_is_cleared_for_pi_under_all_and_kept_everywhere_else() { ); assert_eq!(install::scoped_for(CodingAgent::Pi, all, None), None); } + +#[test] +fn doctor_accepts_a_managed_bundle_without_changing_personal_install_flags() { + const DIGEST: &str = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + let cli = Cli::try_parse_from([ + "nemo-relay", + "doctor", + "--managed-bundle", + "/opt/nvidia/nemo-relay-managed-v1", + "--managed-bundle-sha256", + DIGEST, + "--json", + ]) + .unwrap(); + let command = cli.command.unwrap(); + assert!(command.skips_logging()); + let Command::Doctor(command) = command else { + panic!("expected doctor command"); + }; + assert_eq!( + command.managed_bundle, + Some(PathBuf::from("/opt/nvidia/nemo-relay-managed-v1")) + ); + assert_eq!(command.managed_bundle_sha256.unwrap().to_string(), DIGEST); + assert!(command.json); + + for arguments in [ + vec![ + "nemo-relay", + "doctor", + "--managed-bundle", + "/managed", + "--plugin", + "codex", + "--managed-bundle-sha256", + DIGEST, + ], + vec![ + "nemo-relay", + "doctor", + "--managed-bundle", + "/managed", + "--managed-bundle-sha256", + DIGEST, + "--offline", + ], + vec!["nemo-relay", "doctor", "--managed-bundle", "/managed"], + vec![ + "nemo-relay", + "doctor", + "--managed-bundle", + "/managed", + "--managed-bundle-sha256", + "ABCDEF", + ], + ] { + assert!(Cli::try_parse_from(arguments).is_err()); + } +} diff --git a/crates/cli/tests/coverage/daemon/address_tests.rs b/crates/cli/tests/coverage/daemon/address_tests.rs new file mode 100644 index 000000000..8fa2bbc1d --- /dev/null +++ b/crates/cli/tests/coverage/daemon/address_tests.rs @@ -0,0 +1,84 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn daemon_target_requires_tls_away_from_loopback() { + assert!(daemon_url("http://127.0.0.1:47632").is_ok()); + assert!(daemon_url("https://relay.example.com:443").is_ok()); + assert!(daemon_url("http://relay.example.com:47632").is_err()); + assert!(daemon_url("https://0.0.0.0:47632").is_err()); + assert!(daemon_url("https://[::]:47632").is_err()); + assert!(daemon_url("https://relay.example.com").is_err()); +} + +#[test] +fn worker_port_zero_is_implicit_only() { + assert_eq!( + worker_socket(Ipv4Addr::LOCALHOST, None).unwrap(), + "127.0.0.1:0".parse().unwrap() + ); + assert!(worker_socket(Ipv4Addr::LOCALHOST, Some(0)).is_err()); + assert!(worker_socket(Ipv4Addr::new(10, 0, 0, 1), None).is_err()); +} + +#[test] +fn process_origins_preserve_default_ports_across_revalidation() { + for raw in [ + "https://relay.example.com:443/", + "http://127.0.0.1:80/", + "https://[::1]:443/", + "https://relay.example.com:8443/", + ] { + let origin = explicit_daemon_origin(raw).unwrap(); + assert_eq!(origin, raw.trim_end_matches('/')); + assert_eq!(explicit_daemon_origin(&origin).unwrap(), origin); + assert!(daemon_url(&origin).is_ok()); + } + assert!(explicit_daemon_origin("https://relay.example.com").is_err()); +} + +#[test] +fn unspecified_worker_requires_concrete_advertisement() { + let local: SocketAddr = "0.0.0.0:43210".parse().unwrap(); + assert!(worker_advertised_address(local, None).is_err()); + assert_eq!( + worker_advertised_address(local, Some("worker.example.com")).unwrap(), + "worker.example.com:43210" + ); + assert!(worker_advertised_address(local, Some("::")).is_err()); + for host in [ + "worker.internal:8443", + "https://worker.internal", + "worker/path", + "user@host", + "[::1]:8443", + "-worker", + "worker-", + "valid.-worker", + "worker-.valid", + &"a".repeat(254), + ] { + assert!( + worker_advertised_address(local, Some(host)).is_err(), + "{host}" + ); + } + for host in ["2001:db8::1", "[2001:db8::1]"] { + assert_eq!( + worker_advertised_address(local, Some(host)).unwrap(), + "[2001:db8::1]:43210" + ); + } +} + +#[test] +fn concrete_worker_bind_rejects_advertisement() { + let local: SocketAddr = "127.0.0.1:43210".parse().unwrap(); + assert_eq!( + worker_advertised_address(local, None).unwrap(), + local.to_string() + ); + assert!(worker_advertised_address(local, Some("worker.example.com")).is_err()); +} diff --git a/crates/cli/tests/coverage/daemon/client_tests.rs b/crates/cli/tests/coverage/daemon/client_tests.rs new file mode 100644 index 000000000..776fbb3a6 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/client_tests.rs @@ -0,0 +1,338 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +use std::convert::Infallible; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; + +use axum::body::Body; +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::post; +use axum::{Json, Router}; +use bytes::Bytes; +use serde_json::{Value, json}; +use tokio::net::TcpListener; + +#[test] +fn control_client_has_a_bounded_configuration() { + control_client().expect("control client"); +} + +#[tokio::test] +async fn default_port_process_origins_reach_handshake_transport_for_mcp_and_worker() { + let requests = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + // An explicit local proxy captures both HTTP requests and HTTPS CONNECT attempts without + // binding privileged ports, requiring external DNS, or contacting a real daemon. + let app = Router::new().fallback({ + let requests = Arc::clone(&requests); + move || { + requests.fetch_add(1, Ordering::SeqCst); + async { StatusCode::BAD_GATEWAY } + } + }); + let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let client = reqwest::Client::builder() + .proxy(reqwest::Proxy::all(format!("http://{address}")).unwrap()) + .timeout(Duration::from_secs(5)) + .build() + .unwrap(); + let identity = MachineIdentity::generate().unwrap().identity; + for raw in ["http://127.0.0.1:80/", "https://relay.example:443/"] { + let origin = crate::daemon::common::address::explicit_daemon_origin(raw).unwrap(); + for role in [ComponentRole::Mcp, ComponentRole::Worker] { + let before = requests.load(Ordering::SeqCst); + let result = + begin_handshake(&client, &origin, role, &identity, "default-port-test", None).await; + assert!(matches!( + result, + Err(CliError::Upstream(_) | CliError::Launch(_)) + )); + assert_eq!(requests.load(Ordering::SeqCst), before + 1); + } + } + server.abort(); +} + +#[tokio::test] +async fn rejects_an_oversized_control_response_without_collecting_it() { + let second_polled = Arc::new(AtomicBool::new(false)); + let endpoint_flag = Arc::clone(&second_polled); + let chunks = futures_util::stream::iter([ + Ok::<_, Infallible>(Bytes::from(vec![b'a'; MAX_CONTROL_RESPONSE_BYTES + 1])), + Ok(Bytes::from_static(b"b")), + ]) + .inspect(move |item| { + if item.as_ref().is_ok_and(|bytes| bytes.as_ref() == b"b") { + endpoint_flag.store(true, Ordering::Release); + } + }); + let result = read_bounded_control_chunks(chunks, 0, |never| match never {}).await; + + let error = result.expect_err("oversized response must be rejected"); + assert!( + error + .error + .to_string() + .contains("daemon control response exceeded 262144 bytes") + ); + assert!(!second_polled.load(Ordering::Acquire)); +} + +#[derive(Default)] +struct RetryState { + json_bodies: Mutex>, + empty_bodies: Mutex>, +} + +#[tokio::test] +async fn idempotent_json_retry_reuses_the_exact_encoded_request() { + async fn endpoint(State(state): State>, body: Bytes) -> Response { + let attempt = { + let mut bodies = state.json_bodies.lock().expect("json bodies"); + bodies.push(body); + bodies.len() + }; + if attempt == 1 { + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + Json(json!({"accepted": true})).into_response() + } + + let state = Arc::new(RetryState::default()); + let origin = spawn( + Router::new() + .route("/control", post(endpoint)) + .with_state(Arc::clone(&state)), + ) + .await; + let result: Value = post_json_idempotent( + &control_client().expect("client"), + &format!("{origin}/control"), + &json!({"sequence": 7, "request_id": "same"}), + None, + fast_retry_policy(), + ) + .await + .expect("transient response should be retried"); + + assert_eq!(result, json!({"accepted": true})); + let bodies = state.json_bodies.lock().expect("json bodies"); + assert_eq!(bodies.len(), 2); + assert_eq!(bodies[0], bodies[1]); +} + +#[tokio::test] +async fn idempotent_empty_retry_reuses_the_exact_encoded_request() { + async fn endpoint(State(state): State>, body: Bytes) -> StatusCode { + let attempt = { + let mut bodies = state.empty_bodies.lock().expect("empty bodies"); + bodies.push(body); + bodies.len() + }; + if attempt == 1 { + StatusCode::BAD_GATEWAY + } else { + StatusCode::NO_CONTENT + } + } + + let state = Arc::new(RetryState::default()); + let origin = spawn( + Router::new() + .route("/control", post(endpoint)) + .with_state(Arc::clone(&state)), + ) + .await; + post_empty_idempotent( + &control_client().expect("client"), + &format!("{origin}/control"), + &json!({"sequence": 8, "request_id": "same"}), + fast_retry_policy(), + ) + .await + .expect("transient response should be retried"); + + let bodies = state.empty_bodies.lock().expect("empty bodies"); + assert_eq!(bodies.len(), 2); + assert_eq!(bodies[0], bodies[1]); +} + +#[tokio::test] +async fn control_response_failures_preserve_auth_status_and_json_context() { + async fn invalid_json() -> Response { + (StatusCode::OK, "not-json").into_response() + } + async fn unauthorized() -> Response { + ( + StatusCode::UNAUTHORIZED, + Json(json!({"error": {"message": "bad session"}})), + ) + .into_response() + } + async fn bad_request() -> Response { + (StatusCode::BAD_REQUEST, "opaque rejection").into_response() + } + async fn declared_oversized() -> Response { + Response::new(Body::from(vec![b'x'; MAX_CONTROL_RESPONSE_BYTES + 1])) + } + let origin = spawn( + Router::new() + .route("/invalid-json", post(invalid_json)) + .route("/unauthorized", post(unauthorized)) + .route("/bad-request", post(bad_request)) + .route("/oversized", post(declared_oversized)), + ) + .await; + let client = control_client().unwrap(); + + let invalid: Result = + post_json(&client, &format!("{origin}/invalid-json"), &json!({}), None).await; + assert!( + invalid + .unwrap_err() + .to_string() + .contains("invalid daemon control response") + ); + + let unauthorized: Result = + post_json(&client, &format!("{origin}/unauthorized"), &json!({}), None).await; + assert!( + matches!(unauthorized, Err(CliError::Unauthorized(message)) if message == "bad session") + ); + + let bad_request: Result = + post_json(&client, &format!("{origin}/bad-request"), &json!({}), None).await; + assert!(bad_request.unwrap_err().to_string().contains("HTTP 400")); + + let oversized: Result = + post_json(&client, &format!("{origin}/oversized"), &json!({}), None).await; + assert!( + oversized + .unwrap_err() + .to_string() + .contains("exceeded 262144 bytes") + ); + + let unauthorized = post_empty_idempotent( + &client, + &format!("{origin}/unauthorized"), + &json!({}), + fast_retry_policy(), + ) + .await; + assert!(matches!(unauthorized, Err(CliError::Unauthorized(_)))); +} + +#[tokio::test(start_paused = true)] +async fn bounded_control_retry_stops_on_permanent_error_and_total_deadline() { + let permanent: Result<(), _> = retry_control(fast_retry_policy(), || async { + Err(ControlAttemptError::permanent(CliError::Config( + "permanent".into(), + ))) + }) + .await; + assert!(matches!(permanent, Err(CliError::Config(message)) if message == "permanent")); + + let timed_out: Result<(), _> = retry_control( + ControlRetryPolicy::new( + Duration::from_millis(5), + Duration::from_millis(10), + Duration::ZERO, + ), + std::future::pending, + ) + .await; + assert!( + timed_out + .unwrap_err() + .to_string() + .contains("attempt timed out") + ); + assert!(is_transient_status(StatusCode::TOO_EARLY)); + assert!(!is_transient_status(StatusCode::BAD_REQUEST)); +} + +#[tokio::test(start_paused = true)] +async fn retry_control_honors_retry_after_and_caps_exponential_backoff() { + let started = tokio::time::Instant::now(); + let mut attempts = 0; + let value = retry_control( + ControlRetryPolicy::new( + Duration::from_secs(1), + Duration::from_secs(10), + Duration::from_millis(10), + ), + || { + attempts += 1; + async move { + if attempts == 1 { + Err( + ControlAttemptError::transient(CliError::Launch("retry".into())) + .with_retry_after(Some(Duration::from_secs(2))), + ) + } else { + Ok(attempts) + } + } + }, + ) + .await + .expect("retry succeeds"); + assert_eq!(value, 2); + assert_eq!( + tokio::time::Instant::now() - started, + Duration::from_secs(2) + ); + + assert_eq!(retry_backoff(Duration::ZERO, 10), Duration::ZERO); + for attempt in 0..10 { + let delay = retry_backoff(Duration::from_secs(1), attempt); + assert!(delay >= Duration::from_millis(875)); + assert!(delay < Duration::from_millis(5_625)); + } +} + +#[tokio::test] +async fn a_daemon_role_cannot_initiate_a_client_handshake() { + let identity = MachineIdentity::generate().unwrap().identity; + let result = begin_handshake( + &control_client().unwrap(), + "http://127.0.0.1:1", + ComponentRole::Daemon, + &identity, + "daemon-client", + None, + ) + .await; + let error = match result { + Ok(_) => panic!("daemon role must be rejected before network I/O"), + Err(error) => error, + }; + assert!(matches!( + error, + CliError::Config(message) if message == "a daemon cannot initiate a daemon client handshake" + )); +} + +fn fast_retry_policy() -> ControlRetryPolicy { + ControlRetryPolicy::new( + Duration::from_secs(1), + Duration::from_secs(2), + Duration::ZERO, + ) +} + +async fn spawn(router: Router) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("local address"); + tokio::spawn(async move { + axum::serve(listener, router).await.expect("serve"); + }); + format!("http://{address}") +} diff --git a/crates/cli/tests/coverage/daemon/control_tests.rs b/crates/cli/tests/coverage/daemon/control_tests.rs new file mode 100644 index 000000000..5be05178f --- /dev/null +++ b/crates/cli/tests/coverage/daemon/control_tests.rs @@ -0,0 +1,175 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn daemon_challenge_signature_binds_the_request_before_token_disclosure() { + let daemon = MachineIdentity::generate().expect("daemon").identity; + let mcp = MachineIdentity::generate().expect("mcp").identity; + let request = ChallengeRequest { + initiator: descriptor(ComponentRole::Mcp), + initiator_instance_id: "mcp-one".into(), + initiator_public_identity: mcp.public_identity(), + initiator_fingerprint: mcp.fingerprint(), + initiator_nonce: fresh_nonce().expect("nonce"), + }; + let challenge = super::super::identity::ChallengeRecord::generate(1, 10) + .expect("challenge") + .challenge(); + let mut response = ChallengeResponse { + daemon: descriptor(ComponentRole::Daemon), + daemon_instance_id: "daemon-one".into(), + daemon_public_identity: daemon.public_identity(), + daemon_fingerprint: daemon.fingerprint(), + challenge, + daemon_challenge_proof: daemon.sign(b"placeholder"), + }; + response.daemon_challenge_proof = + daemon.sign(&daemon_challenge_bytes(&request, &response).expect("canonical challenge")); + response + .verify_attestation(&request) + .expect("signed challenge"); + + let mut substituted = request; + substituted.initiator_instance_id = "mcp-two".into(); + assert!(response.verify_attestation(&substituted).is_err()); +} + +#[test] +fn session_request_hash_covers_the_payload_and_sensitive_values_are_redacted() { + let request = SessionRequest::new( + "mcp-1".into(), + SensitiveString::new("session-secret").expect("secret"), + 1, + ActivationFailedPayload { + activation_id: "activation-1".into(), + reason: "bind failed".into(), + }, + ) + .expect("request"); + assert!(request.validate_payload_hash()); + assert!(!format!("{request:?}").contains("session-secret")); + + let mut changed = request; + changed.payload.reason = "different".into(); + assert!(!changed.validate_payload_hash()); +} + +#[test] +fn launch_directive_becomes_worker_bootstrap_without_reencoding_fields() { + let directive = BrokerDirective::LaunchWorker { + activation_id: "activation-1".into(), + activation_token: SensitiveString::new("activation-secret").expect("secret"), + deadline_unix_ms: 42, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }; + let bootstrap = WorkerBootstrap::from_directive(directive).expect("launch directive"); + assert_eq!(bootstrap.activation_id, "activation-1"); + assert_eq!(bootstrap.activation_token.expose(), "activation-secret"); + assert_eq!(bootstrap.deadline_unix_ms, 42); + assert_eq!(bootstrap.bind_ip, Ipv4Addr::LOCALHOST); + assert_eq!(bootstrap.port, 0); +} + +#[test] +fn worker_network_hint_is_signed_and_accepts_concrete_hostnames() { + let identity = MachineIdentity::generate().expect("identity").identity; + let challenge = super::super::identity::ChallengeRecord::generate(1, 10) + .expect("challenge") + .challenge(); + let hint = WorkerNetworkHint::new("Worker.Example.COM", Some(443)).expect("hint"); + assert_eq!(hint.advertised_host, "worker.example.com"); + let proof = WorkerNetworkHintProof::sign( + hint, + "https://daemon.example.com:443", + "mcp-one", + &challenge.id, + &identity.fingerprint(), + &identity, + ) + .expect("signed hint"); + proof + .verify( + "https://daemon.example.com:443", + "mcp-one", + &challenge.id, + &identity.fingerprint(), + &identity.public_identity(), + ) + .expect("valid hint"); + + let mut changed = proof; + changed.hint.advertised_host = "attacker.example.com".into(); + assert!( + changed + .verify( + "https://daemon.example.com:443", + "mcp-one", + &challenge.id, + &identity.fingerprint(), + &identity.public_identity(), + ) + .is_err() + ); + assert!(WorkerNetworkHint::new("https://worker.example.com", None).is_err()); + assert!(WorkerNetworkHint::new("0.0.0.0", None).is_err()); +} + +#[test] +fn worker_generation_grant_binds_endpoint_and_tls_root() { + let daemon = MachineIdentity::generate().expect("daemon").identity; + let worker = MachineIdentity::generate().expect("worker").identity; + let grant = WorkerGenerationGrant::issue( + "worker-one", + worker.fingerprint(), + "https://worker.example.com:9443", + Some("root-certificate"), + &daemon, + ) + .expect("generation grant"); + grant + .verify( + "worker-one", + worker.fingerprint(), + "https://worker.example.com:9443", + Some("root-certificate"), + &daemon.public_identity(), + ) + .expect("valid generation"); + assert!( + grant + .verify( + "worker-one", + worker.fingerprint(), + "https://attacker.example.com:9443", + Some("root-certificate"), + &daemon.public_identity(), + ) + .is_err() + ); + assert!( + grant + .verify( + "worker-one", + worker.fingerprint(), + "https://worker.example.com:9443", + Some("substituted-root"), + &daemon.public_identity(), + ) + .is_err() + ); + assert!( + grant + .verify( + "worker-one", + worker.fingerprint(), + "https://worker.example.com:9443", + None, + &daemon.public_identity(), + ) + .is_err() + ); +} diff --git a/crates/cli/tests/coverage/daemon/daemon_worker_e2e_tests.rs b/crates/cli/tests/coverage/daemon/daemon_worker_e2e_tests.rs new file mode 100644 index 000000000..8cfd84bd5 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/daemon_worker_e2e_tests.rs @@ -0,0 +1,2296 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::VecDeque; +use std::convert::Infallible; +use std::net::Ipv4Addr; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use http::header::{AUTHORIZATION, CONTENT_TYPE, HeaderValue, TE, TRAILER, UPGRADE}; +use http_body_util::{BodyExt as _, Empty, Full}; +use hyper::body::{Frame, Incoming, SizeHint}; +use hyper::server::conn::{http1, http2}; +use hyper::service::service_fn; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use hyper_util::service::TowerToHyperService; +use sha2::{Digest, Sha256}; +use tokio::net::TcpListener; +use tokio::sync::{Barrier, OwnedSemaphorePermit, Semaphore, oneshot}; + +use super::*; +use crate::daemon::common::transport::pooled_worker_h2c_client; +use crate::daemon::worker::{TestWorkerHandle, test_router as worker_test_router}; + +const EVENT_A: &[u8] = b": heartbeat\r\nevent: response.output_text.delta\r\nid: 7\r\nretry: 1000\r\ndata: first\r\ndata: second\r\n\r\n"; +const EVENT_B: &[u8] = b"data: [DONE]\r\n\r\n\x80\xff"; +const WORKER_TOKEN: &str = "test-daemon-to-worker-token"; +const TEST_SEQUENCE_HEADER: &str = "x-test-stream-sequence"; +const SEQUENCE_PARTS: usize = 4; + +#[derive(Clone, Copy, Debug)] +enum TestProtocol { + Http1, + Http2, +} + +#[derive(Clone, Copy, Debug)] +enum ProviderKind { + OpenAi, + Anthropic, +} + +impl ProviderKind { + const fn path(self) -> &'static str { + match self { + Self::OpenAi => "/v1/responses", + Self::Anthropic => "/v1/messages", + } + } +} + +#[derive(Clone, Copy)] +enum DuringStream { + None, + PauseBeyondFormerTotalTimeout, + ControlLoss, + Drain, +} + +struct CausalProviderBody { + phase: u8, + release_second: oneshot::Receiver<()>, + trailers: Option, +} + +impl hyper::body::Body for CausalProviderBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + match this.phase { + 0 => { + this.phase = 1; + Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(EVENT_A))))) + } + 1 => match Pin::new(&mut this.release_second).poll(context) { + Poll::Pending => Poll::Pending, + Poll::Ready(_) => { + this.phase = 2; + Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(EVENT_B))))) + } + }, + 2 => { + this.phase = 3; + Poll::Ready( + this.trailers + .take() + .map(|trailers| Ok(Frame::trailers(trailers))), + ) + } + _ => Poll::Ready(None), + } + } + + fn is_end_stream(&self) -> bool { + self.phase >= 3 && self.trailers.is_none() + } + + fn size_hint(&self) -> SizeHint { + SizeHint::default() + } +} + +struct CountedProviderBody { + remaining: usize, + frame: Bytes, + polls: Arc, +} + +impl hyper::body::Body for CountedProviderBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + self.polls.fetch_add(1, Ordering::SeqCst); + if self.remaining == 0 { + return Poll::Ready(None); + } + self.remaining -= 1; + Poll::Ready(Some(Ok(Frame::data(self.frame.clone())))) + } +} + +struct CancellationProviderBody { + first_sent: bool, + dropped: Option>, +} + +impl hyper::body::Body for CancellationProviderBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if self.first_sent { + Poll::Pending + } else { + self.first_sent = true; + Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(b"first\n\n"))))) + } + } +} + +impl Drop for CancellationProviderBody { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } +} + +struct SequencedProviderBody { + sequence: usize, + next_part: usize, + _response_permit: OwnedSemaphorePermit, +} + +impl hyper::body::Body for SequencedProviderBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if self.next_part == SEQUENCE_PARTS { + return Poll::Ready(None); + } + let part = self.next_part; + self.next_part += 1; + Poll::Ready(Some(Ok(Frame::data(sequence_chunk(self.sequence, part))))) + } +} + +struct FidelityProviderBody { + chunks: VecDeque, + trailers: Option, +} + +impl hyper::body::Body for FidelityProviderBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if let Some(chunk) = self.chunks.pop_front() { + return Poll::Ready(Some(Ok(Frame::data(chunk)))); + } + Poll::Ready( + self.trailers + .take() + .map(|trailers| Ok(Frame::trailers(trailers))), + ) + } +} + +fn fidelity_corpus() -> Bytes { + let mut bytes = Vec::with_capacity(300 * 1024); + bytes.extend_from_slice(b": heartbeat\r\n\r\n"); + bytes.extend_from_slice(b"event: response.output_text.delta\r\nid: 17\r\nretry: 500\r\n"); + bytes.extend_from_slice(b"data: first\r\ndata: second\r\n\r\n"); + bytes.extend_from_slice(b"data: \xff\x00\xfe\r\n\r\n"); + bytes.extend_from_slice(b"event: large\r\ndata: "); + bytes.extend(std::iter::repeat_n(b'L', 256 * 1024)); + bytes.extend_from_slice(b"\r\n\r\ndata: [DONE]\r\n\r\n"); + Bytes::from(bytes) +} + +fn arbitrarily_split_fidelity_corpus(corpus: &Bytes) -> VecDeque { + const WIDTHS: &[usize] = &[1, 2, 3, 7, 31, 257, 4_093, 16_384, 65_521]; + let mut chunks = VecDeque::new(); + let mut offset = 0; + let mut split = 0; + while offset < corpus.len() { + let end = offset + .saturating_add(WIDTHS[split % WIDTHS.len()]) + .min(corpus.len()); + chunks.push_back(corpus.slice(offset..end)); + offset = end; + split += 1; + if split == 4 { + chunks.push_back(Bytes::new()); + } + } + chunks +} + +fn sequence_chunk(sequence: usize, part: usize) -> Bytes { + let prefix = format!("stream={sequence};part={part};"); + let mut chunk = vec![b'x'; 16 * 1024]; + chunk[..prefix.len()].copy_from_slice(prefix.as_bytes()); + *chunk.last_mut().expect("sequence chunk is non-empty") = b'\n'; + Bytes::from(chunk) +} + +#[derive(Debug)] +struct ProviderObservation { + path: String, + authorization: Option, + retained_client_token: bool, + retained_worker_token: bool, + body: Bytes, +} + +fn client_for(protocol: TestProtocol) -> PooledClient { + match protocol { + TestProtocol::Http1 => pooled_client().expect("HTTP/1.1 pooled client"), + TestProtocol::Http2 => { + pooled_worker_h2c_client().expect("HTTP/2 prior-knowledge pooled client") + } + } +} + +async fn spawn_causal_provider( + protocol: TestProtocol, +) -> ( + std::net::SocketAddr, + oneshot::Sender<()>, + oneshot::Receiver, + tokio::task::JoinHandle<()>, +) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind causal provider"); + let address = listener.local_addr().expect("provider address"); + let (release_second, wait_for_release) = oneshot::channel(); + let (observed, observation) = oneshot::channel(); + let wait_for_release = Arc::new(Mutex::new(Some(wait_for_release))); + let observed = Arc::new(Mutex::new(Some(observed))); + let task = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept worker connection"); + stream.set_nodelay(true).expect("provider TCP_NODELAY"); + let service = service_fn(move |request: Request| { + let wait_for_release = wait_for_release + .lock() + .expect("release gate lock") + .take() + .expect("provider receives one request"); + let observed = observed + .lock() + .expect("observation lock") + .take() + .expect("provider observes one request"); + async move { + let (parts, body) = request.into_parts(); + let body = body + .collect() + .await + .expect("provider request body") + .to_bytes(); + let _ = observed.send(ProviderObservation { + path: parts + .uri + .path_and_query() + .map_or("/", |value| value.as_str()) + .to_owned(), + authorization: parts.headers.get(AUTHORIZATION).cloned(), + retained_client_token: parts.headers.contains_key(CLIENT_TOKEN_HEADER), + retained_worker_token: parts.headers.contains_key(WORKER_TOKEN_HEADER), + body, + }); + + let mut trailers = HeaderMap::new(); + trailers.append("x-stream-checksum", HeaderValue::from_static("one")); + trailers.append("x-stream-checksum", HeaderValue::from_static("two")); + let mut response = Response::new(box_body(CausalProviderBody { + phase: 0, + release_second: wait_for_release, + trailers: Some(trailers), + })); + *response.status_mut() = StatusCode::CREATED; + response + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); + response + .headers_mut() + .append("x-provider", HeaderValue::from_static("first")); + response + .headers_mut() + .append("x-provider", HeaderValue::from_static("second")); + response + .headers_mut() + .insert(TRAILER, HeaderValue::from_static("x-stream-checksum")); + Ok::<_, Infallible>(response) + } + }); + + match protocol { + TestProtocol::Http1 => { + let mut builder = http1::Builder::new(); + builder.keep_alive(false); + builder + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve HTTP/1.1 provider") + } + TestProtocol::Http2 => http2::Builder::new(TokioExecutor::new()) + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve HTTP/2 provider"), + } + }); + (address, release_second, observation, task) +} + +async fn spawn_counted_provider( + protocol: TestProtocol, + polls: Arc, + frame_count: usize, + frame_size: usize, +) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind counted provider"); + let address = listener.local_addr().expect("counted provider address"); + let task = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept worker connection"); + stream.set_nodelay(true).expect("provider TCP_NODELAY"); + let service = service_fn(move |_request: Request| { + let polls = Arc::clone(&polls); + async move { + Ok::<_, Infallible>(Response::new(box_body(CountedProviderBody { + remaining: frame_count, + frame: Bytes::from(vec![0x5a; frame_size]), + polls, + }))) + } + }); + match protocol { + TestProtocol::Http1 => { + let mut builder = http1::Builder::new(); + builder.keep_alive(false); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + } + TestProtocol::Http2 => { + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + } + } + }); + (address, task) +} + +async fn spawn_cancellation_provider( + protocol: TestProtocol, +) -> ( + std::net::SocketAddr, + oneshot::Receiver<()>, + tokio::task::JoinHandle<()>, +) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind cancellation provider"); + let address = listener + .local_addr() + .expect("cancellation provider address"); + let (dropped, wait_for_drop) = oneshot::channel(); + let dropped = Arc::new(Mutex::new(Some(dropped))); + let task = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept worker connection"); + stream.set_nodelay(true).expect("provider TCP_NODELAY"); + let service = service_fn(move |_request: Request| { + let dropped = dropped + .lock() + .expect("cancellation signal lock") + .take() + .expect("provider receives exactly one request"); + async move { + Ok::<_, Infallible>(Response::new(box_body(CancellationProviderBody { + first_sent: false, + dropped: Some(dropped), + }))) + } + }); + match protocol { + TestProtocol::Http1 => { + let mut builder = http1::Builder::new(); + builder.keep_alive(false); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + } + TestProtocol::Http2 => { + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + } + } + }); + (address, wait_for_drop, task) +} + +async fn spawn_fidelity_provider( + protocol: TestProtocol, +) -> (std::net::SocketAddr, Bytes, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind fidelity provider"); + let address = listener.local_addr().expect("fidelity provider address"); + let corpus = fidelity_corpus(); + let chunks = Arc::new(Mutex::new(Some(arbitrarily_split_fidelity_corpus(&corpus)))); + let task = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept worker connection"); + stream.set_nodelay(true).expect("provider TCP_NODELAY"); + let service = service_fn(move |_request: Request| { + let chunks = chunks + .lock() + .expect("fidelity chunks lock") + .take() + .expect("fidelity provider receives exactly one request"); + async move { + let mut trailers = HeaderMap::new(); + trailers.append("x-stream-checksum", HeaderValue::from_static("first")); + trailers.append("x-stream-checksum", HeaderValue::from_static("second")); + trailers.append("x-binary-safe", HeaderValue::from_static("yes")); + let mut response = Response::new(box_body(FidelityProviderBody { + chunks, + trailers: Some(trailers), + })); + *response.status_mut() = StatusCode::PARTIAL_CONTENT; + response + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); + response + .headers_mut() + .append("x-provider", HeaderValue::from_static("first")); + response + .headers_mut() + .append("x-provider", HeaderValue::from_static("second")); + response.headers_mut().insert( + TRAILER, + HeaderValue::from_static("x-stream-checksum, x-binary-safe"), + ); + Ok::<_, Infallible>(response) + } + }); + match protocol { + TestProtocol::Http1 => { + let mut builder = http1::Builder::new(); + builder.keep_alive(false); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + } + TestProtocol::Http2 => { + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + } + } + }); + (address, corpus, task) +} + +async fn spawn_sequenced_provider( + protocol: TestProtocol, + streams: usize, + connections: Arc, +) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind sequenced provider"); + let address = listener.local_addr().expect("sequenced provider address"); + let barrier = Arc::new(Barrier::new(streams)); + let response_permits = Arc::new(Semaphore::new(16)); + let task = tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.expect("accept worker connection"); + stream.set_nodelay(true).expect("provider TCP_NODELAY"); + connections.fetch_add(1, Ordering::SeqCst); + let barrier = Arc::clone(&barrier); + let response_permits = Arc::clone(&response_permits); + tokio::spawn(async move { + let service = service_fn(move |request: Request| { + let barrier = Arc::clone(&barrier); + let response_permits = Arc::clone(&response_permits); + async move { + let sequence = request + .headers() + .get(TEST_SEQUENCE_HEADER) + .expect("sequence header reaches provider") + .to_str() + .expect("sequence header is ASCII") + .parse::() + .expect("sequence header is an integer"); + request + .into_body() + .collect() + .await + .expect("provider receives the complete request body"); + barrier.wait().await; + let response_permit = response_permits + .acquire_owned() + .await + .expect("response concurrency semaphore remains open"); + Ok::<_, Infallible>(Response::new(box_body(SequencedProviderBody { + sequence, + next_part: 0, + _response_permit: response_permit, + }))) + } + }); + match protocol { + TestProtocol::Http1 => { + let _ = http1::Builder::new() + .serve_connection(TokioIo::new(stream), service) + .await; + } + TestProtocol::Http2 => { + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + } + } + }); + } + }); + (address, task) +} + +async fn spawn_router( + protocol: TestProtocol, + app: Router, +) -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test router"); + let address = listener.local_addr().expect("router address"); + let task = tokio::spawn(async move { + let mut connections = tokio::task::JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let (stream, _) = accepted.expect("accept test connection"); + stream.set_nodelay(true).expect("router TCP_NODELAY"); + let service = TowerToHyperService::new(app.clone()); + connections.spawn(async move { + match protocol { + TestProtocol::Http1 => { + let mut builder = http1::Builder::new(); + builder.keep_alive(true); + builder + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve HTTP/1.1 router"); + } + TestProtocol::Http2 => { + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + builder + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve HTTP/2 router"); + } + } + }); + } + Some(completed) = connections.join_next(), if !connections.is_empty() => { + completed.expect("test router connection task"); + } + } + } + }); + (address, task) +} + +fn configured_worker_router( + protocol: TestProtocol, + provider: ProviderKind, + provider_address: std::net::SocketAddr, +) -> (Router, TestWorkerHandle) { + configured_worker_router_with_token(protocol, provider, provider_address, WORKER_TOKEN) +} + +fn configured_worker_router_with_token( + protocol: TestProtocol, + provider: ProviderKind, + provider_address: std::net::SocketAddr, + worker_token: &str, +) -> (Router, TestWorkerHandle) { + let mut config = GatewayConfig::default(); + match provider { + ProviderKind::OpenAi => { + config.openai_base_url = format!("http://{provider_address}/v1"); + } + ProviderKind::Anthropic => { + config.anthropic_base_url = format!("http://{provider_address}"); + } + } + worker_test_router(config, client_for(protocol), worker_token.as_bytes()) +} + +fn daemon_router_with_ready_worker( + protocol: TestProtocol, + route_token: &str, + worker_address: std::net::SocketAddr, +) -> (Router, Arc, tempfile::TempDir) { + let machine_identity = MachineIdentity::generate() + .expect("machine identity") + .identity; + let fingerprint = machine_identity.fingerprint(); + let credential = RouteCredential::parse(route_token.to_owned()).expect("route credential"); + let registry = Registry::new(false); + let mcp_session = McpSessionId::new("test-mcp-session").expect("MCP session ID"); + let launch = WorkerLaunch { + activation_id: "test-activation".into(), + activation_token: SensitiveString::new("unused-test-activation-token") + .expect("activation token"), + deadline_unix_ms: now_unix_ms().saturating_add(ACTIVATION_LIFETIME_MS), + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }; + let directive = registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: mcp_session, + lease_expires_at_unix_ms: now_unix_ms().saturating_add(MCP_LEASE_MS), + }, + launch, + ) + .expect("register test MCP"); + assert!(matches!(directive, BrokerDirective::LaunchWorker { .. })); + let target = Arc::new( + WorkerTarget::with_client( + "test-worker", + format!("http://{worker_address}"), + SensitiveString::new(WORKER_TOKEN).expect("worker token"), + client_for(protocol), + ) + .expect("worker target"), + ); + registry + .mark_worker_ready(fingerprint, "test-activation", Arc::clone(&target)) + .expect("publish test worker"); + + let identity = MachineIdentity::generate() + .expect("daemon identity") + .identity; + let generation_state = tempfile::tempdir().expect("generation state directory"); + let active_worker_generations = ActiveWorkerGenerations::load_for_test( + generation_state + .path() + .join("active-worker-generations.json"), + ) + .expect("active generation state"); + let state = Arc::new(DaemonState { + registry, + identity, + descriptor: crate::daemon::common::control::descriptor(ComponentRole::Daemon), + instance_id: "test-daemon".into(), + public_origin: "http://127.0.0.1:1".into(), + config: GatewayConfig::default(), + upstream: pooled_client().expect("daemon pass-through client"), + worker_clients: WorkerClientPool::new().expect("daemon worker clients"), + challenges: Mutex::new(HashMap::new()), + activations: Mutex::new(HashMap::new()), + mcp_sessions: Mutex::new(HashMap::new()), + mcp_heartbeat_serialization: Mutex::new(()), + worker_sessions: Mutex::new(HashMap::new()), + pending_directives: Mutex::new(HashMap::new()), + active_worker_generations, + worker_generation_publication: Mutex::new(()), + }); + (router(state), target, generation_state) +} + +fn daemon_router_with_pass_through( + protocol: TestProtocol, + provider: ProviderKind, + route_token: &str, + provider_address: std::net::SocketAddr, +) -> (Router, tempfile::TempDir) { + let credential = RouteCredential::parse(route_token.to_owned()).expect("route credential"); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let registry = Registry::new(true); + let directive = registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: McpSessionId::new("test-pass-through-mcp").expect("MCP session ID"), + lease_expires_at_unix_ms: now_unix_ms().saturating_add(MCP_LEASE_MS), + }, + WorkerLaunch { + activation_id: "unused-pass-through-activation".into(), + activation_token: SensitiveString::new("unused-pass-through-token") + .expect("activation token"), + deadline_unix_ms: now_unix_ms().saturating_add(ACTIVATION_LIFETIME_MS), + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }, + ) + .expect("register pass-through MCP"); + assert!(matches!(directive, BrokerDirective::UsePassThrough)); + + let mut config = GatewayConfig::default(); + match provider { + ProviderKind::OpenAi => { + config.openai_base_url = format!("http://{provider_address}/v1"); + } + ProviderKind::Anthropic => { + config.anthropic_base_url = format!("http://{provider_address}"); + } + } + let identity = MachineIdentity::generate() + .expect("daemon identity") + .identity; + let generation_state = tempfile::tempdir().expect("generation state directory"); + let active_worker_generations = ActiveWorkerGenerations::load_for_test( + generation_state + .path() + .join("active-worker-generations.json"), + ) + .expect("active generation state"); + let state = Arc::new(DaemonState { + registry, + identity, + descriptor: crate::daemon::common::control::descriptor(ComponentRole::Daemon), + instance_id: "test-pass-through-daemon".into(), + public_origin: "http://127.0.0.1:1".into(), + config, + upstream: client_for(protocol), + worker_clients: WorkerClientPool::new().expect("daemon worker clients"), + challenges: Mutex::new(HashMap::new()), + activations: Mutex::new(HashMap::new()), + mcp_sessions: Mutex::new(HashMap::new()), + mcp_heartbeat_serialization: Mutex::new(()), + worker_sessions: Mutex::new(HashMap::new()), + pending_directives: Mutex::new(HashMap::new()), + active_worker_generations, + worker_generation_publication: Mutex::new(()), + }); + (router(state), generation_state) +} + +fn daemon_router_with_two_ready_workers( + protocol: TestProtocol, + route_tokens: [&str; 2], + worker_addresses: [std::net::SocketAddr; 2], + worker_tokens: [&str; 2], +) -> (Router, Vec>, tempfile::TempDir) { + let registry = Registry::new(false); + let mut targets = Vec::with_capacity(2); + + for index in 0..2 { + let credential = + RouteCredential::parse(route_tokens[index].to_owned()).expect("route credential"); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let activation_id = format!("test-activation-{index}"); + let directive = registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: McpSessionId::new(format!("test-mcp-session-{index}")) + .expect("MCP session ID"), + lease_expires_at_unix_ms: now_unix_ms().saturating_add(MCP_LEASE_MS), + }, + WorkerLaunch { + activation_id: activation_id.clone(), + activation_token: SensitiveString::new(format!("unused-activation-{index}")) + .expect("activation token"), + deadline_unix_ms: now_unix_ms().saturating_add(ACTIVATION_LIFETIME_MS), + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }, + ) + .expect("register test MCP"); + assert!(matches!(directive, BrokerDirective::LaunchWorker { .. })); + let target = Arc::new( + WorkerTarget::with_client( + format!("test-worker-{index}"), + format!("http://{}", worker_addresses[index]), + SensitiveString::new(worker_tokens[index]).expect("worker token"), + client_for(protocol), + ) + .expect("worker target"), + ); + registry + .mark_worker_ready(fingerprint, &activation_id, Arc::clone(&target)) + .expect("publish test worker"); + targets.push(target); + } + + let identity = MachineIdentity::generate() + .expect("daemon identity") + .identity; + let generation_state = tempfile::tempdir().expect("generation state directory"); + let active_worker_generations = ActiveWorkerGenerations::load_for_test( + generation_state + .path() + .join("active-worker-generations.json"), + ) + .expect("active generation state"); + let state = Arc::new(DaemonState { + registry, + identity, + descriptor: crate::daemon::common::control::descriptor(ComponentRole::Daemon), + instance_id: "test-two-route-daemon".into(), + public_origin: "http://127.0.0.1:1".into(), + config: GatewayConfig::default(), + upstream: client_for(protocol), + worker_clients: WorkerClientPool::new().expect("daemon worker clients"), + challenges: Mutex::new(HashMap::new()), + activations: Mutex::new(HashMap::new()), + mcp_sessions: Mutex::new(HashMap::new()), + mcp_heartbeat_serialization: Mutex::new(()), + worker_sessions: Mutex::new(HashMap::new()), + pending_directives: Mutex::new(HashMap::new()), + active_worker_generations, + worker_generation_publication: Mutex::new(()), + }); + (router(state), targets, generation_state) +} + +fn route_token() -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x5a; 32]) +} + +fn route_token_with(byte: u8) -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([byte; 32]) +} + +fn provider_request( + daemon_address: std::net::SocketAddr, + provider: ProviderKind, + token: &str, +) -> Request { + Request::post(format!("http://{daemon_address}{}", provider.path())) + .header(CLIENT_TOKEN_HEADER, token) + .header(AUTHORIZATION, "Bearer caller-provider-token") + .header(CONTENT_TYPE, "application/json") + .header(TE, "trailers") + .body(box_body(Full::new(Bytes::from_static( + br#"{"model":"test","stream":true}"#, + )))) + .expect("provider request") +} + +fn sequenced_provider_request( + daemon_address: std::net::SocketAddr, + provider: ProviderKind, + token: &str, + sequence: usize, +) -> Request { + Request::post(format!("http://{daemon_address}{}", provider.path())) + .header(CLIENT_TOKEN_HEADER, token) + .header( + TEST_SEQUENCE_HEADER, + HeaderValue::from_str(&sequence.to_string()).expect("valid sequence header"), + ) + .body(box_body(Empty::::new())) + .expect("sequenced provider request") +} + +fn worker_provider_request( + worker_address: std::net::SocketAddr, + provider: ProviderKind, + worker_token: &str, +) -> Request { + Request::post(format!("http://{worker_address}{}", provider.path())) + .header(WORKER_TOKEN_HEADER, worker_token) + .header(AUTHORIZATION, "Bearer caller-provider-token") + .header(CONTENT_TYPE, "application/json") + .header(TE, "trailers") + .body(box_body(Full::new(Bytes::from_static( + br#"{"model":"test","stream":true}"#, + )))) + .expect("worker provider request") +} + +async fn assert_fidelity_response(response: Response, expected: &Bytes) { + assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT); + assert_eq!(response.headers()[CONTENT_TYPE], "text/event-stream"); + assert_eq!( + response + .headers() + .get_all("x-provider") + .iter() + .map(|value| value.to_str().expect("ASCII provider header")) + .collect::>(), + ["first", "second"] + ); + let mut actual = Vec::with_capacity(expected.len()); + let mut actual_trailers = None; + let mut body = response.into_body(); + while let Some(frame) = body.frame().await { + let frame = frame.expect("fidelity frame succeeds"); + match frame.into_data() { + Ok(data) => actual.extend_from_slice(&data), + Err(frame) => { + actual_trailers = Some(frame.into_trailers().expect("only data or trailers")); + } + } + } + assert_eq!(Sha256::digest(&actual), Sha256::digest(expected)); + assert_eq!(actual.as_slice(), expected.as_ref()); + let trailers = actual_trailers.expect("fidelity trailers are preserved"); + assert_eq!( + trailers + .get_all("x-stream-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(), + ["first", "second"] + ); + assert_eq!(trailers["x-binary-safe"], "yes"); +} + +async fn wait_for_poll_plateau(polls: &AtomicUsize, frame_count: usize) -> usize { + tokio::time::timeout(Duration::from_secs(3), async { + let mut previous = usize::MAX; + let mut stable_rounds = 0; + loop { + tokio::task::yield_now().await; + let current = polls.load(Ordering::SeqCst); + assert!( + current < frame_count, + "an unread client must stop provider polling before the complete body" + ); + if current > 0 && current == previous { + stable_rounds += 1; + if stable_rounds == 32 { + return current; + } + } else { + previous = current; + stable_rounds = 0; + } + } + }) + .await + .expect("provider polling reaches a bounded plateau") +} + +async fn wait_for_in_flight_zero(worker_target: &WorkerTarget, worker_handle: &TestWorkerHandle) { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if worker_target.in_flight() == 0 && worker_handle.in_flight() == 0 { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("daemon and worker in-flight counters return to zero"); +} + +async fn read_exact_data(body: &mut Incoming, expected: &[u8]) { + let mut actual = Vec::new(); + while actual.len() < expected.len() { + let frame = body + .frame() + .await + .expect("data frame exists") + .expect("data frame succeeds"); + let data = frame.into_data().expect("expected data before trailers"); + actual.extend_from_slice(&data); + } + assert_eq!(actual, expected); +} + +async fn assert_new_request_rejected( + client: &PooledClient, + daemon_address: std::net::SocketAddr, + provider: ProviderKind, + token: &str, +) { + let response = client + .request(provider_request(daemon_address, provider, token)) + .await + .expect("daemon returns worker admission failure"); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); +} + +#[allow(clippy::cognitive_complexity)] +async fn assert_full_hop( + protocol: TestProtocol, + provider: ProviderKind, + during_stream: DuringStream, +) { + let (provider_address, release_second, observed, provider_task) = + spawn_causal_provider(protocol).await; + let (worker_router, worker_handle) = + configured_worker_router(protocol, provider, provider_address); + let (worker_address, worker_task) = spawn_router(protocol, worker_router).await; + let token = route_token(); + let (daemon_router, worker_target, _generation_state) = + daemon_router_with_ready_worker(protocol, &token, worker_address); + let (daemon_address, daemon_task) = spawn_router(protocol, daemon_router).await; + let client = client_for(protocol); + + let response = client + .request(provider_request(daemon_address, provider, &token)) + .await + .expect("full-hop request succeeds"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.headers()[CONTENT_TYPE], "text/event-stream"); + assert_eq!(response.headers()[TRAILER], "x-stream-checksum"); + assert_eq!( + response + .headers() + .get_all("x-provider") + .iter() + .map(|value| value.to_str().expect("ASCII provider header")) + .collect::>(), + ["first", "second"] + ); + let observed = observed.await.expect("provider observation"); + assert_eq!(observed.path, provider.path()); + assert_eq!( + observed.authorization.as_ref().expect("provider auth"), + "Bearer caller-provider-token" + ); + assert!(!observed.retained_client_token); + assert!(!observed.retained_worker_token); + assert_eq!(observed.body.as_ref(), br#"{"model":"test","stream":true}"#); + + let mut body = response.into_body(); + read_exact_data(&mut body, EVENT_A).await; + assert_eq!(worker_target.in_flight(), 1); + assert_eq!(worker_handle.in_flight(), 1); + + match during_stream { + DuringStream::None => {} + DuringStream::PauseBeyondFormerTotalTimeout => { + tokio::time::pause(); + tokio::time::advance(Duration::from_secs(301)).await; + } + DuringStream::ControlLoss => { + worker_handle.control_lost(); + assert_new_request_rejected(&client, daemon_address, provider, &token).await; + } + DuringStream::Drain => { + worker_handle.begin_drain(now_unix_ms().saturating_add(DRAIN_LIFETIME_MS)); + assert_new_request_rejected(&client, daemon_address, provider, &token).await; + } + } + + let next = body.frame(); + tokio::pin!(next); + assert!( + futures_util::poll!(next.as_mut()).is_pending(), + "full-hop delivery must expose event A while the provider still withholds event B" + ); + release_second.send(()).expect("release provider event B"); + let frame = next + .await + .expect("event B frame exists") + .expect("event B frame succeeds"); + let mut event_b = frame + .into_data() + .expect("event B begins in a data frame") + .to_vec(); + while event_b.len() < EVENT_B.len() { + let frame = body + .frame() + .await + .expect("remaining event B data") + .expect("remaining event B frame succeeds"); + event_b.extend_from_slice(frame.data_ref().expect("event B completes before trailers")); + } + assert_eq!(event_b, EVENT_B); + + let trailers = body + .frame() + .await + .expect("trailer frame exists") + .expect("trailer frame succeeds") + .into_trailers() + .expect("last frame contains trailers"); + assert_eq!( + trailers + .get_all("x-stream-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(), + ["one", "two"] + ); + assert!(body.frame().await.is_none()); + assert_eq!(worker_target.in_flight(), 0); + assert_eq!(worker_handle.in_flight(), 0); + + daemon_task.abort(); + worker_task.abort(); + if matches!(protocol, TestProtocol::Http1) { + provider_task.await.expect("HTTP/1.1 provider task"); + } else { + provider_task.abort(); + } +} + +#[tokio::test] +async fn full_hop_is_causally_non_aggregating_for_both_providers_over_http1() { + assert_full_hop( + TestProtocol::Http1, + ProviderKind::OpenAi, + DuringStream::None, + ) + .await; + assert_full_hop( + TestProtocol::Http1, + ProviderKind::Anthropic, + DuringStream::None, + ) + .await; +} + +#[tokio::test] +async fn full_hop_is_causally_non_aggregating_for_both_providers_over_http2() { + assert_full_hop( + TestProtocol::Http2, + ProviderKind::OpenAi, + DuringStream::None, + ) + .await; + assert_full_hop( + TestProtocol::Http2, + ProviderKind::Anthropic, + DuringStream::None, + ) + .await; +} + +#[tokio::test] +async fn full_hop_stream_has_no_former_three_hundred_second_total_timeout() { + assert_full_hop( + TestProtocol::Http1, + ProviderKind::OpenAi, + DuringStream::PauseBeyondFormerTotalTimeout, + ) + .await; +} + +#[tokio::test] +async fn admitted_full_hop_stream_survives_worker_control_loss() { + assert_full_hop( + TestProtocol::Http1, + ProviderKind::OpenAi, + DuringStream::ControlLoss, + ) + .await; +} + +#[tokio::test] +async fn admitted_full_hop_stream_survives_worker_drain() { + assert_full_hop( + TestProtocol::Http2, + ProviderKind::Anthropic, + DuringStream::Drain, + ) + .await; +} + +#[tokio::test] +async fn authenticated_codex_responses_route_preserves_pr994_method_compatibility() { + let (provider_address, release_second, observed, provider_task) = + spawn_causal_provider(TestProtocol::Http1).await; + let (worker_router, _worker_handle) = + configured_worker_router(TestProtocol::Http1, ProviderKind::OpenAi, provider_address); + let (worker_address, worker_task) = spawn_router(TestProtocol::Http1, worker_router).await; + let token = route_token(); + let (daemon_router, _worker_target, _generation_state) = + daemon_router_with_ready_worker(TestProtocol::Http1, &token, worker_address); + let (daemon_address, daemon_task) = spawn_router(TestProtocol::Http1, daemon_router).await; + let client = client_for(TestProtocol::Http1); + for path in [ + "/responses", + "/v1/responses", + "/backend-api/codex/responses", + ] { + let uri = format!("http://{daemon_address}{path}"); + let websocket_probe = Request::get(&uri) + .header(CLIENT_TOKEN_HEADER, &token) + .header(UPGRADE, "websocket") + .body(box_body(Full::new(Bytes::new()))) + .expect("WebSocket probe request"); + let response = client + .request(websocket_probe) + .await + .expect("daemon answers WebSocket probe"); + assert_eq!(response.status(), StatusCode::UPGRADE_REQUIRED, "{path}"); + response + .into_body() + .collect() + .await + .expect("WebSocket probe response body"); + + let ordinary_get = Request::get(&uri) + .header(CLIENT_TOKEN_HEADER, &token) + .body(box_body(Full::new(Bytes::new()))) + .expect("ordinary GET request"); + let response = client + .request(ordinary_get) + .await + .expect("daemon answers ordinary GET"); + assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED, "{path}"); + response + .into_body() + .collect() + .await + .expect("ordinary GET response body"); + } + + let uri = format!("http://{daemon_address}/backend-api/codex/responses"); + let mut post = provider_request(daemon_address, ProviderKind::OpenAi, &token); + *post.uri_mut() = format!("{uri}?client=codex") + .parse() + .expect("canonical Codex response URI"); + let response = client + .request(post) + .await + .expect("POST continues to forward"); + assert_eq!(response.status(), StatusCode::CREATED); + let observed = observed.await.expect("provider observes forwarded POST"); + assert_eq!(observed.path, "/v1/responses?client=codex"); + release_second.send(()).expect("release provider response"); + response + .into_body() + .collect() + .await + .expect("forwarded POST response body"); + + drop(client); + daemon_task.abort(); + worker_task.abort(); + provider_task.await.expect("HTTP/1.1 provider task"); +} + +#[tokio::test] +async fn authenticated_pi_named_endpoint_crosses_the_daemon_and_worker() { + let (provider_address, release_second, observed, provider_task) = + spawn_causal_provider(TestProtocol::Http1).await; + let worker_config = GatewayConfig { + openai_base_url: "http://127.0.0.1:1/v1".into(), + ..GatewayConfig::default() + }; + let (worker_router, _worker_handle) = worker_test_router( + worker_config, + client_for(TestProtocol::Http1), + WORKER_TOKEN.as_bytes(), + ); + let (worker_address, worker_task) = spawn_router(TestProtocol::Http1, worker_router).await; + let token = route_token(); + let (daemon_router, _worker_target, _generation_state) = + daemon_router_with_ready_worker(TestProtocol::Http1, &token, worker_address); + let (daemon_address, daemon_task) = spawn_router(TestProtocol::Http1, daemon_router).await; + let client = client_for(TestProtocol::Http1); + let mut request = provider_request(daemon_address, ProviderKind::OpenAi, &token); + *request.uri_mut() = format!("http://{daemon_address}/responses?client=pi") + .parse() + .expect("Pi Responses URI"); + request.headers_mut().insert( + crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER, + HeaderValue::from_str(&format!("http://{provider_address}/v1")).expect("named Pi endpoint"), + ); + + let response = client + .request(request) + .await + .expect("named Pi request crosses daemon and worker"); + assert_eq!(response.status(), StatusCode::CREATED); + let observed = observed.await.expect("named provider receives Pi request"); + assert_eq!(observed.path, "/v1/responses?client=pi"); + assert!(!observed.retained_client_token); + assert!(!observed.retained_worker_token); + + release_second.send(()).expect("release provider response"); + response + .into_body() + .collect() + .await + .expect("named Pi response body"); + drop(client); + daemon_task.abort(); + worker_task.abort(); + provider_task.await.expect("HTTP/1.1 provider task"); +} + +#[tokio::test] +async fn authenticated_pi_named_endpoint_is_honored_in_pass_through() { + let (provider_address, release_second, observed, provider_task) = + spawn_causal_provider(TestProtocol::Http1).await; + let token = route_token(); + let (daemon_router, _generation_state) = daemon_router_with_pass_through( + TestProtocol::Http1, + ProviderKind::OpenAi, + &token, + provider_address, + ); + let (daemon_address, daemon_task) = spawn_router(TestProtocol::Http1, daemon_router).await; + let client = client_for(TestProtocol::Http1); + let mut request = provider_request(daemon_address, ProviderKind::OpenAi, &token); + *request.uri_mut() = format!("http://{daemon_address}/responses?client=pi") + .parse() + .expect("Pi Responses URI"); + request.headers_mut().insert( + crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER, + HeaderValue::from_str(&format!("http://{provider_address}/custom/v1")) + .expect("named Pi endpoint"), + ); + + let response = client + .request(request) + .await + .expect("named Pi request crosses pass-through daemon"); + assert_eq!(response.status(), StatusCode::CREATED); + let observed = observed.await.expect("named provider receives Pi request"); + assert_eq!(observed.path, "/custom/v1/responses?client=pi"); + assert!(!observed.retained_client_token); + assert!(!observed.retained_worker_token); + + release_second.send(()).expect("release provider response"); + response + .into_body() + .collect() + .await + .expect("named Pi response body"); + drop(client); + daemon_task.abort(); + provider_task.await.expect("HTTP/1.1 provider task"); +} + +#[tokio::test] +async fn dropping_full_hop_http2_client_cancels_provider_and_releases_accounting() { + let (provider_address, provider_dropped, provider_task) = + spawn_cancellation_provider(TestProtocol::Http2).await; + let (worker_router, worker_handle) = + configured_worker_router(TestProtocol::Http2, ProviderKind::OpenAi, provider_address); + let (worker_address, worker_task) = spawn_router(TestProtocol::Http2, worker_router).await; + let token = route_token(); + let (daemon_router, worker_target, _generation_state) = + daemon_router_with_ready_worker(TestProtocol::Http2, &token, worker_address); + let (daemon_address, daemon_task) = spawn_router(TestProtocol::Http2, daemon_router).await; + let client = client_for(TestProtocol::Http2); + + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("full-hop cancellation response head"); + let mut body = response.into_body(); + read_exact_data(&mut body, b"first\n\n").await; + assert_eq!(worker_target.in_flight(), 1); + assert_eq!(worker_handle.in_flight(), 1); + + drop(body); + tokio::time::timeout(Duration::from_secs(2), provider_dropped) + .await + .expect("provider body cancellation must be prompt") + .expect("provider drop signal sent"); + wait_for_in_flight_zero(&worker_target, &worker_handle).await; + + drop(client); + daemon_task.abort(); + worker_task.abort(); + provider_task.abort(); +} + +#[tokio::test] +async fn slow_full_hop_http2_reader_applies_bounded_backpressure_and_resumes() { + const FRAME_COUNT: usize = 512; + const FRAME_SIZE: usize = 64 * 1024; + + let polls = Arc::new(AtomicUsize::new(0)); + let (provider_address, provider_task) = spawn_counted_provider( + TestProtocol::Http2, + Arc::clone(&polls), + FRAME_COUNT, + FRAME_SIZE, + ) + .await; + let (worker_router, worker_handle) = + configured_worker_router(TestProtocol::Http2, ProviderKind::OpenAi, provider_address); + let (worker_address, worker_task) = spawn_router(TestProtocol::Http2, worker_router).await; + let token = route_token(); + let (daemon_router, worker_target, _generation_state) = + daemon_router_with_ready_worker(TestProtocol::Http2, &token, worker_address); + let (daemon_address, daemon_task) = spawn_router(TestProtocol::Http2, daemon_router).await; + let client = client_for(TestProtocol::Http2); + + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("full-hop backpressure response head"); + let polls_while_unread = wait_for_poll_plateau(&polls, FRAME_COUNT).await; + assert!(polls_while_unread > 0, "provider body must begin streaming"); + + let mut received = 0; + let mut body = response.into_body(); + while let Some(frame) = body.frame().await { + received += frame + .expect("backpressure frame succeeds") + .into_data() + .expect("provider emits only data") + .len(); + } + assert_eq!(received, FRAME_COUNT * FRAME_SIZE); + assert_eq!(polls.load(Ordering::SeqCst), FRAME_COUNT + 1); + wait_for_in_flight_zero(&worker_target, &worker_handle).await; + + drop(client); + daemon_task.abort(); + worker_task.abort(); + provider_task.abort(); +} + +#[allow(clippy::cognitive_complexity)] +async fn assert_pass_through_causal(protocol: TestProtocol, provider: ProviderKind) { + let (provider_address, release_second, observed, provider_task) = + spawn_causal_provider(protocol).await; + let token = route_token(); + let (daemon_router, _generation_state) = + daemon_router_with_pass_through(protocol, provider, &token, provider_address); + let (daemon_address, daemon_task) = spawn_router(protocol, daemon_router).await; + let client = client_for(protocol); + let mut request = provider_request(daemon_address, provider, &token); + let expected_authorization = if matches!(provider, ProviderKind::OpenAi) { + *request.uri_mut() = format!("http://{daemon_address}/responses") + .parse() + .expect("generic OpenAI responses URI"); + request.headers_mut().insert( + AUTHORIZATION, + HeaderValue::from_static("Bearer at-caller-controlled-token"), + ); + "Bearer at-caller-controlled-token" + } else { + "Bearer caller-provider-token" + }; + + let response = client + .request(request) + .await + .expect("pass-through response head"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.headers()[CONTENT_TYPE], "text/event-stream"); + let observed = observed + .await + .expect("provider observes pass-through request"); + assert_eq!(observed.path, provider.path()); + assert_eq!( + observed.authorization.as_ref().expect("provider auth"), + expected_authorization + ); + assert!(!observed.retained_client_token); + assert!(!observed.retained_worker_token); + + let mut body = response.into_body(); + read_exact_data(&mut body, EVENT_A).await; + let next = body.frame(); + tokio::pin!(next); + assert!( + futures_util::poll!(next.as_mut()).is_pending(), + "pass-through must expose event A while the provider withholds event B" + ); + release_second.send(()).expect("release provider event B"); + let mut remaining = Vec::new(); + let mut trailers = None; + if let Some(frame) = next.await { + let frame = frame.expect("remaining pass-through frame succeeds"); + match frame.into_data() { + Ok(data) => remaining.extend_from_slice(&data), + Err(frame) => trailers = Some(frame.into_trailers().expect("trailers frame")), + } + } + while let Some(frame) = body.frame().await { + let frame = frame.expect("remaining pass-through frame succeeds"); + match frame.into_data() { + Ok(data) => remaining.extend_from_slice(&data), + Err(frame) => trailers = Some(frame.into_trailers().expect("trailers frame")), + } + } + assert_eq!(remaining, EVENT_B); + assert_eq!( + trailers + .expect("pass-through trailers") + .get_all("x-stream-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(), + ["one", "two"] + ); + + drop(client); + daemon_task.abort(); + if matches!(protocol, TestProtocol::Http1) { + provider_task.await.expect("HTTP/1.1 provider task"); + } else { + provider_task.abort(); + } +} + +async fn assert_pass_through_cancellation(protocol: TestProtocol) { + let (provider_address, provider_dropped, provider_task) = + spawn_cancellation_provider(protocol).await; + let token = route_token(); + let (daemon_router, _generation_state) = + daemon_router_with_pass_through(protocol, ProviderKind::OpenAi, &token, provider_address); + let (daemon_address, daemon_task) = spawn_router(protocol, daemon_router).await; + let client = client_for(protocol); + + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("pass-through cancellation response head"); + let mut body = response.into_body(); + read_exact_data(&mut body, b"first\n\n").await; + drop(body); + tokio::time::timeout(Duration::from_secs(2), provider_dropped) + .await + .expect("pass-through provider cancellation must be prompt") + .expect("provider drop signal sent"); + + drop(client); + daemon_task.abort(); + provider_task.abort(); +} + +async fn assert_pass_through_backpressure(protocol: TestProtocol) { + const FRAME_COUNT: usize = 512; + const FRAME_SIZE: usize = 64 * 1024; + + let polls = Arc::new(AtomicUsize::new(0)); + let (provider_address, provider_task) = + spawn_counted_provider(protocol, Arc::clone(&polls), FRAME_COUNT, FRAME_SIZE).await; + let token = route_token(); + let (daemon_router, _generation_state) = + daemon_router_with_pass_through(protocol, ProviderKind::OpenAi, &token, provider_address); + let (daemon_address, daemon_task) = spawn_router(protocol, daemon_router).await; + let client = client_for(protocol); + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("pass-through backpressure response head"); + + let plateau = wait_for_poll_plateau(&polls, FRAME_COUNT).await; + assert!( + plateau > 0, + "provider body begins before it is backpressured" + ); + let mut received = 0; + let mut body = response.into_body(); + while let Some(frame) = body.frame().await { + received += frame + .expect("pass-through backpressure frame succeeds") + .into_data() + .expect("provider emits only data") + .len(); + } + assert_eq!(received, FRAME_COUNT * FRAME_SIZE); + assert_eq!(polls.load(Ordering::SeqCst), FRAME_COUNT + 1); + + drop(client); + daemon_task.abort(); + provider_task.abort(); +} + +async fn assert_pass_through_fidelity(protocol: TestProtocol) { + let (provider_address, expected, provider_task) = spawn_fidelity_provider(protocol).await; + let token = route_token(); + let (daemon_router, _generation_state) = + daemon_router_with_pass_through(protocol, ProviderKind::OpenAi, &token, provider_address); + let (daemon_address, daemon_task) = spawn_router(protocol, daemon_router).await; + let client = client_for(protocol); + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("pass-through fidelity response"); + assert_fidelity_response(response, &expected).await; + + drop(client); + daemon_task.abort(); + provider_task.abort(); +} + +async fn assert_worker_only_fidelity(protocol: TestProtocol) { + let (provider_address, expected, provider_task) = spawn_fidelity_provider(protocol).await; + let (worker_router, worker_handle) = + configured_worker_router(protocol, ProviderKind::OpenAi, provider_address); + let (worker_address, worker_task) = spawn_router(protocol, worker_router).await; + let client = client_for(protocol); + let response = client + .request(worker_provider_request( + worker_address, + ProviderKind::OpenAi, + WORKER_TOKEN, + )) + .await + .expect("worker-only fidelity response"); + assert_fidelity_response(response, &expected).await; + tokio::time::timeout(Duration::from_secs(2), async { + while worker_handle.in_flight() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("worker-only in-flight count returns to zero"); + + drop(client); + worker_task.abort(); + provider_task.abort(); +} + +async fn assert_full_hop_fidelity(protocol: TestProtocol) { + let (provider_address, expected, provider_task) = spawn_fidelity_provider(protocol).await; + let (worker_router, worker_handle) = + configured_worker_router(protocol, ProviderKind::OpenAi, provider_address); + let (worker_address, worker_task) = spawn_router(protocol, worker_router).await; + let token = route_token(); + let (daemon_router, worker_target, _generation_state) = + daemon_router_with_ready_worker(protocol, &token, worker_address); + let (daemon_address, daemon_task) = spawn_router(protocol, daemon_router).await; + let client = client_for(protocol); + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("full-hop fidelity response"); + assert_fidelity_response(response, &expected).await; + wait_for_in_flight_zero(&worker_target, &worker_handle).await; + + drop(client); + daemon_task.abort(); + worker_task.abort(); + provider_task.abort(); +} + +#[tokio::test] +async fn authenticated_pass_through_is_causal_for_both_providers_over_http1() { + assert_pass_through_causal(TestProtocol::Http1, ProviderKind::OpenAi).await; + assert_pass_through_causal(TestProtocol::Http1, ProviderKind::Anthropic).await; +} + +#[tokio::test] +async fn authenticated_pass_through_is_causal_for_both_providers_over_http2() { + assert_pass_through_causal(TestProtocol::Http2, ProviderKind::OpenAi).await; + assert_pass_through_causal(TestProtocol::Http2, ProviderKind::Anthropic).await; +} + +#[tokio::test] +async fn fidelity_corpus_crosses_pass_through_worker_and_full_hop_over_http1() { + assert_pass_through_fidelity(TestProtocol::Http1).await; + assert_worker_only_fidelity(TestProtocol::Http1).await; + assert_full_hop_fidelity(TestProtocol::Http1).await; +} + +#[tokio::test] +async fn fidelity_corpus_crosses_pass_through_worker_and_full_hop_over_http2() { + assert_pass_through_fidelity(TestProtocol::Http2).await; + assert_worker_only_fidelity(TestProtocol::Http2).await; + assert_full_hop_fidelity(TestProtocol::Http2).await; +} + +#[tokio::test] +async fn pass_through_cancellation_and_backpressure_hold_over_http1() { + assert_pass_through_cancellation(TestProtocol::Http1).await; + assert_pass_through_backpressure(TestProtocol::Http1).await; +} + +#[tokio::test] +async fn pass_through_cancellation_and_backpressure_hold_over_http2() { + assert_pass_through_cancellation(TestProtocol::Http2).await; + assert_pass_through_backpressure(TestProtocol::Http2).await; +} + +async fn assert_128_concurrent_full_hop_streams(protocol: TestProtocol) { + const STREAMS: usize = 128; + const ROUTE_STREAMS: usize = STREAMS / 2; + + let provider_connections = [Arc::new(AtomicUsize::new(0)), Arc::new(AtomicUsize::new(0))]; + let (provider_address_a, provider_task_a) = spawn_sequenced_provider( + protocol, + ROUTE_STREAMS, + Arc::clone(&provider_connections[0]), + ) + .await; + let (provider_address_b, provider_task_b) = spawn_sequenced_provider( + protocol, + ROUTE_STREAMS, + Arc::clone(&provider_connections[1]), + ) + .await; + let worker_tokens = ["worker-route-a-token", "worker-route-b-token"]; + let (worker_router_a, worker_handle_a) = configured_worker_router_with_token( + protocol, + ProviderKind::OpenAi, + provider_address_a, + worker_tokens[0], + ); + let (worker_router_b, worker_handle_b) = configured_worker_router_with_token( + protocol, + ProviderKind::OpenAi, + provider_address_b, + worker_tokens[1], + ); + let (worker_address_a, worker_task_a) = spawn_router(protocol, worker_router_a).await; + let (worker_address_b, worker_task_b) = spawn_router(protocol, worker_router_b).await; + let tokens = [route_token_with(0x5a), route_token_with(0xa5)]; + let (daemon_router, worker_targets, _generation_state) = daemon_router_with_two_ready_workers( + protocol, + [&tokens[0], &tokens[1]], + [worker_address_a, worker_address_b], + worker_tokens, + ); + let (daemon_address, daemon_task) = spawn_router(protocol, daemon_router).await; + let client = client_for(protocol); + + let mut tasks = Vec::with_capacity(STREAMS); + for sequence in 0..STREAMS { + let client = client.clone(); + let token = tokens[sequence % tokens.len()].clone(); + tasks.push(tokio::spawn(async move { + let response = client + .request(sequenced_provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + sequence, + )) + .await + .expect("concurrent full-hop request succeeds"); + assert_eq!(response.status(), StatusCode::OK); + let mut actual = Vec::new(); + let mut body = response.into_body(); + while let Some(frame) = body.frame().await { + actual.extend_from_slice( + &frame + .expect("concurrent full-hop frame succeeds") + .into_data() + .expect("sequenced provider emits only data"), + ); + } + let expected = (0..SEQUENCE_PARTS) + .flat_map(|part| sequence_chunk(sequence, part)) + .collect::>(); + assert_eq!(actual, expected); + })); + } + tokio::time::timeout(Duration::from_secs(30), async { + for task in tasks { + task.await.expect("stream verification task succeeds"); + } + }) + .await + .expect("all 128 requests become concurrent and complete"); + + wait_for_in_flight_zero(&worker_targets[0], &worker_handle_a).await; + wait_for_in_flight_zero(&worker_targets[1], &worker_handle_b).await; + for (route, connections) in provider_connections.iter().enumerate() { + match protocol { + TestProtocol::Http1 => assert_eq!( + connections.load(Ordering::SeqCst), + ROUTE_STREAMS, + "each route's blocked HTTP/1.1 streams require independent pooled connections" + ), + TestProtocol::Http2 => assert_eq!( + connections.load(Ordering::SeqCst), + 1, + "route {route} must multiplex all streams over its own provider connection" + ), + } + } + + drop(client); + daemon_task.abort(); + worker_task_a.abort(); + worker_task_b.abort(); + provider_task_a.abort(); + provider_task_b.abort(); +} + +#[tokio::test] +async fn full_hop_keeps_128_concurrent_http2_streams_isolated() { + assert_128_concurrent_full_hop_streams(TestProtocol::Http2).await; +} + +#[tokio::test] +async fn full_hop_keeps_128_concurrent_http1_streams_isolated() { + assert_128_concurrent_full_hop_streams(TestProtocol::Http1).await; +} + +struct LifecycleDaemonHarness { + state: Arc, + fingerprint: Fingerprint, + mcp_session_id: String, + mcp_secret: SensitiveString, + worker_id: String, + worker_control_secret: SensitiveString, + generation_id: String, + target: Arc, + _generation_state: tempfile::TempDir, +} + +fn lifecycle_daemon_router( + route_token: &str, + worker_address: std::net::SocketAddr, +) -> (Router, LifecycleDaemonHarness) { + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let credential = RouteCredential::parse(route_token.to_owned()).expect("route credential"); + let registry = Registry::new(false); + let mcp_session_id = "lifecycle-mcp-session".to_owned(); + let activation_id = "lifecycle-activation"; + registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: McpSessionId::new(mcp_session_id.clone()).expect("MCP session ID"), + lease_expires_at_unix_ms: now_unix_ms().saturating_add(MCP_LEASE_MS), + }, + WorkerLaunch { + activation_id: activation_id.into(), + activation_token: SensitiveString::new("unused-lifecycle-activation-token") + .expect("activation token"), + deadline_unix_ms: now_unix_ms().saturating_add(ACTIVATION_LIFETIME_MS), + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }, + ) + .expect("register lifecycle MCP"); + + let worker_id = "test-worker".to_owned(); + let endpoint = format!("http://{worker_address}"); + let target = Arc::new( + WorkerTarget::with_client( + worker_id.clone(), + endpoint.clone(), + SensitiveString::new(WORKER_TOKEN).expect("worker data token"), + client_for(TestProtocol::Http1), + ) + .expect("worker target"), + ); + registry + .mark_worker_ready(fingerprint, activation_id, Arc::clone(&target)) + .expect("publish lifecycle worker"); + + let daemon_identity = MachineIdentity::generate() + .expect("daemon identity") + .identity; + let generation_grant = + WorkerGenerationGrant::issue(&worker_id, fingerprint, &endpoint, None, &daemon_identity) + .expect("worker generation grant"); + let generation_id = generation_grant.generation_id.clone(); + let generation_state = tempfile::tempdir().expect("generation state directory"); + let active_worker_generations = ActiveWorkerGenerations::load_for_test( + generation_state + .path() + .join("active-worker-generations.json"), + ) + .expect("active generation state"); + active_worker_generations + .publish(fingerprint, &generation_id) + .expect("publish active generation"); + + let mcp_secret = SensitiveString::new("lifecycle-mcp-control-token").expect("MCP token"); + let worker_control_secret = + SensitiveString::new("unused-test-control-token").expect("worker control token"); + let state = Arc::new(DaemonState { + registry, + identity: daemon_identity, + descriptor: crate::daemon::common::control::descriptor(ComponentRole::Daemon), + instance_id: "lifecycle-daemon".into(), + public_origin: "http://127.0.0.1:1".into(), + config: GatewayConfig::default(), + upstream: pooled_client().expect("daemon pass-through client"), + worker_clients: WorkerClientPool::new().expect("daemon worker clients"), + challenges: Mutex::new(HashMap::new()), + activations: Mutex::new(HashMap::new()), + mcp_sessions: Mutex::new(HashMap::from([( + mcp_session_id.clone(), + McpControlSession { + fingerprint, + token_digest: credential.digest(), + secret: mcp_secret.clone(), + secret_digest: TokenDigest::from_token(mcp_secret.expose().as_bytes()), + lease_expires_at_unix_ms: now_unix_ms().saturating_add(MCP_LEASE_MS), + last_sequence: 0, + last_request_id: String::new(), + last_heartbeat: None, + worker_network: WorkerNetworkHint { + advertised_host: Ipv4Addr::LOCALHOST.to_string(), + port: None, + }, + released: false, + }, + )])), + mcp_heartbeat_serialization: Mutex::new(()), + worker_sessions: Mutex::new(HashMap::from([( + worker_id.clone(), + WorkerControlSession { + fingerprint, + worker_id: worker_id.clone(), + secret: worker_control_secret.clone(), + secret_digest: TokenDigest::from_token(worker_control_secret.expose().as_bytes()), + last_sequence: 0, + last_request_id: String::new(), + next_daemon_sequence: 0, + lease_expires_at_unix_ms: now_unix_ms().saturating_add(WORKER_LEASE_MS), + pending_target: Arc::clone(&target), + publication: WorkerPublication::Activation { + activation_id: activation_id.into(), + }, + published: true, + generation_grant, + }, + )])), + pending_directives: Mutex::new(HashMap::new()), + active_worker_generations, + worker_generation_publication: Mutex::new(()), + }); + ( + router(Arc::clone(&state)), + LifecycleDaemonHarness { + state, + fingerprint, + mcp_session_id, + mcp_secret, + worker_id, + worker_control_secret, + generation_id, + target, + _generation_state: generation_state, + }, + ) +} + +async fn finish_causal_lifecycle_stream(body: &mut Incoming, release_second: oneshot::Sender<()>) { + release_second.send(()).expect("release provider event B"); + let mut data = Vec::new(); + let mut trailers = None; + while let Some(frame) = body.frame().await { + let frame = frame.expect("remaining lifecycle stream frame succeeds"); + match frame.into_data() { + Ok(bytes) => data.extend_from_slice(&bytes), + Err(frame) => { + trailers = Some( + frame + .into_trailers() + .expect("remaining lifecycle frame is trailers"), + ); + } + } + } + assert_eq!(data, EVENT_B); + assert_eq!( + trailers + .expect("lifecycle stream trailers") + .get_all("x-stream-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(), + ["one", "two"] + ); +} + +async fn wait_for_worker_drain_control( + client: &PooledClient, + worker_address: std::net::SocketAddr, +) { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let request = Request::get(format!("http://{worker_address}{WORKER_PROBE_PATH}")) + .header(WORKER_TOKEN_HEADER, WORKER_TOKEN) + .body(box_body(Empty::::new())) + .expect("worker readiness request"); + let response = client + .request(request) + .await + .expect("worker readiness response"); + let status = response.status(); + response + .into_body() + .collect() + .await + .expect("worker readiness body"); + if status == StatusCode::SERVICE_UNAVAILABLE { + return; + } + assert_eq!(status, StatusCode::NO_CONTENT); + tokio::task::yield_now().await; + } + }) + .await + .expect("daemon drain control reaches worker"); +} + +#[tokio::test] +async fn broker_release_enters_draining_while_admitted_stream_finishes() { + let (provider_address, release_second, _observed, provider_task) = + spawn_causal_provider(TestProtocol::Http1).await; + let (worker_router, worker_handle) = + configured_worker_router(TestProtocol::Http1, ProviderKind::OpenAi, provider_address); + let (worker_address, worker_task) = spawn_router(TestProtocol::Http1, worker_router).await; + let token = route_token(); + let (daemon_router, harness) = lifecycle_daemon_router(&token, worker_address); + let (daemon_address, daemon_task) = spawn_router(TestProtocol::Http1, daemon_router).await; + let client = client_for(TestProtocol::Http1); + + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("admitted stream response head"); + assert_eq!(response.status(), StatusCode::CREATED); + let mut body = response.into_body(); + read_exact_data(&mut body, EVENT_A).await; + assert_eq!(harness.target.in_flight(), 1); + assert_eq!(worker_handle.in_flight(), 1); + + let release = SessionRequest::new( + harness.mcp_session_id.clone(), + harness.mcp_secret.clone(), + 1, + EmptyPayload::default(), + ) + .expect("MCP release request"); + let request = Request::post(format!("http://{daemon_address}{MCP_RELEASE_PATH}")) + .header(CONTENT_TYPE, "application/json") + .body(box_body(Full::new(Bytes::from( + serde_json::to_vec(&release).expect("serialize MCP release"), + )))) + .expect("MCP release HTTP request"); + let release_response = client.request(request).await.expect("MCP release response"); + assert_eq!(release_response.status(), StatusCode::NO_CONTENT); + release_response + .into_body() + .collect() + .await + .expect("MCP release body"); + + assert_eq!( + harness + .state + .registry + .snapshot(harness.fingerprint) + .expect("draining route") + .state, + crate::daemon::broker::lifecycle::RouteStateKind::Draining + ); + assert!( + !harness + .state + .active_worker_generations + .matches(harness.fingerprint, &harness.generation_id) + .expect("generation revocation") + ); + assert_new_request_rejected(&client, daemon_address, ProviderKind::OpenAi, &token).await; + wait_for_worker_drain_control(&client, worker_address).await; + + finish_causal_lifecycle_stream(&mut body, release_second).await; + wait_for_in_flight_zero(&harness.target, &worker_handle).await; + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if harness + .state + .registry + .snapshot(harness.fingerprint) + .expect("drain completion route") + .state + == crate::daemon::broker::lifecycle::RouteStateKind::Empty + { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("draining route returns to empty after the accepted stream completes"); + + drop(client); + daemon_task.abort(); + worker_task.abort(); + provider_task.await.expect("HTTP/1.1 provider task"); +} + +#[tokio::test] +async fn broker_worker_heartbeat_expiry_rejects_new_work_but_preserves_admitted_stream() { + let (provider_address, release_second, _observed, provider_task) = + spawn_causal_provider(TestProtocol::Http1).await; + let (worker_router, worker_handle) = + configured_worker_router(TestProtocol::Http1, ProviderKind::OpenAi, provider_address); + let (worker_address, worker_task) = spawn_router(TestProtocol::Http1, worker_router).await; + let token = route_token(); + let (daemon_router, harness) = lifecycle_daemon_router(&token, worker_address); + let (daemon_address, daemon_task) = spawn_router(TestProtocol::Http1, daemon_router).await; + let client = client_for(TestProtocol::Http1); + + let heartbeat = SessionRequest::new( + harness.worker_id.clone(), + harness.worker_control_secret.clone(), + 1, + WorkerHeartbeatPayload { + worker_id: harness.worker_id.clone(), + }, + ) + .expect("worker heartbeat request"); + let request = Request::post(format!("http://{daemon_address}{WORKER_HEARTBEAT_PATH}")) + .header(CONTENT_TYPE, "application/json") + .body(box_body(Full::new(Bytes::from( + serde_json::to_vec(&heartbeat).expect("serialize worker heartbeat"), + )))) + .expect("worker heartbeat HTTP request"); + let heartbeat_response = client + .request(request) + .await + .expect("worker heartbeat response"); + assert_eq!(heartbeat_response.status(), StatusCode::NO_CONTENT); + heartbeat_response + .into_body() + .collect() + .await + .expect("worker heartbeat body"); + + let response = client + .request(provider_request( + daemon_address, + ProviderKind::OpenAi, + &token, + )) + .await + .expect("admitted stream response head"); + assert_eq!(response.status(), StatusCode::CREATED); + let mut body = response.into_body(); + read_exact_data(&mut body, EVENT_A).await; + assert_eq!(harness.target.in_flight(), 1); + assert_eq!(worker_handle.in_flight(), 1); + + lock(&harness.state.worker_sessions) + .get_mut(&harness.worker_id) + .expect("live worker control session") + .lease_expires_at_unix_ms = now_unix_ms().saturating_sub(1); + spawn_maintenance(Arc::clone(&harness.state)); + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let worker_expired = + !lock(&harness.state.worker_sessions).contains_key(&harness.worker_id); + let replacement_activating = harness + .state + .registry + .snapshot(harness.fingerprint) + .is_ok_and(|snapshot| { + snapshot.state == crate::daemon::broker::lifecycle::RouteStateKind::Activating + }); + let relaunch_pending = matches!( + lock(&harness.state.pending_directives).get(&harness.mcp_session_id), + Some(BrokerDirective::LaunchWorker { .. }) + ); + if worker_expired && replacement_activating && relaunch_pending { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("expired worker heartbeat removes the control session"); + + assert_eq!( + harness + .state + .registry + .snapshot(harness.fingerprint) + .expect("recovering route") + .state, + crate::daemon::broker::lifecycle::RouteStateKind::Activating + ); + assert!(matches!( + lock(&harness.state.pending_directives).get(&harness.mcp_session_id), + Some(BrokerDirective::LaunchWorker { .. }) + )); + assert!( + !harness + .state + .active_worker_generations + .matches(harness.fingerprint, &harness.generation_id) + .expect("expired generation revocation") + ); + assert_new_request_rejected(&client, daemon_address, ProviderKind::OpenAi, &token).await; + + finish_causal_lifecycle_stream(&mut body, release_second).await; + wait_for_in_flight_zero(&harness.target, &worker_handle).await; + + drop(client); + daemon_task.abort(); + worker_task.abort(); + provider_task.await.expect("HTTP/1.1 provider task"); +} diff --git a/crates/cli/tests/coverage/daemon/hook_tests.rs b/crates/cli/tests/coverage/daemon/hook_tests.rs new file mode 100644 index 000000000..2f8c169a6 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/hook_tests.rs @@ -0,0 +1,282 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::io::{Read, Write}; +use std::net::TcpListener; +use std::sync::{Arc, Mutex}; + +use base64::Engine; + +use super::*; + +fn valid_token() -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([7_u8; 32]) +} + +fn capture_server(response: Vec) -> (String, Arc>>, std::thread::JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let request = Arc::new(Mutex::new(Vec::new())); + let captured = Arc::clone(&request); + let handle = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut bytes = Vec::new(); + let mut byte = [0_u8; 1]; + while !bytes.ends_with(b"\r\n\r\n") { + stream.read_exact(&mut byte).unwrap(); + bytes.push(byte[0]); + } + let headers = String::from_utf8_lossy(&bytes); + let content_length = headers + .lines() + .find_map(|line| { + line.split_once(':').and_then(|(name, value)| { + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().unwrap()) + }) + }) + .unwrap(); + let mut body = vec![0_u8; content_length]; + stream.read_exact(&mut body).unwrap(); + bytes.extend_from_slice(&body); + *captured.lock().unwrap() = bytes; + if let Err(error) = stream.write_all(&response) { + assert!( + matches!( + error.kind(), + std::io::ErrorKind::BrokenPipe | std::io::ErrorKind::ConnectionReset + ), + "response write failed: {error}" + ); + } + }); + (format!("http://{address}"), request, handle) +} + +#[test] +fn hook_payload_is_bounded_and_empty_input_is_normalized() { + assert_eq!(read_hook_payload(&b" \n\t"[..]).unwrap(), b"{}"); + assert_eq!(read_hook_payload(&b"{\"x\":1}"[..]).unwrap(), b"{\"x\":1}"); + + let oversized = vec![b'x'; crate::configuration::DEFAULT_MAX_HOOK_PAYLOAD_BYTES + 1]; + let error = read_hook_payload(oversized.as_slice()).unwrap_err(); + assert!(error.to_string().contains("exceeds"), "{error}"); +} + +#[test] +fn default_failure_policy_is_event_specific() { + assert!(effective_fail_closed( + HookFailurePolicy::Default, + Some(br#"{"hook_event_name":"PreToolUse"}"#), + )); + assert!(effective_fail_closed( + HookFailurePolicy::Default, + Some(br#"{"hook_event_name":"pre_tool_call"}"#), + )); + for event in ["tool_call", "toolCall", "user_bash", "userBash"] { + let payload = format!(r#"{{"hook_event_name":"{event}"}}"#); + assert!( + effective_fail_closed(HookFailurePolicy::Default, Some(payload.as_bytes())), + "managed Pi policy event must fail closed: {event}" + ); + } + assert!(!effective_fail_closed( + HookFailurePolicy::Default, + Some(br#"{"hook_event_name":"PostToolUse"}"#), + )); + assert!(!effective_fail_closed( + HookFailurePolicy::FailOpen, + Some(br#"{"hook_event_name":"PreToolUse"}"#), + )); + assert!(effective_fail_closed( + HookFailurePolicy::FailClosed, + Some(br#"{"hook_event_name":"Stop"}"#), + )); +} + +#[test] +fn route_token_requires_exactly_256_bits_without_exposing_the_value() { + let secret = "not-a-route-credential"; + let error = route_token(secret).unwrap_err().to_string(); + assert!(error.contains(CLIENT_TOKEN_ENV), "{error}"); + assert!(!error.contains(secret), "{error}"); + assert!(route_token(&valid_token()).is_ok()); +} + +#[test] +fn managed_pi_hook_uses_the_existing_root_path() { + let endpoint = hook_endpoint("https://relay.example.com:443", CodingAgent::Pi).unwrap(); + assert_eq!(endpoint.as_str(), "https://relay.example.com/hooks/pi"); +} + +#[test] +fn managed_hook_never_sends_its_route_credential_over_remote_cleartext() { + assert!(hook_endpoint("http://relay.example.com:47632", CodingAgent::Pi).is_err()); + assert!(hook_endpoint("http://127.0.0.1:47632", CodingAgent::Pi).is_ok()); +} + +#[tokio::test] +async fn hook_forward_uses_the_exact_agent_path_and_one_route_header() { + let response_body = b"{\"continue\":true}\n"; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + response_body.len(), + String::from_utf8_lossy(response_body) + ); + let (daemon_address, captured, server) = capture_server(response.into_bytes()); + let options = Options { + agent: CodingAgent::ClaudeCode, + daemon_address, + failure_policy: HookFailurePolicy::FailClosed, + }; + + let body = forward( + &options, + b"{\"hook_event_name\":\"Stop\"}".to_vec(), + route_token(&valid_token()).unwrap(), + ) + .await + .unwrap(); + server.join().unwrap(); + + assert_eq!(body, response_body); + let request = String::from_utf8(captured.lock().unwrap().clone()).unwrap(); + assert!( + request.starts_with("POST /hooks/claude-code HTTP/1.1\r\n"), + "{request}" + ); + assert_eq!( + request + .lines() + .filter(|line| { + line.split_once(':') + .is_some_and(|(name, _)| name.eq_ignore_ascii_case(CLIENT_TOKEN_HEADER)) + }) + .count(), + 1, + "{request}" + ); + assert!( + request.ends_with("{\"hook_event_name\":\"Stop\"}"), + "{request}" + ); +} + +#[tokio::test] +async fn guardrail_rejections_are_never_failed_open() { + let body = r#"{"error":{"type":"nemo_relay_guardrail_rejected","reason":"blocked"}}"#; + let response = format!( + "HTTP/1.1 403 Forbidden\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ); + let (daemon_address, _captured, server) = capture_server(response.into_bytes()); + let options = Options { + agent: CodingAgent::Codex, + daemon_address, + failure_policy: HookFailurePolicy::FailOpen, + }; + + let error = forward( + &options, + b"{}".to_vec(), + route_token(&valid_token()).unwrap(), + ) + .await + .unwrap_err(); + server.join().unwrap(); + assert_eq!(error.guardrail_rejection_reason(), Some("blocked")); +} + +#[test] +fn hook_validation_covers_malformed_payloads_urls_and_default_events() { + assert!(read_hook_payload(&[0xff][..]).is_err()); + assert!(!effective_fail_closed(HookFailurePolicy::Default, None)); + assert!(!effective_fail_closed( + HookFailurePolicy::Default, + Some(b"not-json") + )); + assert!(!effective_fail_closed( + HookFailurePolicy::Default, + Some(br#"{"unrelated":true}"#) + )); + + for address in [ + "not a URL", + "https://user@relay.example:443", + "https://relay.example:443/path", + "https://relay.example:443?query=true", + "https://relay.example:443/#fragment", + ] { + assert!( + hook_endpoint(address, CodingAgent::Pi).is_err(), + "accepted {address}" + ); + } +} + +#[test] +fn guardrail_error_decoder_accepts_message_fallback_and_rejects_other_shapes() { + assert_eq!( + guardrail_rejection_reason( + br#"{"error":{"type":"nemo_relay_guardrail_rejected","message":"fallback"}}"# + ), + Some("fallback".into()) + ); + for body in [ + &b"not-json"[..], + &br#"{}"#[..], + &br#"{"error":{"type":"other","reason":"no"}}"#[..], + &br#"{"error":{"type":"nemo_relay_guardrail_rejected"}}"#[..], + ] { + assert_eq!(guardrail_rejection_reason(body), None); + } +} + +#[test] +fn delivery_failure_policy_wraps_closed_errors_and_swallows_open_errors() { + assert!(handle_delivery_failure(CliError::Config("closed".into()), true).is_err()); + assert!(handle_delivery_failure(CliError::Config("open".into()), false).is_ok()); +} + +#[tokio::test] +async fn non_guardrail_http_failures_and_oversized_responses_are_rejected() { + let response = + b"HTTP/1.1 503 Service Unavailable\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{}"; + let (daemon_address, _captured, server) = capture_server(response.to_vec()); + let error = forward( + &Options { + agent: CodingAgent::Pi, + daemon_address, + failure_policy: HookFailurePolicy::FailClosed, + }, + b"{}".to_vec(), + route_token(&valid_token()).unwrap(), + ) + .await + .unwrap_err(); + server.join().unwrap(); + assert!(error.to_string().contains("503")); + + let large = vec![b'x'; MAX_HOOK_RESPONSE_BYTES + 1]; + let mut response = format!( + "HTTP/1.1 200 OK\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + large.len() + ) + .into_bytes(); + response.extend_from_slice(&large); + let (daemon_address, _captured, server) = capture_server(response); + let error = forward( + &Options { + agent: CodingAgent::Pi, + daemon_address, + failure_policy: HookFailurePolicy::FailClosed, + }, + b"{}".to_vec(), + route_token(&valid_token()).unwrap(), + ) + .await + .unwrap_err(); + server.join().unwrap(); + assert!(matches!(error, CliError::PayloadTooLarge(_))); +} diff --git a/crates/cli/tests/coverage/daemon/identity_tests.rs b/crates/cli/tests/coverage/daemon/identity_tests.rs new file mode 100644 index 000000000..cea3c9539 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/identity_tests.rs @@ -0,0 +1,109 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn generated_identity_round_trips_and_verifies() { + let generated = MachineIdentity::generate().expect("identity"); + let restored = MachineIdentity::from_pkcs8(&generated.pkcs8).expect("restored identity"); + assert_eq!(generated.identity.fingerprint(), restored.fingerprint()); + + let transcript = b"canonical transcript"; + let signature = generated.identity.sign(transcript); + restored + .public_identity() + .verify(transcript, &signature) + .expect("valid signature"); + assert_eq!( + restored + .public_identity() + .verify(b"changed transcript", &signature), + Err(IdentityError::SignatureVerification) + ); +} + +#[test] +fn public_identity_rejects_the_wrong_length() { + assert_eq!( + PublicIdentity::from_bytes(&[0_u8; 31]), + Err(IdentityError::InvalidPublicKey) + ); +} + +#[test] +fn fingerprint_and_token_digest_are_stable() { + let public = PublicIdentity::from_bytes(&[7_u8; 32]).expect("public key"); + assert_eq!( + public.fingerprint().to_string(), + "4bb06f8e4e3a7715d201d573d0aa423762e55dabd61a2c02278fa56cc6d294e0" + ); + assert_eq!( + TokenDigest::from_token(b"route token").to_string(), + "fdd50053ddd4f9762b19d688e79add7403e4c354bb81430aaf25d3041f5c84e3" + ); +} + +#[test] +fn transcript_encoding_is_domain_separated_and_length_prefixed() { + let encoded = encode_transcript(b"test", &[("a", b"b"), ("cd", b"ef")]).expect("transcript"); + let mut expected = TRANSCRIPT_MAGIC.to_vec(); + expected.extend_from_slice(&4_u64.to_be_bytes()); + expected.extend_from_slice(b"test"); + expected.extend_from_slice(&2_u32.to_be_bytes()); + expected.extend_from_slice(&1_u64.to_be_bytes()); + expected.extend_from_slice(b"a"); + expected.extend_from_slice(&1_u64.to_be_bytes()); + expected.extend_from_slice(b"b"); + expected.extend_from_slice(&2_u64.to_be_bytes()); + expected.extend_from_slice(b"cd"); + expected.extend_from_slice(&2_u64.to_be_bytes()); + expected.extend_from_slice(b"ef"); + assert_eq!(encoded, expected); + + let other_domain = + encode_transcript(b"other", &[("a", b"b"), ("cd", b"ef")]).expect("transcript"); + assert_ne!(encoded, other_domain); +} + +#[test] +fn challenge_is_single_use_and_expires_at_the_boundary() { + let challenge = Challenge { + id: ChallengeId([1; CHALLENGE_ID_BYTES]), + nonce: ChallengeNonce([2; CHALLENGE_NONCE_BYTES]), + issued_at_unix_ms: 100, + expires_at_unix_ms: 200, + }; + let mut record = ChallengeRecord::from_challenge(challenge); + assert_eq!(record.consume(&challenge.id, 199), Ok(challenge)); + assert_eq!( + record.consume(&challenge.id, 199), + Err(ChallengeError::Replay) + ); + + let mut expired = ChallengeRecord::from_challenge(challenge); + assert_eq!( + expired.consume(&challenge.id, 200), + Err(ChallengeError::Expired) + ); + assert_eq!( + expired.consume(&challenge.id, 199), + Err(ChallengeError::Replay) + ); +} + +#[test] +fn mismatched_challenge_does_not_consume_record() { + let challenge = Challenge { + id: ChallengeId([1; CHALLENGE_ID_BYTES]), + nonce: ChallengeNonce([2; CHALLENGE_NONCE_BYTES]), + issued_at_unix_ms: 100, + expires_at_unix_ms: 200, + }; + let mut record = ChallengeRecord::from_challenge(challenge); + assert_eq!( + record.consume(&ChallengeId([3; CHALLENGE_ID_BYTES]), 150), + Err(ChallengeError::IdentifierMismatch) + ); + assert_eq!(record.consume(&challenge.id, 150), Ok(challenge)); +} diff --git a/crates/cli/tests/coverage/daemon/lifecycle_tests.rs b/crates/cli/tests/coverage/daemon/lifecycle_tests.rs new file mode 100644 index 000000000..16743b6e2 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/lifecycle_tests.rs @@ -0,0 +1,58 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::daemon::common::identity::PublicIdentity; + +fn fingerprint() -> Fingerprint { + PublicIdentity::from_bytes(&[9; 32]) + .expect("identity") + .fingerprint() +} + +#[test] +fn worker_request_accounts_for_exact_lifetime() { + let target = Arc::new( + WorkerTarget::new( + "worker-1", + "http://127.0.0.1:41000", + SensitiveString::new("worker-secret").expect("secret"), + ) + .expect("target"), + ); + assert_eq!(target.in_flight(), 0); + let request = target.acquire(fingerprint()); + assert_eq!(target.in_flight(), 1); + assert_eq!(request.fingerprint(), fingerprint()); + assert_eq!(request.target().endpoint(), "http://127.0.0.1:41000"); + assert_eq!(request.session_token(), "worker-secret"); + assert!(!format!("{request:?}").contains("worker-secret")); + drop(request); + assert_eq!(target.in_flight(), 0); +} + +#[test] +fn identifiers_and_targets_reject_empty_values() { + assert_eq!(McpSessionId::new(""), Err(LifecycleError::EmptyIdentifier)); + assert_eq!( + McpSessionId::new("mcp-1").expect("session").as_str(), + "mcp-1" + ); + assert_eq!( + WorkerTarget::new( + "", + "http://127.0.0.1:1", + SensitiveString::new("secret").expect("secret") + ) + .unwrap_err(), + LifecycleError::EmptyIdentifier + ); + assert_eq!( + WorkerTarget::new("worker", "", SensitiveString::new("secret").unwrap()).unwrap_err(), + LifecycleError::EmptyIdentifier + ); + assert!(matches!( + SensitiveString::new(""), + Err(crate::daemon::common::protocol::ProtocolError::MissingSensitiveValue) + )); +} diff --git a/crates/cli/tests/coverage/daemon/managed_tests.rs b/crates/cli/tests/coverage/daemon/managed_tests.rs new file mode 100644 index 000000000..d8a71d162 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/managed_tests.rs @@ -0,0 +1,718 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::ffi::OsStr; + +use base64::Engine; +use tempfile::tempdir; + +use super::*; +use crate::test_support::EnvScope; + +fn credential() -> String { + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x42_u8; 32]) +} + +fn spec(agents: impl IntoIterator) -> ManagedBundleSpec { + ManagedBundleSpec::new( + "https://relay.example.com:443", + "/opt/nvidia/bin/nemo-relay-dispatch", + ManagedPlatform::Linux, + agents, + ) + .unwrap() +} + +fn managed_environment(token: &str) -> EnvScope { + let header = format!("x-enterprise-context: fixed\n{ROUTE_TOKEN_HEADER}: {token}"); + EnvScope::set(&[ + (ROUTE_TOKEN_ENV, Some(OsStr::new(token))), + (CLAUDE_CUSTOM_HEADERS_ENV, Some(OsStr::new(&header))), + ]) +} + +fn make_test_artifact_writable(path: &std::path::Path) { + let mut permissions = std::fs::metadata(path).unwrap().permissions(); + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt as _; + permissions.set_mode(0o600); + } + #[cfg(not(unix))] + permissions.set_readonly(false); + std::fs::set_permissions(path, permissions).unwrap(); +} + +#[test] +fn v1_render_is_deterministic_and_contains_only_deployment_constants() { + let spec = spec([ + ManagedAgent::Codex, + ManagedAgent::ClaudeCode, + ManagedAgent::Pi, + ]); + let first = render_bundle(&spec).unwrap(); + let second = render_bundle(&spec).unwrap(); + assert_eq!(first.manifest, second.manifest); + assert_eq!(first.artifacts.len(), 13); + for (left, right) in first.artifacts.iter().zip(&second.artifacts) { + assert_eq!(left.path, right.path); + assert_eq!(left.bytes, right.bytes); + let text = String::from_utf8(left.bytes.clone()).unwrap(); + if left.path != "pi/extension-v1/package.json" && env!("CARGO_PKG_VERSION") != "1.0.0" { + assert!(!text.contains(env!("CARGO_PKG_VERSION"))); + } + for forbidden in [ + "machine-identity", + "generation_token", + "/Users/", + "C:\\Users\\", + ] { + assert!( + !text.contains(forbidden), + "{} contained {forbidden}", + left.path + ); + } + } + + let codex_settings = first + .artifacts + .iter() + .find(|artifact| artifact.path == "codex/settings-v1/config.toml") + .unwrap(); + let codex_settings = String::from_utf8_lossy(&codex_settings.bytes); + assert!(codex_settings.contains("https://relay.example.com:443/v1")); + assert!(codex_settings.contains(ROUTE_TOKEN_HEADER)); + assert!(codex_settings.contains(ROUTE_TOKEN_ENV)); + assert!(!codex_settings.contains(&credential())); + + for artifact in first + .artifacts + .iter() + .filter(|artifact| artifact.path.ends_with("hooks/hooks.json")) + { + let hooks = String::from_utf8_lossy(&artifact.bytes); + assert!(hooks.contains("/opt/nvidia/bin/nemo-relay-dispatch daemon hook")); + assert!(hooks.contains("--daemon-address https://relay.example.com:443")); + assert!(!hooks.contains("hook-forward")); + } + + assert_pi_bundle_artifacts(&first); +} + +fn assert_pi_bundle_artifacts(bundle: &RenderedBundle) { + let pi_config = bundle + .artifacts + .iter() + .find(|artifact| artifact.path == "pi/extension-v1/managed-config.json") + .unwrap(); + let pi_config: serde_json::Value = serde_json::from_slice(&pi_config.bytes).unwrap(); + assert_eq!(pi_config["daemonAddress"], "https://relay.example.com:443"); + assert_eq!( + pi_config["dispatcherCommand"], + "/opt/nvidia/bin/nemo-relay-dispatch" + ); + + for (path, expected) in [ + ( + "pi/extension-v1/README.md", + canonical_embedded_text(include_str!( + "../../../src/daemon/managed/pi_extension/README.md" + )) + .into_bytes(), + ), + ( + "pi/extension-v1/index.ts", + canonical_embedded_text(include_str!( + "../../../src/daemon/managed/pi_extension/index.ts" + )) + .into_bytes(), + ), + ( + "pi/extension-v1/package.json", + canonical_embedded_text(include_str!( + "../../../src/daemon/managed/pi_extension/package.json" + )) + .into_bytes(), + ), + ( + "pi/extension-v1/tsconfig.json", + canonical_embedded_text(include_str!( + "../../../src/daemon/managed/pi_extension/tsconfig.json" + )) + .into_bytes(), + ), + ] { + let rendered = bundle + .artifacts + .iter() + .find(|artifact| artifact.path == path) + .unwrap(); + assert_eq!(rendered.bytes, expected, "{path}"); + } +} + +#[test] +fn managed_pi_launch_disables_discovered_extensions() { + const ISOLATED_LAUNCH: &str = + "pi --no-extensions -e /srv/nemo-relay/nemo-relay-managed-v1/pi/extension-v1/index.ts"; + const NON_ISOLATED_LAUNCH: &str = + "pi -e /srv/nemo-relay/nemo-relay-managed-v1/pi/extension-v1/index.ts"; + + let rendered = render_bundle(&spec([ManagedAgent::Pi])).unwrap(); + let readme = rendered + .artifacts + .iter() + .find(|artifact| artifact.path == "pi/extension-v1/README.md") + .unwrap(); + let readme = std::str::from_utf8(&readme.bytes).unwrap(); + let daemon_docs = include_str!("../../../../../docs/nemo-relay-cli/daemon.mdx"); + + for (source_name, source) in [ + ("rendered managed Pi README", readme), + ("daemon deployment documentation", daemon_docs), + ] { + assert!( + source.contains(ISOLATED_LAUNCH), + "{source_name} must suppress all discovered Pi extensions" + ); + assert!( + !source.contains(NON_ISOLATED_LAUNCH), + "{source_name} must not document a non-isolated managed Pi launch" + ); + } +} + +#[test] +fn managed_pi_extension_forwards_custom_provider_endpoints_without_route_urls() { + let source = include_str!("../../../src/daemon/managed/pi_extension/index.ts"); + + for contract in [ + "['daemon', 'mcp', '--daemon-address'", + "/hooks/pi", + "pi.registerProvider(model.provider", + "'openai-completions'", + "'openai-responses'", + "'anthropic-messages'", + "pi.on('session_start'", + "pi.on('session_before_compact'", + "pi.on('session_compact'", + "pi.on('tool_call'", + "pi.on('user_bash'", + "typeof toolCall.tool_call_id !== 'string'", + "shapeViolation(current, toolCall.input)", + "'tool_arguments_transformed'", + "code: 'model-registry-unavailable'", + "function toolResultText(content: unknown)", + "function sliceAtCodePointBoundary(value: string", + "const MCP_READY_TIMEOUT_MS = 180_000", + "[CLIENT_TOKEN_HEADER]: active.credential", + "const UPSTREAM_BASE_URL_HEADER = 'x-nemo-relay-upstream-base-url'", + "[UPSTREAM_BASE_URL_HEADER]: decision.upstream", + ] { + assert!(source.contains(contract), "missing Pi contract: {contract}"); + } + + assert_eq!( + source.matches("const CLIENT_TOKEN_HEADER =").count(), + 1, + "the managed Pi extension must define one Relay-specific public header" + ); + for forbidden in [ + "x-nemo-relay-session-id", + "x-nemo-relay-fingerprint", + "x-nemo-relay-generation", + ] { + assert!( + !source.contains(forbidden), + "managed Pi source contains forbidden routing metadata: {forbidden}" + ); + } +} + +#[test] +fn managed_pi_config_json_encodes_cross_platform_deployment_values() { + let spec = ManagedBundleSpec::new( + "https://relay.example.com:443", + "C:\\ProgramData\\NVIDIA\\nemo-relay-dispatch.exe", + ManagedPlatform::Windows, + [ManagedAgent::Pi], + ) + .unwrap(); + let rendered = render_bundle(&spec).unwrap(); + let config = rendered + .artifacts + .iter() + .find(|artifact| artifact.path == "pi/extension-v1/managed-config.json") + .unwrap(); + let text = String::from_utf8(config.bytes.clone()).unwrap(); + assert!(!text.contains("__NEMO_RELAY_")); + let value: serde_json::Value = serde_json::from_str(&text).unwrap(); + assert_eq!( + value["dispatcherCommand"], + "C:\\ProgramData\\NVIDIA\\nemo-relay-dispatch.exe" + ); +} + +#[test] +fn v1_artifact_bytes_do_not_depend_on_binary_or_platform_versioning() { + let source = include_str!("../../../src/daemon/managed/mod.rs"); + assert!(!source.contains("env!(\"CARGO_PKG_VERSION\")")); + + let linux = render_bundle(&spec([ManagedAgent::Codex])).unwrap(); + let macos_spec = ManagedBundleSpec::new( + "https://relay.example.com:443", + "/opt/nvidia/bin/nemo-relay-dispatch", + ManagedPlatform::Macos, + [ManagedAgent::Codex], + ) + .unwrap(); + let macos = render_bundle(&macos_spec).unwrap(); + let linux_artifacts = linux + .artifacts + .into_iter() + .map(|artifact| (artifact.path, artifact.bytes)) + .collect::>(); + let macos_artifacts = macos + .artifacts + .into_iter() + .map(|artifact| (artifact.path, artifact.bytes)) + .collect::>(); + assert_eq!(linux_artifacts, macos_artifacts); +} + +#[test] +fn canonical_v1_bundle_matches_the_release_frozen_golden_digest() { + // This digest pins the release-candidate v1 manifest and every artifact byte for a canonical + // deployment. After v1 is published, behavior changes must use a separately named v2 family. + const GOLDEN_SHA256: &str = "2016bff51f766cd6c9857fa6c5c089ac160c9e3747ac89112e69681c5a2da6a9"; + let rendered = render_bundle(&spec([ + ManagedAgent::Codex, + ManagedAgent::ClaudeCode, + ManagedAgent::Pi, + ])) + .unwrap(); + assert_eq!(rendered_bundle_digest(&rendered).to_string(), GOLDEN_SHA256); +} + +#[test] +fn embedded_managed_text_has_platform_independent_line_endings() { + assert_eq!( + canonical_embedded_text("first\r\nsecond\r\n"), + "first\nsecond\n" + ); + assert_eq!( + canonical_embedded_text("first\nsecond\n"), + "first\nsecond\n" + ); +} + +#[test] +fn write_is_create_only_and_existing_exact_bundle_is_not_rewritten() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let spec = spec([ManagedAgent::Codex]); + let first_digest = write_new_bundle(&root, &spec).unwrap(); + let manifest = root.join(MANIFEST_FILE); + let before = std::fs::metadata(&manifest).unwrap().modified().unwrap(); + let second_digest = write_new_bundle(&root, &spec).unwrap(); + assert_eq!(first_digest, second_digest); + let after = std::fs::metadata(&manifest).unwrap().modified().unwrap(); + assert_eq!(before, after); + + let other_deployment = ManagedBundleSpec::new( + "https://other-relay.example.com:443", + "/opt/nvidia/bin/nemo-relay-dispatch", + ManagedPlatform::Linux, + [ManagedAgent::Codex], + ) + .unwrap(); + let error = write_new_bundle(&root, &other_deployment) + .unwrap_err() + .to_string(); + assert!(error.contains("different deployment bytes"), "{error}"); + assert_eq!( + before, + std::fs::metadata(&manifest).unwrap().modified().unwrap() + ); + + let config = root.join("codex/settings-v1/config.toml"); + make_test_artifact_writable(&config); + std::fs::write(&config, "changed\n").unwrap(); + let error = write_new_bundle(&root, &spec).unwrap_err().to_string(); + assert!(error.contains("exact canonical bytes"), "{error}"); + assert_eq!( + std::fs::read_to_string(root.join("codex/settings-v1/config.toml")).unwrap(), + "changed\n" + ); +} + +#[test] +fn managed_refresh_is_validation_only() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let digest = write_new_bundle(&root, &spec([ManagedAgent::Codex])).unwrap(); + let before = std::fs::read(root.join(MANIFEST_FILE)).unwrap(); + let token = credential(); + let _environment = EnvScope::set(&[(ROUTE_TOKEN_ENV, Some(OsStr::new(&token)))]); + + refresh_bundle(&root, &digest).unwrap(); + + assert_eq!(std::fs::read(root.join(MANIFEST_FILE)).unwrap(), before); +} + +#[test] +fn doctor_validation_checks_exact_bytes_and_managed_environment() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let digest = write_new_bundle( + &root, + &spec([ManagedAgent::Codex, ManagedAgent::ClaudeCode]), + ) + .unwrap(); + let token = credential(); + let _environment = managed_environment(&token); + + let validation = refresh_bundle(&root, &digest).unwrap(); + assert_eq!(validation.artifact_count, 8); + assert_eq!(validation.daemon_address, "https://relay.example.com:443"); + assert_eq!(validation.sha256, digest); + + let hooks = root.join("claude-code/plugin-v1/hooks/hooks.json"); + make_test_artifact_writable(&hooks); + let mut changed = std::fs::read(&hooks).unwrap(); + changed.push(b' '); + std::fs::write(&hooks, changed).unwrap(); + let error = refresh_bundle(&root, &digest).unwrap_err().to_string(); + assert!( + error.contains("expected size limit") || error.contains("exact canonical bytes"), + "{error}" + ); +} + +#[test] +fn doctor_requires_the_separately_provisioned_bundle_digest() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let digest = write_new_bundle(&root, &spec([ManagedAgent::Codex])).unwrap(); + let wrong: ManagedBundleDigest = + "0000000000000000000000000000000000000000000000000000000000000000" + .parse() + .unwrap(); + let _environment = EnvScope::set(&[(ROUTE_TOKEN_ENV, None)]); + + let error = refresh_bundle(&root, &wrong).unwrap_err().to_string(); + assert!(error.contains("SHA-256 mismatch"), "{error}"); + assert!(error.contains(&digest.to_string()), "{error}"); + assert!(!error.contains(ROUTE_TOKEN_ENV), "{error}"); + + assert!(digest.to_string().parse::().is_ok()); + for malformed in [ + "0", + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", + "gggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggggg", + ] { + assert!(malformed.parse::().is_err()); + } +} + +#[test] +fn claude_environment_must_bind_the_custom_header_to_the_route_token() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let digest = write_new_bundle(&root, &spec([ManagedAgent::ClaudeCode])).unwrap(); + let token = credential(); + let wrong = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x24_u8; 32]); + let header = format!("{ROUTE_TOKEN_HEADER}: {wrong}"); + let _environment = EnvScope::set(&[ + (ROUTE_TOKEN_ENV, Some(OsStr::new(&token))), + (CLAUDE_CUSTOM_HEADERS_ENV, Some(OsStr::new(&header))), + ]); + + let error = refresh_bundle(&root, &digest).unwrap_err().to_string(); + assert!(error.contains("must contain exactly one"), "{error}"); + assert!(!error.contains(&token)); + assert!(!error.contains(&wrong)); +} + +#[test] +fn managed_environment_requires_the_enterprise_provisioned_credential() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let digest = write_new_bundle(&root, &spec([ManagedAgent::Codex])).unwrap(); + let _environment = EnvScope::set(&[(ROUTE_TOKEN_ENV, None)]); + + let error = refresh_bundle(&root, &digest).unwrap_err().to_string(); + assert!(error.contains(ROUTE_TOKEN_ENV), "{error}"); + assert!( + error.contains("enterprise") || error.contains("managed"), + "{error}" + ); +} + +#[test] +fn doctor_rejects_extra_files_and_noncanonical_manifest_bytes() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let digest = write_new_bundle(&root, &spec([ManagedAgent::Pi])).unwrap(); + let token = credential(); + let _environment = EnvScope::set(&[(ROUTE_TOKEN_ENV, Some(OsStr::new(&token)))]); + + std::fs::write(root.join("unmanaged.json"), "{}\n").unwrap(); + let error = refresh_bundle(&root, &digest).unwrap_err().to_string(); + assert!(error.contains("unexpected artifact"), "{error}"); + std::fs::remove_file(root.join("unmanaged.json")).unwrap(); + + let manifest = root.join(MANIFEST_FILE); + make_test_artifact_writable(&manifest); + let value: serde_json::Value = + serde_json::from_slice(&std::fs::read(&manifest).unwrap()).unwrap(); + std::fs::write(&manifest, serde_json::to_vec(&value).unwrap()).unwrap(); + let error = refresh_bundle(&root, &digest).unwrap_err().to_string(); + assert!(error.contains("canonical"), "{error}"); +} + +#[test] +fn doctor_rejects_missing_oversized_and_non_file_artifacts() { + let directory = tempdir().unwrap(); + let token = credential(); + let _environment = EnvScope::set(&[(ROUTE_TOKEN_ENV, Some(OsStr::new(&token)))]); + + let missing_root = directory.path().join("missing-root"); + let digest: ManagedBundleDigest = + "0000000000000000000000000000000000000000000000000000000000000000" + .parse() + .unwrap(); + assert!(refresh_bundle(&missing_root, &digest).is_err()); + + let regular_root = directory.path().join("regular-root"); + std::fs::write(®ular_root, b"not a directory").unwrap(); + assert!(refresh_bundle(®ular_root, &digest).is_err()); + + let root = directory.path().join("bundle"); + let digest = write_new_bundle(&root, &spec([ManagedAgent::Codex])).unwrap(); + let artifact = root.join("codex/settings-v1/config.toml"); + std::fs::remove_file(&artifact).unwrap(); + assert!(refresh_bundle(&root, &digest).is_err()); + + write_new_bundle(&root, &spec([ManagedAgent::Codex])).unwrap_err(); + std::fs::create_dir(&artifact).unwrap(); + assert!(refresh_bundle(&root, &digest).is_err()); +} + +#[cfg(unix)] +#[test] +fn doctor_rejects_symlinked_manifest_and_nested_artifact() { + let directory = tempdir().unwrap(); + let token = credential(); + let _environment = EnvScope::set(&[(ROUTE_TOKEN_ENV, Some(OsStr::new(&token)))]); + + for relative in [MANIFEST_FILE, "codex/settings-v1/config.toml"] { + let root = directory.path().join(relative.replace('/', "-")); + let digest = write_new_bundle(&root, &spec([ManagedAgent::Codex])).unwrap(); + let target = root.join(relative); + let original = directory + .path() + .join(format!("{}.original", relative.replace('/', "-"))); + std::fs::rename(&target, &original).unwrap(); + std::os::unix::fs::symlink(&original, &target).unwrap(); + let error = refresh_bundle(&root, &digest).unwrap_err().to_string(); + assert!(error.contains("symlink"), "{relative}: {error}"); + } +} + +#[test] +fn pi_placeholder_replacement_requires_exactly_one_json_string() { + assert!(replace_json_string_value("{}", "missing", "value").is_err()); + assert!(replace_json_string_value(r#"["same","same"]"#, "same", "value").is_err()); + assert_eq!( + replace_json_string_value( + r#"{"key":"placeholder"}"#, + "placeholder", + "quoted \"value\"" + ) + .unwrap(), + r#"{"key":"quoted \"value\""}"# + ); +} + +#[cfg(unix)] +#[test] +fn doctor_rejects_a_symlinked_bundle_root() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let alias = directory.path().join("bundle-alias"); + let digest = write_new_bundle(&root, &spec([ManagedAgent::Pi])).unwrap(); + std::os::unix::fs::symlink(&root, &alias).unwrap(); + + let error = refresh_bundle(&alias, &digest).unwrap_err().to_string(); + assert!( + error.contains("must be a directory, not a symlink"), + "{error}" + ); +} + +#[test] +fn spec_rejects_mutable_or_unsafe_inputs() { + for address in [ + "https://relay.example.com", + "http://relay.example.com:80", + "https://relay.example.com:443/tenant/alice", + ] { + assert!( + ManagedBundleSpec::new( + address, + "C:\\ProgramData\\NVIDIA\\nemo-relay-dispatch.exe", + ManagedPlatform::Windows, + [ManagedAgent::Codex] + ) + .is_err(), + "accepted {address}" + ); + } + for dispatcher in [ + "", + "nemo relay", + "nemo-relay-dispatch", + "nemo-relay;malicious", + "/Users/alice/bin/nemo-relay", + "/home/alice/bin/nemo-relay", + "/root/bin/nemo-relay", + "/tmp/nemo-relay", + "/opt/nvidia/../alice/nemo-relay", + "C:\\ProgramData\\NVIDIA\\nemo-relay.exe", + ] { + assert!( + ManagedBundleSpec::new( + "https://relay.example.com:443", + dispatcher, + ManagedPlatform::Macos, + [ManagedAgent::Codex] + ) + .is_err(), + "accepted {dispatcher}" + ); + } + + for (platform, dispatcher) in [ + (ManagedPlatform::Windows, "nemo-relay-dispatch"), + (ManagedPlatform::Windows, "C:\\Users\\alice\\nemo-relay.exe"), + ( + ManagedPlatform::Windows, + "C:\\Windows\\Temp\\nemo-relay.exe", + ), + (ManagedPlatform::Windows, "/opt/nvidia/nemo-relay"), + ] { + assert!( + ManagedBundleSpec::new( + "https://relay.example.com:443", + dispatcher, + platform, + [ManagedAgent::Codex] + ) + .is_err(), + "accepted {dispatcher}" + ); + } + + for (platform, dispatcher) in [ + ( + ManagedPlatform::Linux, + "/opt/nvidia/bin/nemo-relay-dispatch", + ), + ( + ManagedPlatform::Macos, + "/Library/NVIDIA/bin/nemo-relay-dispatch", + ), + ( + ManagedPlatform::Windows, + "C:\\ProgramData\\NVIDIA\\nemo-relay-dispatch.exe", + ), + ( + ManagedPlatform::Windows, + "\\\\relay.example.com\\nvidia\\nemo-relay-dispatch.exe", + ), + ] { + ManagedBundleSpec::new( + "https://relay.example.com:443", + dispatcher, + platform, + [ManagedAgent::Codex], + ) + .unwrap(); + } +} + +#[test] +fn managed_bundle_value_types_and_empty_agent_sets_are_strict() { + assert_eq!(ManagedPlatform::Linux.as_str(), "linux"); + assert_eq!(ManagedPlatform::Macos.as_str(), "macos"); + assert_eq!(ManagedPlatform::Windows.as_str(), "windows"); + assert_eq!(ManagedAgent::Codex.hook_argument(), "codex"); + assert_eq!(ManagedAgent::ClaudeCode.hook_argument(), "claude"); + assert_eq!(ManagedAgent::Pi.hook_argument(), "pi"); + assert!( + ManagedBundleSpec::new( + "https://relay.example.com:443", + "/opt/nvidia/bin/nemo-relay", + ManagedPlatform::Linux, + [], + ) + .is_err() + ); + for invalid in [ + "", + "ABCDEF0000000000000000000000000000000000000000000000000000000000", + "g000000000000000000000000000000000000000000000000000000000000000", + ] { + assert!(invalid.parse::().is_err()); + } +} + +#[test] +fn managed_bundle_rejects_invalid_manifest_family_and_artifact_capacity() { + let directory = tempdir().unwrap(); + let root = directory.path().join("bundle"); + let _digest = write_new_bundle(&root, &spec([ManagedAgent::Codex])).unwrap(); + let manifest_path = root.join(MANIFEST_FILE); + make_test_artifact_writable(&manifest_path); + let original = std::fs::read(&manifest_path).unwrap(); + let mut manifest: serde_json::Value = serde_json::from_slice(&original).unwrap(); + manifest["schema_version"] = json!(SCHEMA_VERSION + 1); + std::fs::write(&manifest_path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + assert!(validate_bundle_files(&root, false, None).is_err()); + + manifest = serde_json::from_slice(&original).unwrap(); + manifest["artifacts"] = Value::Array( + (0..=MAX_ARTIFACTS) + .map(|index| { + json!({ + "agent": "codex", + "path": format!("artifact-{index}"), + "byte_length": 0, + "sha256": "0".repeat(64), + }) + }) + .collect(), + ); + std::fs::write(&manifest_path, serde_json::to_vec(&manifest).unwrap()).unwrap(); + assert!(validate_bundle_files(&root, false, None).is_err()); +} + +#[test] +fn managed_file_helpers_reject_overwrite_oversize_and_excessive_entries() { + let directory = tempdir().unwrap(); + let file = directory.path().join("artifact"); + write_new_file(&file, b"first").unwrap(); + assert!(write_new_file(&file, b"second").is_err()); + assert_eq!(read_bounded(&file, 5).unwrap(), b"first"); + assert!(read_bounded(&file, 4).is_err()); + assert!(read_bounded(&directory.path().join("missing"), 4).is_err()); + + let root = directory.path().join("bundle"); + std::fs::create_dir(&root).unwrap(); + for index in 0..=(MAX_ARTIFACTS * 4) { + std::fs::create_dir(root.join(format!("empty-{index}"))).unwrap(); + } + assert!(bundle_files(&root).is_err()); +} diff --git a/crates/cli/tests/coverage/daemon/mcp_tests.rs b/crates/cli/tests/coverage/daemon/mcp_tests.rs new file mode 100644 index 000000000..5b34867fc --- /dev/null +++ b/crates/cli/tests/coverage/daemon/mcp_tests.rs @@ -0,0 +1,611 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +use std::sync::{Arc, Mutex}; + +use axum::body::Bytes; +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::post; +use axum::{Json, Router}; +use tokio::net::TcpListener; + +#[test] +fn pending_worker_child_fixture() { + if std::env::var_os("NEMO_RELAY_TEST_PENDING_WORKER_FIXTURE").is_none() { + return; + } + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + println!("PENDING_WORKER_READY {}", listener.local_addr().unwrap()); + use std::io::Write as _; + std::io::stdout().flush().unwrap(); + let mut line = String::new(); + std::io::stdin().read_line(&mut line).unwrap(); +} + +async fn pending_worker_fixture() -> (ActivationChild, SocketAddr) { + use tokio::io::AsyncBufReadExt as _; + let mut child = Command::new(std::env::current_exe().unwrap()) + .args([ + "--exact", + "daemon::mcp::tests::pending_worker_child_fixture", + "--nocapture", + ]) + .env("NEMO_RELAY_TEST_PENDING_WORKER_FIXTURE", "1") + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .spawn() + .unwrap(); + let stdout = child.stdout.take().unwrap(); + let child = ActivationChild { + child, + published: false, + }; + let address = tokio::time::timeout(Duration::from_secs(10), async { + let mut lines = tokio::io::BufReader::new(stdout).lines(); + while let Some(line) = lines.next_line().await.unwrap() { + if let Some(address) = line.strip_prefix("PENDING_WORKER_READY ") { + return address.parse::().unwrap(); + } + } + panic!("worker fixture exited before readiness"); + }) + .await + .unwrap(); + (child, address) +} + +#[tokio::test] +async fn failed_activation_cleanup_reaps_child_and_frees_its_listener() { + let (child, address) = pending_worker_fixture().await; + let mut pending = Some(( + "failed-activation".into(), + child, + tokio::time::Instant::now(), + )); + stop_pending_launch(&mut pending).await.unwrap(); + assert!(pending.is_none()); + let _replacement = TcpListener::bind(address) + .await + .expect("old worker listener released"); + stop_pending_launch(&mut pending).await.unwrap(); +} + +#[tokio::test] +async fn published_worker_survives_guard_drop_but_pending_worker_is_killed() { + for published in [false, true] { + let (mut child, address) = pending_worker_fixture().await; + child.published = published; + let mut stdin = child.child.stdin.take().unwrap(); + drop(child); + if published { + let _connection = tokio::net::TcpStream::connect(address) + .await + .expect("published worker must remain alive"); + stdin.write_all(b"exit\n").await.unwrap(); + } + tokio::time::timeout(Duration::from_secs(10), async { + loop { + if let Ok(listener) = TcpListener::bind(address).await { + break listener; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("worker must terminate"); + drop(stdin); + } +} + +#[test] +fn launch_directive_is_the_only_directive_with_a_worker_bootstrap() { + assert!(WorkerBootstrap::from_directive(BrokerDirective::UsePassThrough).is_none()); + assert!( + WorkerBootstrap::from_directive(BrokerDirective::WaitForWorker { retry_after_ms: 10 }) + .is_none() + ); +} + +#[test] +fn activation_timeout_uses_the_mcp_monotonic_clock() { + let started = tokio::time::Instant::now(); + assert!(!activation_timed_out( + "activation", + "activation", + started, + started + Duration::from_millis(ACTIVATION_LIFETIME_MS - 1), + )); + assert!(activation_timed_out( + "activation", + "activation", + started, + started + Duration::from_millis(ACTIVATION_LIFETIME_MS), + )); + assert!(!activation_timed_out( + "replacement", + "activation", + started, + started + Duration::from_millis(ACTIVATION_LIFETIME_MS), + )); +} + +#[test] +fn daemon_heartbeat_interval_must_leave_a_safe_lease_margin() { + assert!(validate_heartbeat_interval(999).is_err()); + assert_eq!( + validate_heartbeat_interval(1_000).expect("minimum interval"), + Duration::from_secs(1) + ); + assert_eq!( + validate_heartbeat_interval(MCP_LEASE_MS / 3).expect("maximum interval"), + Duration::from_secs(10) + ); + assert!(validate_heartbeat_interval(MCP_LEASE_MS / 3 + 1).is_err()); + assert!(validate_heartbeat_interval(u64::MAX).is_err()); + assert_eq!( + MCP_HEARTBEAT_INTERVAL_MS + HEARTBEAT_RETRY_WINDOW_MS + 5_000, + MCP_LEASE_MS, + "a full retry window must still leave five seconds before lease expiry" + ); +} + +#[test] +fn prescribed_worker_network_accepts_host_or_ipv4_and_rejects_unsafe_values() { + assert_eq!( + parse_worker_network_overrides(Some("Worker.Example.com"), Some("9443")) + .expect("hostname override"), + (Some("worker.example.com".into()), Some(9443)) + ); + assert_eq!( + parse_worker_network_overrides(Some("192.0.2.10"), None).expect("IPv4 override"), + (Some("192.0.2.10".into()), None) + ); + assert!(parse_worker_network_overrides(Some("0.0.0.0"), None).is_err()); + assert!(parse_worker_network_overrides(Some("[::1]"), None).is_err()); + assert!(parse_worker_network_overrides(Some("https://worker.example.com"), None).is_err()); + assert!(parse_worker_network_overrides(None, Some("0")).is_err()); +} + +#[tokio::test] +async fn loopback_daemon_uses_loopback_worker_network_and_environment_overrides() { + let _environment = crate::test_support::EnvScope::set(&[ + (WORKER_ADVERTISE_ENV, None), + (WORKER_PORT_ENV, None), + ]); + let hint = worker_network_hint("http://127.0.0.1:47632") + .await + .expect("loopback daemon network hint"); + assert_eq!(hint.advertised_host, "127.0.0.1"); + assert_eq!(hint.port, None); + + for origin in ["http://127.0.0.1:80", "https://127.0.0.1:443"] { + let hint = worker_network_hint(origin) + .await + .expect("default-port network hint"); + assert_eq!(hint.advertised_host, "127.0.0.1"); + assert_eq!(hint.port, None); + } + + drop(_environment); + let _environment = crate::test_support::EnvScope::set(&[ + ( + WORKER_ADVERTISE_ENV, + Some(std::ffi::OsStr::new("worker.example")), + ), + (WORKER_PORT_ENV, Some(std::ffi::OsStr::new("9443"))), + ]); + let hint = worker_network_hint("http://127.0.0.1:47632") + .await + .expect("explicit network hint"); + assert_eq!(hint.advertised_host, "worker.example"); + assert_eq!(hint.port, Some(9443)); +} + +#[cfg(unix)] +#[test] +fn worker_network_environment_requires_unicode() { + use std::os::unix::ffi::OsStrExt; + + let _environment = crate::test_support::EnvScope::set(&[( + WORKER_ADVERTISE_ENV, + Some(std::ffi::OsStr::from_bytes(b"worker-\xff")), + )]); + assert!(optional_environment(WORKER_ADVERTISE_ENV).is_err()); +} + +#[tokio::test] +async fn ipv6_only_daemon_has_no_supported_worker_route() { + let _environment = crate::test_support::EnvScope::set(&[ + (WORKER_ADVERTISE_ENV, None), + (WORKER_PORT_ENV, None), + ]); + assert!(worker_network_hint("http://[::1]:47632").await.is_err()); +} + +#[tokio::test] +async fn remote_daemon_rejects_an_explicit_loopback_worker_advertisement() { + let _environment = crate::test_support::EnvScope::set(&[ + ( + WORKER_ADVERTISE_ENV, + Some(std::ffi::OsStr::new("127.0.0.1")), + ), + (WORKER_PORT_ENV, None), + ]); + let error = worker_network_hint("https://192.0.2.1:8443") + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("cannot be loopback for a remote daemon"), + "unexpected network validation error: {error}" + ); +} + +#[tokio::test] +async fn remote_daemon_derives_a_concrete_non_loopback_worker_advertisement() { + let _environment = crate::test_support::EnvScope::set(&[ + (WORKER_ADVERTISE_ENV, None), + (WORKER_PORT_ENV, Some(std::ffi::OsStr::new("9443"))), + ]); + let hint = match worker_network_hint("https://192.0.2.1:8443").await { + Ok(hint) => hint, + Err(CliError::Io(error)) + if matches!( + error.kind(), + std::io::ErrorKind::NetworkUnreachable | std::io::ErrorKind::HostUnreachable + ) => + { + return; + } + Err(error) => panic!("unexpected route derivation error: {error}"), + }; + assert!( + !hint + .advertised_host + .parse::() + .unwrap() + .is_loopback() + ); + assert_eq!(hint.port, Some(9443)); +} + +#[test] +fn spawned_worker_explicitly_removes_the_public_route_credential() { + let bootstrap = WorkerBootstrap { + activation_id: "activation".into(), + activation_token: SensitiveString::new("secret").expect("secret"), + deadline_unix_ms: u64::MAX, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }; + let command = worker_command( + std::path::Path::new("nemo-relay"), + "http://127.0.0.1:47632", + &bootstrap, + ); + assert!( + command + .as_std() + .get_envs() + .any(|(name, value)| { name == ROUTE_TOKEN_ENV && value.is_none() }) + ); +} + +#[test] +fn worker_command_keeps_only_documented_network_arguments() { + let bootstrap = WorkerBootstrap { + activation_id: "activation".into(), + activation_token: SensitiveString::new("secret").expect("secret"), + deadline_unix_ms: u64::MAX, + bind_ip: Ipv4Addr::UNSPECIFIED, + port: 9443, + advertise_address: Some("worker.example".into()), + }; + let origin = explicit_daemon_origin("https://daemon.example:443/").unwrap(); + let command = worker_command(std::path::Path::new("nemo-relay"), &origin, &bootstrap); + let arguments = command + .as_std() + .get_args() + .map(|argument| argument.to_string_lossy().into_owned()) + .collect::>(); + assert_eq!( + arguments, + [ + "daemon", + "worker", + "--daemon-address", + "https://daemon.example:443", + "--bind", + "0.0.0.0", + "--port", + "9443", + "--advertise-address", + "worker.example", + ] + ); +} + +#[test] +fn re_registration_preserves_sequence_until_the_daemon_rotates_the_session() { + let mut lease = test_lease("http://127.0.0.1:1".into()); + lease.sequence = 7; + lease.pending_heartbeat = Some( + SessionRequest::new( + lease.session_id.clone(), + lease.session_token.clone(), + lease.sequence, + EmptyPayload::default(), + ) + .expect("pending heartbeat"), + ); + let same_session = Registration { + directive: BrokerDirective::UsePassThrough, + session_token: lease.session_token.clone(), + heartbeat_interval: Duration::from_secs(5), + }; + apply_registration(&mut lease, &same_session); + assert_eq!(lease.sequence, 7); + assert!(lease.pending_heartbeat.is_some()); + + let rotated = Registration { + directive: BrokerDirective::UsePassThrough, + session_token: SensitiveString::new("rotated-session").expect("rotated token"), + heartbeat_interval: Duration::from_secs(4), + }; + apply_registration(&mut lease, &rotated); + assert_eq!(lease.sequence, 0); + assert!(lease.pending_heartbeat.is_none()); +} + +#[derive(Default)] +struct RequestLog { + bodies: Mutex>, +} + +#[tokio::test] +async fn transient_heartbeat_failure_keeps_the_session_and_exact_request() { + async fn heartbeat(State(log): State>, body: Bytes) -> Response { + let attempt = { + let mut bodies = log.bodies.lock().expect("heartbeat bodies"); + bodies.push(body); + bodies.len() + }; + if attempt == 1 { + return StatusCode::SERVICE_UNAVAILABLE.into_response(); + } + Json(McpHeartbeatResponse { + directive: Some(BrokerDirective::UsePassThrough), + }) + .into_response() + } + + let log = Arc::new(RequestLog::default()); + let origin = spawn( + Router::new() + .route(MCP_HEARTBEAT_PATH, post(heartbeat)) + .with_state(Arc::clone(&log)), + ) + .await; + let mut lease = test_lease(origin); + let response = renew_lease_with(&mut lease, fast_retry_policy()) + .await + .expect("brief daemon failure should not end the MCP lease"); + + assert!(matches!( + response.directive, + Some(BrokerDirective::UsePassThrough) + )); + assert_eq!(lease.sequence, 1); + assert!(lease.pending_heartbeat.is_none()); + let bodies = log.bodies.lock().expect("heartbeat bodies"); + assert_eq!(bodies.len(), 2); + assert_eq!(bodies[0], bodies[1]); +} + +#[tokio::test] +async fn release_retries_the_same_session_request() { + async fn release_handler(State(log): State>, body: Bytes) -> StatusCode { + let attempt = { + let mut bodies = log.bodies.lock().expect("release bodies"); + bodies.push(body); + bodies.len() + }; + if attempt == 1 { + StatusCode::BAD_GATEWAY + } else { + StatusCode::NO_CONTENT + } + } + + let log = Arc::new(RequestLog::default()); + let origin = spawn( + Router::new() + .route( + super::super::common::control::MCP_RELEASE_PATH, + post(release_handler), + ) + .with_state(Arc::clone(&log)), + ) + .await; + let mut lease = test_lease(origin); + release(&mut lease).await; + + assert_eq!(lease.sequence, 1); + let bodies = log.bodies.lock().expect("release bodies"); + assert_eq!(bodies.len(), 2); + assert_eq!(bodies[0], bodies[1]); +} + +#[tokio::test] +async fn release_settles_a_pending_heartbeat_before_releasing() { + async fn heartbeat(State(log): State>, body: Bytes) -> Response { + log.bodies.lock().unwrap().push(body); + Json(McpHeartbeatResponse { directive: None }).into_response() + } + async fn released(State(log): State>, body: Bytes) -> StatusCode { + log.bodies.lock().unwrap().push(body); + StatusCode::NO_CONTENT + } + let log = Arc::new(RequestLog::default()); + let origin = spawn( + Router::new() + .route(MCP_HEARTBEAT_PATH, post(heartbeat)) + .route( + super::super::common::control::MCP_RELEASE_PATH, + post(released), + ) + .with_state(Arc::clone(&log)), + ) + .await; + let mut lease = test_lease(origin); + lease.sequence = 1; + lease.pending_heartbeat = Some( + SessionRequest::new( + lease.session_id.clone(), + lease.session_token.clone(), + 1, + EmptyPayload::default(), + ) + .unwrap(), + ); + release(&mut lease).await; + assert_eq!(lease.sequence, 2); + assert!(lease.pending_heartbeat.is_none()); + assert_eq!(log.bodies.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn release_stops_when_a_pending_heartbeat_cannot_be_settled() { + let origin = spawn(Router::new().route( + MCP_HEARTBEAT_PATH, + post(|| async { StatusCode::UNAUTHORIZED }), + )) + .await; + let mut lease = test_lease(origin); + lease.sequence = 1; + lease.pending_heartbeat = Some( + SessionRequest::new( + lease.session_id.clone(), + lease.session_token.clone(), + 1, + EmptyPayload::default(), + ) + .unwrap(), + ); + release(&mut lease).await; + assert_eq!(lease.sequence, 1); + assert!(lease.pending_heartbeat.is_some()); +} + +#[tokio::test] +async fn activation_failure_is_reported_with_an_authenticated_session_request() { + async fn failed(State(log): State>, body: Bytes) -> StatusCode { + log.bodies.lock().expect("failure bodies").push(body); + StatusCode::NO_CONTENT + } + let log = Arc::new(RequestLog::default()); + let origin = spawn( + Router::new() + .route(MCP_ACTIVATION_FAILED_PATH, post(failed)) + .with_state(Arc::clone(&log)), + ) + .await; + let mut lease = test_lease(origin); + report_activation_failed( + &mut lease, + "activation-id", + &CliError::Launch("worker failed safely".into()), + ) + .await + .expect("activation failure report"); + assert_eq!(lease.sequence, 1); + let bodies = log.bodies.lock().expect("failure bodies"); + assert_eq!(bodies.len(), 1); + let value: serde_json::Value = serde_json::from_slice(&bodies[0]).unwrap(); + assert_eq!(value["payload"]["activation_id"], "activation-id"); + assert_eq!( + value["payload"]["reason"], + "launcher error: worker failed safely" + ); +} + +#[tokio::test] +async fn mcp_control_sequence_exhaustion_fails_without_network_io() { + let mut lease = test_lease("http://127.0.0.1:1".into()); + lease.sequence = u64::MAX; + assert!( + renew_lease_with(&mut lease, fast_retry_policy()) + .await + .is_err() + ); + release(&mut lease).await; + assert_eq!(lease.sequence, u64::MAX); +} + +#[tokio::test] +async fn ready_directives_complete_without_spawning_or_polling() { + let mut lease = test_lease("http://127.0.0.1:1".into()); + make_route_ready(&mut lease, BrokerDirective::UsePassThrough) + .await + .unwrap(); + make_route_ready( + &mut lease, + BrokerDirective::ReuseWorker { + endpoint: "http://127.0.0.1:2".into(), + }, + ) + .await + .unwrap(); +} + +#[tokio::test(start_paused = true)] +async fn waiting_for_a_worker_refreshes_registration_and_surfaces_daemon_failure() { + let mut lease = test_lease("http://127.0.0.1:1".into()); + let error = make_route_ready( + &mut lease, + BrokerDirective::WaitForWorker { retry_after_ms: 1 }, + ) + .await + .expect_err("unreachable daemon"); + assert!(matches!(error, CliError::Upstream(_))); +} + +fn test_lease(daemon_origin: String) -> McpLease { + McpLease { + client: control_client().expect("client"), + daemon_origin, + route_credential: RouteCredential::parse( + "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA".into(), + ) + .expect("route credential"), + identity: MachineIdentity::generate().expect("identity").identity, + session_id: "mcp-test-session".into(), + session_token: SensitiveString::new("session-secret").expect("session token"), + heartbeat_interval: Duration::from_secs(10), + sequence: 0, + pending_heartbeat: None, + } +} + +fn fast_retry_policy() -> ControlRetryPolicy { + ControlRetryPolicy::new( + Duration::from_secs(1), + Duration::from_secs(2), + Duration::ZERO, + ) +} + +async fn spawn(router: Router) -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("local address"); + tokio::spawn(async move { + axum::serve(listener, router).await.expect("serve"); + }); + format!("http://{address}") +} diff --git a/crates/cli/tests/coverage/daemon/protocol_tests.rs b/crates/cli/tests/coverage/daemon/protocol_tests.rs new file mode 100644 index 000000000..3625ad3c7 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/protocol_tests.rs @@ -0,0 +1,207 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::daemon::common::control::fresh_nonce; +use crate::daemon::common::identity::{ChallengeRecord, TokenDigest}; + +fn sample_transcript() -> (HandshakeTranscript, MachineIdentity, MachineIdentity) { + let initiator = MachineIdentity::generate().expect("initiator").identity; + let responder = MachineIdentity::generate().expect("responder").identity; + let challenge = ChallengeRecord::generate(10, 100) + .expect("challenge") + .challenge(); + ( + HandshakeTranscript { + daemon_target: "https://relay.example:443".to_owned(), + initiator: ComponentDescriptor::nemo_relay( + ComponentRole::Mcp, + ProtocolRange::default(), + Capabilities::streaming_transport(), + "0.9.0", + ), + responder: ComponentDescriptor::nemo_relay( + ComponentRole::Daemon, + ProtocolRange::default(), + Capabilities::streaming_transport(), + "2.0.0", + ), + initiator_instance_id: "mcp-1".to_owned(), + responder_instance_id: "daemon-1".to_owned(), + selected_protocol: PROTOCOL_V1, + initiator_public_identity: initiator.public_identity(), + responder_public_identity: responder.public_identity(), + initiator_fingerprint: initiator.fingerprint(), + responder_fingerprint: responder.fingerprint(), + challenge_id: challenge.id, + initiator_nonce: fresh_nonce().expect("initiator nonce"), + responder_nonce: challenge.nonce, + route_token_digest: Some(TokenDigest::from_token(b"token")), + }, + initiator, + responder, + ) +} + +#[test] +fn negotiation_selects_highest_overlap_without_using_binary_version() { + assert_eq!( + ProtocolRange::new(1, 4) + .expect("range") + .negotiate(ProtocolRange::new(2, 3).expect("range")), + Ok(3) + ); + assert_eq!( + ProtocolRange::new(1, 2) + .expect("range") + .negotiate(ProtocolRange::new(3, 4).expect("range")), + Err(ProtocolError::NoProtocolOverlap) + ); +} + +#[test] +fn capability_serialization_and_transcript_order_are_deterministic() { + let first = Capabilities::new(["trailers", "http2", "http1"]).expect("capabilities"); + let second = Capabilities::new(["http1", "trailers", "http2"]).expect("capabilities"); + assert_eq!(first, second); + assert_eq!(first.canonical_bytes(), second.canonical_bytes()); + assert!(first.contains("trailers")); + assert!(first.includes(&Capabilities::new(["http1", "http2"]).expect("required"))); + let (mut transcript, _, _) = sample_transcript(); + transcript.initiator.capabilities = first.clone(); + transcript.responder.capabilities = second.clone(); + let mut reordered = transcript.clone(); + reordered.initiator.capabilities = second; + reordered.responder.capabilities = first; + assert_eq!( + transcript.canonical_bytes().unwrap(), + reordered.canonical_bytes().unwrap() + ); +} + +#[test] +fn both_participants_sign_the_same_transcript() { + let (transcript, initiator, responder) = sample_transcript(); + let initiator_proof = transcript + .sign(ComponentRole::Mcp, &initiator) + .expect("initiator proof"); + let responder_proof = transcript + .sign(ComponentRole::Daemon, &responder) + .expect("responder proof"); + transcript.verify(&initiator_proof).expect("initiator"); + transcript.verify(&responder_proof).expect("responder"); +} + +#[test] +fn any_signed_field_mutation_invalidates_the_proof() { + let (transcript, initiator, _) = sample_transcript(); + let proof = transcript + .sign(ComponentRole::Mcp, &initiator) + .expect("proof"); + let assert_rejected = |changed: HandshakeTranscript, field: &str| { + assert!(changed.verify(&proof).is_err(), "{field}"); + }; + macro_rules! reject_mutation { + ($($field:ident).+, $value:expr) => {{ + let mut changed = transcript.clone(); + changed.$($field).+ = $value; + assert_rejected(changed, stringify!($($field).+)); + }}; + } + reject_mutation!(daemon_target, "https://other.example:443".into()); + reject_mutation!(initiator.service, "other".into()); + reject_mutation!(initiator.role, ComponentRole::Worker); + reject_mutation!(initiator.protocol.minimum, 0); + reject_mutation!(initiator.protocol.maximum, 2); + reject_mutation!( + initiator.capabilities, + Capabilities::new(["different"]).unwrap() + ); + reject_mutation!(initiator.binary_version, "other".into()); + reject_mutation!(responder.service, "other".into()); + reject_mutation!(responder.role, ComponentRole::Mcp); + reject_mutation!(responder.protocol.minimum, 0); + reject_mutation!(responder.protocol.maximum, 2); + reject_mutation!( + responder.capabilities, + Capabilities::new(["different"]).unwrap() + ); + reject_mutation!(responder.binary_version, "other".into()); + reject_mutation!(initiator_instance_id, "other".into()); + reject_mutation!(responder_instance_id, "other".into()); + reject_mutation!(selected_protocol, 2); + reject_mutation!( + initiator_public_identity, + transcript.responder_public_identity + ); + reject_mutation!( + responder_public_identity, + transcript.initiator_public_identity + ); + reject_mutation!(initiator_fingerprint, transcript.responder_fingerprint); + reject_mutation!(responder_fingerprint, transcript.initiator_fingerprint); + reject_mutation!( + challenge_id, + ChallengeRecord::generate(10, 100).unwrap().challenge().id + ); + reject_mutation!(initiator_nonce, transcript.responder_nonce); + reject_mutation!(responder_nonce, transcript.initiator_nonce); + reject_mutation!(route_token_digest, Some(TokenDigest::from_token(b"other"))); + reject_mutation!(route_token_digest, None); +} + +#[test] +fn service_and_fingerprint_are_validated_before_signing() { + let (mut wrong_service, initiator, _) = sample_transcript(); + wrong_service.initiator.service = "impostor".to_owned(); + assert_eq!( + wrong_service.sign(ComponentRole::Mcp, &initiator), + Err(ProtocolError::WrongService) + ); + + let (mut wrong_fingerprint, initiator, _) = sample_transcript(); + wrong_fingerprint.initiator_fingerprint = wrong_fingerprint.responder_fingerprint; + assert_eq!( + wrong_fingerprint.sign(ComponentRole::Mcp, &initiator), + Err(ProtocolError::FingerprintMismatch) + ); +} + +#[test] +fn descriptors_reject_oversized_untrusted_fields() { + let oversized_capability = "a".repeat(129); + assert!(Capabilities::new([oversized_capability]).is_err()); + let too_many = (0..65).map(|index| format!("capability-{index}")); + assert!(Capabilities::new(too_many).is_err()); + + let descriptor = ComponentDescriptor::nemo_relay( + ComponentRole::Mcp, + ProtocolRange::default(), + Capabilities::streaming_transport(), + "v".repeat(257), + ); + assert_eq!( + descriptor.validate(), + Err(ProtocolError::BinaryVersionTooLong) + ); +} + +#[test] +fn activation_token_is_redacted_from_debug_but_serialized() { + let directive = WorkerLaunch { + activation_id: "activation-1".to_owned(), + activation_token: SensitiveString::new("secret-value").expect("token"), + deadline_unix_ms: 100, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + } + .into_directive(); + assert!(!format!("{directive:?}").contains("secret-value")); + assert!( + serde_json::to_string(&directive) + .expect("serialize") + .contains("secret-value") + ); + assert!(serde_json::from_str::("\"\"").is_err()); +} diff --git a/crates/cli/tests/coverage/daemon/registry_tests.rs b/crates/cli/tests/coverage/daemon/registry_tests.rs new file mode 100644 index 000000000..145b3b54a --- /dev/null +++ b/crates/cli/tests/coverage/daemon/registry_tests.rs @@ -0,0 +1,1099 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::net::Ipv4Addr; +use std::sync::Barrier; + +use super::*; +use crate::daemon::common::identity::PublicIdentity; +use crate::daemon::common::protocol::SensitiveString; + +fn fingerprint(byte: u8) -> Fingerprint { + PublicIdentity::from_bytes(&[byte; 32]) + .expect("public identity") + .fingerprint() +} + +fn session(name: &str) -> McpSessionId { + McpSessionId::new(name).expect("session") +} + +fn launch(name: &str) -> WorkerLaunch { + WorkerLaunch { + activation_id: name.to_owned(), + activation_token: SensitiveString::new(format!("{name}-secret")).expect("secret"), + deadline_unix_ms: 15_000, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + } +} + +fn registration( + fingerprint: Fingerprint, + token_digest: TokenDigest, + session_id: &str, +) -> McpRegistration { + McpRegistration { + fingerprint, + token_digest, + session_id: session(session_id), + lease_expires_at_unix_ms: 30_000, + } +} + +fn worker(worker_id: &str) -> Arc { + Arc::new( + WorkerTarget::new( + worker_id, + "http://127.0.0.1:41000", + SensitiveString::new("internal-session-token").expect("token"), + ) + .expect("worker target"), + ) +} + +#[test] +fn first_mcp_wins_singleflight_and_retries_idempotently() { + let registry = Registry::new(false).with_retry_after_ms(25); + let fingerprint = fingerprint(1); + let token = TokenDigest::from_token(b"token-1"); + + let first = registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("first")) + .expect("first registration"); + assert!(matches!( + first, + BrokerDirective::LaunchWorker { + ref activation_id, + .. + } if activation_id == "first" + )); + + let retry = registry + .register_mcp( + registration(fingerprint, token, "mcp-a"), + launch("must-not-replace"), + ) + .expect("idempotent retry"); + assert!(matches!( + retry, + BrokerDirective::LaunchWorker { + ref activation_id, + .. + } if activation_id == "first" + )); + + let concurrent = registry + .register_mcp(registration(fingerprint, token, "mcp-b"), launch("second")) + .expect("concurrent registration"); + assert_eq!( + concurrent, + BrokerDirective::WaitForWorker { retry_after_ms: 25 } + ); + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot"), + RouteSnapshot { + state: RouteStateKind::Activating, + reference_count: 2, + launch_owner: Some(session("mcp-a")), + endpoint: None, + in_flight: 0, + } + ); +} + +#[test] +fn concurrent_registrations_issue_exactly_one_launch() { + const MCP_COUNT: usize = 32; + let registry = Arc::new(Registry::new(false)); + let barrier = Arc::new(Barrier::new(MCP_COUNT)); + let fingerprint = fingerprint(13); + let token = TokenDigest::from_token(b"token-13"); + let handles: Vec<_> = (0..MCP_COUNT) + .map(|index| { + let registry = Arc::clone(®istry); + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + registry + .register_mcp( + registration(fingerprint, token, &format!("mcp-{index:02}")), + launch(&format!("launch-{index:02}")), + ) + .expect("registration") + }) + }) + .collect(); + let directives: Vec<_> = handles + .into_iter() + .map(|handle| handle.join().expect("thread")) + .collect(); + + assert_eq!( + directives + .iter() + .filter(|directive| matches!(directive, BrokerDirective::LaunchWorker { .. })) + .count(), + 1 + ); + assert_eq!( + directives + .iter() + .filter(|directive| matches!(directive, BrokerDirective::WaitForWorker { .. })) + .count(), + MCP_COUNT - 1 + ); + assert_eq!( + registry + .snapshot(fingerprint) + .expect("snapshot") + .reference_count, + MCP_COUNT + ); +} + +#[test] +fn ready_worker_is_reused_and_request_guard_counts_in_flight() { + let registry = Registry::new(false); + let fingerprint = fingerprint(2); + let token = TokenDigest::from_token(b"token-2"); + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("launch")) + .expect("registration"); + let target = worker("worker-1"); + registry + .mark_worker_ready(fingerprint, "launch", Arc::clone(&target)) + .expect("worker ready"); + + assert_eq!( + registry + .register_mcp(registration(fingerprint, token, "mcp-b"), launch("unused")) + .expect("reuse"), + BrokerDirective::ReuseWorker { + endpoint: "http://127.0.0.1:41000".to_owned() + } + ); + let request = match registry.resolve_target(&token).expect("resolved") { + ResolvedTarget::Worker(request) => request, + ResolvedTarget::PassThrough => panic!("expected worker"), + }; + assert_eq!(target.in_flight(), 1); + assert_eq!(request.session_token(), "internal-session-token"); + drop(request); + assert_eq!(target.in_flight(), 0); +} + +#[test] +fn token_and_fingerprint_bindings_cannot_be_reassigned() { + let registry = Registry::new(false); + let first_fingerprint = fingerprint(3); + let other_fingerprint = fingerprint(4); + let token = TokenDigest::from_token(b"stable-token"); + registry + .restore_binding(first_fingerprint, token) + .expect("binding"); + assert_eq!( + registry.restore_binding(other_fingerprint, token), + Err(RegistryError::TokenAlreadyBound) + ); + assert_eq!( + registry.restore_binding( + first_fingerprint, + TokenDigest::from_token(b"different-token") + ), + Err(RegistryError::FingerprintTokenMismatch) + ); +} + +#[test] +fn final_reference_enters_non_revivable_drain() { + let registry = Registry::new(false); + let fingerprint = fingerprint(5); + let token = TokenDigest::from_token(b"token-5"); + let first_session = session("mcp-a"); + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("first")) + .expect("register"); + let target = worker("worker-1"); + registry + .mark_worker_ready(fingerprint, "first", Arc::clone(&target)) + .expect("ready"); + let request = match registry.resolve_target(&token).expect("request") { + ResolvedTarget::Worker(request) => request, + ResolvedTarget::PassThrough => panic!("expected worker"), + }; + + assert!(matches!( + registry + .release_mcp(fingerprint, &first_session, 2_000) + .expect("release"), + ReleaseAction::BeginDrain { + deadline_unix_ms: 2_000, + .. + } + )); + assert!(matches!( + registry.resolve_target(&token), + Err(ResolveError::Unavailable(RouteStateKind::Draining)) + )); + assert_eq!( + registry + .register_mcp(registration(fingerprint, token, "mcp-b"), launch("second")) + .expect("wait during drain"), + BrokerDirective::WaitForWorker { + retry_after_ms: DEFAULT_RETRY_AFTER_MS + } + ); + registry + .release_mcp(fingerprint, &session("mcp-b"), 2_000) + .expect("release waiting MCP"); + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot").state, + RouteStateKind::Draining + ); + assert_eq!( + registry + .register_mcp(registration(fingerprint, token, "mcp-c"), launch("second")) + .expect("replacement waits during drain"), + BrokerDirective::WaitForWorker { + retry_after_ms: DEFAULT_RETRY_AFTER_MS + } + ); + assert_eq!( + registry.finish_draining(fingerprint, 1_999), + Err(RegistryError::DrainInProgress) + ); + drop(request); + assert_eq!( + registry.finish_draining(fingerprint, 1_999), + Ok(DrainCompletion::ActivationRequired { + session_id: session("mcp-c") + }) + ); + assert!(matches!( + registry + .register_mcp(registration(fingerprint, token, "mcp-c"), launch("second")) + .expect("new generation"), + BrokerDirective::LaunchWorker { + ref activation_id, + .. + } if activation_id == "second" + )); +} + +#[test] +fn activation_failure_is_shared_pass_through_until_zero_refs() { + let registry = Registry::new(false); + let fingerprint = fingerprint(6); + let token = TokenDigest::from_token(b"token-6"); + let session_id = session("mcp-a"); + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("failed")) + .expect("register"); + registry + .mark_activation_failed(fingerprint, "failed") + .expect("failure"); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); + assert_eq!( + registry + .register_mcp( + registration(fingerprint, token, "mcp-concurrent"), + launch("ignored") + ) + .unwrap(), + BrokerDirective::UsePassThrough + ); + registry + .release_mcp(fingerprint, &session_id, 2_000) + .expect("release"); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); + registry + .release_mcp(fingerprint, &session("mcp-concurrent"), 2_001) + .expect("last release"); + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot").state, + RouteStateKind::Empty + ); + assert!(matches!( + registry + .register_mcp(registration(fingerprint, token, "mcp-b"), launch("retry")) + .expect("retry"), + BrokerDirective::LaunchWorker { + ref activation_id, + .. + } if activation_id == "retry" + )); +} + +#[test] +fn global_pass_through_never_activates_or_accepts_workers() { + let registry = Registry::new(true); + let fingerprint = fingerprint(7); + let token = TokenDigest::from_token(b"token-7"); + assert_eq!( + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("unused")) + .expect("registration"), + BrokerDirective::UsePassThrough + ); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); + assert_eq!( + registry.mark_worker_ready(fingerprint, "unused", worker("worker-1")), + Err(RegistryError::InvalidState { + expected: RouteStateKind::Activating, + actual: RouteStateKind::PassThrough, + }) + ); +} + +#[test] +fn worker_crash_nominates_one_live_mcp_and_relaunches() { + let registry = Registry::new(false); + let fingerprint = fingerprint(8); + let token = TokenDigest::from_token(b"token-8"); + registry + .register_mcp(registration(fingerprint, token, "mcp-b"), launch("first")) + .expect("first"); + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("unused")) + .expect("second"); + registry + .mark_worker_ready(fingerprint, "first", worker("worker-1")) + .expect("ready"); + assert_eq!( + registry + .worker_failed(fingerprint, "worker-1", 10_000) + .expect("failure"), + WorkerFailureAction::NominateMcp { + session_id: session("mcp-a") + } + ); + assert_eq!( + registry.begin_relaunch(fingerprint, &session("mcp-b"), launch("replacement")), + Err(RegistryError::NotLaunchOwner) + ); + assert!(matches!( + registry + .begin_relaunch(fingerprint, &session("mcp-a"), launch("replacement")) + .expect("relaunch"), + BrokerDirective::LaunchWorker { + ref activation_id, + .. + } if activation_id == "replacement" + )); +} + +#[test] +fn expired_launch_owner_is_transferred_idempotently() { + let registry = Registry::new(false); + let fingerprint = fingerprint(9); + let token = TokenDigest::from_token(b"token-9"); + let mut first = registration(fingerprint, token, "mcp-a"); + first.lease_expires_at_unix_ms = 100; + let mut second = registration(fingerprint, token, "mcp-b"); + second.lease_expires_at_unix_ms = 1_000; + registry + .register_mcp(first, launch("launch")) + .expect("first"); + registry + .register_mcp(second, launch("unused")) + .expect("second"); + + let actions = registry.expire_mcp_leases(100, 2_000); + assert_eq!(actions.len(), 1); + assert!(matches!( + &actions[0].1, + ReleaseAction::TransferActivation { + session_id, + directive: BrokerDirective::LaunchWorker { activation_id, .. }, + } if session_id == &session("mcp-b") && activation_id == "launch" + )); + assert_eq!( + registry + .snapshot(fingerprint) + .expect("snapshot") + .reference_count, + 1 + ); + assert!(registry.expire_mcp_leases(100, 2_000).is_empty()); +} + +#[test] +fn simultaneous_lease_expiry_emits_one_terminal_action() { + let registry = Registry::new(false); + let fingerprint = fingerprint(10); + let token = TokenDigest::from_token(b"token-10"); + let mut first = registration(fingerprint, token, "mcp-a"); + first.lease_expires_at_unix_ms = 100; + let mut second = registration(fingerprint, token, "mcp-b"); + second.lease_expires_at_unix_ms = 100; + registry + .register_mcp(first, launch("launch")) + .expect("first"); + registry + .register_mcp(second, launch("unused")) + .expect("second"); + + let actions = registry.expire_mcp_leases(100, 2_000); + assert_eq!(actions.len(), 1); + assert!(matches!( + &actions[0].1, + ReleaseAction::CancelActivation { activation_id } if activation_id == "launch" + )); + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot"), + RouteSnapshot { + state: RouteStateKind::Empty, + reference_count: 0, + launch_owner: None, + endpoint: None, + in_flight: 0, + } + ); +} + +#[test] +fn recovery_waits_for_deadline_then_nominates_a_live_mcp() { + let registry = Registry::new(false); + let fingerprint = fingerprint(11); + let token = TokenDigest::from_token(b"token-11"); + registry + .restore_binding(fingerprint, token) + .expect("persisted binding"); + registry + .begin_recovery(fingerprint, None, 1_000) + .expect("recovery"); + assert_eq!( + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("unused")) + .expect("reconnecting MCP"), + BrokerDirective::WaitForWorker { + retry_after_ms: DEFAULT_RETRY_AFTER_MS + } + ); + assert!(matches!( + registry.finish_recovery(fingerprint, 999, 2_000), + Err(RegistryError::RecoveryInProgress) + )); + assert!(matches!( + registry + .finish_recovery(fingerprint, 1_000, 2_000) + .expect("recovery deadline"), + RecoveryAction::NominateMcp { session_id } if session_id == session("mcp-a") + )); + assert!(matches!( + registry + .begin_relaunch(fingerprint, &session("mcp-a"), launch("replacement")) + .expect("replacement activation"), + BrokerDirective::LaunchWorker { activation_id, .. } if activation_id == "replacement" + )); +} + +#[test] +fn recovered_worker_becomes_ready_when_an_mcp_reconnects() { + let registry = Registry::new(false); + let fingerprint = fingerprint(12); + let token = TokenDigest::from_token(b"token-12"); + registry + .restore_binding(fingerprint, token) + .expect("persisted binding"); + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("restart")) + .expect("reconnecting MCP"); + let permit = registry + .authorize_worker_recovery(fingerprint, "worker-recovered") + .expect("recovery authorization"); + assert_eq!( + registry + .publish_recovered_worker(fingerprint, &permit, worker("worker-recovered")) + .expect("worker registration"), + Some("restart".to_owned()) + ); + assert_eq!( + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("unused")) + .expect("reconnecting MCP"), + BrokerDirective::ReuseWorker { + endpoint: "http://127.0.0.1:41000".to_owned() + } + ); + registry + .renew_mcp(fingerprint, &session("mcp-a"), 50_000) + .expect("renewal"); + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot").state, + RouteStateKind::Ready + ); +} + +#[test] +fn recovered_worker_without_references_is_not_authorized() { + let registry = Registry::new(false); + let fingerprint = fingerprint(18); + let token = TokenDigest::from_token(b"token-18"); + registry + .restore_binding(fingerprint, token) + .expect("persisted binding"); + registry + .begin_recovery(fingerprint, None, 100) + .expect("recovery"); + assert_eq!( + registry.authorize_worker_recovery(fingerprint, "worker-recovered"), + Err(RegistryError::NoLiveMcpReferences) + ); +} + +#[test] +fn expired_activation_enters_transient_pass_through_until_all_references_leave() { + let registry = Registry::new(false); + let fingerprint = fingerprint(14); + let token = TokenDigest::from_token(b"token-14"); + let mut expiring_launch = launch("expiring"); + expiring_launch.deadline_unix_ms = 100; + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), expiring_launch) + .expect("registration"); + registry + .register_mcp(registration(fingerprint, token, "mcp-b"), launch("unused")) + .expect("second registration"); + + assert!(registry.expire_activations(99).is_empty()); + assert_eq!( + registry.expire_activations(100), + vec![ExpiredActivation { + fingerprint, + activation_id: "expiring".to_owned(), + }] + ); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); + assert_eq!( + registry + .register_mcp( + registration(fingerprint, token, "mcp-c"), + launch("must-not-launch"), + ) + .expect("pass-through registration"), + BrokerDirective::UsePassThrough + ); + + for session_id in ["mcp-a", "mcp-b", "mcp-c"] { + registry + .release_mcp(fingerprint, &session(session_id), 1_000) + .expect("release"); + } + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot").state, + RouteStateKind::Empty + ); +} + +#[test] +fn authenticated_worker_communication_failure_is_route_wide_pass_through() { + let registry = Registry::new(false); + let fingerprint = fingerprint(15); + let token = TokenDigest::from_token(b"token-15"); + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("launch")) + .expect("registration"); + registry + .mark_worker_ready(fingerprint, "launch", worker("worker-failed")) + .expect("ready"); + + assert_eq!( + registry + .mark_worker_communication_failed(fingerprint, "worker-failed") + .expect("communication failure"), + None + ); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); + assert_eq!( + registry.mark_worker_communication_failed(fingerprint, "worker-failed"), + Ok(None) + ); + registry + .release_mcp(fingerprint, &session("mcp-a"), 1_000) + .expect("release"); + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot").state, + RouteStateKind::Empty + ); +} + +#[test] +fn delayed_failure_from_old_worker_does_not_displace_new_ready_generation() { + let registry = Registry::new(false); + let fingerprint = fingerprint(16); + let token = TokenDigest::from_token(b"token-16"); + registry + .register_mcp(registration(fingerprint, token, "mcp-a"), launch("launch")) + .expect("registration"); + registry + .mark_worker_ready(fingerprint, "launch", worker("worker-old")) + .expect("ready"); + registry + .worker_failed(fingerprint, "worker-old", 10_000) + .expect("worker failure"); + registry + .begin_relaunch(fingerprint, &session("mcp-a"), launch("replacement")) + .expect("replacement launch"); + registry + .mark_worker_ready(fingerprint, "replacement", worker("worker-new")) + .expect("replacement ready"); + + assert_eq!( + registry.mark_worker_communication_failed(fingerprint, "worker-old"), + Err(RegistryError::WorkerMismatch) + ); + assert_eq!( + registry.snapshot(fingerprint).expect("snapshot").state, + RouteStateKind::Ready + ); +} + +#[test] +fn recovered_worker_supersedes_restart_activation_without_a_second_worker() { + let registry = Registry::new(false); + let fingerprint = fingerprint(17); + let token = TokenDigest::from_token(b"token-17"); + registry + .register_mcp( + registration(fingerprint, token, "mcp-a"), + launch("restart-activation"), + ) + .expect("reconnected MCP"); + + let permit = registry + .authorize_worker_recovery(fingerprint, "worker-survivor") + .expect("recovery authorization"); + assert_eq!( + registry + .publish_recovered_worker(fingerprint, &permit, worker("worker-survivor")) + .expect("recovered worker"), + Some("restart-activation".to_owned()) + ); + assert_eq!( + registry + .register_mcp(registration(fingerprint, token, "mcp-b"), launch("unused"),) + .expect("reuse recovered worker"), + BrokerDirective::ReuseWorker { + endpoint: "http://127.0.0.1:41000".to_owned(), + } + ); +} + +#[test] +fn recovery_requires_a_live_known_route_and_rejects_permanent_pass_through() { + let unknown = Registry::new(false); + assert_eq!( + unknown.authorize_worker_recovery(fingerprint(21), "worker"), + Err(RegistryError::UnknownRoute) + ); + + let pass_through = Registry::new(true); + let fingerprint = fingerprint(22); + let token = TokenDigest::from_token(b"token-22"); + pass_through + .register_mcp(registration(fingerprint, token, "mcp"), launch("unused")) + .expect("pass-through registration"); + assert_eq!( + pass_through.authorize_worker_recovery(fingerprint, "worker"), + Err(RegistryError::RecoveryNotAuthorized) + ); +} + +#[test] +fn pass_through_route_is_not_routable_without_a_live_mcp_reference() { + let registry = Registry::new(true); + let fingerprint = fingerprint(23); + let token = TokenDigest::from_token(b"token-23"); + registry + .register_mcp(registration(fingerprint, token, "mcp"), launch("unused")) + .expect("registration"); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); + registry + .release_mcp(fingerprint, &session("mcp"), 1_000) + .expect("release"); + assert!(matches!( + registry.resolve_target(&token), + Err(ResolveError::Unavailable(RouteStateKind::PassThrough)) + )); +} + +#[test] +fn stable_route_bindings_are_bounded_without_permitting_rebinding() { + let registry = Registry::new(false).with_route_capacity(1); + let first = fingerprint(24); + let second = fingerprint(25); + let token = TokenDigest::from_token(b"bounded-token"); + registry + .register_mcp(registration(first, token, "mcp-a"), launch("first")) + .expect("first route"); + assert_eq!( + registry.register_mcp( + registration(second, TokenDigest::from_token(b"another-token"), "mcp-b"), + launch("second"), + ), + Err(RegistryError::RouteCapacityReached) + ); + assert_eq!( + registry.register_mcp(registration(second, token, "mcp-c"), launch("rebind")), + Err(RegistryError::TokenAlreadyBound) + ); +} + +#[test] +fn capacity_pressure_evicts_only_a_zero_reference_empty_route() { + let registry = Registry::new(false).with_route_capacity(1); + let first = fingerprint(26); + let first_token = TokenDigest::from_token(b"first-token"); + registry + .register_mcp(registration(first, first_token, "mcp-a"), launch("first")) + .expect("first route"); + registry + .release_mcp(first, &session("mcp-a"), 1_000) + .expect("release empty activation"); + + let second = fingerprint(27); + let second_token = TokenDigest::from_token(b"second-token"); + assert!(matches!( + registry + .register_mcp( + registration(second, second_token, "mcp-b"), + launch("second"), + ) + .expect("inactive route should be evicted"), + BrokerDirective::LaunchWorker { .. } + )); + assert!(matches!( + registry.resolve_target(&first_token), + Err(ResolveError::UnknownToken) + )); +} + +#[test] +fn route_wide_pass_through_cancels_activation_and_preserves_permanent_routes() { + let registry = Registry::new(false); + let route_fingerprint = fingerprint(28); + let token = TokenDigest::from_token(b"token-28"); + registry + .register_mcp( + registration(route_fingerprint, token, "mcp"), + launch("cancel-me"), + ) + .unwrap(); + assert_eq!( + registry.mark_route_pass_through(route_fingerprint).unwrap(), + Some("cancel-me".into()) + ); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); + assert_eq!( + registry.mark_route_pass_through(route_fingerprint).unwrap(), + None + ); + + let permanent = Registry::new(true); + let permanent_fingerprint = fingerprint(29); + permanent + .register_mcp( + registration( + permanent_fingerprint, + TokenDigest::from_token(b"token-29"), + "mcp", + ), + launch("unused"), + ) + .unwrap(); + assert_eq!( + permanent + .mark_route_pass_through(permanent_fingerprint) + .unwrap(), + None + ); + assert_eq!( + permanent.snapshot(permanent_fingerprint).unwrap().state, + RouteStateKind::PassThrough + ); +} + +#[test] +fn recovery_completion_covers_live_empty_and_draining_routes() { + let live = Registry::new(false); + let live_fingerprint = fingerprint(30); + let live_token = TokenDigest::from_token(b"token-30"); + live.register_mcp( + registration(live_fingerprint, live_token, "mcp"), + launch("launch"), + ) + .unwrap(); + live.begin_recovery(live_fingerprint, Some(worker("survivor")), 100) + .unwrap(); + assert!(matches!( + live.finish_recovery(live_fingerprint, 100, 200).unwrap(), + RecoveryAction::WorkerRecovered + )); + + let empty = Registry::new(false); + let empty_fingerprint = fingerprint(31); + empty + .restore_binding(empty_fingerprint, TokenDigest::from_token(b"token-31")) + .unwrap(); + empty.begin_recovery(empty_fingerprint, None, 100).unwrap(); + assert!(matches!( + empty.finish_recovery(empty_fingerprint, 100, 200).unwrap(), + RecoveryAction::RouteEmpty + )); + + let draining = Registry::new(false); + let draining_fingerprint = fingerprint(32); + draining + .restore_binding(draining_fingerprint, TokenDigest::from_token(b"token-32")) + .unwrap(); + draining + .begin_recovery(draining_fingerprint, Some(worker("survivor")), 100) + .unwrap(); + assert!(matches!( + draining + .finish_recovery(draining_fingerprint, 100, 200) + .unwrap(), + RecoveryAction::BeginDrain { + deadline_unix_ms: 200, + .. + } + )); + assert_eq!( + draining.finish_draining(draining_fingerprint, 200).unwrap(), + DrainCompletion::RouteEmpty + ); +} + +#[test] +fn registry_rejects_stale_worker_generations_and_invalid_state_transitions() { + let registry = Registry::new(false); + let fingerprint = fingerprint(33); + let token = TokenDigest::from_token(b"token-33"); + registry + .register_mcp(registration(fingerprint, token, "mcp"), launch("active")) + .unwrap(); + assert_eq!( + registry.mark_worker_ready(fingerprint, "stale", worker("worker")), + Err(RegistryError::ActivationMismatch) + ); + assert_eq!( + registry.mark_activation_failed(fingerprint, "stale"), + Err(RegistryError::ActivationMismatch) + ); + assert_eq!( + registry.worker_failed(fingerprint, "worker", 100), + Err(RegistryError::InvalidState { + expected: RouteStateKind::Ready, + actual: RouteStateKind::Activating, + }) + ); + assert_eq!( + registry.begin_relaunch(fingerprint, &session("mcp"), launch("new")), + Err(RegistryError::InvalidState { + expected: RouteStateKind::Recovering, + actual: RouteStateKind::Activating, + }) + ); + assert_eq!( + registry.finish_draining(fingerprint, 100), + Err(RegistryError::InvalidState { + expected: RouteStateKind::Draining, + actual: RouteStateKind::Activating, + }) + ); + assert_eq!( + registry + .mark_worker_communication_failed(fingerprint, "worker") + .unwrap(), + Some("active".into()) + ); + assert!(matches!( + registry.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); +} + +#[test] +fn registry_unknown_route_and_session_errors_are_explicit() { + let registry = Registry::new(false); + let unknown = fingerprint(34); + assert_eq!(registry.snapshot(unknown), Err(RegistryError::UnknownRoute)); + assert_eq!( + registry.renew_mcp(unknown, &session("missing"), 100), + Err(RegistryError::UnknownRoute) + ); + assert!(matches!( + registry.release_mcp(unknown, &session("missing"), 100), + Err(RegistryError::UnknownRoute) + )); + assert_eq!( + registry.mark_route_pass_through(unknown), + Err(RegistryError::UnknownRoute) + ); + assert!(matches!( + registry.finish_recovery(unknown, 100, 200), + Err(RegistryError::UnknownRoute) + )); + + let fingerprint = fingerprint(35); + registry + .register_mcp( + registration(fingerprint, TokenDigest::from_token(b"token-35"), "known"), + launch("launch"), + ) + .unwrap(); + assert_eq!( + registry.renew_mcp(fingerprint, &session("missing"), 100), + Err(RegistryError::UnknownMcpSession) + ); + assert!(matches!( + registry + .release_mcp(fingerprint, &session("missing"), 100) + .unwrap(), + ReleaseAction::NoChange + )); +} + +#[test] +fn recovery_permits_cover_existing_ready_and_recovering_worker_generations() { + let registry = Registry::new(false); + let fingerprint = fingerprint(36); + let token = TokenDigest::from_token(b"token-36"); + registry + .register_mcp(registration(fingerprint, token, "mcp"), launch("launch")) + .unwrap(); + registry + .mark_worker_ready(fingerprint, "launch", worker("survivor")) + .unwrap(); + + let ready = registry + .authorize_worker_recovery(fingerprint, "survivor") + .expect("ready worker permit"); + assert!(matches!( + ready, + RecoveryPermit::ExistingWorker { + recovering: false, + .. + } + )); + assert_eq!( + registry + .publish_recovered_worker(fingerprint, &ready, worker("survivor")) + .unwrap(), + None + ); + + registry + .begin_recovery(fingerprint, Some(worker("survivor")), 1_000) + .unwrap(); + let recovering = registry + .authorize_worker_recovery(fingerprint, "survivor") + .expect("recovering worker permit"); + assert!(matches!( + recovering, + RecoveryPermit::ExistingWorker { + recovering: true, + .. + } + )); + assert_eq!( + registry + .publish_recovered_worker(fingerprint, &recovering, worker("survivor")) + .unwrap(), + None + ); + + registry + .begin_recovery(fingerprint, Some(worker("new-generation")), 2_000) + .unwrap(); + assert_eq!( + registry.publish_recovered_worker(fingerprint, &recovering, worker("survivor")), + Err(RegistryError::RecoveryGenerationChanged) + ); +} + +#[test] +fn communication_failures_preserve_draining_and_mismatched_recovery_generations() { + let registry = Registry::new(false); + let draining_fingerprint = fingerprint(37); + let token = TokenDigest::from_token(b"token-37"); + registry + .register_mcp( + registration(draining_fingerprint, token, "mcp"), + launch("launch"), + ) + .unwrap(); + registry + .mark_worker_ready(draining_fingerprint, "launch", worker("active")) + .unwrap(); + registry + .release_mcp(draining_fingerprint, &session("mcp"), 1_000) + .unwrap(); + assert_eq!( + registry.mark_worker_communication_failed(draining_fingerprint, "active"), + Err(RegistryError::InvalidState { + expected: RouteStateKind::Ready, + actual: RouteStateKind::Draining, + }) + ); + assert_eq!( + registry.mark_worker_communication_failed(draining_fingerprint, "different"), + Err(RegistryError::WorkerMismatch) + ); + + let recovering = Registry::new(false); + let fingerprint = fingerprint(38); + let token = TokenDigest::from_token(b"token-38"); + recovering + .register_mcp(registration(fingerprint, token, "mcp"), launch("launch")) + .unwrap(); + recovering + .begin_recovery(fingerprint, Some(worker("survivor")), 1_000) + .unwrap(); + assert_eq!( + recovering.mark_worker_communication_failed(fingerprint, "different"), + Err(RegistryError::WorkerMismatch) + ); + assert_eq!( + recovering + .snapshot(fingerprint) + .unwrap() + .endpoint + .as_deref(), + Some("http://127.0.0.1:41000") + ); + assert_eq!( + recovering + .mark_worker_communication_failed(fingerprint, "survivor") + .unwrap(), + None + ); + assert!(matches!( + recovering.resolve_target(&token), + Ok(ResolvedTarget::PassThrough) + )); +} diff --git a/crates/cli/tests/coverage/daemon/routes_tests.rs b/crates/cli/tests/coverage/daemon/routes_tests.rs new file mode 100644 index 000000000..9d6044d6b --- /dev/null +++ b/crates/cli/tests/coverage/daemon/routes_tests.rs @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn classifies_only_supported_public_paths() { + assert_eq!( + PublicRoute::from_path("/v1/messages"), + Some(PublicRoute::Provider(ProviderRoute::Anthropic)) + ); + assert_eq!( + PublicRoute::from_path("/hooks/codex"), + Some(PublicRoute::Hook(HookRoute::Codex)) + ); + assert_eq!(PublicRoute::from_path("/admin"), None); +} + +#[test] +fn composes_openai_v1_once() { + let config = GatewayConfig::default(); + assert_eq!( + ProviderRoute::OpenAi.upstream_url(&config, "/v1/responses?x=1"), + "https://api.openai.com/v1/responses?x=1" + ); + assert_eq!( + ProviderRoute::OpenAi.upstream_url(&config, "/responses"), + "https://api.openai.com/v1/responses" + ); + assert_eq!( + ProviderRoute::OpenAi.upstream_url(&config, "/backend-api/codex/responses?client=codex"), + "https://api.openai.com/v1/responses?client=codex" + ); + assert_eq!( + PublicRoute::from_path("/backend-api/codex/responses"), + Some(PublicRoute::Provider(ProviderRoute::OpenAi)) + ); +} + +#[test] +fn composes_anthropic_paths_without_openai_normalization() { + let mut config = GatewayConfig { + anthropic_base_url: "https://api.anthropic.com/custom/".into(), + ..GatewayConfig::default() + }; + assert_eq!( + ProviderRoute::Anthropic.upstream_url(&config, "/v1/messages?beta=true"), + "https://api.anthropic.com/custom/v1/messages?beta=true" + ); + config.anthropic_base_url = "https://api.anthropic.com".into(); + assert_eq!( + ProviderRoute::Anthropic.upstream_url(&config, "/v1/messages"), + "https://api.anthropic.com/v1/messages" + ); + config.anthropic_base_url = "https://api.anthropic.com/v1".into(); + assert_eq!( + ProviderRoute::Anthropic.upstream_url(&config, "/v1/messages"), + "https://api.anthropic.com/v1/messages" + ); +} diff --git a/crates/cli/tests/coverage/daemon/server_tests.rs b/crates/cli/tests/coverage/daemon/server_tests.rs new file mode 100644 index 000000000..bcfed1994 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/server_tests.rs @@ -0,0 +1,3181 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use base64::Engine; +use std::sync::Arc; + +use super::*; +use crate::daemon::common::client::{begin_handshake, control_client, post_json}; +use crate::daemon::common::control::{ + WorkerBootstrap, WorkerNetworkHintProof, WorkerReadyPayload, WorkerRegisterResponse, +}; +use crate::daemon::common::routes::HookRoute; +use crate::daemon::common::state::ROUTE_TOKEN_ENV; +use crate::daemon::common::worker_tls::pooled_worker_tls_client; +use crate::daemon::worker::test_router_with_control_tokens; +use crate::test_support::{EnvScope, PLUGIN_CONFIG_TEST_LOCK}; +use axum::Router; +use axum::extract::State; +use axum::routing::post; +use http_body_util::BodyExt as _; +use tower::ServiceExt as _; + +type CapturedProviderRequest = Arc>>; + +#[tokio::test] +async fn shared_model_catalogs_disable_cache_reuse_between_credentials() { + let provider = Router::new().route( + "/v1/models", + axum::routing::get(|headers: HeaderMap| async move { + Response::builder() + .header(axum::http::header::CACHE_CONTROL, "public, max-age=3600") + .header(axum::http::header::CACHE_CONTROL, "private") + .body(Body::from( + headers[AUTHORIZATION].to_str().unwrap().to_owned(), + )) + .unwrap() + }), + ); + let provider_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let provider_origin = format!("http://{}", provider_listener.local_addr().unwrap()); + let provider_task = + tokio::spawn(async { axum::serve(provider_listener, provider).await.unwrap() }); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("http://{}", listener.local_addr().unwrap()); + let state = test_daemon_state_at( + true, + "", + GatewayConfig { + openai_base_url: provider_origin, + ..GatewayConfig::default() + }, + origin.clone(), + ); + let daemon_task = tokio::spawn({ + let state = Arc::clone(&state); + async move { axum::serve(listener, router(state)).await.unwrap() } + }); + let app = router(Arc::clone(&state)); + for (index, provider_auth) in [(1_u8, "Bearer first-user"), (2, "Bearer second-user")] { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([index; 32]); + let identity = MachineIdentity::generate().unwrap().identity; + assert_eq!( + enroll_test_mcp(&state, &origin, &identity, &token, provider_auth) + .await + .status(), + StatusCode::OK, + ); + let response = app + .clone() + .oneshot( + Request::get("/models") + .header(CLIENT_TOKEN_HEADER, &token) + .header(AUTHORIZATION, provider_auth) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers()[axum::http::header::CACHE_CONTROL], + "no-store" + ); + assert_eq!( + response + .headers() + .get_all(axum::http::header::CACHE_CONTROL) + .iter() + .count(), + 1 + ); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + provider_auth + ); + } + for path in ["/models", "/v1/models"] { + let response = app + .clone() + .oneshot(Request::get(path).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!( + response.headers()[axum::http::header::CACHE_CONTROL], + "no-store" + ); + } + daemon_task.abort(); + provider_task.abort(); +} + +#[tokio::test(start_paused = true)] +async fn challenge_admission_limits_transport_peers_before_polling_bodies() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x91; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let app = router(Arc::clone(&state)); + let identity = MachineIdentity::generate().unwrap().identity; + let mut challenge = ChallengeRequest { + initiator: crate::daemon::common::control::descriptor(ComponentRole::Mcp), + initiator_instance_id: "limited-peer".into(), + initiator_public_identity: identity.public_identity(), + initiator_fingerprint: identity.fingerprint(), + initiator_nonce: ChallengeRecord::generate(1, 1).unwrap().challenge().nonce, + }; + for port in 1..=CHALLENGES_PER_PEER_WINDOW { + let response = app + .clone() + .oneshot( + Request::post(CHALLENGE_PATH) + .extension(ConnectInfo(SocketAddr::from(([192, 0, 2, 1], port as u16)))) + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_vec(&challenge).unwrap())) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + } + let response = app + .clone() + .oneshot( + Request::post(CHALLENGE_PATH) + .extension(ConnectInfo(SocketAddr::from(([192, 0, 2, 1], 65535)))) + .header("x-forwarded-for", "192.0.2.99") + .header(CONTENT_TYPE, "application/json") + .body(Body::from_stream(futures_util::stream::poll_fn( + |_| -> std::task::Poll>> { + panic!("rate-limited challenge body polled") + }, + ))) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + assert_eq!(response.headers()[RETRY_AFTER], "15"); + assert_eq!( + lock(&state.challenges).len(), + CHALLENGES_PER_PEER_WINDOW as usize + ); + + // A worker on another peer still has capacity and needs no MCP route-token header. + challenge.initiator = crate::daemon::common::control::descriptor(ComponentRole::Worker); + let request = |peer| { + Request::post(CHALLENGE_PATH) + .extension(ConnectInfo(SocketAddr::from((peer, 1)))) + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_vec(&challenge).unwrap())) + .unwrap() + }; + assert_eq!( + app.clone() + .oneshot(request([192, 0, 2, 2])) + .await + .unwrap() + .status(), + StatusCode::OK + ); + tokio::time::advance(Duration::from_millis(CHALLENGE_LIFETIME_MS)).await; + assert_eq!( + app.oneshot(request([192, 0, 2, 1])).await.unwrap().status(), + StatusCode::OK + ); +} + +#[tokio::test(start_paused = true)] +async fn challenge_peer_tracking_is_bounded_and_expired_entries_are_reclaimed() { + let now = tokio::time::Instant::now(); + let peers: ChallengePeers = Arc::new(Mutex::new( + (0..MAX_CHALLENGE_PEERS) + .map(|index| (IpAddr::V4(Ipv4Addr::from(index as u32)), (now, 1))) + .collect(), + )); + let app = Router::new() + .route("/", post(|| async { StatusCode::NO_CONTENT })) + .layer(from_fn_with_state(Arc::clone(&peers), limit_challenges)); + let request = || { + Request::post("/") + .extension(ConnectInfo("192.0.2.1:1".parse::().unwrap())) + .body(Body::empty()) + .unwrap() + }; + assert_eq!( + app.clone().oneshot(request()).await.unwrap().status(), + StatusCode::TOO_MANY_REQUESTS + ); + assert_eq!(lock(&peers).len(), MAX_CHALLENGE_PEERS); + tokio::time::advance(Duration::from_millis(CHALLENGE_LIFETIME_MS)).await; + assert_eq!( + app.oneshot(request()).await.unwrap().status(), + StatusCode::NO_CONTENT + ); + assert_eq!(lock(&peers).len(), 1); +} + +async fn enroll_test_mcp( + state: &Arc, + origin: &str, + identity: &MachineIdentity, + token: &str, + session: &str, +) -> Response { + let credential = RouteCredential::parse(token.to_owned()).unwrap(); + let handshake = begin_handshake( + &control_client().unwrap(), + origin, + ComponentRole::Mcp, + identity, + session, + Some(credential.digest()), + ) + .await + .unwrap(); + let worker_network = WorkerNetworkHintProof::sign( + WorkerNetworkHint::new("127.0.0.1", None).unwrap(), + &handshake.proof.transcript.daemon_target, + session, + &handshake.proof.transcript.challenge_id, + &identity.fingerprint(), + identity, + ) + .unwrap(); + register_mcp( + State(Arc::clone(state)), + HeaderMap::from_iter([( + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderValue::from_str(token).unwrap(), + )]), + Json(McpRegisterRequest { + proof: handshake.proof, + worker_network, + }), + ) + .await +} + +#[tokio::test] +async fn open_enrollment_binds_new_tokens_without_allowlist_and_prevents_route_takeover() { + for pass_through in [false, true] { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("http://{}", listener.local_addr().unwrap()); + let first = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0xa1; 32]); + let second = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0xa2; 32]); + let state = test_daemon_state_at( + pass_through, + &first, + GatewayConfig::default(), + origin.clone(), + ); + let server = tokio::spawn({ + let state = Arc::clone(&state); + async move { axum::serve(listener, router(state)).await.unwrap() } + }); + let app = router(Arc::clone(&state)); + let identity = MachineIdentity::generate().unwrap().identity; + let other_identity = MachineIdentity::generate().unwrap().identity; + for (token, machine, session) in [ + (&first, &identity, "first"), + (&second, &other_identity, "second"), + ] { + // Rejection must happen from the request head, before reading even one body frame. + let body = Body::from_stream(futures_util::stream::poll_fn( + |_| -> std::task::Poll>> { + panic!("unregistered request body polled") + }, + )); + let response = app + .clone() + .oneshot( + Request::post("/hooks/pi") + .header(CLIENT_TOKEN_HEADER, token) + .body(body) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let response = enroll_test_mcp(&state, &origin, machine, token, session).await; + assert_eq!(response.status(), StatusCode::OK); + let response: McpRegisterResponse = + serde_json::from_slice(&response.into_body().collect().await.unwrap().to_bytes()) + .unwrap(); + assert_eq!( + matches!(response.directive, BrokerDirective::UsePassThrough), + pass_through + ); + assert_eq!( + matches!(response.directive, BrokerDirective::LaunchWorker { .. }), + !pass_through + ); + let response = app + .clone() + .oneshot( + Request::post("/hooks/pi") + .header(CLIENT_TOKEN_HEADER, token) + .body(Body::from("{}")) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!( + response.status(), + if pass_through { + StatusCode::OK + } else { + StatusCode::SERVICE_UNAVAILABLE + } + ); + } + assert_eq!( + enroll_test_mcp(&state, &origin, &other_identity, &first, "takeover") + .await + .status(), + StatusCode::UNAUTHORIZED, + ); + let third = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0xa3; 32]); + assert_eq!( + enroll_test_mcp(&state, &origin, &identity, &third, "rebind") + .await + .status(), + StatusCode::UNAUTHORIZED, + ); + state + .registry + .release_mcp( + identity.fingerprint(), + &McpSessionId::new("first").unwrap(), + u64::MAX, + ) + .unwrap(); + let response = app + .oneshot( + Request::post("/hooks/pi") + .header(CLIENT_TOKEN_HEADER, &first) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + if pass_through { + assert!(lock(&state.activations).is_empty()); + } + server.abort(); + } +} + +#[test] +fn worker_endpoint_rejects_bind_only_and_non_origin_values() { + assert!(validate_worker_endpoint("http://127.0.0.1:1234", None).is_ok()); + assert!(validate_worker_endpoint("http://0.0.0.0:1234", None).is_err()); + assert!(validate_worker_endpoint("http://127.0.0.1:1234/path", None).is_err()); + assert!(validate_worker_endpoint("http://127.0.0.1", None).is_err()); + assert!(validate_worker_endpoint("http://192.0.2.2:1234", None).is_err()); + assert!(validate_worker_endpoint("https://192.0.2.2:1234", Some("root")).is_ok()); +} + +#[test] +fn public_credential_requires_exactly_one_valid_value() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([7_u8; 32]); + let mut headers = HeaderMap::new(); + assert!(public_credential(&headers).is_err()); + headers.insert( + CLIENT_TOKEN_HEADER, + HeaderValue::from_str(&token).expect("header"), + ); + assert!(public_credential(&headers).is_ok()); + headers.append( + CLIENT_TOKEN_HEADER, + HeaderValue::from_str(&token).expect("header"), + ); + assert!(public_credential(&headers).is_err()); + + let mut non_utf8 = HeaderMap::new(); + non_utf8.insert( + CLIENT_TOKEN_HEADER, + HeaderValue::from_bytes(&[0xff]).expect("opaque header value"), + ); + assert!(public_credential(&non_utf8).is_err()); +} + +#[test] +fn daemon_tls_parsing_rejects_malformed_and_incomplete_pem_material() { + assert!(decode_pem_blocks(&[0xff], "CERTIFICATE").is_err()); + assert!(decode_pem_blocks(b"-----BEGIN CERTIFICATE-----", "CERTIFICATE").is_err()); + assert!( + decode_pem_blocks( + b"-----BEGIN CERTIFICATE-----!-----END CERTIFICATE-----", + "CERTIFICATE" + ) + .is_err() + ); + assert!( + decode_pem_blocks( + b"-----BEGIN CERTIFICATE-----\n-----END CERTIFICATE-----", + "CERTIFICATE" + ) + .is_err() + ); + + let directory = tempfile::tempdir().unwrap(); + let certificate = directory.path().join("certificate.pem"); + let key = directory.path().join("key.pem"); + std::fs::write(&certificate, b"no certificate blocks").unwrap(); + std::fs::write(&key, b"no key blocks").unwrap(); + assert!(load_tls_config(&certificate, &key).is_err()); + std::fs::write( + &certificate, + b"-----BEGIN CERTIFICATE-----AQ==-----END CERTIFICATE-----", + ) + .unwrap(); + assert!(load_tls_config(&certificate, &key).is_err()); + std::fs::write( + &key, + b"-----BEGIN PRIVATE KEY-----AQ==-----END PRIVATE KEY-----\n-----BEGIN PRIVATE KEY-----Ag==-----END PRIVATE KEY-----", + ) + .unwrap(); + assert!(load_tls_config(&certificate, &key).is_err()); +} + +#[test] +fn worker_activation_helpers_cover_remote_and_rejected_endpoint_shapes() { + let loopback = fresh_launch(WorkerNetworkHint::new("localhost", None).unwrap()).unwrap(); + assert_eq!(loopback.bind_ip, Ipv4Addr::LOCALHOST); + assert_eq!(loopback.port, 0); + assert!(loopback.advertise_address.is_none()); + + let remote = fresh_launch(WorkerNetworkHint::new("worker.example", Some(9443)).unwrap()) + .expect("remote worker launch"); + assert_eq!(remote.bind_ip, Ipv4Addr::UNSPECIFIED); + assert_eq!(remote.advertise_address.as_deref(), Some("worker.example")); + let activation = Activation { + fingerprint: MachineIdentity::generate().unwrap().identity.fingerprint(), + secret_digest: TokenDigest::from_token(b"secret"), + deadline_unix_ms: u64::MAX, + bind_ip: remote.bind_ip, + port: remote.port, + advertise_address: remote.advertise_address, + consumed: false, + }; + assert!(activation_endpoint_matches( + "https://worker.example:9443", + &activation + )); + assert!(!activation_endpoint_matches("not a URL", &activation)); + assert!(!activation_endpoint_matches( + "https://worker.example", + &activation + )); + assert!(!activation_endpoint_matches( + "http://worker.example:9443", + &activation + )); + assert!(!reserve_challenge_slot( + &mut HashMap::new(), + now_unix_ms(), + ComponentRole::Daemon + )); +} + +#[tokio::test] +async fn forwarding_rejects_invalid_destinations_and_worker_credentials_before_io() { + let client = pooled_client().unwrap(); + let invalid_destination = forward( + &client, + Request::post("/").body(Body::empty()).unwrap(), + "http://[invalid", + None, + None, + ) + .await; + assert_eq!( + invalid_destination.response.status(), + StatusCode::BAD_GATEWAY + ); + assert!(!invalid_destination.communication_failure); + + let invalid_credential = forward( + &client, + Request::post("/").body(Body::empty()).unwrap(), + "http://127.0.0.1:1", + Some(( + HeaderName::from_static(WORKER_TOKEN_HEADER), + "bad\nvalue".into(), + )), + None, + ) + .await; + assert_eq!( + invalid_credential.response.status(), + StatusCode::INTERNAL_SERVER_ERROR + ); + assert!(!invalid_credential.communication_failure); +} + +#[tokio::test] +async fn worker_response_body_errors_trigger_the_route_callback_once() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let calls = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&calls); + let frames = futures_util::stream::iter([Err::, std::io::Error>( + std::io::Error::other("body failed"), + )]); + let mut body = ErrorObservedBody { + body: http_body_util::StreamBody::new(frames), + on_error: Some(move || { + observed.fetch_add(1, Ordering::SeqCst); + }), + }; + assert!(body.frame().await.unwrap().is_err()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(body.frame().await.is_none()); + assert_eq!(calls.load(Ordering::SeqCst), 1); +} + +#[test] +fn responses_websocket_probe_is_narrow() { + let probe = Request::get("/backend-api/codex/responses") + .header(axum::http::header::UPGRADE, "WebSocket") + .body(Body::empty()) + .expect("probe"); + assert!(responses_websocket_probe(&probe)); + + let ordinary = Request::get("/backend-api/codex/responses") + .body(Body::empty()) + .expect("ordinary GET"); + assert!(!responses_websocket_probe(&ordinary)); + assert!(!public_method_allowed( + ordinary.method(), + ordinary.uri().path() + )); + assert!(public_method_allowed( + &Method::POST, + "/backend-api/codex/responses" + )); + assert!(public_method_allowed(&Method::GET, "/v1/models")); + assert!(!public_method_allowed(&Method::POST, "/v1/models")); +} + +#[test] +fn unavailable_response_includes_retry_after() { + let response = unavailable_response(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(response.headers().get(RETRY_AFTER).unwrap(), "1"); +} + +#[tokio::test] +async fn global_pass_through_authenticates_and_forwards_only_provider_headers() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x11_u8; 32]); + let captured: CapturedProviderRequest = Arc::new(std::sync::Mutex::new(None)); + let provider = Router::new() + .route( + "/v1/responses", + post( + |State(captured): State, + request: Request| async move { + let (parts, body) = request.into_parts(); + let body = axum::body::to_bytes(body, usize::MAX) + .await + .expect("provider request body"); + *captured.lock().expect("capture provider request") = + Some((parts.headers, body)); + Response::builder() + .status(StatusCode::CREATED) + .header("x-provider-result", "preserved") + .body(Body::from("raw provider response")) + .expect("provider response") + }, + ), + ) + .with_state(Arc::clone(&captured)); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind provider"); + let address = listener.local_addr().expect("provider address"); + let provider_task = tokio::spawn(async move { + axum::serve(listener, provider) + .await + .expect("serve provider"); + }); + let state = test_daemon_state( + true, + &token, + GatewayConfig { + openai_base_url: format!("http://{address}"), + openai_auth_header: Some("Bearer configured-provider".into()), + ..GatewayConfig::default() + }, + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("http://{}", listener.local_addr().unwrap()); + // The challenge transcript must name the same origin as the serving daemon. + let mut state = state; + Arc::get_mut(&mut state).unwrap().public_origin = origin.clone(); + let server = tokio::spawn({ + let state = Arc::clone(&state); + async move { axum::serve(listener, router(state)).await.unwrap() } + }); + let identity = MachineIdentity::generate().unwrap().identity; + assert_eq!( + enroll_test_mcp(&state, &origin, &identity, &token, "pass-through") + .await + .status(), + StatusCode::OK + ); + let app = router(state); + + let response = app + .clone() + .oneshot( + Request::post("/v1/responses") + .header(CLIENT_TOKEN_HEADER, &token) + .header("x-nemo-relay-session-id", "untrusted-private-header") + .header("x-provider-feature", "preserved") + .body(Body::from("provider request bytes")) + .expect("public request"), + ) + .await + .expect("public response"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.headers()["x-provider-result"], "preserved"); + assert_eq!( + response + .into_body() + .collect() + .await + .expect("provider response body") + .to_bytes(), + "raw provider response" + ); + let (headers, body) = captured + .lock() + .expect("captured provider request") + .take() + .expect("provider received request"); + assert_eq!(headers[AUTHORIZATION], "Bearer configured-provider"); + assert_eq!(headers["x-provider-feature"], "preserved"); + assert!(!headers.contains_key(CLIENT_TOKEN_HEADER)); + assert!(!headers.contains_key("x-nemo-relay-session-id")); + assert_eq!(body, "provider request bytes"); + + let hook = app + .clone() + .oneshot( + Request::post("/hooks/claude-code") + .header(CLIENT_TOKEN_HEADER, &token) + .body(Body::from("{}")) + .expect("hook request"), + ) + .await + .expect("hook response"); + assert_eq!(hook.status(), StatusCode::OK); + assert_eq!( + hook.into_body() + .collect() + .await + .expect("hook body") + .to_bytes(), + HookRoute::Claude.pass_through_body() + ); + + for request in [ + Request::get("/v1/responses") + .header(CLIENT_TOKEN_HEADER, &token) + .body(Body::empty()) + .expect("disallowed method"), + Request::post("/not-a-route") + .header(CLIENT_TOKEN_HEADER, &token) + .body(Body::empty()) + .expect("unknown route"), + ] { + let response = app.clone().oneshot(request).await.expect("response"); + assert!(matches!( + response.status(), + StatusCode::METHOD_NOT_ALLOWED | StatusCode::NOT_FOUND + )); + } + provider_task.abort(); + server.abort(); +} + +#[tokio::test] +async fn unreachable_worker_marks_its_authenticated_route_pass_through() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x22_u8; 32]); + let credential = RouteCredential::parse(token.clone()).expect("route credential"); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + state + .registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: McpSessionId::new("worker-route-session").expect("session"), + lease_expires_at_unix_ms: u64::MAX, + }, + WorkerLaunch { + activation_id: "worker-route-activation".into(), + activation_token: SensitiveString::new("activation-token").expect("token"), + deadline_unix_ms: u64::MAX, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }, + ) + .expect("register route"); + state + .registry + .mark_worker_ready( + fingerprint, + "worker-route-activation", + Arc::new( + WorkerTarget::new( + "unreachable-worker", + "http://127.0.0.1:9", + SensitiveString::new("worker-session-token").expect("token"), + ) + .expect("worker target"), + ), + ) + .expect("publish route"); + + let response = router(Arc::clone(&state)) + .oneshot( + Request::post("/v1/responses") + .header(CLIENT_TOKEN_HEADER, &token) + .body(Body::from("request bytes")) + .expect("public request"), + ) + .await + .expect("public response"); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + assert!(matches!( + state.registry.resolve_target(&credential.digest()), + Ok(ResolvedTarget::PassThrough) + )); +} + +#[tokio::test] +async fn aborted_public_upload_does_not_demote_a_healthy_worker() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x72_u8; 32]); + let credential = RouteCredential::parse(token.clone()).expect("route credential"); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let activation_id = "aborted-upload-activation"; + state + .registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: McpSessionId::new("aborted-upload-session").expect("session"), + lease_expires_at_unix_ms: u64::MAX, + }, + WorkerLaunch { + activation_id: activation_id.into(), + activation_token: SensitiveString::new("activation-token").expect("token"), + deadline_unix_ms: u64::MAX, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }, + ) + .expect("register route"); + + let listener = TcpListener::bind("127.0.0.1:0").await.expect("worker bind"); + let address = listener.local_addr().expect("worker address"); + let worker_task = tokio::spawn(async move { + let app = Router::new().fallback(|body: Body| async move { + let _ = body.collect().await; + StatusCode::NO_CONTENT + }); + axum::serve(listener, app).await.expect("worker serve"); + }); + state + .registry + .mark_worker_ready( + fingerprint, + activation_id, + Arc::new( + WorkerTarget::new( + "healthy-worker", + format!("http://{address}"), + SensitiveString::new("worker-session-token").expect("token"), + ) + .expect("worker target"), + ), + ) + .expect("publish route"); + + let failed_body = futures_util::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"partial")), + Err(std::io::Error::new( + std::io::ErrorKind::ConnectionReset, + "client upload aborted", + )), + ]); + let response = router(Arc::clone(&state)) + .oneshot( + Request::post("/v1/responses") + .header(CLIENT_TOKEN_HEADER, &token) + .body(Body::from_stream(failed_body)) + .expect("public request"), + ) + .await + .expect("public response"); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + assert!(matches!( + state.registry.resolve_target(&credential.digest()), + Ok(ResolvedTarget::Worker(_)) + )); + worker_task.abort(); +} + +#[tokio::test] +async fn public_ingress_rejects_credentials_methods_websockets_and_unready_routes_early() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x23_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let app = router(Arc::clone(&state)); + + let unknown = app + .clone() + .oneshot(Request::get("/unknown").body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(unknown.status(), StatusCode::NOT_FOUND); + + for request in [ + Request::post("/v1/responses").body(Body::empty()).unwrap(), + Request::post("/v1/responses") + .header(CLIENT_TOKEN_HEADER, "invalid") + .body(Body::empty()) + .unwrap(), + ] { + assert_eq!( + app.clone().oneshot(request).await.unwrap().status(), + StatusCode::UNAUTHORIZED + ); + } + + let websocket = Request::get("/v1/responses") + .header(CLIENT_TOKEN_HEADER, &token) + .header(axum::http::header::UPGRADE, "websocket") + .body(Body::empty()) + .unwrap(); + assert_eq!( + app.clone().oneshot(websocket).await.unwrap().status(), + StatusCode::UPGRADE_REQUIRED + ); + + let fingerprint = MachineIdentity::generate() + .expect("identity") + .identity + .fingerprint(); + let credential = RouteCredential::parse(token.clone()).expect("credential"); + state + .registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: McpSessionId::new("pending-route").expect("session"), + lease_expires_at_unix_ms: u64::MAX, + }, + WorkerLaunch { + activation_id: "pending-activation".into(), + activation_token: SensitiveString::new("activation-secret").expect("secret"), + deadline_unix_ms: u64::MAX, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }, + ) + .expect("register pending route"); + let unavailable = app + .oneshot( + Request::post("/v1/responses") + .header(CLIENT_TOKEN_HEADER, &token) + .body(Body::from("must not be forwarded")) + .expect("request"), + ) + .await + .expect("unavailable response"); + assert_eq!(unavailable.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(unavailable.headers()[RETRY_AFTER], "1"); +} + +#[test] +fn pass_through_auth_injection_supports_environment_and_anthropic_configuration() { + let _environment = EnvScope::set(&[ + ("OPENAI_API_KEY", Some(std::ffi::OsStr::new(" openai-env "))), + ( + "ANTHROPIC_API_KEY", + Some(std::ffi::OsStr::new(" anthropic-env ")), + ), + ]); + let mut openai = HeaderMap::new(); + inject_provider_auth( + &mut openai, + ProviderRoute::OpenAi, + &GatewayConfig::default(), + ); + assert_eq!(openai[AUTHORIZATION], "Bearer openai-env"); + + let mut anthropic = HeaderMap::new(); + inject_provider_auth( + &mut anthropic, + ProviderRoute::Anthropic, + &GatewayConfig::default(), + ); + assert_eq!(anthropic["x-api-key"], "anthropic-env"); + + let mut configured = HeaderMap::new(); + inject_provider_auth( + &mut configured, + ProviderRoute::Anthropic, + &GatewayConfig { + anthropic_auth_header: Some("configured".into()), + ..GatewayConfig::default() + }, + ); + assert_eq!(configured[AUTHORIZATION], "configured"); +} + +fn test_daemon_state(pass_through: bool, token: &str, config: GatewayConfig) -> Arc { + test_daemon_state_at(pass_through, token, config, "http://127.0.0.1:47632".into()) +} + +fn test_daemon_state_at( + pass_through: bool, + _token: &str, + config: GatewayConfig, + public_origin: String, +) -> Arc { + let generation_directory = tempfile::tempdir().expect("generation directory"); + let generation_path = generation_directory + .keep() + .join("active-worker-generations.json"); + Arc::new(DaemonState { + registry: Registry::new(pass_through), + descriptor: crate::daemon::common::control::descriptor(ComponentRole::Daemon), + instance_id: "public-proxy-test-daemon".into(), + public_origin, + config, + upstream: pooled_client().expect("daemon client"), + worker_clients: WorkerClientPool::new().expect("worker clients"), + challenges: Mutex::new(HashMap::new()), + activations: Mutex::new(HashMap::new()), + mcp_sessions: Mutex::new(HashMap::new()), + mcp_heartbeat_serialization: Mutex::new(()), + worker_sessions: Mutex::new(HashMap::new()), + pending_directives: Mutex::new(HashMap::new()), + active_worker_generations: ActiveWorkerGenerations::load_for_test(generation_path) + .expect("generation state"), + worker_generation_publication: Mutex::new(()), + identity: MachineIdentity::generate() + .expect("daemon identity") + .identity, + }) +} + +#[test] +fn registration_requires_the_complete_lossless_transport_capability_set() { + let complete = crate::daemon::common::control::descriptor(ComponentRole::Mcp); + assert!(has_required_transport_capabilities(&complete)); + + let missing_trailers = crate::daemon::common::protocol::ComponentDescriptor::nemo_relay( + ComponentRole::Mcp, + crate::daemon::common::protocol::ProtocolRange::default(), + Capabilities::new(["http1", "http2", "streaming_body_frames", "sse_passthrough"]) + .expect("capabilities"), + "future-version", + ); + assert!(!has_required_transport_capabilities(&missing_trailers)); +} + +#[test] +fn pending_challenge_storage_is_bounded_and_prunes_expired_entries() { + let generated = MachineIdentity::generate().expect("identity"); + let identity = generated.identity; + let descriptor = crate::daemon::common::control::descriptor(ComponentRole::Mcp); + let mut challenges = HashMap::new(); + for index in 0..MAX_PENDING_CHALLENGES { + let record = ChallengeRecord::generate(100, 100).expect("challenge"); + let challenge = record.challenge(); + challenges.insert( + challenge.id, + PendingChallenge { + request: ChallengeRequest { + initiator: descriptor.clone(), + initiator_instance_id: format!("mcp-{index}"), + initiator_public_identity: identity.public_identity(), + initiator_fingerprint: identity.fingerprint(), + initiator_nonce: challenge.nonce, + }, + record, + }, + ); + } + assert_eq!(challenges.len(), MAX_PENDING_CHALLENGES); + assert!(!reserve_challenge_slot( + &mut challenges, + 199, + ComponentRole::Mcp + )); + assert!(reserve_challenge_slot( + &mut challenges, + 200, + ComponentRole::Mcp + )); + assert!(challenges.is_empty()); +} + +#[test] +fn active_mcp_registration_reuses_its_session_credential() { + let identity = MachineIdentity::generate().expect("identity").identity; + let fingerprint = identity.fingerprint(); + let token_digest = TokenDigest::from_token(b"route-token"); + let original = SensitiveString::new("original-session-secret").expect("secret"); + let mut sessions = HashMap::new(); + sessions.insert( + "mcp-session".to_owned(), + McpControlSession { + fingerprint, + token_digest, + secret: original.clone(), + secret_digest: TokenDigest::from_token(original.expose().as_bytes()), + lease_expires_at_unix_ms: 200, + last_sequence: 0, + last_request_id: String::new(), + last_heartbeat: None, + worker_network: worker_network(), + released: false, + }, + ); + + let (selected, reused) = select_mcp_session_token( + &sessions, + "mcp-session", + fingerprint, + token_digest, + worker_network(), + 199, + SensitiveString::new("must-not-rotate").expect("secret"), + ) + .expect("selection"); + assert!(reused); + assert_eq!(selected, original); + + let (selected, reused) = select_mcp_session_token( + &sessions, + "mcp-session", + fingerprint, + token_digest, + worker_network(), + 200, + SensitiveString::new("fresh-after-expiry").expect("secret"), + ) + .expect("expired selection"); + assert!(!reused); + assert_eq!(selected.expose(), "fresh-after-expiry"); +} + +#[test] +fn mcp_session_reuse_rejects_identity_token_and_network_rebinding() { + let identity = MachineIdentity::generate().unwrap().identity; + let fingerprint = identity.fingerprint(); + let token = TokenDigest::from_token(b"route-token"); + let secret = SensitiveString::new("session-secret").unwrap(); + let sessions = HashMap::from([( + "mcp".into(), + McpControlSession { + fingerprint, + token_digest: token, + secret: secret.clone(), + secret_digest: TokenDigest::from_token(secret.expose().as_bytes()), + lease_expires_at_unix_ms: 200, + last_sequence: 0, + last_request_id: String::new(), + last_heartbeat: None, + worker_network: worker_network(), + released: false, + }, + )]); + for (candidate_fingerprint, candidate_token, candidate_network) in [ + ( + MachineIdentity::generate().unwrap().identity.fingerprint(), + token, + worker_network(), + ), + ( + fingerprint, + TokenDigest::from_token(b"different"), + worker_network(), + ), + ( + fingerprint, + token, + WorkerNetworkHint { + advertised_host: "worker.example".into(), + port: Some(443), + }, + ), + ] { + assert!( + select_mcp_session_token( + &sessions, + "mcp", + candidate_fingerprint, + candidate_token, + candidate_network, + 100, + SensitiveString::new("fresh").unwrap(), + ) + .is_err() + ); + } +} + +#[test] +fn authenticated_sequence_rejects_tampering_and_out_of_order_messages() { + let secret = SensitiveString::new("session-secret").unwrap(); + let digest = TokenDigest::from_token(secret.expose().as_bytes()); + let mut sequence = 0; + let mut request_id = String::new(); + let request = + SessionRequest::new("session".into(), secret.clone(), 1, EmptyPayload::default()).unwrap(); + assert!(!authenticate_sequence(digest, &mut sequence, &mut request_id, &request).unwrap()); + assert!(authenticate_sequence(digest, &mut sequence, &mut request_id, &request).unwrap()); + + let wrong = SessionRequest::new( + "session".into(), + SensitiveString::new("wrong").unwrap(), + 2, + EmptyPayload::default(), + ) + .unwrap(); + assert!(authenticate_sequence(digest, &mut sequence, &mut request_id, &wrong).is_err()); + let out_of_order = + SessionRequest::new("session".into(), secret, 3, EmptyPayload::default()).unwrap(); + assert!(authenticate_sequence(digest, &mut sequence, &mut request_id, &out_of_order).is_err()); +} + +#[test] +fn registry_errors_map_to_stable_control_statuses() { + for error in [ + RegistryError::UnknownRoute, + RegistryError::UnknownMcpSession, + RegistryError::TokenAlreadyBound, + RegistryError::FingerprintTokenMismatch, + RegistryError::ActivationMismatch, + RegistryError::WorkerMismatch, + RegistryError::RecoveryNotAuthorized, + RegistryError::RecoveryGenerationChanged, + RegistryError::RouteCapacityReached, + RegistryError::McpReferenceCapacityReached, + ] { + assert!(registry_error(error).status().is_client_error()); + } + assert_eq!( + registry_error(RegistryError::DrainInProgress).status(), + StatusCode::CONFLICT + ); +} + +#[test] +fn duplicate_heartbeat_replays_the_exact_cached_directive() { + let identity = MachineIdentity::generate().expect("identity").identity; + let secret = SensitiveString::new("session-secret").expect("secret"); + let mut request = SessionRequest::new( + "mcp-session".to_owned(), + secret.clone(), + 7, + EmptyPayload::default(), + ) + .expect("request"); + request.request_id = "stable-request-id".to_owned(); + let expected = McpHeartbeatResponse { + directive: Some(BrokerDirective::WaitForWorker { + retry_after_ms: 321, + }), + }; + let session = McpControlSession { + fingerprint: identity.fingerprint(), + token_digest: TokenDigest::from_token(b"route-token"), + secret: secret.clone(), + secret_digest: TokenDigest::from_token(secret.expose().as_bytes()), + lease_expires_at_unix_ms: 1_000, + last_sequence: request.sequence, + last_request_id: request.request_id.clone(), + last_heartbeat: Some(CachedHeartbeat { + sequence: request.sequence, + request_id: request.request_id.clone(), + response: expected, + }), + worker_network: worker_network(), + released: false, + }; + + let replayed = + cached_heartbeat_response(&session, &request, true).expect("duplicate heartbeat response"); + assert!(matches!( + replayed.directive, + Some(BrokerDirective::WaitForWorker { + retry_after_ms: 321 + }) + )); + assert!(cached_heartbeat_response(&session, &request, false).is_none()); + request.request_id = "different-request-id".to_owned(); + assert!(cached_heartbeat_response(&session, &request, true).is_none()); +} + +#[test] +fn expired_mcp_control_sessions_and_pending_directives_are_removed_together() { + let identity = MachineIdentity::generate().expect("identity").identity; + let fingerprint = identity.fingerprint(); + let token_digest = TokenDigest::from_token(b"route-token"); + let session = |lease_expires_at_unix_ms| { + let secret = SensitiveString::new("session-secret").expect("secret"); + McpControlSession { + fingerprint, + token_digest, + secret: secret.clone(), + secret_digest: TokenDigest::from_token(secret.expose().as_bytes()), + lease_expires_at_unix_ms, + last_sequence: 0, + last_request_id: String::new(), + last_heartbeat: None, + worker_network: worker_network(), + released: false, + } + }; + let mut sessions = HashMap::from([ + ("expired".to_owned(), session(100)), + ("live".to_owned(), session(101)), + ]); + let mut pending = HashMap::from([ + ( + "expired".to_owned(), + BrokerDirective::WaitForWorker { + retry_after_ms: 100, + }, + ), + ("live".to_owned(), BrokerDirective::UsePassThrough), + ]); + + prune_expired_mcp_control_state(&mut sessions, &mut pending, 100); + + assert!(!sessions.contains_key("expired")); + assert!(!pending.contains_key("expired")); + assert!(sessions.contains_key("live")); + assert!(pending.contains_key("live")); +} + +fn worker_network() -> WorkerNetworkHint { + WorkerNetworkHint { + advertised_host: Ipv4Addr::LOCALHOST.to_string(), + port: None, + } +} + +#[test] +fn advertised_https_is_valid_behind_a_reverse_proxy_without_native_tls() { + let options = crate::daemon::ServerOptions { + bind: Ipv4Addr::LOCALHOST, + port: 8080, + advertise_address: Some("https://relay.example.com:443".into()), + pass_through: false, + gateway: crate::server::GatewayOverrides::default(), + tls_cert: None, + tls_key: None, + }; + assert_eq!( + daemon_origin(&options, "127.0.0.1:8080".parse().unwrap()).expect("proxy origin"), + "https://relay.example.com" + ); + + let native_http = crate::daemon::ServerOptions { + tls_cert: Some("cert.pem".into()), + tls_key: Some("key.pem".into()), + advertise_address: Some("http://127.0.0.1:8080".into()), + ..options + }; + assert!(daemon_origin(&native_http, "127.0.0.1:8080".parse().unwrap()).is_err()); +} + +#[test] +fn daemon_origin_enforces_bind_and_tls_advertisement_contracts() { + let local: SocketAddr = "127.0.0.1:47632".parse().unwrap(); + let loopback = ServerOptions { + bind: Ipv4Addr::LOCALHOST, + port: 47632, + advertise_address: None, + pass_through: false, + gateway: crate::server::GatewayOverrides::default(), + tls_cert: None, + tls_key: None, + }; + assert_eq!( + daemon_origin(&loopback, local).unwrap(), + "http://127.0.0.1:47632" + ); + assert_eq!( + daemon_origin( + &ServerOptions { + tls_cert: Some("cert".into()), + ..loopback.clone() + }, + local + ) + .unwrap(), + "https://127.0.0.1:47632" + ); + for (address, tls, expected) in [ + ("127.0.0.1:80", false, "http://127.0.0.1"), + ("127.0.0.1:443", true, "https://127.0.0.1"), + ] { + let options = ServerOptions { + tls_cert: tls.then(|| "cert".into()), + tls_key: tls.then(|| "key".into()), + ..loopback.clone() + }; + assert_eq!( + daemon_origin(&options, address.parse().unwrap()).unwrap(), + expected + ); + } + let unspecified = ServerOptions { + bind: Ipv4Addr::UNSPECIFIED, + ..loopback + }; + assert!(daemon_origin(&unspecified, local).is_err()); + for origin in ["http://relay.example.com", "http://[::1]", "https://[::1]"] { + assert!( + daemon_origin( + &ServerOptions { + advertise_address: Some(origin.into()), + ..unspecified.clone() + }, + local + ) + .is_err(), + "public bind accepted invalid origin {origin}" + ); + } + assert!( + daemon_origin( + &ServerOptions { + advertise_address: Some("https://127.0.0.1:443".into()), + ..unspecified.clone() + }, + local + ) + .is_err() + ); + assert_eq!( + daemon_origin( + &ServerOptions { + advertise_address: Some("https://relay.example.com:443/".into()), + ..unspecified + }, + local + ) + .unwrap(), + "https://relay.example.com" + ); +} + +#[tokio::test] +async fn daemon_startup_rejects_an_unpaired_tls_identity_after_initializing_state() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let state_directory = tempfile::tempdir().unwrap(); + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x46_u8; 32]); + let _environment = EnvScope::set(&[ + (ROUTE_TOKEN_ENV, Some(std::ffi::OsStr::new(&token))), + ("XDG_CONFIG_HOME", Some(state_directory.path().as_os_str())), + ]); + let error = serve(ServerOptions { + bind: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + pass_through: false, + gateway: crate::server::GatewayOverrides::default(), + tls_cert: Some(state_directory.path().join("certificate.pem")), + tls_key: None, + }) + .await + .expect_err("unpaired TLS configuration"); + assert!(error.to_string().contains("must be supplied together")); +} + +#[tokio::test] +async fn daemon_plain_listener_starts_and_reports_address_conflicts() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let state_directory = tempfile::tempdir().unwrap(); + let _environment = EnvScope::set(&[ + (ROUTE_TOKEN_ENV, None), + ("XDG_CONFIG_HOME", Some(state_directory.path().as_os_str())), + ]); + let reserved = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let running_port = reserved.local_addr().unwrap().port(); + drop(reserved); + let options = ServerOptions { + bind: Ipv4Addr::LOCALHOST, + port: running_port, + advertise_address: None, + pass_through: true, + gateway: crate::server::GatewayOverrides::default(), + tls_cert: None, + tls_key: None, + }; + let running = tokio::spawn(serve(options)); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + if let Ok(response) = + reqwest::get(format!("http://127.0.0.1:{running_port}/unknown")).await + { + assert_eq!(response.status(), StatusCode::NOT_FOUND); + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("daemon starts"); + running.abort(); + let _ = running.await; + + let occupied = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = occupied.local_addr().unwrap().port(); + let error = serve(ServerOptions { + port, + ..ServerOptions { + bind: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + pass_through: true, + gateway: crate::server::GatewayOverrides::default(), + tls_cert: None, + tls_key: None, + } + }) + .await + .expect_err("occupied daemon port"); + assert!(error.to_string().contains("failed to bind daemon listener")); +} + +#[tokio::test] +async fn challenge_endpoint_rejects_incompatible_and_forged_components() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x41_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let identity = MachineIdentity::generate().unwrap().identity; + let valid = ChallengeRequest { + initiator: crate::daemon::common::control::descriptor(ComponentRole::Mcp), + initiator_instance_id: "mcp-instance".into(), + initiator_public_identity: identity.public_identity(), + initiator_fingerprint: identity.fingerprint(), + initiator_nonce: ChallengeRecord::generate(1, 1).unwrap().challenge().nonce, + }; + + let mut incompatible = valid.clone(); + incompatible.initiator = crate::daemon::common::protocol::ComponentDescriptor::nemo_relay( + ComponentRole::Mcp, + crate::daemon::common::protocol::ProtocolRange::default(), + Capabilities::new(["http1"]).unwrap(), + "test", + ); + assert_eq!( + issue_challenge(State(Arc::clone(&state)), Json(incompatible)) + .await + .status(), + StatusCode::UPGRADE_REQUIRED + ); + + let mut daemon_role = valid.clone(); + daemon_role.initiator = crate::daemon::common::control::descriptor(ComponentRole::Daemon); + assert_eq!( + issue_challenge(State(Arc::clone(&state)), Json(daemon_role)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let mut forged = valid; + forged.initiator_fingerprint = MachineIdentity::generate().unwrap().identity.fingerprint(); + assert_eq!( + issue_challenge(State(state), Json(forged)).await.status(), + StatusCode::UNAUTHORIZED + ); +} + +#[tokio::test] +async fn control_handlers_reject_unknown_sessions_and_oversized_failure_payloads() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x43_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let secret = SensitiveString::new("unknown-session-secret").unwrap(); + let empty = + SessionRequest::new("unknown".into(), secret.clone(), 1, EmptyPayload::default()).unwrap(); + assert_eq!( + heartbeat_mcp(State(Arc::clone(&state)), Json(empty.clone())) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + assert_eq!( + release_mcp(State(Arc::clone(&state)), Json(empty)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let ready = SessionRequest::new( + "unknown".into(), + secret.clone(), + 1, + WorkerReadyPayload { + worker_id: "unknown".into(), + }, + ) + .unwrap(); + assert_eq!( + ready_worker(State(Arc::clone(&state)), Json(ready)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + let heartbeat = SessionRequest::new( + "unknown".into(), + secret.clone(), + 1, + WorkerHeartbeatPayload { + worker_id: "unknown".into(), + }, + ) + .unwrap(); + assert_eq!( + heartbeat_worker(State(Arc::clone(&state)), Json(heartbeat)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let oversized = SessionRequest::new( + "unknown".into(), + secret, + 1, + ActivationFailedPayload { + activation_id: "x".repeat(129), + reason: "reason".into(), + }, + ) + .unwrap(); + assert_eq!( + activation_failed(State(state), Json(oversized)) + .await + .status(), + StatusCode::BAD_REQUEST + ); +} + +#[tokio::test] +async fn worker_control_handlers_bind_payload_identity_before_sequence_authentication() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x44_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + lock(&state.worker_sessions).insert("staged".into(), staged_worker_session("staged", u64::MAX)); + let secret = SensitiveString::new("control-secret").unwrap(); + let ready = SessionRequest::new( + "staged".into(), + secret.clone(), + 1, + WorkerReadyPayload { + worker_id: "different".into(), + }, + ) + .unwrap(); + assert_eq!( + ready_worker(State(Arc::clone(&state)), Json(ready)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + let heartbeat = SessionRequest::new( + "staged".into(), + secret, + 1, + WorkerHeartbeatPayload { + worker_id: "different".into(), + }, + ) + .unwrap(); + assert_eq!( + heartbeat_worker(State(state), Json(heartbeat)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); +} + +#[tokio::test] +async fn mcp_handlers_cover_registry_loss_duplicate_cache_and_released_sessions() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x48_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let fingerprint = MachineIdentity::generate().unwrap().identity.fingerprint(); + let token_digest = TokenDigest::from_token(token.as_bytes()); + let secret = SensitiveString::new("mcp-handler-secret").unwrap(); + let session = |released, last_sequence, last_request_id: String| McpControlSession { + fingerprint, + token_digest, + secret: secret.clone(), + secret_digest: TokenDigest::from_token(secret.expose().as_bytes()), + lease_expires_at_unix_ms: u64::MAX, + last_sequence, + last_request_id, + last_heartbeat: None, + worker_network: worker_network(), + released, + }; + + lock(&state.mcp_sessions).insert("lost-route".into(), session(false, 0, String::new())); + let lost = SessionRequest::new( + "lost-route".into(), + secret.clone(), + 1, + EmptyPayload::default(), + ) + .unwrap(); + assert_eq!( + heartbeat_mcp(State(Arc::clone(&state)), Json(lost)) + .await + .status(), + StatusCode::NOT_FOUND + ); + assert!(!lock(&state.mcp_sessions).contains_key("lost-route")); + + let duplicate = SessionRequest::new( + "duplicate".into(), + secret.clone(), + 1, + EmptyPayload::default(), + ) + .unwrap(); + lock(&state.mcp_sessions).insert( + "duplicate".into(), + session(false, 1, duplicate.request_id.clone()), + ); + assert_eq!( + heartbeat_mcp(State(Arc::clone(&state)), Json(duplicate)) + .await + .status(), + StatusCode::CONFLICT + ); + + let released_duplicate = SessionRequest::new( + "released".into(), + secret.clone(), + 1, + EmptyPayload::default(), + ) + .unwrap(); + lock(&state.mcp_sessions).insert( + "released".into(), + session(true, 1, released_duplicate.request_id.clone()), + ); + assert_eq!( + release_mcp(State(Arc::clone(&state)), Json(released_duplicate.clone())) + .await + .status(), + StatusCode::NO_CONTENT + ); + let released_new = SessionRequest::new( + "released".into(), + secret.clone(), + 2, + EmptyPayload::default(), + ) + .unwrap(); + assert_eq!( + release_mcp(State(Arc::clone(&state)), Json(released_new)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + lock(&state.mcp_sessions).insert("activation".into(), session(false, 0, String::new())); + let failed = SessionRequest::new( + "activation".into(), + secret, + 1, + ActivationFailedPayload { + activation_id: "unknown".into(), + reason: "failed".into(), + }, + ) + .unwrap(); + assert_eq!( + activation_failed(State(state), Json(failed)).await.status(), + StatusCode::NOT_FOUND + ); +} + +#[tokio::test] +async fn worker_handlers_cover_bad_sequence_and_already_published_readiness() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x49_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let mut session = staged_worker_session("staged", u64::MAX); + session.published = true; + lock(&state.worker_sessions).insert("staged".into(), session); + + let wrong = SessionRequest::new( + "staged".into(), + SensitiveString::new("wrong-secret").unwrap(), + 1, + WorkerHeartbeatPayload { + worker_id: "staged".into(), + }, + ) + .unwrap(); + assert_eq!( + heartbeat_worker(State(Arc::clone(&state)), Json(wrong)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let ready = SessionRequest::new( + "staged".into(), + SensitiveString::new("control-secret").unwrap(), + 1, + WorkerReadyPayload { + worker_id: "staged".into(), + }, + ) + .unwrap(); + assert_eq!( + ready_worker(State(state), Json(ready)).await.status(), + StatusCode::NO_CONTENT + ); +} + +#[tokio::test] +async fn worker_registration_and_recovery_fail_closed_at_each_authority_boundary() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4a_u8; 32]); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("http://{}", listener.local_addr().unwrap()); + let state = test_daemon_state_at(false, &token, GatewayConfig::default(), origin.clone()); + let server = tokio::spawn({ + let state = Arc::clone(&state); + async move { axum::serve(listener, router(state)).await.unwrap() } + }); + let client = control_client().unwrap(); + let worker = MachineIdentity::generate().unwrap().identity; + + let proof = begin_handshake( + &client, + &origin, + ComponentRole::Worker, + &worker, + "unknown-activation", + None, + ) + .await + .unwrap() + .proof; + let unknown = WorkerRegisterRequest { + proof, + worker_id: "unknown-activation".into(), + endpoint: "http://127.0.0.1:41000".into(), + activation_id: "missing".into(), + activation_token: SensitiveString::new("missing-token").unwrap(), + tls_root_certificate: None, + }; + assert_eq!( + register_worker(State(Arc::clone(&state)), Json(unknown)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let proof = begin_handshake( + &client, + &origin, + ComponentRole::Worker, + &worker, + "wrong-secret", + None, + ) + .await + .unwrap() + .proof; + lock(&state.activations).insert( + "wrong-secret".into(), + Activation { + fingerprint: worker.fingerprint(), + secret_digest: TokenDigest::from_token(b"correct-token"), + deadline_unix_ms: u64::MAX, + consumed: false, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }, + ); + let wrong_secret = WorkerRegisterRequest { + proof, + worker_id: "wrong-secret".into(), + endpoint: "http://127.0.0.1:41000".into(), + activation_id: "wrong-secret".into(), + activation_token: SensitiveString::new("incorrect-token").unwrap(), + tls_root_certificate: None, + }; + assert_eq!( + register_worker(State(Arc::clone(&state)), Json(wrong_secret)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let proof = begin_handshake( + &client, + &origin, + ComponentRole::Worker, + &worker, + "endpoint-mismatch", + None, + ) + .await + .unwrap() + .proof; + lock(&state.activations).insert( + "endpoint-mismatch".into(), + Activation { + fingerprint: worker.fingerprint(), + secret_digest: TokenDigest::from_token(b"activation-token"), + deadline_unix_ms: u64::MAX, + consumed: false, + bind_ip: Ipv4Addr::LOCALHOST, + port: 41000, + advertise_address: None, + }, + ); + let mismatched = WorkerRegisterRequest { + proof, + worker_id: "endpoint-mismatch".into(), + endpoint: "http://127.0.0.1:41001".into(), + activation_id: "endpoint-mismatch".into(), + activation_token: SensitiveString::new("activation-token").unwrap(), + tls_root_certificate: None, + }; + assert_eq!( + register_worker(State(Arc::clone(&state)), Json(mismatched)) + .await + .status(), + StatusCode::BAD_REQUEST + ); + + lock(&state.activations).insert( + "replay-activation".into(), + Activation { + fingerprint: worker.fingerprint(), + secret_digest: TokenDigest::from_token(b"replay-token"), + deadline_unix_ms: u64::MAX, + consumed: false, + bind_ip: Ipv4Addr::LOCALHOST, + port: 41002, + advertise_address: None, + }, + ); + for worker_id in ["replay-worker", "replay-worker"] { + let proof = begin_handshake( + &client, + &origin, + ComponentRole::Worker, + &worker, + worker_id, + None, + ) + .await + .unwrap() + .proof; + let replay = WorkerRegisterRequest { + proof, + worker_id: worker_id.into(), + endpoint: "http://127.0.0.1:41002".into(), + activation_id: "replay-activation".into(), + activation_token: SensitiveString::new("replay-token").unwrap(), + tls_root_certificate: None, + }; + assert!( + register_worker(State(Arc::clone(&state)), Json(replay)) + .await + .status() + .is_success() + ); + } + let proof = begin_handshake( + &client, + &origin, + ComponentRole::Worker, + &worker, + "activation-thief", + None, + ) + .await + .unwrap() + .proof; + let stolen = WorkerRegisterRequest { + proof, + worker_id: "activation-thief".into(), + endpoint: "http://127.0.0.1:41002".into(), + activation_id: "replay-activation".into(), + activation_token: SensitiveString::new("replay-token").unwrap(), + tls_root_certificate: None, + }; + assert_eq!( + register_worker(State(Arc::clone(&state)), Json(stolen)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let proof = begin_handshake( + &client, + &origin, + ComponentRole::Worker, + &worker, + "invalid-recovery", + None, + ) + .await + .unwrap() + .proof; + let foreign = MachineIdentity::generate().unwrap().identity; + let invalid_grant = WorkerGenerationGrant::issue( + "invalid-recovery", + worker.fingerprint(), + "http://127.0.0.1:41000", + None, + &foreign, + ) + .unwrap(); + let invalid_recovery = WorkerRecoverRequest { + proof, + worker_id: "invalid-recovery".into(), + endpoint: "http://127.0.0.1:41000".into(), + tls_root_certificate: None, + generation_grant: invalid_grant, + }; + assert_eq!( + recover_worker(State(Arc::clone(&state)), Json(invalid_recovery)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let proof = begin_handshake( + &client, + &origin, + ComponentRole::Worker, + &worker, + "revoked-recovery", + None, + ) + .await + .unwrap() + .proof; + let revoked_grant = WorkerGenerationGrant::issue( + "revoked-recovery", + worker.fingerprint(), + "http://127.0.0.1:41000", + None, + &state.identity, + ) + .unwrap(); + let revoked = WorkerRecoverRequest { + proof, + worker_id: "revoked-recovery".into(), + endpoint: "http://127.0.0.1:41000".into(), + tls_root_certificate: None, + generation_grant: revoked_grant, + }; + assert_eq!( + recover_worker(State(state), Json(revoked)).await.status(), + StatusCode::UNAUTHORIZED + ); + server.abort(); +} + +#[tokio::test] +async fn mcp_registration_rejects_credential_and_network_proof_mismatches_before_routing() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4d_u8; 32]); + let credential = RouteCredential::parse(token.clone()).unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let origin = format!("http://{}", listener.local_addr().unwrap()); + let state = test_daemon_state_at(false, &token, GatewayConfig::default(), origin.clone()); + let server = tokio::spawn({ + let state = Arc::clone(&state); + async move { axum::serve(listener, router(state)).await.unwrap() } + }); + let client = control_client().unwrap(); + let identity = MachineIdentity::generate().unwrap().identity; + let handshake = begin_handshake( + &client, + &origin, + ComponentRole::Mcp, + &identity, + "mcp-boundary", + Some(credential.digest()), + ) + .await + .unwrap(); + let hint = WorkerNetworkHint::new("127.0.0.1", None).unwrap(); + let hint_proof = WorkerNetworkHintProof::sign( + hint, + &handshake.proof.transcript.daemon_target, + "mcp-boundary", + &handshake.proof.transcript.challenge_id, + &identity.fingerprint(), + &identity, + ) + .unwrap(); + let request = McpRegisterRequest { + proof: handshake.proof.clone(), + worker_network: hint_proof, + }; + assert_eq!( + register_mcp( + State(Arc::clone(&state)), + HeaderMap::new(), + Json(request.clone()), + ) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + let other = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4e_u8; 32]); + let headers = HeaderMap::from_iter([( + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderValue::from_str(&other).unwrap(), + )]); + assert_eq!( + register_mcp(State(Arc::clone(&state)), headers, Json(request.clone()),) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let headers = HeaderMap::from_iter([( + HeaderName::from_static(CLIENT_TOKEN_HEADER), + HeaderValue::from_str(&token).unwrap(), + )]); + let mut mismatched = request.clone(); + mismatched.proof.transcript.route_token_digest = None; + assert_eq!( + register_mcp(State(Arc::clone(&state)), headers.clone(), Json(mismatched),) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let mut bad_network = request; + bad_network.worker_network.hint.advertised_host = "localhost".into(); + assert_eq!( + register_mcp(State(Arc::clone(&state)), headers, Json(bad_network)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + server.abort(); +} + +#[tokio::test] +async fn worker_readiness_probe_failure_removes_the_staged_session() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4b_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let worker = MachineIdentity::generate().unwrap().identity; + let fingerprint = worker.fingerprint(); + let daemon_proof = HandshakeProof { + signer: ComponentRole::Daemon, + signature: state.identity.sign(b"readiness-failure"), + }; + let response = stage_worker( + &state, + fingerprint, + "unreachable-worker".into(), + "http://127.0.0.1:1".into(), + None, + None, + WorkerPublication::Activation { + activation_id: "unreachable-activation".into(), + }, + daemon_proof, + ); + assert_eq!(response.status(), StatusCode::OK); + let registration: WorkerRegisterResponse = + serde_json::from_slice(&response.into_body().collect().await.unwrap().to_bytes()).unwrap(); + let ready = SessionRequest::new( + "unreachable-worker".into(), + registration.session_token, + 1, + WorkerReadyPayload { + worker_id: "unreachable-worker".into(), + }, + ) + .unwrap(); + assert_eq!( + ready_worker(State(Arc::clone(&state)), Json(ready)) + .await + .status(), + StatusCode::BAD_GATEWAY + ); + assert!(!lock(&state.worker_sessions).contains_key("unreachable-worker")); +} + +#[tokio::test] +async fn native_tls_configuration_serves_a_daemon_request_with_pinned_trust() { + let rcgen::CertifiedKey { cert, key_pair } = + rcgen::generate_simple_self_signed(vec!["127.0.0.1".into()]).expect("certificate"); + let directory = tempfile::tempdir().expect("temporary TLS directory"); + let certificate_path = directory.path().join("daemon.crt"); + let key_path = directory.path().join("daemon.pk8"); + std::fs::write(&certificate_path, cert.pem()).expect("write certificate"); + std::fs::write(&key_path, key_pair.serialize_pem()).expect("write key"); + let config = load_tls_config(&certificate_path, &key_path).expect("daemon TLS config"); + assert_eq!( + config.alpn_protocols, + [b"h2".to_vec(), b"http/1.1".to_vec()] + ); + + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind daemon TLS listener"); + let address = listener.local_addr().expect("daemon TLS address"); + let app = Router::new().route( + "/probe", + axum::routing::get(|| async { StatusCode::NO_CONTENT }), + ); + let server = tokio::spawn(serve_tls(listener, app, config)); + + let root = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(cert.der()); + let client = pooled_worker_tls_client(&root).expect("pinned TLS client"); + let request = Request::get(format!("https://127.0.0.1:{}/probe", address.port())) + .body(box_body(http_body_util::Empty::::new())) + .expect("probe request"); + let response = client.request(request).await.expect("daemon TLS response"); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + server.abort(); + assert!( + server + .await + .expect_err("TLS server is stopped") + .is_cancelled() + ); +} + +#[test] +fn public_ingress_keeps_only_authenticated_provider_routing_metadata() { + let mut headers = HeaderMap::new(); + headers.insert(CLIENT_TOKEN_HEADER, HeaderValue::from_static("route")); + headers.insert( + "x-nemo-relay-internal-dispatch-url", + HeaderValue::from_static("http://attacker.invalid"), + ); + headers.insert(WORKER_TOKEN_HEADER, HeaderValue::from_static("attacker")); + headers.insert( + "x-nemo-relay-bootstrap-proof", + HeaderValue::from_static("attacker"), + ); + headers.insert( + crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER, + HeaderValue::from_static("https://custom.example/v1"), + ); + let mut hook_headers = headers.clone(); + strip_public_relay_headers(&mut headers, PublicRoute::Provider(ProviderRoute::OpenAi)); + assert!(headers.contains_key(CLIENT_TOKEN_HEADER)); + assert!(headers.contains_key(crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER)); + assert!(!headers.contains_key("x-nemo-relay-internal-dispatch-url")); + assert!(!headers.contains_key(WORKER_TOKEN_HEADER)); + assert!(!headers.contains_key("x-nemo-relay-bootstrap-proof")); + + strip_public_relay_headers( + &mut hook_headers, + PublicRoute::Hook(crate::daemon::common::routes::HookRoute::Pi), + ); + assert!(!hook_headers.contains_key(crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER)); +} + +#[test] +fn pass_through_provider_auth_preserves_callers_and_fills_missing_configured_auth() { + let config = GatewayConfig { + openai_auth_header: Some("Bearer configured".into()), + ..GatewayConfig::default() + }; + let mut headers = HeaderMap::new(); + headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer caller")); + inject_provider_auth(&mut headers, ProviderRoute::OpenAi, &config); + assert_eq!( + headers.get(AUTHORIZATION).expect("caller auth"), + "Bearer caller" + ); + + headers.remove(AUTHORIZATION); + inject_provider_auth(&mut headers, ProviderRoute::OpenAi, &config); + assert_eq!( + headers.get(AUTHORIZATION).expect("configured auth"), + "Bearer configured" + ); +} + +#[test] +fn worker_route_failure_signal_is_consumed_before_the_public_response() { + let mut response = Response::new(Body::empty()); + response.headers_mut().insert( + WORKER_ROUTE_FAILURE_HEADER, + HeaderValue::from_static("pass-through"), + ); + + assert!(take_worker_route_failure(&mut response)); + assert!(!response.headers().contains_key(WORKER_ROUTE_FAILURE_HEADER)); + assert!(!take_worker_route_failure(&mut response)); +} + +#[test] +fn activation_endpoint_is_bound_to_signed_worker_network_policy() { + let activation = Activation { + fingerprint: MachineIdentity::generate().unwrap().identity.fingerprint(), + secret_digest: TokenDigest::from_token(b"secret"), + deadline_unix_ms: u64::MAX, + consumed: false, + bind_ip: Ipv4Addr::UNSPECIFIED, + port: 9443, + advertise_address: Some("worker.example.com".into()), + }; + assert!(activation_endpoint_matches( + "https://worker.example.com:9443", + &activation + )); + assert!(!activation_endpoint_matches( + "http://worker.example.com:9443", + &activation + )); + assert!(!activation_endpoint_matches( + "https://attacker.example.com:9443", + &activation + )); + let mut activation = activation; + activation.port = 443; + assert!(activation_endpoint_matches( + "https://worker.example.com:443", + &activation + )); + assert!(!activation_endpoint_matches( + "https://worker.example.com", + &activation + )); + assert!(!activation_endpoint_matches( + "https://worker.example.com:80", + &activation + )); + activation.bind_ip = Ipv4Addr::LOCALHOST; + activation.advertise_address = None; + activation.port = 80; + assert!(activation_endpoint_matches( + "http://127.0.0.1:80", + &activation + )); + assert!(!activation_endpoint_matches( + "http://127.0.0.1", + &activation + )); + assert!(!activation_endpoint_matches( + "http://127.0.0.1:443", + &activation + )); +} + +#[test] +fn released_mcp_session_credential_is_never_reused() { + let identity = MachineIdentity::generate().expect("identity").identity; + let fingerprint = identity.fingerprint(); + let token_digest = TokenDigest::from_token(b"route-token"); + let secret = SensitiveString::new("released-secret").expect("secret"); + let sessions = HashMap::from([( + "released".to_owned(), + McpControlSession { + fingerprint, + token_digest, + secret: secret.clone(), + secret_digest: TokenDigest::from_token(secret.expose().as_bytes()), + lease_expires_at_unix_ms: 1_000, + last_sequence: 1, + last_request_id: "release-request".into(), + last_heartbeat: None, + worker_network: worker_network(), + released: true, + }, + )]); + let fresh = SensitiveString::new("fresh-secret").expect("fresh"); + let (selected, reused) = select_mcp_session_token( + &sessions, + "released", + fingerprint, + token_digest, + worker_network(), + 999, + fresh.clone(), + ) + .expect("fresh selection"); + assert!(!reused); + assert_eq!(selected, fresh); +} + +#[test] +fn staged_worker_sessions_are_bounded_pruned_and_collision_safe() { + let mut sessions = HashMap::from([("staged".into(), staged_worker_session("staged", 100))]); + assert!(!reserve_worker_session_slot(&mut sessions, 99, "other", 1)); + assert!(!reserve_worker_session_slot(&mut sessions, 99, "staged", 2)); + assert!(reserve_worker_session_slot(&mut sessions, 100, "other", 1)); + assert!(sessions.is_empty()); +} + +#[test] +fn worker_staging_and_replay_are_bound_to_endpoint_and_publication() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x45_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let fingerprint = MachineIdentity::generate().unwrap().identity.fingerprint(); + let proof = HandshakeProof { + signer: ComponentRole::Daemon, + signature: state.identity.sign(b"test-daemon-proof"), + }; + let publication = WorkerPublication::Activation { + activation_id: "activation".into(), + }; + assert_eq!( + stage_worker( + &state, + fingerprint, + String::new(), + "http://127.0.0.1:41000".into(), + None, + None, + publication.clone(), + proof.clone(), + ) + .status(), + StatusCode::BAD_REQUEST + ); + assert!( + stage_worker( + &state, + fingerprint, + "worker".into(), + "http://127.0.0.1:41000".into(), + None, + None, + publication.clone(), + proof.clone(), + ) + .status() + .is_success() + ); + assert!( + replay_worker_registration( + &state, + fingerprint, + "worker", + "http://127.0.0.1:41000", + None, + Some("activation"), + None, + proof.clone(), + ) + .is_some() + ); + assert!( + replay_worker_registration( + &state, + fingerprint, + "worker", + "http://127.0.0.1:41001", + None, + Some("activation"), + None, + proof.clone(), + ) + .is_none() + ); + assert_eq!( + stage_worker( + &state, + fingerprint, + "worker".into(), + "http://127.0.0.1:41000".into(), + None, + None, + publication, + proof, + ) + .status(), + StatusCode::TOO_MANY_REQUESTS + ); +} + +fn staged_worker_session(worker_id: &str, lease_expires_at_unix_ms: u64) -> WorkerControlSession { + let worker = MachineIdentity::generate() + .expect("worker identity") + .identity; + let daemon = MachineIdentity::generate() + .expect("daemon identity") + .identity; + let endpoint = "http://127.0.0.1:41000"; + let secret = SensitiveString::new("control-secret").expect("control secret"); + let data = SensitiveString::new("data-secret").expect("data secret"); + WorkerControlSession { + fingerprint: worker.fingerprint(), + worker_id: worker_id.into(), + secret: secret.clone(), + secret_digest: TokenDigest::from_token(secret.expose().as_bytes()), + last_sequence: 0, + last_request_id: String::new(), + next_daemon_sequence: 0, + lease_expires_at_unix_ms, + pending_target: Arc::new( + WorkerTarget::new(worker_id, endpoint, data).expect("worker target"), + ), + publication: WorkerPublication::Activation { + activation_id: "activation".into(), + }, + published: false, + generation_grant: WorkerGenerationGrant::issue( + worker_id, + worker.fingerprint(), + endpoint, + None, + &daemon, + ) + .expect("generation grant"), + } +} + +#[tokio::test] +async fn daemon_drain_control_is_authenticated_retried_and_sequence_bounded() { + async fn drain(headers: HeaderMap, body: Bytes) -> Response { + assert_eq!(headers[WORKER_TOKEN_HEADER], "data-secret"); + let request: SessionRequest = serde_json::from_slice(&body).unwrap(); + assert_eq!(request.session_id, "drain-worker"); + assert_eq!(request.sequence, 1); + StatusCode::NO_CONTENT.into_response() + } + + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4b_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve( + listener, + Router::new().route(WORKER_DRAIN_PATH, post(drain)), + ) + .await + .unwrap(); + }); + let endpoint = format!("http://{address}"); + let target = Arc::new( + WorkerTarget::new( + "drain-worker", + endpoint, + SensitiveString::new("data-secret").unwrap(), + ) + .unwrap(), + ); + let mut session = staged_worker_session("drain-worker", u64::MAX); + session.pending_target = Arc::clone(&target); + lock(&state.worker_sessions).insert("drain-worker".into(), session); + + request_worker_drain(&state, &target, now_unix_ms().saturating_add(1_000)).await; + assert_eq!( + lock(&state.worker_sessions)["drain-worker"].next_daemon_sequence, + 1 + ); + + lock(&state.worker_sessions) + .get_mut("drain-worker") + .unwrap() + .next_daemon_sequence = u64::MAX; + request_worker_drain(&state, &target, u64::MAX).await; + lock(&state.worker_sessions).remove("drain-worker"); + request_worker_drain(&state, &target, u64::MAX).await; + server.abort(); +} + +#[test] +fn worker_generation_snapshot_releases_the_session_lock() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x73_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let session = staged_worker_session("snapshot-worker", u64::MAX); + let expected = ( + session.fingerprint, + session.generation_grant.generation_id.clone(), + ); + lock(&state.worker_sessions).insert("snapshot-worker".into(), session); + + let snapshot = worker_generation(&state, "snapshot-worker").unwrap(); + let mut sessions = state + .worker_sessions + .try_lock() + .expect("snapshot released lock"); + sessions.remove("snapshot-worker"); + assert_eq!( + snapshot, expected, + "snapshot owns its generation independently" + ); +} + +#[test] +fn communication_failure_invalidates_route_without_waiting_for_durable_revocation() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x74_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let session = staged_worker_session("failed-worker", u64::MAX); + let fingerprint = session.fingerprint; + let generation = session.generation_grant.generation_id.clone(); + let digest = RouteCredential::parse(token).unwrap().digest(); + let launch = fresh_launch(WorkerNetworkHint::new("127.0.0.1", None).unwrap()).unwrap(); + let activation_id = launch.activation_id.clone(); + state + .registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: digest, + session_id: McpSessionId::new("failure-mcp").unwrap(), + lease_expires_at_unix_ms: u64::MAX, + }, + launch, + ) + .unwrap(); + state + .registry + .mark_worker_ready( + fingerprint, + &activation_id, + Arc::clone(&session.pending_target), + ) + .unwrap(); + state + .active_worker_generations + .publish(fingerprint, &generation) + .unwrap(); + lock(&state.worker_sessions).insert("failed-worker".into(), session); + + // A publication transaction may be stalled on disk. Failure handling on a body-polling + // thread must nevertheless return and invalidate the route before that transaction ends. + let publication = lock(&state.worker_generation_publication); + let (sent, received) = std::sync::mpsc::channel(); + let failure_state = Arc::clone(&state); + let task = runtime.spawn_blocking(move || { + handle_worker_communication_failure(&failure_state, fingerprint, "failed-worker"); + sent.send(()).unwrap(); + }); + let completed = received.recv_timeout(Duration::from_secs(5)); + let route = state.registry.resolve_target(&digest); + let session_removed = state + .worker_sessions + .try_lock() + .is_ok_and(|sessions| !sessions.contains_key("failed-worker")); + // Always release contention before asserting, so a regression cannot strand runtime shutdown. + drop(publication); + runtime.block_on(task).unwrap(); + completed.expect("failure handling waited for the publication lock"); + assert!(matches!(route, Ok(ResolvedTarget::PassThrough))); + assert!(session_removed); + runtime.block_on(async { + tokio::time::timeout(Duration::from_secs(5), async { + while state + .active_worker_generations + .matches(fingerprint, &generation) + .unwrap() + { + tokio::task::yield_now().await; + } + }) + .await + .expect("generation was not durably revoked"); + }); +} + +#[tokio::test] +async fn release_actions_revoke_activation_transfer_directives_and_ignore_absent_nominees() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4d_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let fingerprint = MachineIdentity::generate().unwrap().identity.fingerprint(); + let launch = fresh_launch(WorkerNetworkHint::new("127.0.0.1", None).unwrap()).unwrap(); + let activation_id = launch.activation_id.clone(); + let directive = launch.into_directive(); + remember_activation(&state, fingerprint, &directive); + lock(&state.pending_directives).insert("launch-owner".into(), directive); + + handle_release_action( + Arc::clone(&state), + fingerprint, + ReleaseAction::CancelActivation { + activation_id: activation_id.clone(), + }, + ); + assert!(!lock(&state.activations).contains_key(&activation_id)); + assert!(!lock(&state.pending_directives).contains_key("launch-owner")); + + let session_id = McpSessionId::new("transfer-owner").unwrap(); + handle_release_action( + Arc::clone(&state), + fingerprint, + ReleaseAction::TransferActivation { + session_id: session_id.clone(), + directive: BrokerDirective::UsePassThrough, + }, + ); + assert!(matches!( + lock(&state.pending_directives).get(session_id.as_str()), + Some(BrokerDirective::UsePassThrough) + )); + + handle_release_action( + Arc::clone(&state), + fingerprint, + ReleaseAction::NominateMcp { + session_id: McpSessionId::new("missing-owner").unwrap(), + }, + ); + handle_release_action(state, fingerprint, ReleaseAction::NoChange); +} + +#[tokio::test] +async fn nominated_live_mcp_receives_a_fresh_relaunch_activation() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4e_u8; 32]); + let credential = RouteCredential::parse(token.clone()).unwrap(); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let fingerprint = MachineIdentity::generate().unwrap().identity.fingerprint(); + let session_id = McpSessionId::new("relaunch-owner").unwrap(); + let activation_id = match state + .registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id: session_id.clone(), + lease_expires_at_unix_ms: u64::MAX, + }, + fresh_launch(worker_network()).unwrap(), + ) + .unwrap() + { + BrokerDirective::LaunchWorker { activation_id, .. } => activation_id, + directive => panic!("unexpected directive: {directive:?}"), + }; + let target = Arc::new( + WorkerTarget::new( + "failed-worker", + "http://127.0.0.1:41001", + SensitiveString::new("worker-secret").unwrap(), + ) + .unwrap(), + ); + state + .registry + .mark_worker_ready(fingerprint, &activation_id, target) + .unwrap(); + let action = state + .registry + .worker_failed(fingerprint, "failed-worker", u64::MAX) + .unwrap(); + let nominee = match action { + WorkerFailureAction::NominateMcp { session_id } => session_id, + WorkerFailureAction::RouteEmpty => panic!("route unexpectedly empty"), + }; + let secret = SensitiveString::new("mcp-secret").unwrap(); + lock(&state.mcp_sessions).insert( + nominee.as_str().to_owned(), + McpControlSession { + fingerprint, + token_digest: credential.digest(), + secret: secret.clone(), + secret_digest: TokenDigest::from_token(secret.expose().as_bytes()), + lease_expires_at_unix_ms: u64::MAX, + last_sequence: 0, + last_request_id: String::new(), + last_heartbeat: None, + worker_network: worker_network(), + released: false, + }, + ); + + handle_release_action( + Arc::clone(&state), + fingerprint, + ReleaseAction::NominateMcp { + session_id: nominee.clone(), + }, + ); + assert!(matches!( + lock(&state.pending_directives).get(nominee.as_str()), + Some(BrokerDirective::LaunchWorker { .. }) + )); + assert_eq!(lock(&state.activations).len(), 1); +} + +#[tokio::test] +async fn worker_probe_reports_a_reachable_but_unready_endpoint() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + axum::serve( + listener, + Router::new().route( + WORKER_PROBE_PATH, + axum::routing::get(|| async { StatusCode::OK }), + ), + ) + .await + .unwrap(); + }); + let target = Arc::new( + WorkerTarget::new( + "unready-worker", + endpoint, + SensitiveString::new("data-secret").unwrap(), + ) + .unwrap(), + ); + let error = probe_worker(&target).await.unwrap_err(); + assert!(error.to_string().contains("returned HTTP 200 OK")); + server.abort(); +} + +#[test] +fn publication_failures_revoke_activation_and_fail_recovery_closed() { + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4f_u8; 32]); + let credential = RouteCredential::parse(token.clone()).unwrap(); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + let fingerprint = MachineIdentity::generate().unwrap().identity.fingerprint(); + let session_id = McpSessionId::new("publication-owner").unwrap(); + let launch = WorkerLaunch { + activation_id: "publication-activation".into(), + activation_token: SensitiveString::new("publication-secret").unwrap(), + deadline_unix_ms: u64::MAX, + bind_ip: Ipv4Addr::LOCALHOST, + port: 0, + advertise_address: None, + }; + let directive = state + .registry + .register_mcp( + McpRegistration { + fingerprint, + token_digest: credential.digest(), + session_id, + lease_expires_at_unix_ms: u64::MAX, + }, + launch, + ) + .unwrap(); + remember_activation(&state, fingerprint, &directive); + fail_worker_publication( + &state, + fingerprint, + &WorkerPublication::Activation { + activation_id: "publication-activation".into(), + }, + ); + assert!(lock(&state.activations).is_empty()); + assert!(matches!( + state.registry.resolve_target(&credential.digest()), + Ok(ResolvedTarget::PassThrough) + )); + + let unrelated = MachineIdentity::generate().unwrap().identity.fingerprint(); + fail_worker_publication( + &state, + unrelated, + &WorkerPublication::Recovery { + permit: RecoveryPermit::ExistingWorker { + worker_id: "missing-worker".into(), + recovering: true, + }, + }, + ); +} + +#[tokio::test] +async fn readiness_rechecks_activation_and_recovery_authority_after_the_probe() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + axum::serve( + listener, + Router::new().route( + WORKER_PROBE_PATH, + axum::routing::get(|| async { StatusCode::NO_CONTENT }), + ), + ) + .await + .unwrap(); + }); + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x4c_u8; 32]); + let state = test_daemon_state(false, &token, GatewayConfig::default()); + + let mut activation = staged_worker_session("stale-activation", u64::MAX); + activation.pending_target = Arc::new( + WorkerTarget::new( + "stale-activation", + endpoint.clone(), + SensitiveString::new("data-secret").unwrap(), + ) + .unwrap(), + ); + activation.publication = WorkerPublication::Activation { + activation_id: "already-revoked".into(), + }; + lock(&state.worker_sessions).insert("stale-activation".into(), activation); + let ready = SessionRequest::new( + "stale-activation".into(), + SensitiveString::new("control-secret").unwrap(), + 1, + WorkerReadyPayload { + worker_id: "stale-activation".into(), + }, + ) + .unwrap(); + assert_eq!( + ready_worker(State(Arc::clone(&state)), Json(ready)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + + let mut recovery = staged_worker_session("stale-recovery", u64::MAX); + recovery.pending_target = Arc::new( + WorkerTarget::new( + "stale-recovery", + endpoint, + SensitiveString::new("data-secret").unwrap(), + ) + .unwrap(), + ); + recovery.publication = WorkerPublication::Recovery { + permit: RecoveryPermit::ExistingWorker { + worker_id: "stale-recovery".into(), + recovering: true, + }, + }; + lock(&state.worker_sessions).insert("stale-recovery".into(), recovery); + let ready = SessionRequest::new( + "stale-recovery".into(), + SensitiveString::new("control-secret").unwrap(), + 1, + WorkerReadyPayload { + worker_id: "stale-recovery".into(), + }, + ) + .unwrap(); + assert_eq!( + ready_worker(State(Arc::clone(&state)), Json(ready)) + .await + .status(), + StatusCode::UNAUTHORIZED + ); + server.abort(); +} + +#[tokio::test] +async fn authenticated_control_plane_activates_heartbeats_and_drains_a_worker() { + let route_token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x5a_u8; 32]); + let credential = RouteCredential::parse(route_token.clone()).expect("route credential"); + let daemon_listener = TcpListener::bind("127.0.0.1:0").await.expect("daemon bind"); + let daemon_address = daemon_listener.local_addr().expect("daemon address"); + let daemon_origin = format!("http://{daemon_address}"); + let generation_directory = tempfile::tempdir().expect("generation directory"); + let daemon_identity = MachineIdentity::generate() + .expect("daemon identity") + .identity; + let state = Arc::new(DaemonState { + registry: Registry::new(false), + descriptor: crate::daemon::common::control::descriptor(ComponentRole::Daemon), + instance_id: "control-plane-test-daemon".into(), + public_origin: daemon_origin.clone(), + config: GatewayConfig::default(), + upstream: pooled_client().expect("daemon client"), + worker_clients: WorkerClientPool::new().expect("worker clients"), + challenges: Mutex::new(HashMap::new()), + activations: Mutex::new(HashMap::new()), + mcp_sessions: Mutex::new(HashMap::new()), + mcp_heartbeat_serialization: Mutex::new(()), + worker_sessions: Mutex::new(HashMap::new()), + pending_directives: Mutex::new(HashMap::new()), + active_worker_generations: ActiveWorkerGenerations::load_for_test( + generation_directory.path().join("active-workers.json"), + ) + .expect("generation state"), + worker_generation_publication: Mutex::new(()), + identity: daemon_identity, + }); + let daemon_task = tokio::spawn({ + let state = Arc::clone(&state); + async move { + axum::serve(daemon_listener, router(state)) + .await + .expect("daemon serve"); + } + }); + + let client = control_client().expect("control client"); + let machine_identity = MachineIdentity::generate() + .expect("machine identity") + .identity; + let mcp_session_id = "control-plane-test-mcp"; + let mcp_handshake = begin_handshake( + &client, + &daemon_origin, + ComponentRole::Mcp, + &machine_identity, + mcp_session_id, + Some(credential.digest()), + ) + .await + .expect("MCP handshake"); + let hint = + WorkerNetworkHint::new(Ipv4Addr::LOCALHOST.to_string(), None).expect("worker network hint"); + let hint_proof = WorkerNetworkHintProof::sign( + hint, + &mcp_handshake.proof.transcript.daemon_target, + mcp_session_id, + &mcp_handshake.proof.transcript.challenge_id, + &machine_identity.fingerprint(), + &machine_identity, + ) + .expect("worker network proof"); + let mcp_registration: McpRegisterResponse = post_json( + &client, + &format!("{daemon_origin}{MCP_REGISTER_PATH}"), + &McpRegisterRequest { + proof: mcp_handshake.proof.clone(), + worker_network: hint_proof, + }, + Some(&route_token), + ) + .await + .expect("MCP registration"); + mcp_handshake + .authenticate_daemon(&mcp_registration.daemon_proof) + .expect("daemon MCP proof"); + let bootstrap = WorkerBootstrap::from_directive(mcp_registration.directive.clone()) + .expect("launch directive"); + + let worker_listener = TcpListener::bind("127.0.0.1:0").await.expect("worker bind"); + let worker_address = worker_listener.local_addr().expect("worker address"); + let worker_endpoint = format!("http://{worker_address}"); + let worker_id = "test-worker"; + let worker_handshake = begin_handshake( + &client, + &daemon_origin, + ComponentRole::Worker, + &machine_identity, + worker_id, + None, + ) + .await + .expect("worker handshake"); + let worker_registration: WorkerRegisterResponse = post_json( + &client, + &format!("{daemon_origin}{WORKER_REGISTER_PATH}"), + &WorkerRegisterRequest { + proof: worker_handshake.proof.clone(), + worker_id: worker_id.into(), + endpoint: worker_endpoint, + activation_id: bootstrap.activation_id, + activation_token: bootstrap.activation_token, + tls_root_certificate: None, + }, + None, + ) + .await + .expect("worker registration"); + worker_handshake + .authenticate_daemon(&worker_registration.daemon_proof) + .expect("daemon worker proof"); + + let upstream_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("provider bind"); + let upstream_address = upstream_listener.local_addr().expect("provider address"); + let upstream_task = tokio::spawn(async move { + axum::serve( + upstream_listener, + Router::new().route( + "/v1/responses", + post(|| async { + Response::builder() + .status(StatusCode::CREATED) + .header("x-worker-provider", "reached") + .body(Body::from("worker provider bytes")) + .expect("provider response") + }), + ), + ) + .await + .expect("serve provider"); + }); + let worker_config = GatewayConfig { + openai_base_url: format!("http://{upstream_address}"), + ..GatewayConfig::default() + }; + let (worker_router, worker_handle) = test_router_with_control_tokens( + worker_config, + pooled_client().expect("worker upstream client"), + worker_registration.data_token.expose().as_bytes(), + worker_registration.session_token.expose().as_bytes(), + ); + let worker_task = tokio::spawn(async move { + axum::serve(worker_listener, worker_router) + .await + .expect("worker serve"); + }); + + let ready = SessionRequest::new( + worker_id.into(), + worker_registration.session_token.clone(), + 1, + WorkerReadyPayload { + worker_id: worker_id.into(), + }, + ) + .expect("ready request"); + let response = client + .post(format!("{daemon_origin}{WORKER_READY_PATH}")) + .json(&ready) + .send() + .await + .expect("ready response"); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + assert!(state.registry.resolve_target(&credential.digest()).is_ok()); + + let public_response = router(Arc::clone(&state)) + .oneshot( + Request::post("/v1/responses") + .header(CLIENT_TOKEN_HEADER, &route_token) + .body(Body::from(r#"{"model":"test","stream":true}"#)) + .expect("public worker request"), + ) + .await + .expect("public worker response"); + assert_eq!(public_response.status(), StatusCode::CREATED); + assert_eq!(public_response.headers()["x-worker-provider"], "reached"); + assert_eq!( + public_response + .into_body() + .collect() + .await + .expect("public worker body") + .to_bytes(), + "worker provider bytes" + ); + + let worker_heartbeat = SessionRequest::new( + worker_id.into(), + worker_registration.session_token, + 2, + WorkerHeartbeatPayload { + worker_id: worker_id.into(), + }, + ) + .expect("worker heartbeat"); + let response = client + .post(format!("{daemon_origin}{WORKER_HEARTBEAT_PATH}")) + .json(&worker_heartbeat) + .send() + .await + .expect("worker heartbeat response"); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + lock(&state.worker_sessions).remove(worker_id); + let (recovery_data, recovery_session, _generation_grant) = + crate::daemon::worker::test_recover_control_session( + &daemon_origin, + &machine_identity, + worker_id, + &format!("http://{worker_address}"), + worker_registration.generation_grant, + ) + .await + .expect("worker recovery registration"); + worker_handle.stage_recovery_tokens(recovery_data.as_bytes(), recovery_session.as_bytes()); + let recovery_ready = SessionRequest::new( + worker_id.into(), + SensitiveString::new(recovery_session).expect("recovery session token"), + 1, + WorkerReadyPayload { + worker_id: worker_id.into(), + }, + ) + .expect("recovery ready request"); + let response = client + .post(format!("{daemon_origin}{WORKER_READY_PATH}")) + .json(&recovery_ready) + .send() + .await + .expect("recovery ready response"); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + let mcp_heartbeat = SessionRequest::new( + mcp_session_id.into(), + mcp_registration.session_token.clone(), + 1, + EmptyPayload::default(), + ) + .expect("MCP heartbeat"); + let heartbeat: McpHeartbeatResponse = post_json( + &client, + &format!("{daemon_origin}{MCP_HEARTBEAT_PATH}"), + &mcp_heartbeat, + None, + ) + .await + .expect("MCP heartbeat response"); + assert!(matches!( + heartbeat.directive, + Some(BrokerDirective::ReuseWorker { .. }) + )); + + let release = SessionRequest::new( + mcp_session_id.into(), + mcp_registration.session_token, + 2, + EmptyPayload::default(), + ) + .expect("MCP release"); + let response = client + .post(format!("{daemon_origin}{MCP_RELEASE_PATH}")) + .json(&release) + .send() + .await + .expect("MCP release response"); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + tokio::time::timeout(Duration::from_secs(2), async { + while !worker_handle.is_draining() { + tokio::task::yield_now().await; + } + }) + .await + .expect("daemon sent authenticated drain request"); + worker_task.abort(); + upstream_task.abort(); + daemon_task.abort(); +} diff --git a/crates/cli/tests/coverage/daemon/state_tests.rs b/crates/cli/tests/coverage/daemon/state_tests.rs new file mode 100644 index 000000000..03d831721 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/state_tests.rs @@ -0,0 +1,356 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn route_credential_is_exactly_256_bits() { + let value = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([7_u8; 32]); + let credential = RouteCredential::parse(value.clone()).unwrap(); + assert_eq!(credential.expose(), value); + assert!(RouteCredential::parse("short".into()).is_err()); + assert!(RouteCredential::parse(format!(" {value}")).is_err()); + assert!(RouteCredential::parse("!invalid!".into()).is_err()); + assert!(RouteCredential::parse(format!("{value}\n")).is_err()); +} + +#[test] +fn route_credential_loads_from_the_environment_and_reports_absence() { + let value = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([8_u8; 32]); + let environment = crate::test_support::EnvScope::set(&[( + ROUTE_TOKEN_ENV, + Some(std::ffi::OsStr::new(&value)), + )]); + let credential = RouteCredential::from_environment().expect("environment credential"); + assert_eq!(credential.expose(), value); + assert_eq!( + credential.digest(), + TokenDigest::from_token(value.as_bytes()) + ); + drop(environment); + + let _environment = crate::test_support::EnvScope::set(&[(ROUTE_TOKEN_ENV, None)]); + assert!(RouteCredential::from_environment().is_err()); +} + +#[cfg(unix)] +#[test] +fn identity_and_lock_files_reject_symlinks_and_repair_owner_private_modes() { + use std::os::unix::fs::{PermissionsExt, symlink}; + + let directory = tempfile::tempdir().expect("tempdir"); + let identity = directory.path().join("identity.pk8"); + load_or_create_identity(&identity).expect("create identity"); + std::fs::set_permissions(&identity, std::fs::Permissions::from_mode(0o644)) + .expect("loosen identity mode"); + load_or_create_identity(&identity).expect("repair identity mode"); + assert_eq!( + std::fs::metadata(&identity) + .expect("identity metadata") + .permissions() + .mode() + & 0o777, + 0o600 + ); + + let target = directory.path().join("target"); + std::fs::write(&target, b"not-an-identity").expect("target"); + let linked_identity = directory.path().join("linked.pk8"); + symlink(&target, &linked_identity).expect("identity symlink"); + assert!(load_or_create_identity(&linked_identity).is_err()); + + let lock_target = directory.path().join("lock-target"); + std::fs::write(&lock_target, b"").expect("lock target"); + let lock_identity = directory.path().join("lock-linked.pk8"); + symlink(&lock_target, lock_identity.with_extension("lock")).expect("lock symlink"); + assert!(load_or_create_identity(&lock_identity).is_err()); +} + +#[test] +fn active_generation_survives_restart_but_revoked_generation_does_not() { + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join(ACTIVE_WORKER_GENERATIONS_FILENAME); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let generations = + ActiveWorkerGenerations::load_for_test(path.clone()).expect("load empty state"); + + assert!(!generations.matches(fingerprint, "generation-one").unwrap()); + assert_eq!( + generations.publish(fingerprint, "generation-one").unwrap(), + None + ); + assert!(generations.matches(fingerprint, "generation-one").unwrap()); + assert!(!generations.matches(fingerprint, "generation-two").unwrap()); + + let reloaded = ActiveWorkerGenerations::load_for_test(path.clone()).expect("reload state"); + assert!(reloaded.matches(fingerprint, "generation-one").unwrap()); + assert_eq!( + reloaded.publish(fingerprint, "generation-two").unwrap(), + Some("generation-one".into()) + ); + assert!( + !reloaded + .revoke_if_matches(fingerprint, "generation-one") + .unwrap() + ); + assert!(reloaded.matches(fingerprint, "generation-two").unwrap()); + assert!( + reloaded + .revoke_if_matches(fingerprint, "generation-two") + .unwrap() + ); + + let after_revoke = ActiveWorkerGenerations::load_for_test(path).expect("reload revoked state"); + assert!(!after_revoke.matches(fingerprint, "generation-two").unwrap()); +} + +#[test] +fn active_worker_generation_restore_is_compare_and_set() { + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join(ACTIVE_WORKER_GENERATIONS_FILENAME); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let generations = ActiveWorkerGenerations::load_for_test(path).expect("load state"); + generations + .publish(fingerprint, "generation-old") + .expect("publish old generation"); + let previous = generations + .publish(fingerprint, "generation-candidate") + .expect("publish candidate"); + + assert!( + !generations + .restore_if_matches(fingerprint, "different-candidate", previous.as_deref(),) + .unwrap() + ); + assert!( + generations + .restore_if_matches(fingerprint, "generation-candidate", previous.as_deref(),) + .unwrap() + ); + assert!(generations.matches(fingerprint, "generation-old").unwrap()); + assert!( + generations + .restore_if_matches(fingerprint, "generation-old", None) + .unwrap() + ); + assert!(!generations.matches(fingerprint, "generation-old").unwrap()); +} + +#[test] +fn active_worker_generation_ids_are_validated_on_every_public_operation() { + let directory = tempfile::tempdir().expect("tempdir"); + let generations = ActiveWorkerGenerations::load_for_test( + directory.path().join(ACTIVE_WORKER_GENERATIONS_FILENAME), + ) + .expect("load generation state"); + let fingerprint = MachineIdentity::generate().unwrap().identity.fingerprint(); + for invalid in ["", &"x".repeat(MAX_GENERATION_ID_BYTES + 1)] { + assert!(generations.matches(fingerprint, invalid).is_err()); + assert!(generations.publish(fingerprint, invalid).is_err()); + assert!(generations.revoke_if_matches(fingerprint, invalid).is_err()); + assert!( + generations + .restore_if_matches(fingerprint, invalid, None) + .is_err() + ); + } + generations.publish(fingerprint, "current").unwrap(); + assert!( + generations + .restore_if_matches( + fingerprint, + "current", + Some(&"x".repeat(MAX_GENERATION_ID_BYTES + 1)), + ) + .is_err() + ); +} + +#[test] +fn identity_and_daemon_pin_state_fail_closed_on_corruption_or_replacement() { + let directory = tempfile::tempdir().expect("tempdir"); + let invalid = directory.path().join("invalid.pk8"); + std::fs::write(&invalid, b"not a private key").unwrap(); + assert!(load_or_create_identity(&invalid).is_err()); + + let oversized = directory.path().join("oversized.pk8"); + std::fs::File::create(&oversized) + .unwrap() + .set_len(MAX_IDENTITY_BYTES + 1) + .unwrap(); + assert!(load_or_create_identity(&oversized).is_err()); + + let _environment = crate::test_support::EnvScope::set(&[( + "XDG_CONFIG_HOME", + Some(directory.path().as_os_str()), + )]); + let first = MachineIdentity::generate() + .unwrap() + .identity + .public_identity(); + let second = MachineIdentity::generate() + .unwrap() + .identity + .public_identity(); + verify_or_store_daemon_pin("https://daemon.example:443", first).unwrap(); + verify_or_store_daemon_pin("https://daemon.example:443", first).unwrap(); + assert!(verify_or_store_daemon_pin("https://daemon.example:443", second).is_err()); + load_or_create_machine_identity().unwrap(); + load_or_create_daemon_identity().unwrap(); +} + +#[cfg(unix)] +#[test] +fn private_state_helpers_reject_non_directories_and_non_regular_files() { + use std::os::unix::fs::symlink; + + let directory = tempfile::tempdir().expect("tempdir"); + let ordinary = directory.path().join("ordinary"); + std::fs::write(&ordinary, b"value").unwrap(); + assert!(create_private_directory(&ordinary).is_err()); + + let linked = directory.path().join("linked-directory"); + symlink(directory.path(), &linked).unwrap(); + assert!(create_private_directory(&linked).is_err()); + + let opened_directory = std::fs::File::open(directory.path()).unwrap(); + assert!(validate_private_file(&opened_directory, directory.path()).is_err()); + assert!( + read_bounded(&directory.path().join("missing"), 8, "missing") + .unwrap() + .is_none() + ); + assert!(read_bounded(directory.path(), 8, "directory").is_err()); +} + +#[test] +fn active_worker_generation_state_corruption_fails_closed() { + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join(ACTIVE_WORKER_GENERATIONS_FILENAME); + std::fs::write(&path, b"{").expect("write corrupt state"); + + let error = ActiveWorkerGenerations::load_for_test(path) + .expect_err("corrupt state must fail") + .to_string(); + + assert!(error.contains("corrupt"), "{error}"); +} + +#[test] +fn active_worker_generation_state_enforces_its_file_bound() { + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join(ACTIVE_WORKER_GENERATIONS_FILENAME); + let file = std::fs::File::create(&path).expect("create oversized state"); + file.set_len(MAX_ACTIVE_WORKER_GENERATIONS_BYTES + 1) + .expect("extend oversized state"); + + let error = ActiveWorkerGenerations::load_for_test(path) + .expect_err("oversized state must fail") + .to_string(); + + assert!(error.contains("exceeds"), "{error}"); +} + +#[test] +fn active_worker_generation_state_rejects_duplicate_routes_and_unknown_schema() { + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join(ACTIVE_WORKER_GENERATIONS_FILENAME); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let duplicate = serde_json::json!({ + "schema_version": ACTIVE_WORKER_GENERATIONS_SCHEMA_VERSION, + "generations": [ + {"fingerprint": fingerprint, "generation_id": "generation-one"}, + {"fingerprint": fingerprint, "generation_id": "generation-two"} + ] + }); + std::fs::write(&path, serde_json::to_vec(&duplicate).unwrap()).expect("write duplicate state"); + assert!( + ActiveWorkerGenerations::load_for_test(path.clone()) + .unwrap_err() + .to_string() + .contains("duplicate fingerprint") + ); + + let too_many = serde_json::json!({ + "schema_version": ACTIVE_WORKER_GENERATIONS_SCHEMA_VERSION, + "generations": (0..=MAX_ACTIVE_WORKER_GENERATIONS) + .map(|index| { + let mut distinct = [0_u8; 32]; + distinct[..std::mem::size_of::()] + .copy_from_slice(&index.to_le_bytes()); + serde_json::json!({ + "fingerprint": distinct, + "generation_id": format!("generation-{index}") + }) + }) + .collect::>() + }); + std::fs::write(&path, serde_json::to_vec(&too_many).unwrap()).expect("write oversized map"); + assert!( + ActiveWorkerGenerations::load_for_test(path.clone()) + .unwrap_err() + .to_string() + .contains("routes") + ); + + let unknown = serde_json::json!({"schema_version": 2, "generations": []}); + std::fs::write(&path, serde_json::to_vec(&unknown).unwrap()).expect("write unknown state"); + assert!( + ActiveWorkerGenerations::load_for_test(path) + .unwrap_err() + .to_string() + .contains("unsupported schema version") + ); +} + +#[cfg(unix)] +#[test] +fn active_worker_generation_state_is_owner_private_and_rejects_symlinks() { + use std::os::unix::fs::{PermissionsExt, symlink}; + + let directory = tempfile::tempdir().expect("tempdir"); + let path = directory.path().join(ACTIVE_WORKER_GENERATIONS_FILENAME); + let fingerprint = MachineIdentity::generate() + .expect("machine identity") + .identity + .fingerprint(); + let generations = + ActiveWorkerGenerations::load_for_test(path.clone()).expect("load generation state"); + generations + .publish(fingerprint, "generation") + .expect("publish generation"); + assert_eq!( + std::fs::metadata(&path) + .expect("state metadata") + .permissions() + .mode() + & 0o777, + 0o600 + ); + std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o644)) + .expect("loosen state mode"); + ActiveWorkerGenerations::load_for_test(path.clone()).expect("repair state mode"); + assert_eq!( + std::fs::metadata(&path) + .expect("repaired state metadata") + .permissions() + .mode() + & 0o777, + 0o600 + ); + + let target = directory.path().join("target.json"); + std::fs::write(&target, b"{}").expect("write target"); + let linked = directory.path().join("linked.json"); + symlink(&target, &linked).expect("create state symlink"); + assert!(ActiveWorkerGenerations::load_for_test(linked).is_err()); +} diff --git a/crates/cli/tests/coverage/daemon/transport_tests.rs b/crates/cli/tests/coverage/daemon/transport_tests.rs new file mode 100644 index 000000000..bac986262 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/transport_tests.rs @@ -0,0 +1,1116 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::VecDeque; +use std::convert::Infallible; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use http::header::{HeaderValue, SET_COOKIE, TE, TRAILER}; +use http_body_util::{BodyExt, Empty}; +use hyper::body::{Frame, Incoming, SizeHint}; +use hyper::server::conn::{http1, http2}; +use hyper::service::service_fn; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use tokio::net::TcpListener; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; + +use super::*; + +#[test] +fn forwarding_rejects_remote_cleartext_destinations() { + let request = Request::get("/") + .body(box_body(Empty::::new())) + .expect("request"); + let remote = "http://provider.example/v1/responses" + .parse() + .expect("remote URI"); + assert!(prepare_forward_request(request, remote, &[]).is_err()); + + let request = Request::get("/") + .body(box_body(Empty::::new())) + .expect("request"); + let loopback = "http://127.0.0.1:8080/v1/responses" + .parse() + .expect("loopback URI"); + assert!(prepare_forward_request(request, loopback, &[]).is_ok()); +} + +struct CausalBody { + phase: u8, + release_second: oneshot::Receiver<()>, + trailers: Option, +} + +impl CausalBody { + fn new(release_second: oneshot::Receiver<()>, trailers: HeaderMap) -> Self { + Self { + phase: 0, + release_second, + trailers: Some(trailers), + } + } +} + +impl Body for CausalBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + match this.phase { + 0 => { + this.phase = 1; + Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(b"event-a\n\n"))))) + } + 1 => match Pin::new(&mut this.release_second).poll(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(_) => { + this.phase = 2; + Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(b"event-b\n\n"))))) + } + }, + 2 => { + this.phase = 3; + Poll::Ready( + this.trailers + .take() + .map(|trailers| Ok(Frame::trailers(trailers))), + ) + } + _ => Poll::Ready(None), + } + } + + fn is_end_stream(&self) -> bool { + self.phase >= 3 && self.trailers.is_none() + } + + fn size_hint(&self) -> SizeHint { + SizeHint::default() + } +} + +#[derive(Clone, Copy)] +enum TestProtocol { + Http1, + Http2, +} + +fn client_for(protocol: TestProtocol) -> PooledHttpClient { + match protocol { + TestProtocol::Http1 => pooled_http_client(), + TestProtocol::Http2 => pooled_h2c_client(), + } +} + +fn request_with_empty_body(uri: String, protocol: TestProtocol) -> Request { + let mut request = Request::get(uri) + .body(box_body(Empty::::new())) + .expect("valid request"); + if matches!(protocol, TestProtocol::Http1) { + request + .headers_mut() + .insert(TE, HeaderValue::from_static("trailers")); + } + request +} + +async fn serve_one_connection(listener: TcpListener, protocol: TestProtocol, service: S) +where + S: hyper::service::Service< + Request, + Response = Response, + Error = Infallible, + > + Send + + 'static, + S::Future: Send, +{ + let (stream, _) = listener.accept().await.expect("accept test client"); + match protocol { + TestProtocol::Http1 => http1::Builder::new() + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve HTTP/1.1 test connection"), + TestProtocol::Http2 => { + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + builder + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve HTTP/2 test connection") + } + } +} + +async fn spawn_causal_provider( + protocol: TestProtocol, + release_second: oneshot::Receiver<()>, +) -> (std::net::SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind causal provider"); + let address = listener.local_addr().expect("bound provider address"); + let release_second = Arc::new(Mutex::new(Some(release_second))); + let service = service_fn(move |_request: Request| { + let release_second = release_second + .lock() + .expect("release gate lock") + .take() + .expect("causal provider receives exactly one request"); + async move { + let mut trailers = HeaderMap::new(); + trailers.append("x-checksum", HeaderValue::from_static("one")); + trailers.append("x-checksum", HeaderValue::from_static("two")); + let response = Response::builder() + .status(StatusCode::CREATED) + .header(TRAILER, "x-checksum") + .body(box_body(CausalBody::new(release_second, trailers))) + .expect("valid causal response"); + Ok::<_, Infallible>(response) + } + }); + let task = tokio::spawn(serve_one_connection(listener, protocol, service)); + (address, task) +} + +async fn spawn_relay( + protocol: TestProtocol, + provider: std::net::SocketAddr, +) -> (std::net::SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind relay"); + let address = listener.local_addr().expect("bound relay address"); + let upstream = client_for(protocol); + let service = service_fn(move |request: Request| { + let upstream = upstream.clone(); + async move { + let path = request + .uri() + .path_and_query() + .map(|value| value.as_str()) + .unwrap_or("/"); + let destination = format!("http://{provider}{path}") + .parse() + .expect("valid provider URI"); + let request = request.map(box_body); + let request = prepare_forward_request(request, destination, &[]) + .expect("relay request head is valid"); + let response = upstream.request(request).await.expect("provider responds"); + let response = prepare_forward_response(response, &[]) + .expect("relay response head is valid") + .map(box_body); + Ok::<_, Infallible>(response) + } + }); + let task = tokio::spawn(serve_one_connection(listener, protocol, service)); + (address, task) +} + +#[allow(clippy::cognitive_complexity)] +async fn assert_causal_relay(protocol: TestProtocol, path: &str) { + let (release_second, wait_for_release) = oneshot::channel(); + let (provider, provider_task) = spawn_causal_provider(protocol, wait_for_release).await; + let (relay, relay_task) = spawn_relay(protocol, provider).await; + + let client = client_for(protocol); + let request = request_with_empty_body(format!("http://{relay}{path}"), protocol); + let response = client.request(request).await.expect("relay responds"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.headers()[TRAILER], "x-checksum"); + + let mut body = box_body(response.into_body()); + let first = body + .frame() + .await + .expect("first frame exists") + .expect("first frame succeeds") + .into_data() + .expect("first frame is data"); + assert_eq!(first, "event-a\n\n"); + + let second = body.frame(); + tokio::pin!(second); + assert!( + futures_util::poll!(second.as_mut()).is_pending(), + "relay must expose event A without waiting for event B" + ); + release_second.send(()).expect("release provider event B"); + + assert_eq!( + second + .await + .expect("second frame exists") + .expect("second frame succeeds") + .into_data() + .expect("second frame is data"), + "event-b\n\n" + ); + let trailers = body + .frame() + .await + .expect("trailer frame exists") + .expect("trailer frame succeeds") + .into_trailers() + .expect("last frame contains trailers"); + assert_eq!( + trailers + .get_all("x-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(), + ["one", "two"] + ); + assert!(body.frame().await.is_none()); + + drop(client); + match protocol { + TestProtocol::Http1 => { + relay_task.await.expect("relay task succeeds"); + provider_task.await.expect("provider task succeeds"); + } + TestProtocol::Http2 => { + relay_task.abort(); + provider_task.abort(); + } + } +} + +struct FramesBody { + frames: VecDeque>, +} + +impl FramesBody { + fn new(frames: impl IntoIterator>) -> Self { + Self { + frames: frames.into_iter().collect(), + } + } +} + +impl Body for FramesBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + Poll::Ready(self.frames.pop_front().map(Ok)) + } + + fn is_end_stream(&self) -> bool { + self.frames.is_empty() + } +} + +struct CountedBody { + remaining: usize, + frame: Bytes, + polls: Arc, +} + +impl Body for CountedBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + self.polls.fetch_add(1, Ordering::SeqCst); + if self.remaining == 0 { + return Poll::Ready(None); + } + self.remaining -= 1; + Poll::Ready(Some(Ok(Frame::data(self.frame.clone())))) + } +} + +struct CancellationBody { + first_sent: bool, + dropped: Option>, +} + +impl Body for CancellationBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + if self.first_sent { + Poll::Pending + } else { + self.first_sent = true; + Poll::Ready(Some(Ok(Frame::data(Bytes::from_static(b"first\n\n"))))) + } + } +} + +impl Drop for CancellationBody { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } +} + +#[test] +fn strips_connection_scoped_headers_and_preserves_trailer_declaration() { + let mut headers = HeaderMap::new(); + headers.append( + CONNECTION, + HeaderValue::from_static("keep-alive, x-private"), + ); + headers.append(CONNECTION, HeaderValue::from_static("proxy-connection")); + headers.insert("keep-alive", HeaderValue::from_static("timeout=5")); + headers.insert("proxy-connection", HeaderValue::from_static("keep-alive")); + headers.insert("x-private", HeaderValue::from_static("secret")); + headers.append(TRAILER, HeaderValue::from_static("x-checksum")); + headers.append(TRAILER, HeaderValue::from_static("x-signature")); + headers.append("x-end-to-end", HeaderValue::from_static("one")); + headers.append("x-end-to-end", HeaderValue::from_static("two")); + + strip_hop_by_hop_headers(&mut headers).expect("valid headers"); + + assert!(!headers.contains_key(CONNECTION)); + assert!(!headers.contains_key("keep-alive")); + assert!(!headers.contains_key("proxy-connection")); + assert!(!headers.contains_key("x-private")); + assert_eq!( + headers + .get_all(TRAILER) + .iter() + .map(|value| value.to_str().expect("ASCII trailer declaration")) + .collect::>(), + ["x-checksum", "x-signature"] + ); + assert_eq!( + headers + .get_all("x-end-to-end") + .iter() + .map(|value| value.to_str().expect("ASCII test header")) + .collect::>(), + ["one", "two"] + ); +} + +#[test] +fn removes_trailer_declaration_when_connection_nominates_it() { + let mut headers = HeaderMap::new(); + headers.insert(CONNECTION, HeaderValue::from_static("trailer")); + headers.insert(TRAILER, HeaderValue::from_static("x-checksum")); + + strip_hop_by_hop_headers(&mut headers).expect("valid headers"); + + assert!(!headers.contains_key(TRAILER)); +} + +#[test] +fn rewrites_only_the_request_head() { + let (release, receiver) = oneshot::channel(); + let body = CausalBody::new(receiver, HeaderMap::new()); + let mut request = Request::post("http://old.example/v1/responses") + .header(HOST, "old.example") + .header("x-route-token", "private") + .body(body) + .expect("valid request"); + *request.version_mut() = http::Version::HTTP_2; + request + .headers_mut() + .append("x-preserved", HeaderValue::from_static("first")); + request + .headers_mut() + .append("x-preserved", HeaderValue::from_static("second")); + + let destination = "https://worker.example:8443/v1/responses?stream=true" + .parse() + .expect("valid destination"); + let request = prepare_forward_request( + request, + destination, + &[HeaderName::from_static("x-route-token")], + ) + .expect("request can be forwarded"); + + assert_eq!( + request.uri(), + &"https://worker.example:8443/v1/responses?stream=true" + .parse::() + .expect("valid expected URI") + ); + assert_eq!(request.headers()[HOST], "worker.example:8443"); + assert_eq!(request.version(), http::Version::HTTP_11); + assert!(!request.headers().contains_key("x-route-token")); + assert_eq!(request.headers().get_all("x-preserved").iter().count(), 2); + + drop(request); + assert!( + release.send(()).is_err(), + "the unchanged body owns the receiver" + ); +} + +#[tokio::test] +async fn forwards_first_frame_before_source_releases_second_and_preserves_trailers() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test server"); + let address = listener.local_addr().expect("bound address"); + let (release_second, wait_for_release) = oneshot::channel(); + + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept test client"); + let wait_for_release = Arc::new(Mutex::new(Some(wait_for_release))); + let service = service_fn(move |_request: Request| { + let wait_for_release = wait_for_release + .lock() + .expect("release gate lock") + .take() + .expect("test server receives exactly one request"); + async move { + let mut trailers = HeaderMap::new(); + trailers.append("x-checksum", HeaderValue::from_static("one")); + trailers.append("x-checksum", HeaderValue::from_static("two")); + let body = box_body(CausalBody::new(wait_for_release, trailers)); + let response = Response::builder() + .header(TRAILER, "x-checksum") + .body(body) + .expect("valid response"); + Ok::<_, Infallible>(response) + } + }); + + let mut connection = http1::Builder::new(); + connection.keep_alive(false); + connection + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve causal response"); + }); + + let client = pooled_http_client(); + let request = Request::get(format!("http://{address}/v1/responses")) + .header(TE, "trailers") + .body(box_body(Empty::::new())) + .expect("valid request"); + let response = client.request(request).await.expect("request succeeds"); + assert_eq!(response.headers()[TRAILER], "x-checksum"); + + let mut body = box_body(response.into_body()); + let first = body + .frame() + .await + .expect("first frame exists") + .expect("first frame succeeds") + .into_data() + .expect("first frame is data"); + assert_eq!(first, "event-a\n\n"); + + let second = body.frame(); + tokio::pin!(second); + assert!(futures_util::poll!(second.as_mut()).is_pending()); + release_second.send(()).expect("release source"); + + let second = second + .await + .expect("second frame exists") + .expect("second frame succeeds") + .into_data() + .expect("second frame is data"); + assert_eq!(second, "event-b\n\n"); + + let trailers = body + .frame() + .await + .expect("trailer frame exists") + .expect("trailer frame succeeds") + .into_trailers() + .expect("last frame contains trailers"); + let values = trailers + .get_all("x-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(); + assert_eq!(values, ["one", "two"]); + assert!(body.frame().await.is_none()); + + drop(client); + server.await.expect("server task succeeds"); +} + +#[tokio::test] +async fn h2_preserves_duplicate_trailer_multimap_values() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind test server"); + let address = listener.local_addr().expect("bound address"); + let (release_second, wait_for_release) = oneshot::channel(); + + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept test client"); + let release = Arc::new(Mutex::new(Some(wait_for_release))); + let service = service_fn(move |_request: Request| { + let wait_for_release = release + .lock() + .expect("release gate lock") + .take() + .expect("test server receives exactly one request"); + async move { + let mut trailers = HeaderMap::new(); + trailers.append("x-checksum", HeaderValue::from_static("one")); + trailers.append("x-checksum", HeaderValue::from_static("two")); + Ok::<_, Infallible>(Response::new(box_body(CausalBody::new( + wait_for_release, + trailers, + )))) + } + }); + + http2::Builder::new(TokioExecutor::new()) + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve causal HTTP/2 response"); + }); + + let client = pooled_h2c_client(); + let request = Request::get(format!("http://{address}/v1/messages")) + .body(box_body(Empty::::new())) + .expect("valid request"); + let response = client.request(request).await.expect("request succeeds"); + let mut body = box_body(response.into_body()); + assert_eq!( + body.frame() + .await + .expect("first frame") + .expect("first frame succeeds") + .into_data() + .expect("data frame"), + "event-a\n\n" + ); + release_second.send(()).expect("release source"); + assert_eq!( + body.frame() + .await + .expect("second frame") + .expect("second frame succeeds") + .into_data() + .expect("data frame"), + "event-b\n\n" + ); + let trailers = body + .frame() + .await + .expect("trailer frame") + .expect("trailer succeeds") + .into_trailers() + .expect("trailers"); + let values = trailers + .get_all("x-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(); + assert_eq!(values, ["one", "two"]); + assert!(body.frame().await.is_none()); + + drop(client); + server.abort(); +} + +#[tokio::test] +async fn relay_is_causally_non_aggregating_for_openai_and_anthropic_over_http1() { + assert_causal_relay(TestProtocol::Http1, "/v1/responses").await; + assert_causal_relay(TestProtocol::Http1, "/v1/messages").await; +} + +#[tokio::test] +async fn relay_is_causally_non_aggregating_for_openai_and_anthropic_over_http2() { + assert_causal_relay(TestProtocol::Http2, "/v1/responses").await; + assert_causal_relay(TestProtocol::Http2, "/v1/messages").await; +} + +async fn assert_exact_relay_fidelity(protocol: TestProtocol) { + let chunks = [ + Bytes::from_static(b": heartbeat\r\n\r\n"), + Bytes::from_static(b"event: delta\r\nid: 17\r\nretry: 500\r\n"), + Bytes::new(), + Bytes::from_static(b"data: first\r\ndata: second\r\n\r\n"), + Bytes::from_static(b"data: \xff\x00\xfe\r\n\r\n"), + Bytes::from_static(b"data: [DONE]\r\n\r\n"), + ]; + let expected = chunks + .iter() + .flat_map(|chunk| chunk.iter().copied()) + .collect::>(); + + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind fidelity provider"); + let provider = listener.local_addr().expect("bound provider address"); + let service = service_fn(move |_request: Request| { + let chunks = chunks.clone(); + async move { + let mut trailers = HeaderMap::new(); + trailers.append("x-checksum", HeaderValue::from_static("first")); + trailers.append("x-checksum", HeaderValue::from_static("second")); + trailers.append("x-binary-safe", HeaderValue::from_static("yes")); + let frames = chunks + .into_iter() + .map(Frame::data) + .chain(std::iter::once(Frame::trailers(trailers))); + let mut response = Response::builder() + .status(StatusCode::PARTIAL_CONTENT) + .header(TRAILER, "x-checksum, x-binary-safe") + .body(box_body(FramesBody::new(frames))) + .expect("valid fidelity response"); + response + .headers_mut() + .append(SET_COOKIE, HeaderValue::from_static("a=1")); + response + .headers_mut() + .append(SET_COOKIE, HeaderValue::from_static("b=2")); + Ok::<_, Infallible>(response) + } + }); + let provider_task = tokio::spawn(serve_one_connection(listener, protocol, service)); + let (relay, relay_task) = spawn_relay(protocol, provider).await; + + let client = client_for(protocol); + let response = client + .request(request_with_empty_body( + format!("http://{relay}/v1/responses"), + protocol, + )) + .await + .expect("relay responds"); + assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT); + assert_eq!( + response + .headers() + .get_all(SET_COOKIE) + .iter() + .map(|value| value.to_str().expect("ASCII cookie")) + .collect::>(), + ["a=1", "b=2"] + ); + + let mut actual = Vec::new(); + let mut actual_trailers = None; + let mut body = box_body(response.into_body()); + while let Some(frame) = body.frame().await { + let frame = frame.expect("fidelity frame succeeds"); + match frame.into_data() { + Ok(data) => actual.extend_from_slice(&data), + Err(frame) => { + actual_trailers = Some(frame.into_trailers().expect("only data or trailers")); + } + } + } + assert_eq!(actual, expected); + let trailers = actual_trailers.expect("trailers preserved"); + assert_eq!( + trailers + .get_all("x-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII checksum")) + .collect::>(), + ["first", "second"] + ); + assert_eq!(trailers["x-binary-safe"], "yes"); + + drop(client); + match protocol { + TestProtocol::Http1 => { + relay_task.await.expect("relay task succeeds"); + provider_task.await.expect("provider task succeeds"); + } + TestProtocol::Http2 => { + relay_task.abort(); + provider_task.abort(); + } + } +} + +#[tokio::test] +async fn relay_preserves_exact_bytes_duplicate_headers_and_trailers_over_http1() { + assert_exact_relay_fidelity(TestProtocol::Http1).await; +} + +#[tokio::test] +async fn relay_preserves_exact_bytes_duplicate_headers_and_trailers_over_http2() { + assert_exact_relay_fidelity(TestProtocol::Http2).await; +} + +#[tokio::test] +async fn slow_http2_reader_applies_bounded_backpressure_and_then_resumes() { + const FRAME_COUNT: usize = 128; + const FRAME_SIZE: usize = 64 * 1024; + + let polls = Arc::new(AtomicUsize::new(0)); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind backpressure provider"); + let provider = listener.local_addr().expect("bound provider address"); + let provider_polls = polls.clone(); + let service = service_fn(move |_request: Request| { + let polls = provider_polls.clone(); + async move { + Ok::<_, Infallible>(Response::new(box_body(CountedBody { + remaining: FRAME_COUNT, + frame: Bytes::from(vec![0x5a; FRAME_SIZE]), + polls, + }))) + } + }); + let provider_task = tokio::spawn(serve_one_connection(listener, TestProtocol::Http2, service)); + let (relay, relay_task) = spawn_relay(TestProtocol::Http2, provider).await; + + let client = pooled_h2c_client(); + let response = client + .request(request_with_empty_body( + format!("http://{relay}/v1/responses"), + TestProtocol::Http2, + )) + .await + .expect("relay returns response head"); + assert!( + polls.load(Ordering::SeqCst) < FRAME_COUNT, + "an unread downstream must stop the provider before the full 8 MiB body is polled" + ); + + let mut received = 0; + let mut body = box_body(response.into_body()); + while let Some(frame) = body.frame().await { + received += frame + .expect("backpressure frame succeeds") + .into_data() + .expect("provider emits only data") + .len(); + } + assert_eq!(received, FRAME_COUNT * FRAME_SIZE); + assert_eq!(polls.load(Ordering::SeqCst), FRAME_COUNT + 1); + + drop(client); + relay_task.abort(); + provider_task.abort(); +} + +#[tokio::test] +async fn dropping_http2_client_body_promptly_cancels_provider_body() { + let (dropped, wait_for_drop) = oneshot::channel(); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind cancellation provider"); + let provider = listener.local_addr().expect("bound provider address"); + let dropped = Arc::new(Mutex::new(Some(dropped))); + let service = service_fn(move |_request: Request| { + let dropped = dropped + .lock() + .expect("cancellation signal lock") + .take() + .expect("provider receives exactly one request"); + async move { + Ok::<_, Infallible>(Response::new(box_body(CancellationBody { + first_sent: false, + dropped: Some(dropped), + }))) + } + }); + let provider_task = tokio::spawn(serve_one_connection(listener, TestProtocol::Http2, service)); + let (relay, relay_task) = spawn_relay(TestProtocol::Http2, provider).await; + + let client = pooled_h2c_client(); + let response = client + .request(request_with_empty_body( + format!("http://{relay}/v1/responses"), + TestProtocol::Http2, + )) + .await + .expect("relay returns response head"); + let mut body = box_body(response.into_body()); + assert_eq!( + body.frame() + .await + .expect("first frame exists") + .expect("first frame succeeds") + .into_data() + .expect("first frame is data"), + "first\n\n" + ); + drop(body); + drop(client); + + tokio::time::timeout(Duration::from_secs(2), wait_for_drop) + .await + .expect("provider body cancellation must be prompt") + .expect("provider drop signal sent"); + relay_task.abort(); + provider_task.abort(); +} + +async fn spawn_multiplexed_provider( + connections: Arc, +) -> (std::net::SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind multiplexed provider"); + let address = listener.local_addr().expect("bound provider address"); + let task = tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.expect("accept provider client"); + connections.fetch_add(1, Ordering::SeqCst); + tokio::spawn(async move { + let service = service_fn(|request: Request| async move { + let sequence = request.uri().path().trim_start_matches('/').to_owned(); + let frames = (0..4).map(|part| Frame::data(multiplexed_chunk(&sequence, part))); + Ok::<_, Infallible>(Response::new(box_body(FramesBody::new(frames)))) + }); + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + }); + } + }); + (address, task) +} + +fn multiplexed_chunk(sequence: &str, part: usize) -> Bytes { + let prefix = format!("stream={sequence};part={part};"); + let mut chunk = vec![b'x'; 512]; + chunk[..prefix.len()].copy_from_slice(prefix.as_bytes()); + chunk[511] = b'\n'; + Bytes::from(chunk) +} + +async fn spawn_multiplexed_relay( + provider: std::net::SocketAddr, + connections: Arc, +) -> (std::net::SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind multiplexed relay"); + let address = listener.local_addr().expect("bound relay address"); + let upstream = pooled_h2c_client(); + let task = tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.expect("accept relay client"); + connections.fetch_add(1, Ordering::SeqCst); + let upstream = upstream.clone(); + tokio::spawn(async move { + let service = service_fn(move |request: Request| { + let upstream = upstream.clone(); + async move { + let path = request + .uri() + .path_and_query() + .map(|value| value.as_str()) + .unwrap_or("/"); + let destination = format!("http://{provider}{path}") + .parse() + .expect("valid provider URI"); + let request = + prepare_forward_request(request.map(box_body), destination, &[]) + .expect("valid multiplexed request head"); + let response = upstream.request(request).await.expect("provider responds"); + Ok::<_, Infallible>( + prepare_forward_response(response, &[]) + .expect("valid multiplexed response head") + .map(box_body), + ) + } + }); + let mut builder = http2::Builder::new(TokioExecutor::new()); + builder.max_concurrent_streams(256); + builder.max_pending_accept_reset_streams(256); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + }); + } + }); + (address, task) +} + +#[tokio::test] +async fn multiplexed_http2_keeps_128_concurrent_streams_isolated() { + const STREAMS: usize = 128; + + let provider_connections = Arc::new(AtomicUsize::new(0)); + let relay_connections = Arc::new(AtomicUsize::new(0)); + let (provider, provider_task) = spawn_multiplexed_provider(provider_connections.clone()).await; + let (relay, relay_task) = spawn_multiplexed_relay(provider, relay_connections.clone()).await; + let client = pooled_h2c_client(); + + // Warm both pools before introducing concurrency so all work multiplexes over established + // HTTP/2 connections rather than racing connection establishment. + let warm = client + .request(request_with_empty_body( + format!("http://{relay}/warm"), + TestProtocol::Http2, + )) + .await + .expect("warm request succeeds"); + let mut warm_body = box_body(warm.into_body()); + while warm_body.frame().await.is_some() {} + + let mut tasks = Vec::with_capacity(STREAMS); + for sequence in 0..STREAMS { + let client = client.clone(); + tasks.push(tokio::spawn(async move { + let response = client + .request(request_with_empty_body( + format!("http://{relay}/{sequence}"), + TestProtocol::Http2, + )) + .await + .expect("concurrent request succeeds"); + let mut body = box_body(response.into_body()); + let mut actual = Vec::new(); + while let Some(frame) = body.frame().await { + actual.extend_from_slice( + &frame + .expect("concurrent frame succeeds") + .into_data() + .expect("concurrent provider emits data"), + ); + } + let expected = (0..4) + .flat_map(|part| multiplexed_chunk(&sequence.to_string(), part)) + .collect::>(); + assert_eq!(actual, expected); + })); + } + for task in tasks { + task.await.expect("stream verification task succeeds"); + } + + assert_eq!(relay_connections.load(Ordering::SeqCst), 1); + assert_eq!(provider_connections.load(Ordering::SeqCst), 1); + + drop(client); + relay_task.abort(); + provider_task.abort(); +} + +async fn spawn_pooled_http1_provider() -> (std::net::SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind pooled HTTP/1.1 provider"); + let address = listener.local_addr().expect("bound provider address"); + let task = tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.expect("accept provider client"); + tokio::spawn(async move { + let service = service_fn(|request: Request| async move { + let sequence = request.uri().path().trim_start_matches('/').to_owned(); + let frames = (0..4).map(|part| Frame::data(multiplexed_chunk(&sequence, part))); + Ok::<_, Infallible>(Response::new(box_body(FramesBody::new(frames)))) + }); + let _ = http1::Builder::new() + .serve_connection(TokioIo::new(stream), service) + .await; + }); + } + }); + (address, task) +} + +async fn spawn_pooled_http1_relay( + provider: std::net::SocketAddr, +) -> (std::net::SocketAddr, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind pooled HTTP/1.1 relay"); + let address = listener.local_addr().expect("bound relay address"); + let upstream = pooled_http_client(); + let task = tokio::spawn(async move { + loop { + let (stream, _) = listener.accept().await.expect("accept relay client"); + let upstream = upstream.clone(); + tokio::spawn(async move { + let service = service_fn(move |request: Request| { + let upstream = upstream.clone(); + async move { + let path = request + .uri() + .path_and_query() + .map(|value| value.as_str()) + .unwrap_or("/"); + let destination = format!("http://{provider}{path}") + .parse() + .expect("valid provider URI"); + let request = + prepare_forward_request(request.map(box_body), destination, &[]) + .expect("valid pooled request head"); + let response = upstream.request(request).await.expect("provider responds"); + Ok::<_, Infallible>( + prepare_forward_response(response, &[]) + .expect("valid pooled response head") + .map(box_body), + ) + } + }); + let _ = http1::Builder::new() + .serve_connection(TokioIo::new(stream), service) + .await; + }); + } + }); + (address, task) +} + +#[tokio::test] +async fn pooled_http1_keeps_128_concurrent_streams_isolated() { + const STREAMS: usize = 128; + + let (provider, provider_task) = spawn_pooled_http1_provider().await; + let (relay, relay_task) = spawn_pooled_http1_relay(provider).await; + let client = pooled_http_client(); + let mut tasks = Vec::with_capacity(STREAMS); + for sequence in 0..STREAMS { + let client = client.clone(); + tasks.push(tokio::spawn(async move { + let response = client + .request(request_with_empty_body( + format!("http://{relay}/{sequence}"), + TestProtocol::Http1, + )) + .await + .expect("concurrent request succeeds"); + let mut body = box_body(response.into_body()); + let mut actual = Vec::new(); + while let Some(frame) = body.frame().await { + actual.extend_from_slice( + &frame + .expect("concurrent frame succeeds") + .into_data() + .expect("concurrent provider emits data"), + ); + } + let sequence = sequence.to_string(); + let expected = (0..4) + .flat_map(|part| multiplexed_chunk(&sequence, part)) + .collect::>(); + assert_eq!(actual, expected); + })); + } + for task in tasks { + task.await.expect("stream verification task succeeds"); + } + + drop(client); + relay_task.abort(); + provider_task.abort(); +} diff --git a/crates/cli/tests/coverage/daemon/worker_control_tests.rs b/crates/cli/tests/coverage/daemon/worker_control_tests.rs new file mode 100644 index 000000000..44dc0669a --- /dev/null +++ b/crates/cli/tests/coverage/daemon/worker_control_tests.rs @@ -0,0 +1,134 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +use std::sync::{Arc, Mutex}; + +use bytes::Bytes; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; + +#[test] +fn daemon_heartbeat_interval_is_bounded() { + assert!(validate_heartbeat_interval(99).is_err()); + assert_eq!( + validate_heartbeat_interval(5_000).expect("normal interval"), + Duration::from_secs(5) + ); + assert!(validate_heartbeat_interval(20_001).is_err()); +} + +#[test] +fn registration_exposes_scoped_digests_grant_interval_and_sequence_exhaustion() { + let mut registration = test_registration("data-token", "session-token"); + assert_eq!( + registration.data_token_digest(), + TokenDigest::from_token(b"data-token") + ); + assert_eq!( + registration.session_token_digest(), + TokenDigest::from_token(b"session-token") + ); + assert_eq!(registration.heartbeat_interval(), Duration::from_secs(5)); + assert_eq!(registration.generation_grant().worker_id, "worker-one"); + + registration.next_sequence = u64::MAX; + let error = registration + .advance_sequence() + .expect_err("sequence exhaustion must be fatal"); + assert!(error.to_string().contains("sequence was exhausted")); +} + +#[tokio::test] +async fn readiness_retries_the_exact_request_after_a_lost_response() { + let (origin, received) = lost_response_server().await; + let mut registration = test_registration("data", "session"); + + registration + .ready(&origin, "worker-one") + .await + .expect("readiness retry"); + + assert_exact_retry(received); + assert_eq!(registration.next_sequence, 2); + assert!(registration.pending_ready.is_none()); +} + +#[tokio::test] +async fn heartbeat_retries_the_exact_request_after_a_lost_response() { + let (origin, received) = lost_response_server().await; + let mut registration = test_registration("data", "session"); + + registration + .heartbeat(&origin, "worker-one") + .await + .expect("heartbeat retry"); + + assert_exact_retry(received); + assert_eq!(registration.next_sequence, 2); + assert!(registration.pending_heartbeat.is_none()); +} + +fn assert_exact_retry(received: Arc>>) { + let received = received.lock().expect("received request bodies"); + assert_eq!(received.len(), 2); + assert_eq!(received[0], received[1]); +} + +async fn lost_response_server() -> (String, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let address = listener.local_addr().expect("local address"); + let received = Arc::new(Mutex::new(Vec::new())); + let server_received = Arc::clone(&received); + tokio::spawn(async move { + for attempt in 0..2 { + let (mut stream, _) = listener.accept().await.expect("accept request"); + let body = read_http_body(&mut stream).await; + server_received + .lock() + .expect("received request bodies") + .push(body); + if attempt == 0 { + // The daemon applied the request but its response was lost. Closing the socket + // makes the client retry the same authenticated envelope on a new connection. + continue; + } + stream + .write_all(b"HTTP/1.1 204 No Content\r\nConnection: close\r\n\r\n") + .await + .expect("write response"); + } + }); + (format!("http://{address}"), received) +} + +async fn read_http_body(stream: &mut TcpStream) -> Bytes { + let mut request = Vec::new(); + let (body_offset, content_length) = loop { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.expect("read request"); + assert_ne!(count, 0, "request ended before its headers"); + request.extend_from_slice(&chunk[..count]); + if let Some(offset) = request.windows(4).position(|window| window == b"\r\n\r\n") { + let body_offset = offset + 4; + let headers = std::str::from_utf8(&request[..offset]).expect("HTTP headers"); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().expect("content length")) + }) + .expect("content length header"); + break (body_offset, content_length); + } + }; + while request.len() < body_offset + content_length { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.expect("read body"); + assert_ne!(count, 0, "request ended before its body"); + request.extend_from_slice(&chunk[..count]); + } + Bytes::copy_from_slice(&request[body_offset..body_offset + content_length]) +} diff --git a/crates/cli/tests/coverage/daemon/worker_managed_tests.rs b/crates/cli/tests/coverage/daemon/worker_managed_tests.rs new file mode 100644 index 000000000..21b138c06 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/worker_managed_tests.rs @@ -0,0 +1,1798 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use std::convert::Infallible; +use std::ffi::OsStr; +use std::io; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; + +use axum::Router; +use axum::extract::State; +use axum::routing::post; +use http_body_util::{BodyExt as _, Full, StreamBody}; +use nemo_relay::api::llm::LlmRequestInterceptOutcome; +use nemo_relay::api::registry::{RuntimeRegistrationOwner, RuntimeRegistrationOwnerKind}; +use nemo_relay::api::registry::{ + deregister_llm_execution_intercept, deregister_llm_request_intercept, + register_llm_execution_intercept, register_llm_request_intercept, +}; + +use crate::daemon::common::transport::pooled_client; +use crate::test_support::{EnvScope, PLUGIN_CONFIG_TEST_LOCK}; + +type CapturedProviderRequest = Arc>>; +type ProviderRequests = Arc>>; + +#[tokio::test] +async fn observation_preserves_delivery_while_capturing_json() { + let expected = Bytes::from_static(br#"{"ok":true}"#); + let (body, observation) = + observe_body(Full::new(expected.clone()), StatusCode::OK, expected.len()); + let delivered = body.collect().await.expect("delivered body").to_bytes(); + let observed = observation + .finish(ProviderSurface::OpenAIResponses, false) + .await; + assert_eq!(delivered, expected); + assert_eq!(observed.value, Some(json!({ "ok": true }))); + assert!(!observed.truncated); + assert_eq!(observed.terminal, OBSERVATION_COMPLETE); +} + +#[tokio::test] +async fn capture_limit_truncates_observation_without_truncating_delivery() { + let expected = Bytes::from_static(br#"{"too":"large"}"#); + let (body, observation) = observe_body(Full::new(expected.clone()), StatusCode::OK, 3); + let delivered = body.collect().await.expect("delivered body").to_bytes(); + let observed = observation + .finish(ProviderSurface::OpenAIResponses, false) + .await; + assert_eq!(delivered, expected); + assert!(observed.value.is_none()); + assert!(observed.truncated); + assert_eq!(observed.terminal, OBSERVATION_COMPLETE); +} + +#[tokio::test] +async fn saturated_observation_queue_never_blocks_or_truncates_delivery() { + let expected = (0..OBSERVATION_QUEUE_FRAMES + 8) + .map(|index| Bytes::from(vec![u8::try_from(index).expect("test byte")])) + .collect::>(); + let frames = expected + .clone() + .into_iter() + .map(|bytes| Ok::<_, Infallible>(Frame::data(bytes))); + let (mut body, observation) = observe_body( + StreamBody::new(futures_util::stream::iter(frames)), + StatusCode::OK, + usize::MAX, + ); + + let mut delivered = Vec::new(); + while let Some(frame) = body.frame().await { + let frame = frame.expect("delivery frame"); + if let Some(data) = frame.data_ref() { + delivered.extend_from_slice(data); + } + } + let observed = observation + .finish(ProviderSurface::OpenAIResponses, false) + .await; + + let expected_delivery = expected + .iter() + .flat_map(|bytes| bytes.iter().copied()) + .collect::>(); + assert_eq!(delivered, expected_delivery); + assert!(observed.truncated); + assert_eq!(observed.terminal, OBSERVATION_COMPLETE); +} + +#[tokio::test] +async fn dropping_delivery_marks_observation_cancelled_and_terminates_it() { + let frames = [ + Ok::<_, Infallible>(Frame::data(Bytes::from_static(b"first"))), + Ok(Frame::data(Bytes::from_static(b"second"))), + ]; + let (mut body, observation) = observe_body( + StreamBody::new(futures_util::stream::iter(frames)), + StatusCode::OK, + usize::MAX, + ); + let first = body + .frame() + .await + .expect("first frame") + .expect("first delivery frame") + .into_data() + .expect("first data"); + assert_eq!(first, "first"); + drop(body); + + let observed = tokio::time::timeout( + Duration::from_secs(1), + observation.finish(ProviderSurface::OpenAIResponses, false), + ) + .await + .expect("observation task must terminate after cancellation"); + assert_eq!(observed.terminal, OBSERVATION_CANCELLED); + assert!( + observed + .failure + .as_deref() + .is_some_and(|failure| failure.contains("cancelled")) + ); +} + +#[tokio::test(start_paused = true)] +async fn observation_body_deadline_terminates_a_stalled_observer() { + let started = tokio::time::Instant::now(); + let (_sender, receiver) = tokio::sync::mpsc::channel(1); + let observation = ObservationReceiver { + receiver, + signal: Arc::new(ObservationSignal::new()), + status: StatusCode::OK, + }; + let observed = observation + .finish(ProviderSurface::OpenAIResponses, false) + .await; + assert!(observed.truncated); + assert_eq!(observed.terminal, OBSERVATION_ACTIVE); + assert_eq!( + observed.failure.as_deref(), + Some("provider response observation timed out") + ); + assert_eq!(started.elapsed(), OBSERVATION_COMPLETION_TIMEOUT); +} + +#[tokio::test(start_paused = true)] +async fn successful_stream_observation_can_outlive_the_response_head_deadline() { + let (sender, receiver) = tokio::sync::mpsc::channel(4); + let signal = Arc::new(ObservationSignal::new()); + let observation = ObservationReceiver { + receiver, + signal: Arc::clone(&signal), + status: StatusCode::OK, + }; + let task = tokio::spawn(observation.finish(ProviderSurface::OpenAIChat, true)); + tokio::task::yield_now().await; + sender.send(Bytes::from_static(b"data: {\"id\":\"long\",\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"hello\"}}]}\n\n")).await.unwrap(); + tokio::time::advance(RESPONSE_HEAD_TIMEOUT + Duration::from_secs(1)).await; + tokio::task::yield_now().await; + assert!( + !task.is_finished(), + "response-head deadline must not truncate observation" + ); + sender.send(Bytes::from_static(b"data: {\"id\":\"long\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n")).await.unwrap(); + drop(sender); + signal.finish(OBSERVATION_COMPLETE); + let observed = task.await.unwrap(); + assert_eq!(observed.terminal, OBSERVATION_COMPLETE); + assert!(!observed.truncated); + assert!(observed.failure.is_none()); + assert!(observed.value.is_some()); +} + +#[test] +fn internal_worker_headers_are_not_forwarded_to_providers() { + let mut headers = HeaderMap::new(); + headers.insert( + WORKER_TOKEN_HEADER, + HeaderValue::from_static("worker-secret"), + ); + headers.insert( + CLIENT_TOKEN_HEADER, + HeaderValue::from_static("client-secret"), + ); + headers.insert( + "x-nemo-relay-session-id", + HeaderValue::from_static("session"), + ); + headers.insert("x-provider-header", HeaderValue::from_static("kept")); + strip_worker_headers(&mut headers); + strip_internal_headers(&mut headers); + assert!(!headers.contains_key(WORKER_TOKEN_HEADER)); + assert!(!headers.contains_key(CLIENT_TOKEN_HEADER)); + assert!(!headers.contains_key("x-nemo-relay-session-id")); + assert_eq!(headers["x-provider-header"], "kept"); +} + +#[tokio::test] +async fn daemon_hop_cannot_supply_worker_local_dispatch_overrides() { + let request = Request::post("/v1/responses") + .header( + INTERNAL_DISPATCH_URL_HEADER, + "https://attacker.invalid/v1/responses", + ) + .header(INTERNAL_DISPATCH_ROUTE_HEADER, "anthropic_messages") + .header(INTERNAL_DISPATCH_BACKEND_HEADER, "attacker") + .header(INTERNAL_RETRY_AWARE_HEADER, "true") + .header("x-nemo-relay-session-id", "session-kept-for-correlation") + .body(Body::from(r#"{"model":"test","stream":true}"#)) + .expect("provider request"); + let prepared = PreparedProviderRequest::read(request, &GatewayConfig::default()) + .await + .expect("prepared request"); + + assert!(!prepared.headers.contains_key(INTERNAL_DISPATCH_URL_HEADER)); + assert!( + !prepared + .headers + .contains_key(INTERNAL_DISPATCH_ROUTE_HEADER) + ); + assert!( + !prepared + .headers + .contains_key(INTERNAL_DISPATCH_BACKEND_HEADER) + ); + assert!(!prepared.headers.contains_key(INTERNAL_RETRY_AWARE_HEADER)); + assert_eq!( + prepared.headers["x-nemo-relay-session-id"], + "session-kept-for-correlation" + ); +} + +#[test] +fn execution_middleware_is_explicitly_incompatible_with_raw_delivery() { + let owner = RuntimeRegistrationOwner { + kind: RuntimeRegistrationOwnerKind::GlobalApi, + plugin_kind: None, + component_ordinal: None, + }; + let registrations = [ + RuntimeRegistrationIdentity { + kind: RuntimeRegistrationKind::LlmExecutionIntercept, + local_name: "buffered".into(), + effective_name: "plugin.buffered".into(), + owner: owner.clone(), + }, + RuntimeRegistrationIdentity { + kind: RuntimeRegistrationKind::LlmStreamExecutionIntercept, + local_name: "streaming".into(), + effective_name: "plugin.streaming".into(), + owner, + }, + RuntimeRegistrationIdentity { + kind: RuntimeRegistrationKind::LlmRequestIntercept, + local_name: "request".into(), + effective_name: "plugin.request".into(), + owner: RuntimeRegistrationOwner { + kind: RuntimeRegistrationOwnerKind::GlobalApi, + plugin_kind: None, + component_ordinal: None, + }, + }, + ]; + + assert_eq!( + incompatible_registration_names(®istrations), + ["plugin.buffered", "plugin.streaming"] + ); +} + +#[test] +fn only_request_middleware_requires_request_body_decoding() { + let owner = RuntimeRegistrationOwner { + kind: RuntimeRegistrationOwnerKind::GlobalApi, + plugin_kind: None, + component_ordinal: None, + }; + for kind in [ + RuntimeRegistrationKind::LlmSanitizeRequestGuardrail, + RuntimeRegistrationKind::LlmConditionalExecutionGuardrail, + RuntimeRegistrationKind::LlmRequestIntercept, + ] { + assert!(registration_reads_request_body( + &RuntimeRegistrationIdentity { + kind, + local_name: "request-reader".into(), + effective_name: "request-reader".into(), + owner: owner.clone(), + } + )); + } + for kind in [ + RuntimeRegistrationKind::Subscriber, + RuntimeRegistrationKind::LlmSanitizeResponseGuardrail, + ] { + assert!(!registration_reads_request_body( + &RuntimeRegistrationIdentity { + kind, + local_name: "response-only".into(), + effective_name: "response-only".into(), + owner: owner.clone(), + } + )); + } +} + +struct PendingRequestBody { + polls: Arc, +} + +impl HttpBody for PendingRequestBody { + type Data = Bytes; + type Error = Infallible; + + fn poll_frame( + self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + self.polls.fetch_add(1, AtomicOrdering::Relaxed); + Poll::Pending + } + + fn is_end_stream(&self) -> bool { + false + } + + fn size_hint(&self) -> SizeHint { + SizeHint::default() + } +} + +#[tokio::test] +async fn unbuffered_dispatch_returns_response_head_without_collecting_request_body() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind provider"); + let address = listener.local_addr().expect("provider address"); + let app = Router::new().route( + "/v1/responses", + post(|| async { + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream; charset=utf-8") + .body(Body::from("data: [DONE]\n\n")) + .expect("provider response") + }), + ); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve provider"); + }); + let polls = Arc::new(AtomicUsize::new(0)); + let request = Request::post("/v1/responses") + .header(ACCEPT, "text/event-stream") + .body(Body::new(PendingRequestBody { + polls: Arc::clone(&polls), + })) + .expect("streaming request"); + let config = GatewayConfig { + openai_base_url: format!("http://{address}"), + ..GatewayConfig::default() + }; + + let (response, observation, streaming) = tokio::time::timeout( + Duration::from_secs(1), + dispatch_unbuffered_observed( + crate::daemon::common::transport::pooled_client().expect("provider client"), + request, + ProviderRoute::OpenAi, + &config, + DEFAULT_OBSERVATION_CAPTURE_BYTES, + ), + ) + .await + .expect("response head must not wait for request completion") + .expect("provider response"); + + assert_eq!(response.status(), StatusCode::OK); + assert!(streaming); + assert!(polls.load(AtomicOrdering::Relaxed) > 0); + let delivered = response + .into_body() + .collect() + .await + .expect("delivered response") + .to_bytes(); + assert_eq!(delivered, "data: [DONE]\n\n"); + let observed = observation + .finish(ProviderSurface::OpenAIResponses, true) + .await; + assert_eq!(observed.terminal, OBSERVATION_COMPLETE); + server.abort(); +} + +#[test] +fn changing_stream_mode_is_a_route_wide_transport_incompatibility() { + let incompatible = CliError::Flow(FlowError::InvalidArgument( + STREAM_MODE_MUTATION_ERROR.into(), + )); + assert!(requires_route_pass_through(&incompatible)); + assert!(!requires_route_pass_through(&CliError::Flow( + FlowError::InvalidArgument("some other request error".into()) + ))); +} + +#[test] +fn chatgpt_shaped_responses_requests_use_the_managed_responses_pipeline() { + assert_eq!( + provider_surface("/backend-api/codex/responses"), + Some(ProviderSurface::OpenAIResponses) + ); +} + +#[test] +fn managed_worker_canonicalizes_chatgpt_responses_before_alignment() { + let _environment = EnvScope::set(&[("OPENAI_API_KEY", None)]); + let mut headers = HeaderMap::new(); + headers.insert( + AUTHORIZATION, + HeaderValue::from_static("Bearer at-managed-chatgpt-token"), + ); + let prepared = PreparedProviderRequest { + method: Method::POST, + version: http::Version::HTTP_11, + headers, + path: "/backend-api/codex/responses".into(), + path_and_query: "/backend-api/codex/responses?client=codex".into(), + body: Bytes::from_static(br#"{"model":"test","stream":true}"#), + request_json: json!({"model": "test", "stream": true}), + streaming: true, + }; + let destination = effective_destination( + &prepared, + ProviderRoute::OpenAi, + None, + &GatewayConfig::default(), + ) + .expect("ChatGPT destination"); + + assert_eq!( + destination, + "https://chatgpt.com/backend-api/codex/responses?client=codex" + .parse::() + .unwrap() + ); +} + +#[test] +fn managed_worker_does_not_infer_upstream_authority_from_generic_bearer_tokens() { + let _environment = EnvScope::set(&[("OPENAI_API_KEY", None)]); + let mut headers = HeaderMap::new(); + headers.insert( + AUTHORIZATION, + HeaderValue::from_static("Bearer at-caller-controlled-token"), + ); + let prepared = PreparedProviderRequest { + method: Method::POST, + version: http::Version::HTTP_11, + headers, + path: "/responses".into(), + path_and_query: "/responses?client=pi".into(), + body: Bytes::from_static(br#"{"model":"test","stream":true}"#), + request_json: json!({"model": "test", "stream": true}), + streaming: true, + }; + let config = GatewayConfig { + openai_base_url: "https://administrator.example/v1".into(), + ..GatewayConfig::default() + }; + + let destination = effective_destination(&prepared, ProviderRoute::OpenAi, None, &config) + .expect("administrator-selected destination"); + + assert_eq!( + destination, + "https://administrator.example/v1/responses?client=pi" + .parse::() + .unwrap() + ); +} + +#[test] +fn managed_worker_uses_pi_named_provider_endpoint_and_strips_the_routing_header() { + let mut headers = HeaderMap::new(); + headers.insert( + crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER, + HeaderValue::from_static("https://custom.example/inference/v1"), + ); + headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer caller-provider-token"), + ); + let prepared = PreparedProviderRequest { + method: Method::POST, + version: http::Version::HTTP_11, + headers, + path: "/chat/completions".into(), + path_and_query: "/chat/completions?client=pi".into(), + body: Bytes::from_static(br#"{"model":"custom","stream":true}"#), + request_json: json!({"model": "custom", "stream": true}), + streaming: true, + }; + + let destination = effective_destination( + &prepared, + ProviderRoute::OpenAi, + None, + &GatewayConfig::default(), + ) + .expect("Pi-selected destination"); + assert_eq!( + destination, + "https://custom.example/inference/v1/chat/completions?client=pi" + .parse::() + .unwrap() + ); + + let (forwarded, _, _) = effective_request(&prepared, None).expect("forwarded request"); + assert!(!forwarded.contains_key(crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER)); +} + +#[test] +fn unchanged_middleware_headers_preserve_credentials_and_duplicate_values() { + let mut headers = HeaderMap::new(); + headers.append("x-provider-feature", HeaderValue::from_static("first")); + headers.append("x-provider-feature", HeaderValue::from_static("second")); + headers.insert( + http::header::COOKIE, + HeaderValue::from_static("session=secret"), + ); + headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer provider")); + headers.insert("x-api-key", HeaderValue::from_static("provider-key")); + let prepared = prepared_request(headers); + let effective = LlmRequest { + headers: crate::gateway::daemon_observable_headers(&prepared.headers), + content: prepared.request_json.clone(), + }; + + let (forwarded, body, explicit_target) = + effective_request(&prepared, Some(&effective)).expect("effective request"); + + let values = forwarded + .get_all("x-provider-feature") + .iter() + .map(|value| value.to_str().unwrap()) + .collect::>(); + assert_eq!(values, ["first", "second"]); + assert_eq!(forwarded[http::header::COOKIE], "session=secret"); + assert_eq!(forwarded[AUTHORIZATION], "Bearer provider"); + assert_eq!(forwarded["x-api-key"], "provider-key"); + assert_eq!(body, prepared.body); + assert!(!explicit_target); +} + +#[test] +fn middleware_header_diff_changes_only_the_named_observable_header() { + let mut headers = HeaderMap::new(); + headers.append("x-unchanged", HeaderValue::from_static("first")); + headers.append("x-unchanged", HeaderValue::from_static("second")); + headers.insert("x-changed", HeaderValue::from_static("before")); + headers.insert( + http::header::COOKIE, + HeaderValue::from_static("session=secret"), + ); + headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer provider")); + let prepared = prepared_request(headers); + let mut effective_headers = crate::gateway::daemon_observable_headers(&prepared.headers); + effective_headers.insert("x-changed".into(), json!("after")); + let effective = LlmRequest { + headers: effective_headers, + content: prepared.request_json.clone(), + }; + + let (forwarded, _, explicit_target) = + effective_request(&prepared, Some(&effective)).expect("effective request"); + + let unchanged = forwarded + .get_all("x-unchanged") + .iter() + .map(|value| value.to_str().unwrap()) + .collect::>(); + assert_eq!(unchanged, ["first", "second"]); + assert_eq!(forwarded["x-changed"], "after"); + assert_eq!(forwarded[http::header::COOKIE], "session=secret"); + assert_eq!(forwarded[AUTHORIZATION], "Bearer provider"); + assert!(!explicit_target); +} + +#[test] +fn explicit_target_removes_hidden_provider_credentials() { + let mut headers = HeaderMap::new(); + headers.insert( + http::header::COOKIE, + HeaderValue::from_static("session=secret"), + ); + headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer provider")); + headers.insert("x-api-key", HeaderValue::from_static("provider-key")); + let prepared = prepared_request(headers); + let mut effective_headers = crate::gateway::daemon_observable_headers(&prepared.headers); + effective_headers.insert( + INTERNAL_DISPATCH_URL_HEADER.into(), + json!("https://selected.example/v1/responses"), + ); + effective_headers.insert("authorization".into(), json!("Bearer selected-provider")); + let effective = LlmRequest { + headers: effective_headers, + content: prepared.request_json.clone(), + }; + + let (forwarded, _, explicit_target) = + effective_request(&prepared, Some(&effective)).expect("effective request"); + + assert!(explicit_target); + assert!(!forwarded.contains_key(http::header::COOKIE)); + assert_eq!(forwarded[AUTHORIZATION], "Bearer selected-provider"); + assert!(!forwarded.contains_key("x-api-key")); +} + +fn prepared_request(headers: HeaderMap) -> PreparedProviderRequest { + PreparedProviderRequest { + method: Method::POST, + version: http::Version::HTTP_11, + headers, + path: "/v1/responses".into(), + path_and_query: "/v1/responses".into(), + body: Bytes::from_static(br#"{"model":"test","stream":true}"#), + request_json: json!({"model": "test", "stream": true}), + streaming: true, + } +} + +#[test] +fn observation_capture_limit_defaults_and_accepts_a_positive_override() { + { + let _environment = EnvScope::set(&[(OBSERVATION_CAPTURE_BYTES_ENV, None)]); + assert_eq!( + observation_capture_limit_from_environment().unwrap(), + DEFAULT_OBSERVATION_CAPTURE_BYTES + ); + } + { + let _environment = + EnvScope::set(&[(OBSERVATION_CAPTURE_BYTES_ENV, Some(OsStr::new("65536")))]); + assert_eq!( + observation_capture_limit_from_environment().unwrap(), + 65_536 + ); + } +} + +#[test] +fn observation_capture_limit_rejects_zero_and_invalid_values() { + for value in ["0", "not-a-number"] { + let _environment = + EnvScope::set(&[(OBSERVATION_CAPTURE_BYTES_ENV, Some(OsStr::new(value)))]); + let error = observation_capture_limit_from_environment() + .unwrap_err() + .to_string(); + assert!(error.contains("positive integer"), "{error}"); + } +} + +#[tokio::test] +async fn prepared_requests_and_both_dispatch_paths_preserve_provider_contracts() { + let received: ProviderRequests = Arc::new(std::sync::Mutex::new(Vec::new())); + let app = Router::new() + .route( + "/v1/responses", + post( + |State(received): State, request: Request| async move { + let (parts, body) = request.into_parts(); + let body = axum::body::to_bytes(body, usize::MAX) + .await + .expect("provider request body"); + received + .lock() + .expect("capture provider request") + .push((parts.headers, body)); + Response::builder() + .status(StatusCode::CREATED) + .header(CONTENT_TYPE, "text/event-stream; charset=utf-8") + .header("x-provider-response", "preserved") + .body(Body::from("event: message\ndata: raw\n\n")) + .expect("provider response") + }, + ), + ) + .with_state(Arc::clone(&received)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind provider"); + let address = listener.local_addr().expect("provider address"); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve provider"); + }); + let config = GatewayConfig { + openai_base_url: format!("http://{address}"), + openai_auth_header: Some("Bearer configured-secret".into()), + ..GatewayConfig::default() + }; + + let source = Request::post("/v1/responses?trace=one") + .header(WORKER_TOKEN_HEADER, "worker-secret") + .header(CLIENT_TOKEN_HEADER, "client-secret") + .header(INTERNAL_DISPATCH_URL_HEADER, "https://attacker.invalid") + .header("x-provider-request", "preserved") + .body(Body::from(br#"{"model":"test","stream":true}"#.as_slice())) + .expect("provider request"); + let prepared = PreparedProviderRequest::read(source, &config) + .await + .expect("prepared provider request"); + assert_eq!(prepared.path, "/v1/responses"); + assert_eq!(prepared.path_and_query, "/v1/responses?trace=one"); + assert!(prepared.streaming); + assert!(!prepared.headers.contains_key(WORKER_TOKEN_HEADER)); + assert!(!prepared.headers.contains_key(INTERNAL_DISPATCH_URL_HEADER)); + + let response = dispatch_observed( + pooled_client().expect("provider client"), + prepared, + ProviderRoute::OpenAi, + None, + &config, + DEFAULT_OBSERVATION_CAPTURE_BYTES, + ) + .await + .expect("observed provider dispatch"); + assert_eq!(response.0.status(), StatusCode::CREATED); + assert_eq!(response.0.headers()["x-provider-response"], "preserved"); + let delivered = response + .0 + .into_body() + .collect() + .await + .expect("observed delivery") + .to_bytes(); + assert_eq!(delivered, "event: message\ndata: raw\n\n"); + let observed = response + .1 + .finish(ProviderSurface::OpenAIResponses, true) + .await; + assert_eq!(observed.terminal, OBSERVATION_COMPLETE); + + let raw = Request::post("/v1/responses") + .header(WORKER_TOKEN_HEADER, "worker-secret") + .header(CLIENT_TOKEN_HEADER, "client-secret") + .body(Body::from("raw provider payload")) + .expect("unmanaged request"); + let response = dispatch_unmanaged( + pooled_client().expect("provider client"), + raw, + ProviderRoute::OpenAi, + &config, + ) + .await + .expect("unmanaged provider dispatch"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!( + response + .into_body() + .collect() + .await + .expect("unmanaged delivery") + .to_bytes(), + "event: message\ndata: raw\n\n" + ); + + let captured = received.lock().expect("read provider requests"); + assert_eq!(captured.len(), 2); + for (headers, _) in captured.iter() { + assert_eq!(headers[AUTHORIZATION], "Bearer configured-secret"); + assert!(!headers.contains_key(WORKER_TOKEN_HEADER)); + assert!(!headers.contains_key(CLIENT_TOKEN_HEADER)); + } + assert_eq!( + captured[0].1, + br#"{"model":"test","stream":true}"#.as_slice() + ); + assert_eq!(captured[1].1, "raw provider payload"); + server.abort(); +} + +#[tokio::test] +async fn prepared_request_enforces_body_limit_and_normalizes_non_json_payloads() { + let limited = GatewayConfig { + max_passthrough_body_bytes: 2, + ..GatewayConfig::default() + }; + let too_large = PreparedProviderRequest::read( + Request::post("/v1/responses") + .body(Body::from("abc")) + .expect("request"), + &limited, + ) + .await; + let too_large = match too_large { + Ok(_) => panic!("request must honor configured body limit"), + Err(error) => error, + }; + assert!(matches!(too_large, CliError::PayloadTooLarge(_))); + + let prepared = PreparedProviderRequest::read( + Request::post("/v1/responses") + .body(Body::from("not-json")) + .expect("request"), + &GatewayConfig::default(), + ) + .await + .expect("non-json requests remain forwardable"); + assert_eq!(prepared.request_json, Value::Null); + assert!(!prepared.streaming); +} + +#[test] +fn managed_request_helpers_cover_overrides_streaming_and_provider_credentials() { + let _environment = EnvScope::set(&[ + ("OPENAI_API_KEY", Some(OsStr::new(" openai-env-key "))), + ("ANTHROPIC_API_KEY", Some(OsStr::new("anthropic-env-key"))), + ]); + let prepared = prepared_request(HeaderMap::from_iter([( + CONTENT_LENGTH, + HeaderValue::from_static("123"), + )])); + let (headers, body, explicit) = effective_request(&prepared, None).expect("raw request"); + assert!(!explicit); + assert!(!headers.contains_key(CONTENT_LENGTH)); + assert_eq!(body, prepared.body); + + let mut effective_headers = serde_json::Map::new(); + effective_headers.insert( + INTERNAL_DISPATCH_URL_HEADER.into(), + json!("http://127.0.0.1:9999/v1/responses"), + ); + let effective = LlmRequest { + headers: effective_headers, + content: json!({"model":"replacement","stream":false}), + }; + assert!(has_explicit_target(&effective)); + assert!(!stream_mode(&effective)); + assert!(!prepared_streaming(&effective)); + let destination = effective_destination( + &prepared, + ProviderRoute::OpenAi, + Some(&effective), + &GatewayConfig::default(), + ) + .expect("explicit destination"); + assert_eq!( + destination, + "http://127.0.0.1:9999/v1/responses".parse::().unwrap() + ); + let (_, body, explicit) = + effective_request(&prepared, Some(&effective)).expect("rewritten request"); + assert!(explicit); + assert_eq!( + body, + Bytes::from_static(br#"{"model":"replacement","stream":false}"#) + ); + + let mut headers = HeaderMap::new(); + inject_provider_auth( + &mut headers, + ProviderRoute::OpenAi, + &GatewayConfig::default(), + ); + assert_eq!(headers[AUTHORIZATION], "Bearer openai-env-key"); + headers.clear(); + inject_provider_auth( + &mut headers, + ProviderRoute::Anthropic, + &GatewayConfig::default(), + ); + assert_eq!(headers["x-api-key"], "anthropic-env-key"); + headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer caller")); + inject_provider_auth( + &mut headers, + ProviderRoute::OpenAi, + &GatewayConfig::default(), + ); + assert_eq!(headers[AUTHORIZATION], "Bearer caller"); + + let accepts_sse = HeaderMap::from_iter([( + ACCEPT, + HeaderValue::from_static("application/json, text/event-stream; charset=utf-8"), + )]); + assert!(request_streaming_hint(&accepts_sse)); + let content_type = HeaderMap::from_iter([( + CONTENT_TYPE, + HeaderValue::from_static("text/event-stream; charset=utf-8"), + )]); + assert!(response_streaming(&content_type)); + assert_eq!( + provider_surface("/v1/messages"), + Some(ProviderSurface::AnthropicMessages) + ); + assert_eq!(provider_surface("/unsupported"), None); +} + +#[tokio::test] +async fn observation_reports_provider_body_errors_and_metadata() { + let frames = futures_util::stream::iter(vec![Err::, io::Error>( + io::Error::other("provider body failed"), + )]); + let (mut body, observation) = observe_body( + StreamBody::new(frames), + StatusCode::BAD_GATEWAY, + DEFAULT_OBSERVATION_CAPTURE_BYTES, + ); + assert!(body.frame().await.expect("body frame").is_err()); + let observed = observation + .finish(ProviderSurface::OpenAIResponses, false) + .await; + assert_eq!(observed.terminal, OBSERVATION_BODY_ERROR); + assert_eq!(observed.status, StatusCode::BAD_GATEWAY); + assert!(observed.failure.is_some()); + assert_eq!( + observed.metadata()["daemon_worker_observation"]["terminal"], + "body_error" + ); +} + +#[tokio::test] +async fn managed_runtime_forwards_all_managed_hook_shapes_and_closes_cleanly() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let runtime = + ManagedRuntime::initialize(GatewayConfig::default(), Vec::new(), "machine-owner".into()) + .await + .expect("initialize managed runtime"); + + for (route, payload, expected) in [ + (HookRoute::Codex, json!({}), json!({})), + (HookRoute::Claude, json!({}), json!({"continue": true})), + (HookRoute::Pi, json!({}), json!({})), + ] { + let response = runtime + .handle_hook( + route, + Request::post("/hook") + .header(WORKER_TOKEN_HEADER, "internal-only") + .body(Body::from(payload.to_string())) + .expect("hook request"), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = response + .into_body() + .collect() + .await + .expect("hook response") + .to_bytes(); + assert_eq!( + serde_json::from_slice::(&body).expect("hook JSON response"), + expected + ); + } + + runtime + .ensure_streaming_transport_compatible() + .expect("no incompatible middleware"); + runtime.close().await.expect("close managed runtime"); +} + +#[tokio::test] +async fn managed_runtime_uses_unmanaged_provider_fallback_without_buffering_response() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind provider"); + let address = listener.local_addr().expect("provider address"); + let app = Router::new().route( + "/v1/custom", + post(|| async { + Response::builder() + .status(StatusCode::ACCEPTED) + .header(CONTENT_TYPE, "application/octet-stream") + .body(Body::from(Bytes::from_static(b"opaque provider bytes"))) + .expect("provider response") + }), + ); + let provider = tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve provider"); + }); + let config = GatewayConfig { + openai_base_url: format!("http://{address}"), + ..GatewayConfig::default() + }; + let runtime = ManagedRuntime::initialize(config, Vec::new(), "machine-owner".into()) + .await + .expect("initialize managed runtime"); + let response = runtime + .proxy_provider( + pooled_client().expect("provider client"), + Request::post("/custom") + .body(Body::from("unbuffered request bytes")) + .expect("provider request"), + ProviderRoute::OpenAi, + ) + .await + .expect("unmanaged dispatch"); + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!( + response + .into_body() + .collect() + .await + .expect("provider body") + .to_bytes(), + "opaque provider bytes" + ); + runtime.close().await.expect("close managed runtime"); + provider.abort(); +} + +#[tokio::test] +async fn cancelled_provider_dispatch_releases_sessions_in_both_worker_paths() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const INTERCEPT: &str = "daemon-cancel-before-head"; + let _registration = RequestInterceptRegistration(INTERCEPT); + for buffered in [false, true] { + if buffered { + register_llm_request_intercept( + INTERCEPT, + 1, + false, + Arc::new(|_, request, annotated| { + Box::pin(async move { Ok(LlmRequestInterceptOutcome::new(request, annotated)) }) + }), + ) + .unwrap(); + } + let entered = Arc::new(Notify::new()); + let release = Arc::new(Notify::new()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let app = Router::new().route( + "/v1/responses", + post({ + let entered = Arc::clone(&entered); + let release = Arc::clone(&release); + move || { + let entered = Arc::clone(&entered); + let release = Arc::clone(&release); + async move { + entered.notify_one(); + release.notified().await; + StatusCode::OK + } + } + }), + ); + let provider = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + let runtime = Arc::new( + ManagedRuntime::initialize( + GatewayConfig { + openai_base_url: format!("http://{address}"), + ..GatewayConfig::default() + }, + Vec::new(), + "cancel-owner".into(), + ) + .await + .unwrap(), + ); + let dispatch_runtime = Arc::clone(&runtime); + let dispatch = tokio::spawn(async move { + dispatch_runtime + .proxy_provider( + pooled_client().unwrap(), + Request::post("/v1/responses") + .body(Body::from( + r#"{"model":"test","input":"test","stream":true}"#, + )) + .unwrap(), + ProviderRoute::OpenAi, + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), entered.notified()) + .await + .unwrap(); + assert!(runtime.sessions.has_open_sessions().await); + dispatch.abort(); + assert!(dispatch.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(5), async { + while runtime.sessions.has_open_sessions().await { + runtime + .sessions + .close_idle_sessions_at( + std::time::Instant::now(), + Duration::ZERO, + "cancelled_test_call", + ) + .await + .unwrap(); + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled request retained its gateway session"); + release.notify_one(); + runtime.close().await.unwrap(); + provider.abort(); + } +} + +#[tokio::test] +async fn cancelled_request_middleware_releases_its_gateway_session() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const INTERCEPT: &str = "daemon-cancel-in-middleware"; + let _registration = RequestInterceptRegistration(INTERCEPT); + let entered = Arc::new(Notify::new()); + register_llm_request_intercept( + INTERCEPT, + 1, + false, + Arc::new({ + let entered = Arc::clone(&entered); + move |_, _, _| { + let entered = Arc::clone(&entered); + Box::pin(async move { + entered.notify_one(); + std::future::pending().await + }) + } + }), + ) + .unwrap(); + let runtime = Arc::new( + ManagedRuntime::initialize( + GatewayConfig::default(), + Vec::new(), + "cancel-middleware-owner".into(), + ) + .await + .unwrap(), + ); + let dispatch_runtime = Arc::clone(&runtime); + let dispatch = tokio::spawn(async move { + dispatch_runtime + .proxy_provider( + pooled_client().unwrap(), + Request::post("/v1/responses") + .body(Body::from( + r#"{"model":"test","input":"test","stream":true}"#, + )) + .unwrap(), + ProviderRoute::OpenAi, + ) + .await + }); + tokio::time::timeout(Duration::from_secs(5), entered.notified()) + .await + .unwrap(); + assert!(runtime.sessions.has_open_sessions().await); + dispatch.abort(); + assert!(dispatch.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(5), async { + while runtime.sessions.has_open_sessions().await { + runtime + .sessions + .close_idle_sessions_at( + std::time::Instant::now(), + Duration::ZERO, + "cancelled_test_call", + ) + .await + .unwrap(); + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled middleware retained its gateway session"); + runtime.close().await.unwrap(); +} + +#[tokio::test] +async fn managed_runtime_bypasses_middleware_for_a_claude_startup_probe() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const INTERCEPT: &str = "daemon-worker-startup-probe-bypass-coverage"; + let _ = deregister_llm_request_intercept(INTERCEPT); + let _registration = RequestInterceptRegistration(INTERCEPT); + register_llm_request_intercept( + INTERCEPT, + 1, + false, + Arc::new(|_name, request, annotated| { + Box::pin(async move { Ok(LlmRequestInterceptOutcome::new(request, annotated)) }) + }), + ) + .unwrap(); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let provider = tokio::spawn(async move { + axum::serve( + listener, + Router::new().route( + "/v1/messages", + post(|| async { (StatusCode::ACCEPTED, "startup probe forwarded") }), + ), + ) + .await + .unwrap(); + }); + let runtime = ManagedRuntime::initialize( + GatewayConfig { + anthropic_base_url: format!("http://{address}"), + ..GatewayConfig::default() + }, + Vec::new(), + "machine-owner".into(), + ) + .await + .unwrap(); + let response = runtime + .proxy_provider( + pooled_client().unwrap(), + Request::post("/v1/messages") + .header("x-claude-code-session-id", "startup-probe") + .body(Body::from( + r#"{"model":"claude-opus-4-8[1m]","max_tokens":1,"messages":[{"role":"user","content":"test"}]}"#, + )) + .unwrap(), + ProviderRoute::Anthropic, + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "startup probe forwarded" + ); + drop(_registration); + let response = runtime + .proxy_provider( + pooled_client().unwrap(), + Request::post("/v1/messages") + .header("x-claude-code-session-id", "unbuffered-startup-probe") + .body(Body::from( + r#"{"model":"claude-opus-4-8[1m]","max_tokens":1,"messages":[{"role":"user","content":"test"}]}"#, + )) + .unwrap(), + ProviderRoute::Anthropic, + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!( + response.into_body().collect().await.unwrap().to_bytes(), + "startup probe forwarded" + ); + runtime.close().await.unwrap(); + provider.abort(); +} + +#[tokio::test] +async fn managed_runtime_observes_a_supported_stream_without_rewriting_its_bytes() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind provider"); + let address = listener.local_addr().expect("provider address"); + let app = Router::new().route( + "/v1/responses", + post(|| async { + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .body(Body::from( + "event: response.output_text.delta\ndata: hello\n\ndata: [DONE]\n\n", + )) + .expect("provider response") + }), + ); + let provider = tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve provider"); + }); + let runtime = ManagedRuntime::initialize( + GatewayConfig { + openai_base_url: format!("http://{address}"), + ..GatewayConfig::default() + }, + Vec::new(), + "machine-owner".into(), + ) + .await + .expect("initialize managed runtime"); + let response = runtime + .proxy_provider( + pooled_client().expect("provider client"), + Request::post("/v1/responses") + .header(ACCEPT, "text/event-stream") + .body(Body::from( + r#"{"model":"test","input":"hello","stream":true}"#, + )) + .expect("provider request"), + ProviderRoute::OpenAi, + ) + .await + .expect("managed stream dispatch"); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response + .into_body() + .collect() + .await + .expect("managed stream body") + .to_bytes(), + "event: response.output_text.delta\ndata: hello\n\ndata: [DONE]\n\n" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + runtime.close().await.expect("close managed runtime"); + provider.abort(); +} + +#[tokio::test] +async fn managed_runtime_closes_the_llm_lifecycle_when_provider_dispatch_fails() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let runtime = ManagedRuntime::initialize( + GatewayConfig { + openai_base_url: "http://127.0.0.1:9".into(), + ..GatewayConfig::default() + }, + Vec::new(), + "machine-owner".into(), + ) + .await + .expect("initialize managed runtime"); + let error = runtime + .proxy_provider( + pooled_client().expect("provider client"), + Request::post("/v1/responses") + .body(Body::from(r#"{"model":"test","input":"hello"}"#)) + .expect("provider request"), + ProviderRoute::OpenAi, + ) + .await + .expect_err("unreachable provider must fail"); + assert!( + matches!(error, CliError::Launch(_)), + "unexpected dispatch error: {error:?}" + ); + runtime.close().await.expect("close managed runtime"); +} + +struct RequestInterceptRegistration(&'static str); + +impl Drop for RequestInterceptRegistration { + fn drop(&mut self) { + let _ = deregister_llm_request_intercept(self.0); + } +} + +struct ExecutionInterceptRegistration(&'static str); + +impl Drop for ExecutionInterceptRegistration { + fn drop(&mut self) { + let _ = deregister_llm_execution_intercept(self.0); + } +} + +#[tokio::test] +async fn managed_runtime_rejects_response_mutating_execution_middleware() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const INTERCEPT: &str = "daemon-worker-incompatible-execution-coverage"; + let _ = deregister_llm_execution_intercept(INTERCEPT); + let _registration = ExecutionInterceptRegistration(INTERCEPT); + register_llm_execution_intercept( + INTERCEPT, + 1, + Arc::new(|_name, _request, _next| Box::pin(async { Ok(json!({})) })), + ) + .expect("register execution middleware"); + + let result = + ManagedRuntime::initialize(GatewayConfig::default(), Vec::new(), "machine-owner".into()) + .await; + let error = match result { + Ok(_) => panic!("raw worker delivery must reject execution middleware"), + Err(error) => error, + }; + assert!(error.to_string().contains(INTERCEPT)); +} + +#[tokio::test] +async fn managed_runtime_applies_request_middleware_before_raw_provider_delivery() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const INTERCEPT: &str = "daemon-worker-managed-request-coverage"; + let _ = deregister_llm_request_intercept(INTERCEPT); + let _registration = RequestInterceptRegistration(INTERCEPT); + register_llm_request_intercept( + INTERCEPT, + 1, + false, + Arc::new(|_name, mut request, annotated| { + request + .headers + .insert("x-worker-intercept".into(), json!("applied")); + request.content["input"] = json!("rewritten by middleware"); + Box::pin(async move { Ok(LlmRequestInterceptOutcome::new(request, annotated)) }) + }), + ) + .expect("register request middleware"); + + let captured: CapturedProviderRequest = Arc::new(std::sync::Mutex::new(None)); + let app = Router::new() + .route( + "/v1/responses", + post( + |State(captured): State, + request: Request| async move { + let (parts, body) = request.into_parts(); + let body = axum::body::to_bytes(body, usize::MAX) + .await + .expect("provider body"); + *captured.lock().expect("capture request") = Some((parts.headers, body)); + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .body(Body::from("data: raw provider result\n\n")) + .expect("provider response") + }, + ), + ) + .with_state(Arc::clone(&captured)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind provider"); + let address = listener.local_addr().expect("provider address"); + let provider = tokio::spawn(async move { + axum::serve(listener, app).await.expect("serve provider"); + }); + let runtime = ManagedRuntime::initialize( + GatewayConfig { + openai_base_url: format!("http://{address}"), + ..GatewayConfig::default() + }, + Vec::new(), + "machine-owner".into(), + ) + .await + .expect("initialize managed runtime"); + let response = runtime + .proxy_provider( + pooled_client().expect("provider client"), + Request::post("/v1/responses") + .header(ACCEPT, "text/event-stream") + .body(Body::from( + r#"{"model":"test","input":"original","stream":true}"#, + )) + .expect("provider request"), + ProviderRoute::OpenAi, + ) + .await + .expect("managed provider dispatch"); + assert_eq!( + response + .into_body() + .collect() + .await + .expect("provider response body") + .to_bytes(), + "data: raw provider result\n\n" + ); + let (headers, body) = captured + .lock() + .expect("captured request") + .take() + .expect("provider saw request"); + assert_eq!(headers["x-worker-intercept"], "applied"); + assert_eq!( + serde_json::from_slice::(&body).expect("rewritten request JSON")["input"], + "rewritten by middleware" + ); + runtime.close().await.expect("close managed runtime"); + provider.abort(); +} + +#[tokio::test] +async fn managed_runtime_rejects_request_middleware_that_changes_stream_mode() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + const INTERCEPT: &str = "daemon-worker-stream-mode-mutation-coverage"; + let _ = deregister_llm_request_intercept(INTERCEPT); + let _registration = RequestInterceptRegistration(INTERCEPT); + register_llm_request_intercept( + INTERCEPT, + 1, + false, + Arc::new(|_name, mut request, annotated| { + request.content["stream"] = json!(false); + Box::pin(async move { Ok(LlmRequestInterceptOutcome::new(request, annotated)) }) + }), + ) + .expect("register request middleware"); + + let runtime = + ManagedRuntime::initialize(GatewayConfig::default(), Vec::new(), "machine-owner".into()) + .await + .expect("initialize managed runtime"); + let error = runtime + .proxy_provider( + pooled_client().expect("provider client"), + Request::post("/v1/responses") + .header(ACCEPT, "text/event-stream") + .body(Body::from( + r#"{"model":"test","input":"value","stream":true}"#, + )) + .expect("provider request"), + ProviderRoute::OpenAi, + ) + .await + .expect_err("stream-mode mutation must fail closed"); + assert!(error.to_string().contains(STREAM_MODE_MUTATION_ERROR)); + runtime.close().await.expect("close managed runtime"); +} + +#[tokio::test] +async fn managed_helper_error_and_metadata_paths_are_explicit() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let prepared = prepared_request(HeaderMap::new()); + let invalid_destination = LlmRequest { + headers: serde_json::Map::from_iter([( + INTERNAL_DISPATCH_URL_HEADER.into(), + json!("http://[invalid"), + )]), + content: Value::Null, + }; + assert!( + effective_destination( + &prepared, + ProviderRoute::OpenAi, + Some(&invalid_destination), + &GatewayConfig::default(), + ) + .is_err() + ); + assert!(!has_explicit_target(&LlmRequest { + headers: serde_json::Map::from_iter([( + INTERNAL_DISPATCH_ROUTE_HEADER.into(), + json!(" "), + )]), + content: Value::Null, + })); + assert_eq!(json_header_value(&json!(42)).unwrap(), "42"); + assert!(json_header_value(&json!("\n")).is_none()); + assert_eq!( + merge_object(json!("not-an-object"), json!({"added": true})), + json!({"added": true}) + ); + let mut metadata = Value::Null; + insert_metadata(&mut metadata, "key", json!("value")); + assert_eq!(metadata, json!({"key": "value"})); + assert!(!request_body_decode_required().expect("empty middleware registry")); + assert!(reject_incompatible_execution_middleware().is_ok()); +} + +#[test] +fn managed_header_and_observation_helpers_cover_ignored_and_invalid_values() { + let mut source = HeaderMap::new(); + source.insert("x-remove-me", HeaderValue::from_static("before")); + source.insert(CONTENT_ENCODING, HeaderValue::from_static("gzip")); + source.insert( + INTERNAL_DISPATCH_BACKEND_HEADER, + HeaderValue::from_static("untrusted"), + ); + let prepared = prepared_request(source); + let effective = LlmRequest { + headers: serde_json::Map::from_iter([ + ("bad header".into(), json!("ignored")), + (INTERNAL_RETRY_AWARE_HEADER.into(), json!(true)), + ("x-json".into(), json!({"nested": true})), + ]), + content: json!({"model": "rewritten", "stream": true}), + }; + let (headers, body, explicit) = effective_request(&prepared, Some(&effective)).unwrap(); + assert!(!explicit); + assert!(!headers.contains_key("x-remove-me")); + assert!(!headers.contains_key(CONTENT_ENCODING)); + assert!(!headers.contains_key(INTERNAL_DISPATCH_BACKEND_HEADER)); + assert!(!headers.contains_key(INTERNAL_RETRY_AWARE_HEADER)); + assert_eq!(headers["x-json"], r#"{"nested":true}"#); + assert_ne!(body, prepared.body); + + let null_content = LlmRequest { + headers: serde_json::Map::new(), + content: Value::Null, + }; + assert_eq!( + effective_request(&prepared, Some(&null_content)).unwrap().1, + prepared.body + ); + assert!(!request_streaming_hint(&HeaderMap::from_iter([( + ACCEPT, + HeaderValue::from_static("application/json"), + )]))); + assert!(!response_streaming(&HeaderMap::from_iter([( + CONTENT_TYPE, + HeaderValue::from_static("application/json"), + )]))); + assert_eq!( + provider_surface("/responses"), + Some(ProviderSurface::OpenAIResponses) + ); + assert_eq!( + provider_surface("/chat/completions"), + Some(ProviderSurface::OpenAIChat) + ); + assert_eq!( + provider_surface("/v1/chat/completions"), + Some(ProviderSurface::OpenAIChat) + ); + assert_eq!( + provider_surface("/backend-api/codex/responses"), + Some(ProviderSurface::OpenAIResponses) + ); + + let unknown = ObservedResponse { + value: None, + truncated: false, + terminal: 99, + status: StatusCode::OK, + failure: Some("unexpected".into()), + }; + assert_eq!( + unknown.metadata()["daemon_worker_observation"]["terminal"], + "unknown" + ); +} + +#[cfg(unix)] +#[test] +fn managed_observation_and_provider_environment_values_reject_invalid_inputs() { + use std::os::unix::ffi::OsStrExt; + + let _environment = EnvScope::set(&[ + ( + OBSERVATION_CAPTURE_BYTES_ENV, + Some(std::ffi::OsStr::from_bytes(b"\xff")), + ), + ("OPENAI_API_KEY", Some(OsStr::new(" "))), + ("ANTHROPIC_API_KEY", None), + ]); + assert!(observation_capture_limit_from_environment().is_err()); + assert!(environment_value("OPENAI_API_KEY").is_none()); + assert!(environment_value("ANTHROPIC_API_KEY").is_none()); + let mut headers = HeaderMap::new(); + inject_provider_auth( + &mut headers, + ProviderRoute::OpenAi, + &GatewayConfig { + openai_auth_header: Some("bad\nvalue".into()), + ..GatewayConfig::default() + }, + ); + assert!(headers.is_empty()); +} + +#[tokio::test] +async fn observation_signal_waits_for_terminal_result_and_marks_truncation() { + let signal = Arc::new(ObservationSignal::new()); + signal.truncate(); + let waiter = { + let signal = Arc::clone(&signal); + tokio::spawn(async move { signal.wait().await }) + }; + tokio::task::yield_now().await; + signal.finish(OBSERVATION_COMPLETE); + assert_eq!(waiter.await.expect("waiter task"), OBSERVATION_COMPLETE); + assert!(signal.truncated.load(AtomicOrdering::Acquire)); + signal.finish(OBSERVATION_CANCELLED); + assert_eq!( + signal.terminal.load(AtomicOrdering::Acquire), + OBSERVATION_COMPLETE + ); +} + +#[tokio::test] +async fn managed_hook_rejects_invalid_and_oversized_json_before_adapter_dispatch() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let runtime = ManagedRuntime::initialize( + GatewayConfig { + max_hook_payload_bytes: 8, + ..GatewayConfig::default() + }, + Vec::new(), + "machine-owner".into(), + ) + .await + .expect("initialize managed runtime"); + + for body in ["not-json", r#"{"payload":"too-large"}"#] { + let response = runtime + .handle_hook( + HookRoute::Pi, + Request::post("/hook") + .body(Body::from(body)) + .expect("hook request"), + ) + .await; + assert!(response.status().is_client_error()); + } + runtime.close().await.expect("close managed runtime"); +} + +#[tokio::test] +async fn observation_handles_empty_invalid_and_unsuccessful_provider_responses() { + let (body, observation) = observe_body( + http_body_util::Empty::::new(), + StatusCode::BAD_REQUEST, + DEFAULT_OBSERVATION_CAPTURE_BYTES, + ); + assert!(body.is_end_stream()); + let observed = observation + .finish(ProviderSurface::OpenAIResponses, false) + .await; + assert_eq!(observed.terminal, OBSERVATION_COMPLETE); + assert!(observed.value.is_none()); + assert!(observed.failure.unwrap().contains("HTTP 400")); + + let malformed = Bytes::from_static(b"data: {not-json}\n\n"); + let (body, observation) = observe_body( + Full::new(malformed.clone()), + StatusCode::OK, + DEFAULT_OBSERVATION_CAPTURE_BYTES, + ); + assert_eq!(body.collect().await.unwrap().to_bytes(), malformed); + let observed = observation + .finish(ProviderSurface::OpenAIResponses, true) + .await; + assert!(observed.truncated); + assert!(observed.value.is_none()); + assert!(observed.failure.unwrap().contains("truncated")); +} + +#[tokio::test] +async fn malformed_permission_hooks_fail_closed_in_each_native_response_shape() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let runtime = + ManagedRuntime::initialize(GatewayConfig::default(), Vec::new(), "machine-owner".into()) + .await + .expect("initialize managed runtime"); + + let codex = runtime + .handle_hook( + HookRoute::Codex, + Request::post("/hook") + .body(Body::from( + json!({ + "session_id": "codex-session", + "hook_event_name": "PermissionRequest", + "tool_name": "shell", + "arguments": {"cmd": "pwd"} + }) + .to_string(), + )) + .unwrap(), + ) + .await; + assert_eq!(codex.status(), StatusCode::OK); + let codex: Value = + serde_json::from_slice(&codex.into_body().collect().await.unwrap().to_bytes()).unwrap(); + assert_eq!(codex["decision"], "deny"); + assert!( + codex["reason"] + .as_str() + .is_some_and(|reason| !reason.is_empty()) + ); + + let claude = runtime + .handle_hook( + HookRoute::Claude, + Request::post("/hook") + .body(Body::from( + json!({ + "session_id": "claude-session", + "hook_event_name": "PermissionRequest", + "tool_name": "Write", + "tool_input": {"file_path": "README.md"} + }) + .to_string(), + )) + .unwrap(), + ) + .await; + assert_eq!(claude.status(), StatusCode::OK); + let claude: Value = + serde_json::from_slice(&claude.into_body().collect().await.unwrap().to_bytes()).unwrap(); + assert_eq!(claude["hookSpecificOutput"]["decision"]["behavior"], "deny"); + assert!( + claude["hookSpecificOutput"]["decision"]["message"] + .as_str() + .is_some_and(|message| !message.is_empty()) + ); + + runtime.close().await.expect("close managed runtime"); +} diff --git a/crates/cli/tests/coverage/daemon/worker_runtime_tests.rs b/crates/cli/tests/coverage/daemon/worker_runtime_tests.rs new file mode 100644 index 000000000..db8bde1f9 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/worker_runtime_tests.rs @@ -0,0 +1,858 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::daemon::common::protocol::SensitiveString; +use base64::Engine as _; +use std::convert::Infallible; +use std::pin::Pin; +use std::sync::Mutex; +use std::task::{Context, Poll}; + +use http::header::TRAILER; +use http_body_util::BodyExt as _; +use hyper::body::{Frame, Incoming, SizeHint}; +use hyper::server::conn::http1; +use hyper::service::service_fn; +use hyper_util::rt::TokioIo; +use tokio::sync::oneshot; +use tower::ServiceExt as _; + +use crate::test_support::EnvScope; + +fn state() -> Arc { + state_with_config(GatewayConfig::default()) +} + +fn state_with_config(config: GatewayConfig) -> Arc { + Arc::new(WorkerState { + worker_id: "worker-one".into(), + config, + upstream: pooled_client().expect("pooled client"), + managed: None, + auth: RwLock::new(AuthTokens { + data: TokenDigest::from_token(b"data-secret"), + pending_data: None, + readiness_data: None, + control: TokenDigest::from_token(b"control-secret"), + last_control_sequence: 0, + last_control_request_id: String::new(), + }), + accepting: AtomicBool::new(true), + draining: AtomicBool::new(false), + exiting: AtomicBool::new(false), + in_flight: AtomicUsize::new(0), + drain_deadline: RwLock::new(None), + lifecycle: Notify::new(), + }) +} + +#[test] +fn relative_drain_timeout_does_not_depend_on_the_daemon_wall_clock() { + let request = WorkerDrainRequest { + worker_id: "worker-one".into(), + deadline_unix_ms: 0, + timeout_ms: Some(321), + }; + assert_eq!(drain_timeout_ms(&request), 321); +} + +#[tokio::test] +async fn authenticated_readiness_probe_opens_admission_before_publication() { + let state = state(); + state.accepting.store(false, Ordering::Release); + write_lock(&state.auth).readiness_data = Some(TokenDigest::from_token(b"data-secret")); + let headers = HeaderMap::from_iter([( + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderValue::from_static("data-secret"), + )]); + let response = readiness_probe(State(Arc::clone(&state)), headers).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + assert!(state.admit().is_some()); +} + +#[tokio::test] +async fn worker_stops_immediately_after_exit_or_an_empty_drain() { + let exited = state(); + exited.request_exit(); + tokio::time::timeout(Duration::from_millis(50), exited.wait_until_stopped()) + .await + .expect("exit stops the worker"); + + let drained = state(); + drained.begin_drain(0); + tokio::time::timeout(Duration::from_millis(50), drained.wait_until_stopped()) + .await + .expect("empty drain stops the worker"); +} + +#[tokio::test] +async fn unmanaged_worker_forwards_hooks_without_contacting_a_provider() { + let response = router(state()) + .oneshot( + Request::post("/hooks/pi") + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(Body::from("native hook payload")) + .expect("hook request"), + ) + .await + .expect("hook response"); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response + .into_body() + .collect() + .await + .expect("hook body") + .to_bytes(), + crate::daemon::common::routes::HookRoute::Pi.pass_through_body() + ); +} + +#[test] +fn drain_control_requires_scoped_sequence_hash_and_exact_replay() { + let state = state(); + let request = SessionRequest::new( + "worker-one".into(), + SensitiveString::new("control-secret").expect("secret"), + 1, + WorkerDrainRequest { + worker_id: "worker-one".into(), + deadline_unix_ms: 100, + timeout_ms: Some(100), + }, + ) + .expect("drain request"); + assert!(state.authenticate_control(&request)); + assert!(state.authenticate_control(&request)); + + let mut replay_mutation = request.clone(); + replay_mutation.request_id = "different-request".into(); + assert!(!state.authenticate_control(&replay_mutation)); + + let out_of_order = SessionRequest::new( + "worker-one".into(), + SensitiveString::new("control-secret").expect("secret"), + 3, + WorkerDrainRequest { + worker_id: "worker-one".into(), + deadline_unix_ms: 100, + timeout_ms: Some(100), + }, + ) + .expect("out of order"); + assert!(!state.authenticate_control(&out_of_order)); +} + +#[tokio::test] +async fn control_and_readiness_reject_mismatched_credentials_without_admission() { + let state = state(); + state.accepting.store(false, Ordering::Release); + write_lock(&state.auth).readiness_data = Some(TokenDigest::from_token(b"data-secret")); + let mut duplicate = HeaderMap::new(); + duplicate.append( + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderValue::from_static("data-secret"), + ); + duplicate.append( + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderValue::from_static("data-secret"), + ); + assert_eq!( + readiness_probe(State(Arc::clone(&state)), duplicate) + .await + .status(), + StatusCode::NO_CONTENT + ); + assert!(state.admit().is_none()); + + let wrong_worker = SessionRequest::new( + "wrong-worker".into(), + SensitiveString::new("control-secret").expect("secret"), + 1, + WorkerDrainRequest { + worker_id: "wrong-worker".into(), + deadline_unix_ms: 1, + timeout_ms: None, + }, + ) + .expect("control request"); + assert!(!state.authenticate_control(&wrong_worker)); +} + +struct CausalBody { + phase: u8, + release_second: oneshot::Receiver<()>, + trailers: Option, +} + +struct PanicBody; + +impl hyper::body::Body for PanicBody { + type Data = bytes::Bytes; + type Error = Infallible; + + fn poll_frame( + self: Pin<&mut Self>, + _context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + panic!("an unauthenticated request body must not be polled") + } +} + +impl hyper::body::Body for CausalBody { + type Data = bytes::Bytes; + type Error = Infallible; + + fn poll_frame( + self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll, Self::Error>>> { + let this = self.get_mut(); + match this.phase { + 0 => { + this.phase = 1; + Poll::Ready(Some(Ok(Frame::data(bytes::Bytes::from_static( + b"event: first\r\ndata: A\r\n\r\n", + ))))) + } + 1 => match Pin::new(&mut this.release_second).poll(context) { + Poll::Pending => Poll::Pending, + Poll::Ready(_) => { + this.phase = 2; + Poll::Ready(Some(Ok(Frame::data(bytes::Bytes::from_static( + b": heartbeat\r\ndata: [DONE]\r\n\r\n", + ))))) + } + }, + 2 => { + this.phase = 3; + Poll::Ready( + this.trailers + .take() + .map(|trailers| Ok(Frame::trailers(trailers))), + ) + } + _ => Poll::Ready(None), + } + } + + fn is_end_stream(&self) -> bool { + self.phase >= 3 && self.trailers.is_none() + } + + fn size_hint(&self) -> SizeHint { + SizeHint::default() + } +} + +#[test] +fn data_request_requires_exactly_one_matching_header() { + let state = state(); + let mut headers = HeaderMap::new(); + assert!(!state.authenticate_data(&headers)); + headers.insert(WORKER_TOKEN_HEADER, HeaderValue::from_static("wrong")); + assert!(!state.authenticate_data(&headers)); + headers.insert(WORKER_TOKEN_HEADER, HeaderValue::from_static("data-secret")); + assert!(state.authenticate_data(&headers)); + headers.append(WORKER_TOKEN_HEADER, HeaderValue::from_static("data-secret")); + assert!(!state.authenticate_data(&headers)); +} + +#[test] +fn recovery_probe_accepts_staged_token_only_until_commit_or_discard() { + let state = state(); + let registration = control::test_registration("new-data-secret", "new-control-secret"); + state.stage_recovery_data_token(®istration); + + let mut headers = HeaderMap::new(); + headers.insert( + WORKER_TOKEN_HEADER, + HeaderValue::from_static("new-data-secret"), + ); + assert!(state.authenticate_data(&headers)); + headers.insert(WORKER_TOKEN_HEADER, HeaderValue::from_static("data-secret")); + assert!(state.authenticate_data(&headers)); + + state.discard_recovery_data_token(); + headers.insert( + WORKER_TOKEN_HEADER, + HeaderValue::from_static("new-data-secret"), + ); + assert!(!state.authenticate_data(&headers)); + + state.stage_recovery_data_token(®istration); + state.control_restored(®istration); + assert!(state.authenticate_data(&headers)); + headers.insert(WORKER_TOKEN_HEADER, HeaderValue::from_static("data-secret")); + assert!(!state.authenticate_data(&headers)); +} + +#[tokio::test] +async fn only_the_staged_registration_token_can_reopen_readiness() { + let state = state(); + state.control_lost(); + let registration = control::test_registration("new-data-secret", "new-control-secret"); + state.stage_recovery_data_token(®istration); + + let old_headers = HeaderMap::from_iter([( + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderValue::from_static("data-secret"), + )]); + assert_eq!( + readiness_probe(State(Arc::clone(&state)), old_headers) + .await + .status(), + StatusCode::NO_CONTENT + ); + assert!(state.admit().is_none()); + + let staged_headers = HeaderMap::from_iter([( + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderValue::from_static("new-data-secret"), + )]); + assert_eq!( + readiness_probe(State(Arc::clone(&state)), staged_headers) + .await + .status(), + StatusCode::NO_CONTENT + ); + assert!(state.admit().is_some()); +} + +#[tokio::test] +async fn unauthenticated_request_is_rejected_before_its_body_is_polled() { + let request = Request::post("/v1/responses") + .body(Body::new(PanicBody)) + .expect("worker request"); + let response = router(state()) + .oneshot(request) + .await + .expect("worker response"); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); +} + +#[tokio::test] +async fn authenticated_readiness_probe_works_before_public_requests_are_admitted() { + let state = state(); + state.control_lost(); + let probe = Request::get(WORKER_PROBE_PATH) + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(Body::empty()) + .expect("probe request"); + assert_eq!( + router(Arc::clone(&state)) + .oneshot(probe) + .await + .expect("probe response") + .status(), + StatusCode::NO_CONTENT + ); + + let provider = Request::post("/v1/responses") + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(Body::new(PanicBody)) + .expect("provider request"); + let response = router(state) + .oneshot(provider) + .await + .expect("provider response"); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + response + .headers() + .get(WORKER_ROUTE_FAILURE_HEADER) + .expect("route failure signal"), + "pass-through" + ); +} + +#[test] +fn route_failure_responses_are_explicitly_signaled_to_the_daemon() { + let response = route_failure_response(CliError::Config("incompatible middleware".into())); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!( + response + .headers() + .get(WORKER_ROUTE_FAILURE_HEADER) + .expect("route failure signal"), + "pass-through" + ); +} + +#[test] +fn control_loss_rejects_new_admissions_without_touching_existing_one() { + let state = state(); + let accepted = state.admit().expect("request admitted"); + assert_eq!(state.in_flight.load(Ordering::Acquire), 1); + state.control_lost(); + assert!(state.admit().is_none()); + assert_eq!(state.in_flight.load(Ordering::Acquire), 1); + drop(accepted); + assert_eq!(state.in_flight.load(Ordering::Acquire), 0); +} + +#[test] +fn configured_provider_auth_never_replaces_caller_auth() { + let config = GatewayConfig { + openai_auth_header: Some("Bearer configured".into()), + ..GatewayConfig::default() + }; + let mut headers = HeaderMap::new(); + headers.insert(AUTHORIZATION, HeaderValue::from_static("Bearer caller")); + inject_provider_auth(&mut headers, ProviderRoute::OpenAi, &config); + assert_eq!(headers.get(AUTHORIZATION).expect("auth"), "Bearer caller"); + + headers.remove(AUTHORIZATION); + inject_provider_auth(&mut headers, ProviderRoute::OpenAi, &config); + assert_eq!( + headers.get(AUTHORIZATION).expect("configured auth"), + "Bearer configured" + ); +} + +#[tokio::test] +async fn provider_frames_are_forwarded_causally_with_status_headers_and_trailers() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind provider"); + let provider_address = listener.local_addr().expect("provider address"); + let (release_second, wait_for_release) = oneshot::channel(); + let provider = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept worker"); + let release = Arc::new(Mutex::new(Some(wait_for_release))); + let service = service_fn(move |_request: Request| { + let release = release + .lock() + .expect("release lock") + .take() + .expect("one request"); + async move { + let mut trailers = HeaderMap::new(); + trailers.append("x-stream-checksum", HeaderValue::from_static("one")); + trailers.append("x-stream-checksum", HeaderValue::from_static("two")); + let mut response = Response::new(box_body(CausalBody { + phase: 0, + release_second: release, + trailers: Some(trailers), + })); + *response.status_mut() = StatusCode::CREATED; + response + .headers_mut() + .append("x-provider", HeaderValue::from_static("first")); + response + .headers_mut() + .append("x-provider", HeaderValue::from_static("second")); + response.headers_mut().insert( + WORKER_ROUTE_FAILURE_HEADER, + HeaderValue::from_static("provider-spoof"), + ); + response + .headers_mut() + .insert(TRAILER, HeaderValue::from_static("x-stream-checksum")); + Ok::<_, Infallible>(response) + } + }); + let mut connection = http1::Builder::new(); + connection.keep_alive(false); + connection + .serve_connection(TokioIo::new(stream), service) + .await + .expect("serve provider response"); + }); + + let config = GatewayConfig { + openai_base_url: format!("http://{provider_address}/v1"), + ..GatewayConfig::default() + }; + let state = state_with_config(config); + let request = Request::post("/v1/responses") + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(Body::empty()) + .expect("worker request"); + let response = router(Arc::clone(&state)) + .oneshot(request) + .await + .expect("worker response"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.headers().get_all("x-provider").iter().count(), 2); + assert!(!response.headers().contains_key(WORKER_ROUTE_FAILURE_HEADER)); + assert_eq!(response.headers()[TRAILER], "x-stream-checksum"); + assert_eq!(state.in_flight.load(Ordering::Acquire), 1); + + let mut body = response.into_body(); + let first = body + .frame() + .await + .expect("first frame") + .expect("first frame succeeds") + .into_data() + .expect("first is data"); + assert_eq!(first, "event: first\r\ndata: A\r\n\r\n"); + let second = body.frame(); + tokio::pin!(second); + assert!(futures_util::poll!(second.as_mut()).is_pending()); + release_second.send(()).expect("release provider"); + assert_eq!( + second + .await + .expect("second frame") + .expect("second frame succeeds") + .into_data() + .expect("second is data"), + ": heartbeat\r\ndata: [DONE]\r\n\r\n" + ); + let trailers = body + .frame() + .await + .expect("trailer frame") + .expect("trailer succeeds") + .into_trailers() + .expect("trailers"); + assert_eq!( + trailers + .get_all("x-stream-checksum") + .iter() + .map(|value| value.to_str().expect("ASCII trailer")) + .collect::>(), + ["one", "two"] + ); + assert!(body.frame().await.is_none()); + assert_eq!(state.in_flight.load(Ordering::Acquire), 0); + provider.await.expect("provider task"); +} + +#[tokio::test] +async fn draining_waits_for_accepted_work_and_honors_the_deadline() { + let draining_state = state(); + let accepted = draining_state + .admit() + .expect("request admitted before drain"); + draining_state.begin_drain(10_000); + + let waiting = draining_state.wait_until_stopped(); + tokio::pin!(waiting); + assert!(futures_util::poll!(waiting.as_mut()).is_pending()); + drop(accepted); + tokio::time::timeout(Duration::from_millis(50), waiting) + .await + .expect("releasing the last request completes drain"); + + let expired = state(); + let _accepted = expired.admit().expect("request admitted"); + expired.begin_drain(0); + tokio::time::timeout(Duration::from_millis(50), expired.wait_until_stopped()) + .await + .expect("expired drain deadline stops waiting"); +} + +#[tokio::test] +async fn stopping_readiness_unknown_routes_and_control_auth_have_stable_responses() { + let stopping_state = state(); + stopping_state.begin_drain(0); + let response = readiness_probe(State(Arc::clone(&stopping_state)), HeaderMap::new()).await; + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + + let unknown = router(state()) + .oneshot( + Request::get("/not-a-public-route") + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(Body::empty()) + .expect("unknown request"), + ) + .await + .expect("unknown response"); + assert_eq!(unknown.status(), StatusCode::NOT_FOUND); + + let invalid = SessionRequest::new( + "worker-one".into(), + SensitiveString::new("wrong-control-secret").expect("secret"), + 1, + WorkerDrainRequest { + worker_id: "worker-one".into(), + deadline_unix_ms: now_unix_ms() + 1_000, + timeout_ms: None, + }, + ) + .expect("drain request"); + let response = drain(State(state()), Json(invalid)).await; + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + + let valid = SessionRequest::new( + "worker-one".into(), + SensitiveString::new("control-secret").expect("secret"), + 1, + WorkerDrainRequest { + worker_id: "worker-one".into(), + deadline_unix_ms: now_unix_ms() + 1_000, + timeout_ms: None, + }, + ) + .expect("drain request"); + let state = state(); + let response = drain(State(Arc::clone(&state)), Json(valid)).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + assert!(state.draining.load(Ordering::Acquire)); +} + +#[tokio::test] +async fn provider_routing_reports_invalid_destinations_and_connection_failures() { + let invalid = state_with_config(GatewayConfig { + openai_base_url: "not a URI".into(), + ..GatewayConfig::default() + }); + let response = router(invalid) + .oneshot( + Request::post("/v1/responses") + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(Body::empty()) + .expect("provider request"), + ) + .await + .expect("invalid destination response"); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + + let unavailable = state_with_config(GatewayConfig { + openai_base_url: "http://127.0.0.1:9".into(), + ..GatewayConfig::default() + }); + let response = router(unavailable) + .oneshot( + Request::post("/v1/responses") + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(Body::empty()) + .expect("provider request"), + ) + .await + .expect("connection failure response"); + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); +} + +#[test] +fn provider_auth_supports_both_configured_and_environment_credentials() { + let _environment = EnvScope::set(&[ + ("OPENAI_API_KEY", Some(std::ffi::OsStr::new(" openai-env "))), + ( + "ANTHROPIC_API_KEY", + Some(std::ffi::OsStr::new(" anthropic-env ")), + ), + ]); + let mut openai = HeaderMap::new(); + inject_provider_auth( + &mut openai, + ProviderRoute::OpenAi, + &GatewayConfig::default(), + ); + assert_eq!(openai[AUTHORIZATION], "Bearer openai-env"); + + let mut anthropic = HeaderMap::new(); + inject_provider_auth( + &mut anthropic, + ProviderRoute::Anthropic, + &GatewayConfig::default(), + ); + assert_eq!(anthropic["x-api-key"], "anthropic-env"); + + let mut configured = HeaderMap::new(); + inject_provider_auth( + &mut configured, + ProviderRoute::Anthropic, + &GatewayConfig { + anthropic_auth_header: Some("configured-anthropic".into()), + ..GatewayConfig::default() + }, + ); + assert_eq!(configured[AUTHORIZATION], "configured-anthropic"); + assert!(header_value("bad\nheader").is_none()); +} + +#[test] +fn runtime_state_constructor_stages_registration_credentials_until_readiness() { + let registration = control::test_registration("fresh-data", "fresh-control"); + let state = Arc::new( + WorkerState::new( + "worker-one".into(), + GatewayConfig::default(), + None, + ®istration, + ) + .expect("worker state"), + ); + assert!(state.admit().is_none()); + let headers = HeaderMap::from_iter([( + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderValue::from_static("fresh-data"), + )]); + assert!(state.authenticate_data(&headers)); + state.activate_pending_readiness(&headers); + assert!(state.admit().is_some()); +} + +#[tokio::test] +async fn runtime_serve_reaches_readiness_and_exposes_the_authenticated_listener() { + let _environment = EnvScope::set(&[("NEMO_RELAY_DAEMON_OBSERVATION_CAPTURE_BYTES", None)]); + let daemon_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind daemon fixture"); + let daemon_address = daemon_listener.local_addr().expect("daemon address"); + let daemon = tokio::spawn(async move { + axum::serve( + daemon_listener, + Router::new().fallback(|| async { StatusCode::NO_CONTENT }), + ) + .await + .expect("serve daemon fixture"); + }); + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind worker"); + let worker_address = listener.local_addr().expect("worker address"); + let identity = MachineIdentity::generate().expect("identity").identity; + let worker = tokio::spawn(serve( + listener, + RuntimeOptions { + daemon_origin: format!("http://{daemon_address}"), + identity, + worker_id: "worker-one".into(), + endpoint: format!("http://{worker_address}"), + worker_tls_root: None, + tls_config: None, + config: GatewayConfig::default(), + dynamic_plugins: Vec::new(), + registration: control::test_registration("runtime-data", "runtime-control"), + }, + )); + + let client = reqwest::Client::new(); + let probe_url = format!("http://{worker_address}{WORKER_PROBE_PATH}"); + let mut response = None; + for _ in 0..50 { + if let Ok(candidate) = client + .get(&probe_url) + .header(WORKER_TOKEN_HEADER, "runtime-data") + .send() + .await + { + response = Some(candidate); + break; + } + tokio::task::yield_now().await; + } + assert_eq!( + response.expect("worker became reachable").status(), + StatusCode::NO_CONTENT + ); + worker.abort(); + daemon.abort(); +} + +#[tokio::test] +async fn tls_worker_listener_discards_bad_handshakes_and_serves_authenticated_requests() { + let _ = rustls::crypto::ring::default_provider().install_default(); + let rcgen::CertifiedKey { cert, key_pair } = + rcgen::generate_simple_self_signed(vec!["127.0.0.1".into()]).expect("certificate"); + let key = rustls::pki_types::PrivateKeyDer::Pkcs8(rustls::pki_types::PrivatePkcs8KeyDer::from( + key_pair.serialize_der(), + )); + let mut tls = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![cert.der().clone()], key) + .expect("TLS server config"); + tls.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind TLS worker"); + let address = listener.local_addr().expect("worker address"); + let server = tokio::spawn(serve_tls(listener, router(state()), Arc::new(tls))); + + // A non-TLS peer is ignored without terminating the listener. + let mut plain = tokio::net::TcpStream::connect(address) + .await + .expect("plain connection"); + use tokio::io::AsyncWriteExt as _; + plain.write_all(b"not tls").await.expect("plain write"); + plain.shutdown().await.expect("plain shutdown"); + + let root = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(cert.der()); + let client = crate::daemon::common::worker_tls::pooled_worker_tls_client(&root) + .expect("pinned worker client"); + let request = Request::get(format!( + "https://127.0.0.1:{}{WORKER_PROBE_PATH}", + address.port() + )) + .header(WORKER_TOKEN_HEADER, "data-secret") + .body(box_body(http_body_util::Empty::::new())) + .expect("TLS probe"); + let response = client.request(request).await.expect("TLS response"); + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + server.abort(); + assert!(server.await.expect_err("server aborted").is_cancelled()); +} + +#[tokio::test] +async fn heartbeat_attempt_accepts_only_a_successful_daemon_acknowledgement() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind daemon fixture"); + let address = listener.local_addr().expect("daemon address"); + let daemon = tokio::spawn(async move { + axum::serve( + listener, + Router::new().fallback(|| async { StatusCode::NO_CONTENT }), + ) + .await + .expect("serve daemon fixture"); + }); + let mut registration = control::test_registration("data", "control"); + assert!( + heartbeat_attempt( + &mut registration, + &format!("http://{address}"), + "worker-one" + ) + .await + ); + daemon.abort(); +} + +#[tokio::test(start_paused = true)] +async fn monitor_control_stops_without_contacting_daemon_after_drain_begins() { + let state = state(); + state.begin_drain(10_000); + let identity = MachineIdentity::generate().expect("identity").identity; + let monitor = monitor_control( + Arc::clone(&state), + "http://127.0.0.1:9".into(), + identity, + "worker-one".into(), + "http://127.0.0.1:1".into(), + None, + control::test_registration("data", "control"), + ); + tokio::time::timeout(Duration::from_secs(6), monitor) + .await + .expect("monitor exits on its first lifecycle check"); +} + +#[tokio::test(start_paused = true)] +async fn monitor_control_rejects_work_and_exits_after_recovery_deadline() { + let state = state(); + state.accepting.store(true, Ordering::Release); + let identity = MachineIdentity::generate().expect("identity").identity; + tokio::time::timeout( + Duration::from_millis(RECOVERY_LIFETIME_MS + 30_000), + monitor_control( + Arc::clone(&state), + "http://127.0.0.1:9".into(), + identity, + "worker-one".into(), + "http://127.0.0.1:1".into(), + None, + control::test_registration("data", "control"), + ), + ) + .await + .expect("monitor exits after its bounded recovery window"); + assert!(!state.accepting.load(Ordering::Acquire)); + assert!(state.exiting.load(Ordering::Acquire)); +} diff --git a/crates/cli/tests/coverage/daemon/worker_tests.rs b/crates/cli/tests/coverage/daemon/worker_tests.rs new file mode 100644 index 000000000..52e5bcc0e --- /dev/null +++ b/crates/cli/tests/coverage/daemon/worker_tests.rs @@ -0,0 +1,74 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::daemon::common::control::now_unix_ms; +use crate::daemon::common::protocol::SensitiveString; + +fn bootstrap(bind_ip: Ipv4Addr, port: u16) -> WorkerBootstrap { + WorkerBootstrap { + activation_id: "activation".into(), + activation_token: SensitiveString::new("secret").expect("secret"), + deadline_unix_ms: now_unix_ms().saturating_add(10_000), + bind_ip, + port, + advertise_address: None, + } +} + +#[test] +fn default_worker_network_matches_loopback_ephemeral_grant() { + let options = Options { + daemon_address: "http://127.0.0.1:47632".into(), + bind: Ipv4Addr::LOCALHOST, + port: None, + advertise_address: None, + }; + assert_eq!( + validate_bootstrap(&options, &bootstrap(Ipv4Addr::LOCALHOST, 0)).expect("valid grant"), + "127.0.0.1:0".parse().expect("socket") + ); +} + +#[test] +fn worker_network_must_match_activation_grant() { + let options = Options { + daemon_address: "http://127.0.0.1:47632".into(), + bind: Ipv4Addr::LOCALHOST, + port: Some(4444), + advertise_address: None, + }; + assert!(validate_bootstrap(&options, &bootstrap(Ipv4Addr::LOCALHOST, 0)).is_err()); + + let bind_mismatch = Options { + port: None, + ..options.clone() + }; + assert!(validate_bootstrap(&bind_mismatch, &bootstrap(Ipv4Addr::UNSPECIFIED, 0)).is_err()); + + let advertised = Options { + bind: Ipv4Addr::UNSPECIFIED, + port: None, + advertise_address: Some("worker.example.com".into()), + ..options + }; + let mut advertised_grant = bootstrap(Ipv4Addr::UNSPECIFIED, 0); + advertised_grant.advertise_address = Some("worker.example.com".into()); + assert_eq!( + validate_bootstrap(&advertised, &advertised_grant).unwrap(), + "0.0.0.0:0".parse().unwrap() + ); +} + +#[test] +fn daemon_wall_clock_deadline_is_not_rejected_by_the_worker_clock() { + let options = Options { + daemon_address: "http://127.0.0.1:47632".into(), + bind: Ipv4Addr::LOCALHOST, + port: None, + advertise_address: None, + }; + let mut expired = bootstrap(Ipv4Addr::LOCALHOST, 0); + expired.deadline_unix_ms = now_unix_ms(); + assert!(validate_bootstrap(&options, &expired).is_ok()); +} diff --git a/crates/cli/tests/coverage/daemon/worker_tls_tests.rs b/crates/cli/tests/coverage/daemon/worker_tls_tests.rs new file mode 100644 index 000000000..90bdca306 --- /dev/null +++ b/crates/cli/tests/coverage/daemon/worker_tls_tests.rs @@ -0,0 +1,198 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::convert::Infallible; +use std::sync::Arc; + +use bytes::Bytes; +use http::{Request, Response, StatusCode}; +use http_body_util::{BodyExt as _, Empty, Full}; +use hyper::body::Incoming; +use hyper::service::service_fn; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use hyper_util::server::conn::auto::Builder; +use tokio::net::TcpListener; +use tokio_rustls::TlsAcceptor; + +use super::{WorkerClientPool, WorkerTlsIdentity, pooled_worker_tls_client}; +use crate::daemon::common::transport::box_body; + +#[test] +fn generates_a_daemon_pinnable_worker_identity() { + let identity = WorkerTlsIdentity::generate("127.0.0.1").expect("generate worker identity"); + assert!(!identity.root_certificate().is_empty()); + pooled_worker_tls_client(identity.root_certificate()).expect("build pinned worker client"); + WorkerTlsIdentity::generate("worker.example.com").expect("hostname SAN identity"); +} + +#[test] +fn rejects_invalid_or_non_concrete_worker_roots_and_hosts() { + assert!(WorkerTlsIdentity::generate("").is_err()); + assert!(WorkerTlsIdentity::generate("0.0.0.0").is_err()); + assert!(WorkerTlsIdentity::generate("::").is_err()); + assert!(WorkerTlsIdentity::generate("224.0.0.1").is_err()); + assert!(WorkerTlsIdentity::generate("ff02::1").is_err()); + assert!(pooled_worker_tls_client("").is_err()); + assert!(pooled_worker_tls_client("not_base64!").is_err()); +} + +#[test] +fn process_wide_worker_pool_reuses_cleartext_and_matching_tls_trust() { + let pool = WorkerClientPool::with_tls_capacity(2).expect("worker client pool"); + let first_cleartext = pool.client(None).expect("cleartext client"); + let second_cleartext = pool.client(None).expect("reused cleartext client"); + assert!(Arc::ptr_eq(&first_cleartext, &second_cleartext)); + + let first_identity = WorkerTlsIdentity::generate("127.0.0.1").expect("first identity"); + let second_identity = WorkerTlsIdentity::generate("127.0.0.1").expect("second identity"); + let first_tls = pool + .client(Some(first_identity.root_certificate())) + .expect("first TLS client"); + let first_tls_reused = pool + .client(Some(first_identity.root_certificate())) + .expect("reused first TLS client"); + let second_tls = pool + .client(Some(second_identity.root_certificate())) + .expect("isolated second TLS client"); + + assert!(Arc::ptr_eq(&first_tls, &first_tls_reused)); + assert!(!Arc::ptr_eq(&first_tls, &second_tls)); +} + +#[test] +fn worker_tls_pool_cache_is_bounded_and_evicts_least_recently_used_root() { + let pool = WorkerClientPool::with_tls_capacity(2).expect("worker client pool"); + let first_identity = WorkerTlsIdentity::generate("127.0.0.1").expect("first identity"); + let second_identity = WorkerTlsIdentity::generate("127.0.0.1").expect("second identity"); + let third_identity = WorkerTlsIdentity::generate("127.0.0.1").expect("third identity"); + + let first = pool + .client(Some(first_identity.root_certificate())) + .expect("first TLS client"); + let second = pool + .client(Some(second_identity.root_certificate())) + .expect("second TLS client"); + let first_reused = pool + .client(Some(first_identity.root_certificate())) + .expect("refresh first TLS client recency"); + assert!(Arc::ptr_eq(&first, &first_reused)); + + pool.client(Some(third_identity.root_certificate())) + .expect("third TLS client"); + assert_eq!(super::lock(&pool.tls).entries.len(), 2); + + let second_after_eviction = pool + .client(Some(second_identity.root_certificate())) + .expect("recreated second TLS client"); + assert!(!Arc::ptr_eq(&second, &second_after_eviction)); + assert_eq!(super::lock(&pool.tls).entries.len(), 2); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_matching_tls_lookups_share_one_pool() { + const LOOKUPS: usize = 32; + + let pool = Arc::new(WorkerClientPool::with_tls_capacity(2).expect("worker client pool")); + let identity = WorkerTlsIdentity::generate("127.0.0.1").expect("worker identity"); + let root = Arc::new(identity.root_certificate().to_owned()); + let barrier = Arc::new(tokio::sync::Barrier::new(LOOKUPS)); + let mut lookups = Vec::with_capacity(LOOKUPS); + + for _ in 0..LOOKUPS { + let pool = Arc::clone(&pool); + let root = Arc::clone(&root); + let barrier = Arc::clone(&barrier); + lookups.push(tokio::spawn(async move { + barrier.wait().await; + pool.client(Some(root.as_str())).expect("TLS client") + })); + } + + let first = lookups.remove(0).await.expect("first lookup task"); + for lookup in lookups { + let client = lookup.await.expect("lookup task"); + assert!(Arc::ptr_eq(&first, &client)); + } + assert_eq!(super::lock(&pool.tls).entries.len(), 1); +} + +#[tokio::test] +async fn daemon_pinned_client_completes_a_real_tls_worker_round_trip() { + let identity = WorkerTlsIdentity::generate("127.0.0.1").expect("generate worker identity"); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind TLS worker"); + let address = listener.local_addr().expect("worker address"); + let acceptor = TlsAcceptor::from(identity.server_config()); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept daemon"); + let stream = acceptor.accept(stream).await.expect("authenticate TLS"); + Builder::new(TokioExecutor::new()) + .serve_connection( + TokioIo::new(stream), + service_fn(|_request: Request| async { + let mut response = Response::new(Full::new(Bytes::from_static(b"ready"))); + *response.status_mut() = StatusCode::CREATED; + Ok::<_, Infallible>(response) + }), + ) + .await + .expect("serve pinned daemon request"); + }); + + let pool = WorkerClientPool::with_tls_capacity(2).expect("worker client pool"); + let client = pool + .client(Some(identity.root_certificate())) + .expect("pinned client"); + let request = Request::get(format!("https://127.0.0.1:{}/probe", address.port())) + .body(box_body(Empty::::new())) + .expect("request"); + let response = client.request(request).await.expect("TLS worker response"); + assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!( + response + .into_body() + .collect() + .await + .expect("response body") + .to_bytes(), + "ready" + ); + server.abort(); + assert!( + server + .await + .expect_err("server is stopped after the round trip") + .is_cancelled() + ); +} + +#[tokio::test] +async fn tls_pool_never_reuses_trust_across_different_roots() { + let serving_identity = + WorkerTlsIdentity::generate("127.0.0.1").expect("serving worker identity"); + let unrelated_identity = + WorkerTlsIdentity::generate("127.0.0.1").expect("unrelated worker identity"); + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind TLS worker"); + let address = listener.local_addr().expect("worker address"); + let acceptor = TlsAcceptor::from(serving_identity.server_config()); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.expect("accept daemon"); + // The client rejects this server before an HTTP request can be delivered. + let _ = acceptor.accept(stream).await; + }); + + let pool = WorkerClientPool::with_tls_capacity(2).expect("worker client pool"); + let unrelated_client = pool + .client(Some(unrelated_identity.root_certificate())) + .expect("unrelated pinned client"); + let request = Request::get(format!("https://127.0.0.1:{}/probe", address.port())) + .body(box_body(Empty::::new())) + .expect("request"); + assert!(unrelated_client.request(request).await.is_err()); + + server.abort(); + let _ = server.await; +} diff --git a/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs b/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs index abdf9ea5b..ca90bb90c 100644 --- a/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs +++ b/crates/cli/tests/coverage/shared/bootstrap_state_tests.rs @@ -23,6 +23,27 @@ fn owner_records_are_versioned_endpoint_scoped_and_round_trip() { assert_eq!(lock_name("not a url/with spaces"), "not_a_url_with_spaces"); } +#[test] +fn bootstrap_state_reports_malformed_recovery_and_non_directory_state_paths() { + let dir = tempfile::tempdir().unwrap(); + let url = "http://127.0.0.1:47632"; + std::fs::write(recovery_path(dir.path(), url), b"not-json").unwrap(); + let error = read_recovery(dir.path(), url).unwrap_err(); + assert!( + error.contains("failed to parse gateway recovery"), + "{error}" + ); + std::fs::remove_file(recovery_path(dir.path(), url)).unwrap(); + std::fs::create_dir(recovery_path(dir.path(), url)).unwrap(); + let error = read_recovery(dir.path(), url).unwrap_err(); + assert!(error.contains("failed to read gateway recovery"), "{error}"); + + let file = dir.path().join("not-a-directory"); + std::fs::write(&file, b"occupied").unwrap(); + let error = create_private_dir(&file).unwrap_err(); + assert!(error.contains("failed to create"), "{error}"); +} + #[cfg(any(target_os = "linux", target_os = "macos", windows))] #[test] fn live_owner_record_uses_a_process_instance_identity() { diff --git a/crates/cli/tests/coverage/shared/config_tests.rs b/crates/cli/tests/coverage/shared/config_tests.rs index 82d01d9d4..cf2aa0d34 100644 --- a/crates/cli/tests/coverage/shared/config_tests.rs +++ b/crates/cli/tests/coverage/shared/config_tests.rs @@ -361,6 +361,167 @@ impl Drop for PluginConfigDiscoveryScope { } } +#[tokio::test] +async fn managed_worker_configuration_ignores_user_files_and_non_secret_environment() { + let _guard = crate::test_support::PLUGIN_CONFIG_TEST_LOCK.lock().await; + let temp = tempfile::tempdir().unwrap(); + let user_home = temp.path().join("user"); + let user_directory = user_home.join("nemo-relay"); + let system_directory = temp.path().join("system"); + std::fs::create_dir_all(&user_directory).unwrap(); + std::fs::create_dir_all(&system_directory).unwrap(); + std::fs::write( + user_directory.join("config.toml"), + "invalid user config = [", + ) + .unwrap(); + std::fs::write( + user_directory.join(PLUGINS_TOML), + "invalid user plugins = [", + ) + .unwrap(); + + let system_config = system_directory.join("config.toml"); + let system_plugins = system_directory.join(PLUGINS_TOML); + std::fs::write( + &system_config, + r#" +[gateway] +max_hook_payload_bytes = 1234 +max_passthrough_body_bytes = 5678 + +[upstream] +openai_base_url = "https://admin.example/openai" +openai_auth_header = "Bearer admin-file-openai" +anthropic_base_url = "https://admin.example/anthropic" +anthropic_auth_header = "Basic admin-file-anthropic" +"#, + ) + .unwrap(); + std::fs::write(&system_plugins, "version = 1\ncomponents = []\n").unwrap(); + + let _environment = crate::test_support::EnvScope::set(&[ + ("XDG_CONFIG_HOME", Some(user_home.as_os_str())), + ("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", None), + ( + "NEMO_RELAY_GATEWAY_BIND", + Some(std::ffi::OsStr::new("0.0.0.0:9999")), + ), + ( + "NEMO_RELAY_OPENAI_BASE_URL", + Some(std::ffi::OsStr::new("https://user.example/openai")), + ), + ( + "NEMO_RELAY_ANTHROPIC_BASE_URL", + Some(std::ffi::OsStr::new("https://user.example/anthropic")), + ), + ( + "NEMO_RELAY_MAX_HOOK_PAYLOAD_BYTES", + Some(std::ffi::OsStr::new("1")), + ), + ( + "NEMO_RELAY_MAX_PASSTHROUGH_BODY_BYTES", + Some(std::ffi::OsStr::new("2")), + ), + ( + "NEMO_RELAY_OPENAI_AUTH_HEADER", + Some(std::ffi::OsStr::new("Bearer managed-env-openai")), + ), + ( + "NEMO_RELAY_ANTHROPIC_AUTH_HEADER", + Some(std::ffi::OsStr::new("Basic managed-env-anthropic")), + ), + ]); + + let managed = + resolve_managed_worker_config_from_paths(system_config, system_plugins.clone()).unwrap(); + + assert_eq!(managed.plugin_config_path, system_plugins); + assert_eq!( + managed.resolved.gateway.openai_base_url, + "https://admin.example/openai" + ); + assert_eq!( + managed.resolved.gateway.anthropic_base_url, + "https://admin.example/anthropic" + ); + assert_eq!( + managed.resolved.gateway.openai_auth_header.as_deref(), + Some("Bearer managed-env-openai") + ); + assert_eq!( + managed.resolved.gateway.anthropic_auth_header.as_deref(), + Some("Basic managed-env-anthropic") + ); + assert_eq!( + managed.resolved.gateway.bind, + "127.0.0.1:4040".parse().unwrap() + ); + assert_eq!(managed.resolved.gateway.max_hook_payload_bytes, 1234); + assert_eq!(managed.resolved.gateway.max_passthrough_body_bytes, 5678); + assert_eq!( + managed.resolved.gateway.plugin_config, + Some(json!({ "components": [], "version": 1 })) + ); + assert!(managed.resolved.dynamic_plugins.is_empty()); +} + +#[tokio::test] +async fn managed_worker_dynamic_plugin_loading_ignores_user_lifecycle_state() { + let _guard = crate::test_support::PLUGIN_CONFIG_TEST_LOCK.lock().await; + let temp = tempfile::tempdir().unwrap(); + let user_home = temp.path().join("user"); + let user_directory = user_home.join("nemo-relay"); + let system_directory = temp.path().join("system"); + std::fs::create_dir_all(&user_directory).unwrap(); + std::fs::create_dir_all(&system_directory).unwrap(); + let _environment = crate::test_support::EnvScope::set(&[ + ("XDG_CONFIG_HOME", Some(user_home.as_os_str())), + ("NEMO_RELAY_TEST_SKIP_IMPLICIT_CONFIG", None), + ("NEMO_RELAY_OPENAI_AUTH_HEADER", None), + ("NEMO_RELAY_ANTHROPIC_AUTH_HEADER", None), + ]); + std::fs::write( + user_directory.join(".dynamic-plugins.json"), + "invalid user lifecycle state", + ) + .unwrap(); + + let system_config = system_directory.join("config.toml"); + let system_plugins = system_directory.join(PLUGINS_TOML); + let plugin_directory = system_directory.join("plugins/worker"); + std::fs::create_dir_all(&plugin_directory).unwrap(); + let manifest_path = write_dynamic_manifest(&plugin_directory, "acme.managed-worker"); + let (manifest, manifest_ref) = DynamicPluginManifest::load_from_path(&manifest_path).unwrap(); + let mut record = manifest.into_record(Some(manifest_ref)).unwrap(); + record.spec.enabled = true; + let environment = system_directory.join("plugin-environment"); + write_attested_python_environment(&environment, &manifest_path); + record.source.environment_ref = Some(environment.to_string_lossy().into_owned()); + std::fs::write( + system_directory.join(".dynamic-plugins.json"), + serde_json::to_vec_pretty(&json!({ + "schema_version": 1, + "records": [record], + })) + .unwrap(), + ) + .unwrap(); + std::fs::write(&system_config, "").unwrap(); + std::fs::write( + &system_plugins, + "version = 1\ncomponents = []\n\n[[plugins.dynamic]]\nmanifest = \"plugins/worker/relay-plugin.toml\"\n\n[plugins.policy.overrides.\"acme.managed-worker\"]\nattestation = \"integrity_only\"\n", + ) + .unwrap(); + let managed = resolve_managed_worker_config_from_paths(system_config, system_plugins).unwrap(); + let active = + active_dynamic_plugin_components(Some(&managed.plugin_config_path), &managed.resolved) + .unwrap(); + + assert_eq!(active.len(), 1); + assert_eq!(active[0].plugin_id, "acme.managed-worker"); +} + #[cfg(feature = "__skip-implicit-config")] #[test] fn test_hook_skips_implicit_config_but_retains_explicit_config_and_environment() { diff --git a/crates/cli/tests/coverage/shared/doctor_tests.rs b/crates/cli/tests/coverage/shared/doctor_tests.rs index 5eb015ca1..96fc77a8c 100644 --- a/crates/cli/tests/coverage/shared/doctor_tests.rs +++ b/crates/cli/tests/coverage/shared/doctor_tests.rs @@ -2,6 +2,416 @@ // SPDX-License-Identifier: Apache-2.0 use super::*; + +#[tokio::test] +async fn pi_managed_install_diagnostics_distinguish_current_and_stale_versions() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let agent_dir = tempfile::tempdir().unwrap(); + let _environment = crate::test_support::EnvScope::set(&[ + ( + crate::agents::pi::doctor::PI_AGENT_DIR_ENV, + Some(agent_dir.path().as_os_str()), + ), + (crate::agents::pi::launch::PI_EXTENSION_PATH_ENV, None), + ]); + crate::agents::pi::install::install(crate::installation::InstallRequest { + install_dir: None, + force: false, + dry_run: false, + skip_doctor: true, + }) + .unwrap(); + + let current = pi_managed_install_check().unwrap(); + assert_eq!(current.status, Status::Pass); + assert!(current.details.contains(env!("CARGO_PKG_VERSION"))); + let root = crate::agents::pi::install::install_root().unwrap(); + let state_path = root.join(".nemo-relay-install.json"); + let mut state: serde_json::Value = + serde_json::from_slice(&std::fs::read(&state_path).unwrap()).unwrap(); + state["relay_version"] = serde_json::json!("0.0.0-stale"); + std::fs::write(&state_path, serde_json::to_vec(&state).unwrap()).unwrap(); + + let stale = pi_managed_install_check().unwrap(); + assert_eq!(stale.status, Status::Warn); + assert!(stale.details.contains("0.0.0-stale")); + assert!(stale.details.contains("nemo-relay install pi")); +} + +#[tokio::test] +async fn pi_extension_diagnostics_warn_for_a_project_scoped_copy() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let agent_dir = tempfile::tempdir().unwrap(); + let project = tempfile::tempdir().unwrap(); + let _environment = crate::test_support::EnvScope::set(&[ + ( + crate::agents::pi::doctor::PI_AGENT_DIR_ENV, + Some(agent_dir.path().as_os_str()), + ), + (crate::agents::pi::launch::PI_EXTENSION_PATH_ENV, None), + ]); + let extension = project.path().join(".pi/extensions/nemo-relay"); + std::fs::create_dir_all(&extension).unwrap(); + std::fs::write( + extension.join("package.json"), + r#"{"name":"nemo-relay-pi","pi":{"extensions":["./index.ts"]}}"#, + ) + .unwrap(); + std::fs::write(extension.join("index.ts"), "export default 1").unwrap(); + + let check = pi_extension_trust_check(project.path()); + assert_eq!(check.status, Status::Warn); + assert!( + check.details.contains("project-scoped"), + "{}", + check.details + ); + assert!( + check.details.contains("silently skipped"), + "{}", + check.details + ); +} + +#[tokio::test] +async fn pi_extension_diagnostics_distinguish_disabled_and_undecidable_settings_filters() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let temp = tempfile::tempdir().unwrap(); + let agent_dir = temp.path().join("agent"); + let checkout = temp.path().join("checkout"); + std::fs::create_dir_all(&agent_dir).unwrap(); + std::fs::create_dir_all(&checkout).unwrap(); + std::fs::write( + checkout.join("package.json"), + r#"{"name":"nemo-relay-pi","pi":{"extensions":["./index.ts"]}}"#, + ) + .unwrap(); + std::fs::write(checkout.join("index.ts"), "export default 1").unwrap(); + let _environment = crate::test_support::EnvScope::set(&[ + ( + crate::agents::pi::doctor::PI_AGENT_DIR_ENV, + Some(agent_dir.as_os_str()), + ), + (crate::agents::pi::launch::PI_EXTENSION_PATH_ENV, None), + ]); + + std::fs::write( + agent_dir.join("settings.json"), + r#"{"packages":[{"source":"../checkout","autoload":false}]}"#, + ) + .unwrap(); + let disabled = pi_extension_trust_check(temp.path()); + assert_eq!(disabled.status, Status::Warn); + assert!( + disabled.details.contains("filtered off"), + "{}", + disabled.details + ); + + std::fs::write( + agent_dir.join("settings.json"), + r#"{"packages":[{"source":"../checkout","extensions":["src/*.ts"]}]}"#, + ) + .unwrap(); + let undecidable = pi_extension_trust_check(temp.path()); + assert_eq!(undecidable.status, Status::Warn); + assert!( + undecidable.details.contains("cannot be decided"), + "{}", + undecidable.details + ); +} + +#[test] +fn loopback_detection_accepts_local_hosts_and_rejects_remote_or_invalid_urls() { + for url in [ + "http://localhost:8080", + "http://127.0.0.1:8080", + "http://[::1]:8080", + ] { + assert!(is_loopback(url), "expected loopback: {url}"); + } + for url in [ + "https://relay.example.com:443", + "not-a-url", + "http://192.0.2.1:80", + ] { + assert!(!is_loopback(url), "expected non-loopback: {url}"); + } +} + +#[tokio::test] +async fn offline_pi_gateway_check_is_informational_without_network_io() { + let check = + pi_gateway_reachability_check(DoctorProbeMode::Offline, &ResolvedConfig::default()).await; + assert_eq!(check.name, "pi gateway reachability"); + assert_eq!(check.status, Status::Info); + assert!(check.details.contains("live reachability probe skipped")); +} + +#[tokio::test] +async fn live_pi_gateway_check_distinguishes_foreign_and_unavailable_listeners() { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let server = std::thread::spawn(move || { + let mut stream = accept_bounded(&listener); + let _ = read_headers(&mut stream); + stream + .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}") + .unwrap(); + }); + let mut resolved = ResolvedConfig::default(); + resolved.gateway.bind = address; + let foreign = pi_gateway_reachability_check(DoctorProbeMode::Live, &resolved).await; + assert_eq!(foreign.status, Status::Warn); + assert!(foreign.details.contains("not a NeMo Relay gateway")); + server.join().unwrap(); + + let unavailable = TcpListener::bind("127.0.0.1:0").unwrap(); + resolved.gateway.bind = unavailable.local_addr().unwrap(); + drop(unavailable); + let unavailable = pi_gateway_reachability_check(DoctorProbeMode::Live, &resolved).await; + assert_eq!(unavailable.status, Status::Warn); + assert!(unavailable.details.contains("not answering")); +} + +#[tokio::test] +async fn live_pi_gateway_check_distinguishes_compatible_and_incompatible_relay_versions() { + for (status, body_status, expected) in [ + ("200 OK", "ok", Status::Pass), + ("409 Conflict", "incompatible", Status::Warn), + ] { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let address = listener.local_addr().unwrap(); + let server = std::thread::spawn(move || { + let mut stream = accept_bounded(&listener); + let _ = read_headers(&mut stream); + let body = json!({ + "status": body_status, + "service": "nemo-relay", + "bootstrap_protocol": crate::bootstrap::BOOTSTRAP_PROTOCOL_VERSION, + "instance_id": "doctor-fixture" + }) + .to_string(); + stream + .write_all( + format!( + "HTTP/1.1 {status}\r\nContent-Length: {}\r\n\r\n{body}", + body.len() + ) + .as_bytes(), + ) + .unwrap(); + }); + let mut resolved = ResolvedConfig::default(); + resolved.gateway.bind = address; + let check = pi_gateway_reachability_check(DoctorProbeMode::Live, &resolved).await; + assert_eq!(check.status, expected, "{}", check.details); + server.join().unwrap(); + } +} + +#[test] +fn agent_version_diagnostics_separate_unverified_optional_and_required_failures() { + let mut status = Status::Pass; + let mut details = Vec::new(); + apply_agent_version_status( + CodingAgent::Pi, + Some("0.85.0"), + true, + false, + &mut status, + &mut details, + ); + assert_eq!(status, Status::Warn); + assert!(details.iter().any(|detail| detail.contains("0.84"))); + + let mut status = Status::Pass; + let mut details = Vec::new(); + apply_agent_version_status( + CodingAgent::Pi, + Some("invalid"), + true, + true, + &mut status, + &mut details, + ); + assert_eq!(status, Status::Fail); + assert!( + details + .iter() + .any(|detail| detail.contains("could not parse")) + ); +} + +#[tokio::test] +async fn remote_pi_gateway_probe_reports_success_and_http_failure() { + for (status_line, expected_status, expected_detail) in [ + ("200 OK", Status::Pass, "answered /healthz"), + ("503 Service Unavailable", Status::Warn, "HTTP 503"), + ] { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + let server = std::thread::spawn(move || { + let mut stream = accept_bounded(&listener); + let _ = read_headers(&mut stream); + stream + .write_all( + format!("HTTP/1.1 {status_line}\r\nContent-Length: 0\r\n\r\n").as_bytes(), + ) + .unwrap(); + }); + let gateway_url = format!("http://127.0.0.1:{port}"); + let check = remote_pi_gateway_reachability_check(&gateway_url).await; + assert_eq!(check.status, expected_status, "{}", check.details); + assert!(check.details.contains(expected_detail), "{}", check.details); + server.join().unwrap(); + } +} + +#[tokio::test] +async fn atof_header_validation_covers_all_authored_shape_errors() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let _environment = EnvScope::set(&[ + ("NEMO_RELAY_MISSING_ATOF_HEADER", None), + ( + "NEMO_RELAY_BLANK_ATOF_HEADER", + Some(std::ffi::OsStr::new(" ")), + ), + ( + "NEMO_RELAY_VALID_ATOF_HEADER", + Some(std::ffi::OsStr::new("secret")), + ), + ]); + for (endpoint, expected) in [ + ( + serde_json::json!({"headers": {"bad header": "value"}}), + "invalid HTTP header name", + ), + ( + serde_json::json!({"headers": {"x-test": "bad\nvalue"}}), + "headers.x-test invalid", + ), + ( + serde_json::json!({"header_env": []}), + "header_env must be an object", + ), + ( + serde_json::json!({"header_env": {"x-test": 1}}), + "must be a string", + ), + ( + serde_json::json!({"header_env": {"bad header": "PATH"}}), + "invalid HTTP header name", + ), + ( + serde_json::json!({"header_env": {"x-test": "NEMO_RELAY_MISSING_ATOF_HEADER"}}), + "is not set", + ), + ( + serde_json::json!({"header_env": {"x-test": "NEMO_RELAY_BLANK_ATOF_HEADER"}}), + "is blank", + ), + ] { + let error = endpoint_headers(&endpoint).expect_err("invalid endpoint headers"); + assert!(error.contains(expected), "{error}"); + } + assert_eq!( + endpoint_headers(&serde_json::json!({ + "headers": {"x-static": "literal"}, + "header_env": {"x-secret": "NEMO_RELAY_VALID_ATOF_HEADER"} + })) + .unwrap(), + [ + ("x-static".into(), "literal".into()), + ("x-secret".into(), "secret".into()) + ] + ); + let duplicate = endpoint_headers(&serde_json::json!({ + "headers": {"x-secret": "literal"}, + "header_env": {"x-secret": "NEMO_RELAY_VALID_ATOF_HEADER"} + })) + .expect_err("duplicate static and environment headers must be rejected"); + assert_eq!( + duplicate, + "header \"x-secret\" cannot appear in both headers and header_env" + ); +} + +#[test] +fn atof_probe_target_validation_rejects_malformed_hosts_schemes_and_transports() { + assert!(validate_atof_stream_probe_target(1, "http_post", "not a URL").is_err()); + assert!(validate_atof_stream_probe_target(2, "http_post", "file:///tmp/events").is_err()); + assert!(validate_atof_stream_probe_target(3, "http_post", "ws://example.com/events").is_err()); + assert!( + validate_atof_stream_probe_target(4, "websocket", "http://example.com/events").is_err() + ); + assert!(validate_atof_stream_probe_target(5, "grpc", "https://example.com/events").is_err()); + assert!(validate_atof_stream_probe_target(6, "ndjson", "https://example.com/events").is_ok()); +} + +#[tokio::test] +async fn pi_extension_load_path_reports_missing_user_and_project_scopes() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let directory = tempfile::tempdir().unwrap(); + let agent_dir = directory.path().join("agent"); + std::fs::create_dir_all(&agent_dir).unwrap(); + let _environment = EnvScope::set(&[ + (crate::agents::pi::launch::PI_EXTENSION_PATH_ENV, None), + ( + crate::agents::pi::doctor::PI_AGENT_DIR_ENV, + Some(agent_dir.as_os_str()), + ), + ]); + let missing = pi_extension_trust_check(directory.path()); + assert_eq!(missing.status, Status::Info); + + let user = agent_dir.join("extensions/nemo-relay"); + write_doctor_pi_package(&user); + let installed = pi_extension_trust_check(directory.path()); + assert_eq!(installed.status, Status::Pass); + assert!(installed.details.contains("user scope")); + + std::fs::remove_dir_all(&user).unwrap(); + let project = directory.path().join(".pi/extensions/nemo-relay"); + write_doctor_pi_package(&project); + let project = pi_extension_trust_check(directory.path()); + assert_eq!(project.status, Status::Warn); + assert!(project.details.contains("project-scoped")); +} + +fn write_doctor_pi_package(path: &std::path::Path) { + std::fs::create_dir_all(path).unwrap(); + std::fs::write( + path.join("package.json"), + r#"{"name":"nemo-relay-pi","pi":{"extensions":["./index.ts"]}}"#, + ) + .unwrap(); + std::fs::write(path.join("index.ts"), "export default 1").unwrap(); +} + +#[tokio::test] +async fn atof_websocket_probe_sends_payload_and_reports_success() { + use futures_util::StreamExt as _; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("ws://{}/events", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + socket.next().await.unwrap().unwrap().into_text().unwrap() + }); + let check = probe_atof_websocket( + &url, + vec![("x-observer".into(), "doctor".into())], + "{\"probe\":true}".into(), + std::time::Duration::from_secs(2), + 42, + ) + .await; + assert_eq!(check.status, Status::Pass); + assert_eq!(server.await.unwrap(), "{\"probe\":true}"); +} use std::io::{Read, Write}; use std::net::TcpListener; use std::path::PathBuf; @@ -9,7 +419,7 @@ use std::sync::{Arc, Mutex}; use crate::configuration::{GatewayConfig, ResolvedConfig, ResolvedDynamicPluginConfig}; use crate::server::GatewayOverrides; -use crate::test_support::{EnvScope, accept_bounded, read_headers}; +use crate::test_support::{EnvScope, PLUGIN_CONFIG_TEST_LOCK, accept_bounded, read_headers}; fn start_doctor_http_capture_server() -> (String, Arc>, std::thread::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); @@ -199,6 +609,90 @@ fn exit_code_fails_when_agent_readiness_fails() { assert_eq!(exit_code(&report), 1); } +#[tokio::test] +async fn managed_bundle_doctor_report_is_managed_only() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + use base64::Engine; + use std::ffi::OsStr; + + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().join("bundle"); + let spec = crate::daemon::managed::ManagedBundleSpec::new( + "https://relay.example.com:443", + "/opt/nvidia/bin/nemo-relay-dispatch", + crate::daemon::managed::ManagedPlatform::Linux, + [crate::daemon::managed::ManagedAgent::Pi], + ) + .unwrap(); + let digest = crate::daemon::managed::write_new_bundle(&root, &spec).unwrap(); + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x42_u8; 32]); + let _environment = EnvScope::set(&[( + crate::daemon::common::state::ROUTE_TOKEN_ENV, + Some(OsStr::new(&token)), + )]); + + let report = collect_managed_bundle_report(&root, &digest); + assert_eq!(report.managed_bundle.status, Status::Pass); + let human = format_managed_bundle_human(&report); + assert!(human.contains("Managed bundle validation passed")); + let json: serde_json::Value = + serde_json::from_str(&format_managed_bundle_json(&report).unwrap()).unwrap(); + assert_eq!(json["managed_bundle"]["status"], "pass"); + for personal_section in [ + "environment", + "configuration", + "agents", + "host_plugins", + "observability", + "completions", + ] { + assert!(json.get(personal_section).is_none(), "{personal_section}"); + } +} + +#[tokio::test] +async fn managed_bundle_doctor_reports_digest_failures_and_both_output_modes() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + use base64::Engine; + use std::ffi::OsStr; + + let token = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([0x43_u8; 32]); + let _environment = EnvScope::set(&[( + crate::daemon::common::state::ROUTE_TOKEN_ENV, + Some(OsStr::new(&token)), + )]); + let directory = tempfile::tempdir().unwrap(); + let root = directory.path().join("bundle"); + let spec = crate::daemon::managed::ManagedBundleSpec::new( + "https://relay.example.com:443", + "/opt/nvidia/bin/nemo-relay-dispatch", + crate::daemon::managed::ManagedPlatform::Linux, + [crate::daemon::managed::ManagedAgent::Codex], + ) + .unwrap(); + let digest = crate::daemon::managed::write_new_bundle(&root, &spec).unwrap(); + assert_eq!( + run_managed_bundle_doctor(&root, &digest, true).unwrap(), + std::process::ExitCode::SUCCESS + ); + assert_eq!( + run_managed_bundle_doctor(&root, &digest, false).unwrap(), + std::process::ExitCode::SUCCESS + ); + + let wrong: crate::daemon::managed::ManagedBundleDigest = + "0000000000000000000000000000000000000000000000000000000000000000" + .parse() + .unwrap(); + let report = collect_managed_bundle_report(&root, &wrong); + assert_eq!(report.managed_bundle.status, Status::Fail); + assert!(report.managed_bundle.artifact_count.is_none()); + assert_eq!( + run_managed_bundle_doctor(&root, &wrong, true).unwrap(), + std::process::ExitCode::FAILURE + ); +} + #[test] fn exit_code_fails_when_an_installed_host_plugin_is_unready() { let mut report = empty_report(); @@ -476,6 +970,155 @@ fn format_json_reports_discovered_dynamic_plugin_fields() { assert_eq!(plugin["host_config_status"], "present"); } +#[test] +fn dynamic_plugin_diagnostic_checks_distinguish_registry_and_host_config_sources() { + let registry = DynamicPluginReferenceInfo { + plugin_id: "acme.registry".into(), + manifest_ref: "/opt/acme/relay-plugin.toml".into(), + source: PathBuf::from("/opt/acme/plugins.toml"), + host_config_status: DynamicPluginHostConfigStatus::Absent, + }; + let check = dynamic_plugin_reference_check(®istry); + assert_eq!(check.status, Status::Pass); + assert!(check.details.contains("acme.registry")); + assert!( + dynamic_plugin_host_config_check(®istry) + .details + .contains("host config only") + ); + + let present = DynamicPluginReferenceInfo { + host_config_status: DynamicPluginHostConfigStatus::Present, + ..registry + }; + let check = dynamic_plugin_host_config_check(&present); + assert_eq!(check.status, Status::Info); + assert!(check.details.contains("host-owned config present")); +} + +#[tokio::test] +async fn collect_report_preserves_configuration_and_plugin_resolution_failures() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let directory = tempfile::tempdir().unwrap(); + let missing_config = directory.path().join("missing-config.toml"); + let report = collect_report( + Some(CodingAgent::Codex), + DoctorProbeMode::Offline, + &GatewayOverrides { + config: Some(missing_config), + ..GatewayOverrides::default() + }, + ) + .await + .expect("diagnostic report"); + assert_eq!(report.target_agent.as_deref(), Some("codex")); + assert_eq!(report.configuration.resolution.status, Status::Fail); + assert_eq!( + report.configuration.upstream_auth.openai, + SecretPresence::Unknown + ); + + let valid_config = directory.path().join("config.toml"); + std::fs::write(&valid_config, "[upstream]\n").unwrap(); + let invalid_plugins = directory.path().join("invalid-plugins.toml"); + std::fs::write(&invalid_plugins, "[[invalid").unwrap(); + let report = collect_report( + None, + DoctorProbeMode::Offline, + &GatewayOverrides { + config: Some(valid_config), + plugin_config_path: Some(invalid_plugins), + ..GatewayOverrides::default() + }, + ) + .await + .expect("plugin diagnostic report"); + assert_eq!(report.configuration.plugin_resolution.status, Status::Fail); +} + +#[tokio::test] +async fn component_diagnostics_cover_disabled_malformed_and_explicit_sink_configuration() { + let disabled_cache: PluginConfig = serde_json::from_value(json!({ + "version": 1, + "components": [{ + "kind": "adaptive", + "enabled": false, + "config": {"response_cache": {"ttl_seconds": 60}} + }] + })) + .unwrap(); + let mut checks = Vec::new(); + collect_response_cache_component_checks( + &mut checks, + &disabled_cache, + false, + DoctorProbeMode::Offline, + ) + .await; + assert_eq!(checks[0].status, Status::Info); + assert!(checks[0].details.contains("adaptive plugin disabled")); + + let malformed_cache: PluginConfig = serde_json::from_value(json!({ + "version": 1, + "components": [{ + "kind": "adaptive", + "config": {"response_cache": "not-an-object"} + }] + })) + .unwrap(); + let mut checks = Vec::new(); + collect_response_cache_component_checks( + &mut checks, + &malformed_cache, + false, + DoctorProbeMode::Offline, + ) + .await; + assert_eq!(checks[0].status, Status::Fail); + assert!(checks[0].details.contains("invalid response_cache config")); + + let malformed_pricing: PluginConfig = serde_json::from_value(json!({ + "version": 1, + "components": [{"kind": "pricing", "config": {"sources": "invalid"}}] + })) + .unwrap(); + let mut checks = Vec::new(); + collect_pricing_component_checks(&mut checks, &malformed_pricing); + assert_eq!(checks[0].status, Status::Fail); + assert!(checks[0].details.contains("invalid config")); + + let directory = tempfile::tempdir().unwrap(); + let file_checks = observability_atof_file_checks(&json!({ + "atof": { + "enabled": true, + "sinks": [{"type": "file", "output_directory": directory.path()}] + } + })); + assert_eq!(file_checks[0].status, Status::Pass); + assert!(file_checks[0].details.starts_with("sinks[0]")); + + let bad_scheme = probe_atof_websocket( + "http://127.0.0.1/events", + Vec::new(), + "{}".into(), + Duration::from_millis(10), + 11, + ) + .await; + assert_eq!(bad_scheme.status, Status::Fail); + assert!(bad_scheme.details.contains("must be ws or wss")); +} + +#[test] +fn diagnostic_layer_status_rejects_toml_syntax_before_shape_validation() { + let directory = tempfile::tempdir().unwrap(); + let invalid = directory.path().join("invalid.toml"); + std::fs::write(&invalid, "[[broken").unwrap(); + let layer = layer_status(&invalid); + assert_eq!(layer.status, Status::Fail); + assert!(layer.details.contains("invalid TOML")); +} + #[test] fn check_dir_writable_does_not_create_missing_dir() { let temp = tempfile::tempdir().unwrap(); diff --git a/crates/cli/tests/coverage/shared/filesystem_snapshot_tests.rs b/crates/cli/tests/coverage/shared/filesystem_snapshot_tests.rs new file mode 100644 index 000000000..c24d25b2f --- /dev/null +++ b/crates/cli/tests/coverage/shared/filesystem_snapshot_tests.rs @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +#[test] +fn backup_is_create_only_and_removal_is_idempotent() { + let directory = tempfile::tempdir().unwrap(); + let path = directory.path().join("settings.json"); + assert_eq!( + backup_path(&path), + directory.path().join("settings.json.nemo-relay.bak") + ); + assert_eq!( + backup_path(&directory.path().join("settings")), + directory.path().join("settings.nemo-relay.bak") + ); + backup(&path).unwrap(); + assert!(!backup_path(&path).exists()); + + std::fs::write(&path, b"original").unwrap(); + backup(&path).unwrap(); + std::fs::write(&path, b"changed").unwrap(); + backup(&path).unwrap(); + assert_eq!(std::fs::read(backup_path(&path)).unwrap(), b"original"); + remove_backup(&path).unwrap(); + remove_backup(&path).unwrap(); +} + +#[test] +fn optional_snapshots_restore_existing_and_missing_files() { + let directory = tempfile::tempdir().unwrap(); + let existing = directory.path().join("existing"); + std::fs::write(&existing, b"before").unwrap(); + let snapshot = snapshot_optional_file(&existing).unwrap(); + std::fs::write(&existing, b"after").unwrap(); + restore_file_snapshot(&snapshot).unwrap(); + assert_eq!(std::fs::read(&existing).unwrap(), b"before"); + + let missing = directory.path().join("missing"); + let snapshot = snapshot_optional_file(&missing).unwrap(); + std::fs::write(&missing, b"created later").unwrap(); + restore_file_snapshot(&snapshot).unwrap(); + assert!(!missing.exists()); + restore_file_snapshot(&snapshot).unwrap(); +} + +#[cfg(unix)] +#[test] +fn symlink_preserving_operations_update_targets_and_restore_links() { + let directory = tempfile::tempdir().unwrap(); + let target = directory.path().join("target"); + let link = directory.path().join("nested/link"); + std::fs::write(&target, b"initial").unwrap(); + ensure_symlink_path(&link, &target).unwrap(); + assert_eq!(std::fs::read_link(&link).unwrap(), target); + + atomic_write_preserving_symlink(&link, b"updated").unwrap(); + assert_eq!(std::fs::read(&target).unwrap(), b"updated"); + let snapshot = snapshot_optional_file(&link).unwrap(); + std::fs::remove_file(&link).unwrap(); + std::fs::write(&link, b"replacement file").unwrap(); + restore_file_snapshot(&snapshot).unwrap(); + assert!( + std::fs::symlink_metadata(&link) + .unwrap() + .file_type() + .is_symlink() + ); + assert_eq!(std::fs::read(&target).unwrap(), b"updated"); + + remove_file_preserving_symlink(&link).unwrap(); + assert!(!target.exists()); + assert!(link.is_symlink()); + remove_file_preserving_symlink(&link).unwrap(); +} + +#[cfg(unix)] +#[test] +fn relative_symlink_chains_are_resolved_from_each_parent() { + let directory = tempfile::tempdir().unwrap(); + let target = directory.path().join("target"); + let middle = directory.path().join("middle"); + let link = directory.path().join("link"); + std::os::unix::fs::symlink("target", &middle).unwrap(); + std::os::unix::fs::symlink("middle", &link).unwrap(); + atomic_write_preserving_symlink(&link, b"chain").unwrap(); + assert_eq!(std::fs::read(target).unwrap(), b"chain"); +} diff --git a/crates/cli/tests/coverage/shared/gateway_tests.rs b/crates/cli/tests/coverage/shared/gateway_tests.rs index f96515eb4..33776cfea 100644 --- a/crates/cli/tests/coverage/shared/gateway_tests.rs +++ b/crates/cli/tests/coverage/shared/gateway_tests.rs @@ -1931,6 +1931,90 @@ fn chatgpt_jwt_routes_to_chatgpt_backend_when_no_api_key() { ); } +#[test] +fn daemon_generic_openai_paths_keep_the_administrator_upstream() { + let mut headers = HeaderMap::new(); + headers.insert( + "authorization", + HeaderValue::from_static("Bearer at-caller-controlled-token"), + ); + let config = GatewayConfig { + openai_base_url: "https://administrator.example/v1".into(), + ..GatewayConfig::default() + }; + + assert_eq!( + daemon_provider_upstream_url(&headers, "/responses?client=pi", &config) + .unwrap() + .as_deref(), + Some("https://administrator.example/v1/responses?client=pi") + ); + assert_eq!( + daemon_provider_upstream_url( + &headers, + "/backend-api/codex/responses?client=codex", + &config, + ) + .unwrap() + .as_deref(), + Some("https://chatgpt.com/backend-api/codex/responses?client=codex") + ); +} + +#[test] +fn authenticated_daemon_pi_route_uses_the_exact_named_provider_endpoint() { + let mut headers = HeaderMap::new(); + headers.insert( + crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER, + HeaderValue::from_static("https://custom.example/inference/v1"), + ); + assert!( + daemon_provider_upstream_url( + &headers, + "/chat/completions?client=pi", + &GatewayConfig::default(), + ) + .unwrap_err() + .to_string() + .contains("caller-supplied provider credential") + ); + headers.insert( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer caller-provider-token"), + ); + + assert_eq!( + daemon_provider_upstream_url( + &headers, + "/chat/completions?client=pi", + &GatewayConfig::default(), + ) + .unwrap() + .as_deref(), + Some("https://custom.example/inference/v1/chat/completions?client=pi") + ); + + headers.insert( + crate::agents::pi::alignment::UPSTREAM_BASE_URL_HEADER, + HeaderValue::from_static("http://remote.example/v1"), + ); + assert!( + daemon_provider_upstream_url(&headers, "/chat/completions", &GatewayConfig::default(),) + .is_err() + ); + + let config = GatewayConfig { + openai_auth_header: Some("Bearer administrator-token".into()), + ..GatewayConfig::default() + }; + let forwarded = daemon_provider_forward_headers(&headers, "/chat/completions", &config) + .expect("provider route"); + assert_eq!( + forwarded.get(http::header::AUTHORIZATION), + Some(&HeaderValue::from_static("Bearer caller-provider-token")) + ); +} + #[test] fn provider_key_does_not_trigger_chatgpt_backend() { let mut headers = HeaderMap::new(); diff --git a/crates/cli/tests/coverage/shared/installer_tests.rs b/crates/cli/tests/coverage/shared/installer_tests.rs index f7ae88d74..d3ebfccd9 100644 --- a/crates/cli/tests/coverage/shared/installer_tests.rs +++ b/crates/cli/tests/coverage/shared/installer_tests.rs @@ -11,6 +11,14 @@ use serde_json::Value; use crate::agents::CodingAgent; +#[test] +fn null_hook_configuration_normalizes_to_an_empty_object() { + assert_eq!( + super::merging::hook_config_root(serde_json::Value::Null).unwrap(), + json!({}) + ); +} + struct BootstrapConfigHome { _guard: std::sync::MutexGuard<'static, ()>, previous: Option, diff --git a/crates/cli/tests/coverage/shared/marketplace_state_tests.rs b/crates/cli/tests/coverage/shared/marketplace_state_tests.rs new file mode 100644 index 000000000..980976fff --- /dev/null +++ b/crates/cli/tests/coverage/shared/marketplace_state_tests.rs @@ -0,0 +1,76 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::agents::CodingAgent; +use crate::test_support::{EnvScope, PLUGIN_CONFIG_TEST_LOCK}; + +#[tokio::test] +async fn managed_integration_registry_round_trips_and_deduplicates_install_directories() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let config = tempfile::tempdir().unwrap(); + let install = tempfile::tempdir().unwrap(); + let _environment = EnvScope::set(&[("XDG_CONFIG_HOME", Some(config.path().as_os_str()))]); + assert!( + registered_install_dirs(CodingAgent::Codex) + .unwrap() + .is_empty() + ); + register_managed_integration(CodingAgent::Codex, install.path()).unwrap(); + register_managed_integration(CodingAgent::Codex, install.path()).unwrap(); + assert_eq!( + registered_install_dirs(CodingAgent::Codex).unwrap(), + [install.path().canonicalize().unwrap()] + ); + assert!( + registered_install_dirs(CodingAgent::ClaudeCode) + .unwrap() + .is_empty() + ); + unregister_managed_integration(CodingAgent::Codex, install.path()).unwrap(); + unregister_managed_integration(CodingAgent::Codex, install.path()).unwrap(); + assert!( + registered_install_dirs(CodingAgent::Codex) + .unwrap() + .is_empty() + ); +} + +#[tokio::test] +async fn malformed_managed_integration_registry_is_reported() { + let _guard = PLUGIN_CONFIG_TEST_LOCK.lock().await; + let config = tempfile::tempdir().unwrap(); + let _environment = EnvScope::set(&[("XDG_CONFIG_HOME", Some(config.path().as_os_str()))]); + let path = managed_integrations_registry_path().unwrap(); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + std::fs::write(&path, b"not-json").unwrap(); + assert!(read_managed_integrations_registry().is_err()); + assert!(registered_install_dirs(CodingAgent::Codex).is_err()); +} + +#[test] +fn persisted_state_must_match_the_selected_layout() { + let directory = tempfile::tempdir().unwrap(); + let layout = PluginLayout::new(CodingAgent::Codex, directory.path()); + let valid = PluginState { + marketplace_root: layout.marketplace_root.clone(), + plugin_root: layout.plugin_root.clone(), + host_plugin_removed: false, + host_marketplace_removed: false, + plugin_setup_installed: false, + }; + layout.validate_persisted_state(&valid).unwrap(); + let invalid = PluginState { + plugin_root: directory.path().join("outside"), + ..valid + }; + assert!(layout.validate_persisted_state(&invalid).is_err()); + assert!(!HostRegistrationProgress::default().any_added()); + assert!( + HostRegistrationProgress { + host_plugin_added: true, + host_marketplace_added: false, + } + .any_added() + ); +} diff --git a/crates/cli/tests/coverage/shared/mcp_tests.rs b/crates/cli/tests/coverage/shared/mcp_tests.rs index cfee096d6..36e3a0aa8 100644 --- a/crates/cli/tests/coverage/shared/mcp_tests.rs +++ b/crates/cli/tests/coverage/shared/mcp_tests.rs @@ -210,6 +210,10 @@ fn invalid_and_unknown_requests_return_jsonrpc_errors() { response_for(&json!({"jsonrpc":"2.0", "id":4, "method":"resources/list"})), Some(jsonrpc_error(json!(4), -32601, "Method not found")) ); + assert_eq!( + response_for(&json!({"jsonrpc":"2.0", "id":6, "method":null})), + Some(jsonrpc_error(json!(6), -32600, "Invalid Request")) + ); assert_eq!( response_for(&json!({"jsonrpc":"2.0", "id":5, "method":"initialize", "params":{}})), Some(jsonrpc_error(json!(5), -32602, "Missing protocolVersion")) diff --git a/crates/cli/tests/coverage/shared/plugin_prompt_tests.rs b/crates/cli/tests/coverage/shared/plugin_prompt_tests.rs new file mode 100644 index 000000000..47f25e884 --- /dev/null +++ b/crates/cli/tests/coverage/shared/plugin_prompt_tests.rs @@ -0,0 +1,269 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use nemo_relay::config_editor::{ + EditorFieldKind, EditorFieldSpec, EditorListItemSpec, EditorSchema, EditorTaggedUnionSpec, + EditorVariantSpec, +}; + +fn field(kind: EditorFieldKind) -> EditorFieldSpec { + EditorFieldSpec { + name: "value", + label: "Value", + kind, + enum_values: &[], + optional: true, + nested_schema: None, + nested_default: None, + list_item: None, + tagged_union: None, + } +} + +fn item(kind: EditorFieldKind) -> EditorListItemSpec { + EditorListItemSpec { + kind, + schema: None, + default: None, + tagged_union: None, + list_item: None, + } +} + +fn empty_schema() -> &'static EditorSchema { + static SCHEMA: EditorSchema = EditorSchema { fields: &[] }; + &SCHEMA +} + +fn string_schema() -> &'static EditorSchema { + static FIELDS: [EditorFieldSpec; 1] = [EditorFieldSpec { + name: "value", + label: "Value", + kind: EditorFieldKind::String, + enum_values: &[], + optional: true, + nested_schema: None, + nested_default: None, + list_item: None, + tagged_union: None, + }]; + static SCHEMA: EditorSchema = EditorSchema { fields: &FIELDS }; + &SCHEMA +} + +fn known_variant_default() -> Value { + json!({"kind": "known"}) +} + +static KNOWN_VARIANTS: [EditorVariantSpec; 1] = [EditorVariantSpec { + label: "Known", + tag: "known", + schema: empty_schema, + default: known_variant_default, +}]; + +static KNOWN_UNION: EditorTaggedUnionSpec = EditorTaggedUnionSpec { + discriminator: "kind", + variants: &KNOWN_VARIANTS, +}; + +fn assert_config_error(result: Result, expected: &str) { + let error = result + .err() + .expect("operation should reject malformed metadata"); + assert!(error.to_string().contains(expected), "{error}"); +} + +#[test] +fn malformed_structured_field_metadata_fails_before_terminal_input() { + let theme = ColorfulTheme::default(); + let mut config = json!({}); + + assert_config_error( + edit_section(&theme, &mut config, field(EditorFieldKind::Section)), + "not an editable section", + ); + for (kind, expected) in [ + (EditorFieldKind::List, "does not describe its list entries"), + (EditorFieldKind::Map, "does not describe its map values"), + ( + EditorFieldKind::TaggedUnion, + "does not describe its variants", + ), + ] { + assert_config_error( + edit_config_field(&theme, &mut config, field(kind)), + expected, + ); + } + assert_config_error( + edit_config_field(&theme, &mut config, field(EditorFieldKind::Section)), + "not an editable section", + ); +} + +#[test] +fn malformed_nested_value_metadata_fails_before_terminal_input() { + let theme = ColorfulTheme::default(); + let mut value = json!({}); + let schema = empty_schema(); + + for (kind, expected) in [ + (EditorFieldKind::Section, "not an editable section"), + (EditorFieldKind::List, "does not describe its list entries"), + (EditorFieldKind::Map, "does not describe its map values"), + ( + EditorFieldKind::TaggedUnion, + "does not describe its variants", + ), + ( + EditorFieldKind::DiscriminatedSection, + "does not describe its variants", + ), + ] { + assert_config_error( + edit_value_field(&theme, "root", &mut value, schema, field(kind), None), + expected, + ); + } +} + +#[test] +fn malformed_collection_items_and_tagged_values_are_rejected_without_prompting() { + let theme = ColorfulTheme::default(); + for (kind, expected) in [ + (EditorFieldKind::Section, "list item has no schema"), + ( + EditorFieldKind::List, + "nested list item has no entry description", + ), + ( + EditorFieldKind::Map, + "nested map item has no entry description", + ), + ] { + let mut value = json!({}); + assert_config_error( + edit_editor_item(&theme, "item", &mut value, &item(kind)), + expected, + ); + } + + static EMPTY_UNION: EditorTaggedUnionSpec = EditorTaggedUnionSpec { + discriminator: "kind", + variants: &[], + }; + assert_config_error( + select_tagged_union_variant(&theme, &EMPTY_UNION), + "tagged union has no variants", + ); + assert_config_error( + new_tagged_union_value(&theme, &EMPTY_UNION), + "tagged union has no variants", + ); + + let mut missing = json!({}); + assert_config_error( + edit_tagged_union_payload(&theme, "backend", &mut missing, &KNOWN_UNION), + "tagged union has no discriminator value", + ); + let mut unknown = json!({"kind": "unknown"}); + assert_config_error( + edit_tagged_union_payload(&theme, "backend", &mut unknown, &KNOWN_UNION), + "unknown tagged union type", + ); +} + +#[test] +fn selection_dispatch_resets_values_and_ignores_out_of_range_actions() { + let theme = ColorfulTheme::default(); + let mut value = json!({"value": "changed"}); + let schema = string_schema(); + let default = json!({"value": "default"}); + + assert!( + edit_selected_value_item( + &theme, + "root", + &mut value, + schema, + Some(&default), + schema.fields.len(), + ) + .unwrap() + ); + assert_eq!(value, default); + assert!( + edit_selected_value_item( + &theme, + "root", + &mut value, + schema, + None, + schema.fields.len(), + ) + .unwrap() + ); + assert_eq!(value, json!({})); + assert!( + !edit_selected_value_item( + &theme, + "root", + &mut value, + schema, + None, + schema.fields.len() + 1, + ) + .unwrap() + ); +} + +#[test] +fn scalar_prompt_rejects_structured_kinds_without_reading_the_terminal() { + let theme = ColorfulTheme::default(); + for kind in [ + EditorFieldKind::Section, + EditorFieldKind::List, + EditorFieldKind::Map, + EditorFieldKind::TaggedUnion, + EditorFieldKind::DiscriminatedSection, + ] { + assert_config_error( + prompt_value(&theme, &field(kind), None), + if kind == EditorFieldKind::Section { + "nested section" + } else { + "structured value" + }, + ); + } +} + +#[test] +fn menu_and_editor_io_errors_preserve_cancellation_and_failure_context() { + for kind in [ + std::io::ErrorKind::Interrupted, + std::io::ErrorKind::UnexpectedEof, + ] { + for error in [ + menu_error(std::io::Error::from(kind)), + editor_error(dialoguer::Error::IO(std::io::Error::from(kind))), + ] { + assert!(matches!( + error, + CliError::Config(ref message) if message == PLUGIN_EDIT_CANCELLED_MESSAGE + )); + } + } + assert!( + menu_error(std::io::Error::other("terminal failed")) + .to_string() + .contains("terminal failed") + ); + assert!( + editor_error(dialoguer::Error::IO(std::io::Error::other("editor failed"))) + .to_string() + .contains("editor failed") + ); +} diff --git a/crates/cli/tests/coverage/shared/plugins_lifecycle_tests.rs b/crates/cli/tests/coverage/shared/plugins_lifecycle_tests.rs index 4321f0380..bd0fdde2e 100644 --- a/crates/cli/tests/coverage/shared/plugins_lifecycle_tests.rs +++ b/crates/cli/tests/coverage/shared/plugins_lifecycle_tests.rs @@ -29,6 +29,19 @@ fn allow_unsigned_test_plugins(resolved: &mut ResolvedConfig) { Some(nemo_relay::plugin::dynamic::DynamicPluginAttestationMode::IntegrityOnly); } +#[cfg(unix)] +#[test] +fn process_python_environment_runner_reports_launch_success_and_exit_failures() { + let runner = ProcessPythonEnvironmentCommandRunner; + assert!(runner.run(OsStr::new("true"), &[]).is_ok()); + let failure = runner.run(OsStr::new("false"), &[]).unwrap_err(); + assert!(failure.contains("exited with status"), "{failure}"); + let missing = runner + .run(OsStr::new("/definitely/missing/nemo-relay-python"), &[]) + .unwrap_err(); + assert!(missing.contains("failed to start"), "{missing}"); +} + #[test] fn activation_snapshots_use_a_short_directory_prefix() { let _env = EnvScope::set(&[(ACTIVATION_SNAPSHOT_DIR_ENV, None)]); diff --git a/crates/cli/tests/coverage/shared/plugins_tests.rs b/crates/cli/tests/coverage/shared/plugins_tests.rs index e5811cf9d..de23a01a3 100644 --- a/crates/cli/tests/coverage/shared/plugins_tests.rs +++ b/crates/cli/tests/coverage/shared/plugins_tests.rs @@ -2623,6 +2623,17 @@ fn print_preview_renders_default_plugin_config() { print_preview(&PluginConfig::default()).unwrap(); } +#[test] +fn document_preview_applies_typed_components_without_persisting() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("plugins.toml"); + let document = PluginConfigDocument::read(&path).unwrap(); + let components = editable_components(document.config()).unwrap(); + + preview_document(&document, &components, &[]).unwrap(); + assert!(!path.exists()); +} + #[test] fn validate_config_reports_plugin_diagnostics() { let config = PluginConfig { diff --git a/crates/cli/tests/coverage/shared/server_tests.rs b/crates/cli/tests/coverage/shared/server_tests.rs index f6a7cf72f..2f16dcded 100644 --- a/crates/cli/tests/coverage/shared/server_tests.rs +++ b/crates/cli/tests/coverage/shared/server_tests.rs @@ -301,18 +301,24 @@ async fn responses_websocket_upgrades_request_http_fallback() { #[tokio::test] async fn responses_plain_get_remains_method_not_allowed() { let app = router_with_state(AppState::new(test_config())); - let response = app - .oneshot( - Request::builder() - .method("GET") - .uri("/v1/responses") - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); - - assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED); + for path in [ + "/responses", + "/v1/responses", + "/backend-api/codex/responses", + ] { + let response = app + .clone() + .oneshot( + Request::builder() + .method("GET") + .uri(path) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED, "{path}"); + } } #[test] diff --git a/crates/cli/tests/managed_pi_extension_tests.mjs b/crates/cli/tests/managed_pi_extension_tests.mjs new file mode 100644 index 000000000..65f322827 --- /dev/null +++ b/crates/cli/tests/managed_pi_extension_tests.mjs @@ -0,0 +1,259 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +import assert from 'node:assert/strict'; +import { chmod, mkdtemp, readFile, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { pathToFileURL } from 'node:url'; +import test from 'node:test'; + +import { + decideManagedProviderRedirect, + decideManagedToolTransform, + summarizeManagedToolResult, +} from '../src/daemon/managed/pi_extension/index.ts'; + +const CREDENTIAL = Buffer.alloc(32, 7).toString('base64url'); + +function requireReplace(source, pattern, replacement) { + assert.equal(source.split(pattern).length, 2, `expected exactly one lifecycle rewrite: ${pattern}`); + return source.replace(pattern, replacement); +} + +async function loadLifecycleTestApi() { + const sourceUrl = new URL('../src/daemon/managed/pi_extension/index.ts', import.meta.url); + let source = await readFile(sourceUrl, 'utf8'); + const rewrites = [ + [ + 'export default function managedNemoRelayPi(pi: ExtensionAPI): void {', + 'export default function managedNemoRelayPi(pi: ExtensionAPI, initialize = initializeRuntime): void {', + ], + ['runtimePromise ??= initializeRuntime();', 'runtimePromise ??= initialize();'], + ['function createSharedLease(', 'export function createSharedLease('], + [ + 'const launched = spawn(config.dispatcherCommand, [', + 'const launched = spawn(process.execPath, [config.dispatcherCommand, ', + ], + ]; + for (const [pattern, replacement] of rewrites) source = requireReplace(source, pattern, replacement); + source += '\nexport { managedNemoRelayPi as installManagedNemoRelayPi };\n'; + const directory = await mkdtemp(join(tmpdir(), 'nemo-relay-pi-module-')); + const modulePath = join(directory, 'index.ts'); + await writeFile(modulePath, source); + return import(pathToFileURL(modulePath).href); +} + +test('managed Pi MCP lease becomes ready, restarts, and releases without another launch', async () => { + const directory = await mkdtemp(join(tmpdir(), 'nemo-relay-pi-lease-')); + const launches = join(directory, 'launches'); + const dispatcher = join(directory, 'dispatcher.mjs'); + await writeFile( + dispatcher, + `#!/usr/bin/env node +import { appendFileSync, readFileSync } from 'node:fs'; +import { createInterface } from 'node:readline'; +appendFileSync(${JSON.stringify(launches)}, 'launch\\n'); +const launchCount = readFileSync(${JSON.stringify(launches)}, 'utf8').trim().split('\\n').length; +const lines = createInterface({ input: process.stdin, crlfDelay: Infinity }); +lines.on('line', (line) => { + const request = JSON.parse(line); + if (request.method === 'initialize') { + process.stdout.write(JSON.stringify({ jsonrpc: '2.0', id: request.id, result: { + protocolVersion: '2025-11-25', capabilities: {}, serverInfo: { name: 'nemo-relay', version: 'test' } + } }) + '\\n'); + } else if (request.method === 'notifications/initialized' && launchCount === 1) { + setTimeout(() => process.exit(0), 10); + } +}); +`, + ); + await chmod(dispatcher, 0o755); + + let removed = 0; + const { createSharedLease } = await loadLifecycleTestApi(); + const lease = createSharedLease( + { schema: 'nemo-relay-managed-pi-v1', daemonAddress: 'http://127.0.0.1:47632', dispatcherCommand: dispatcher }, + CREDENTIAL, + () => { removed += 1; }, + ); + await lease.ensureReady(); + const deadline = Date.now() + 10_000; + while ((await readFile(launches, 'utf8')).trim().split('\n').length < 2) { + assert.ok(Date.now() < deadline, 'MCP child did not restart'); + await new Promise((resolve) => setTimeout(resolve, 10)); + await lease.ensureReady(); + } + await lease.release(); + assert.equal((await readFile(launches, 'utf8')).trim().split('\n').length, 2); + assert.equal(removed, 1); + await assert.rejects(lease.ensureReady(), /released/); + assert.equal((await readFile(launches, 'utf8')).trim().split('\n').length, 2); +}); + +test('managed Pi shutdown branches preserve ordering and quit releases without reinitializing', async () => { + const handlers = new Map(); + const pi = { + on(name, handler) { handlers.set(name, handler); }, + registerProvider() {}, + }; + let ensureReady = 0; + let releases = 0; + const runtime = { + config: { schema: 'nemo-relay-managed-pi-v1', daemonAddress: 'https://relay.example.com', dispatcherCommand: '/opt/nemo-relay' }, + credential: CREDENTIAL, + lease: { + ensureReady: async () => { ensureReady += 1; }, + release: async () => { releases += 1; }, + }, + }; + const bodies = []; + const originalFetch = globalThis.fetch; + globalThis.fetch = async (_url, request) => { + bodies.push(JSON.parse(request.body)); + return new Response('{}', { status: 200 }); + }; + try { + const { installManagedNemoRelayPi } = await loadLifecycleTestApi(); + installManagedNemoRelayPi(pi, async () => runtime); + const context = { cwd: '/workspace', sessionManager: { getSessionId: () => 'session' } }; + await handlers.get('session_start')({ type: 'session_start', reason: 'startup' }, context); + await Promise.all([ + handlers.get('turn_start')({ type: 'turn_start', turnIndex: 0 }, context), + handlers.get('turn_start')({ type: 'turn_start', turnIndex: 1 }, context), + ]); + await handlers.get('session_shutdown')({ type: 'session_shutdown', reason: 'reload' }, context); + await handlers.get('session_shutdown')({ type: 'session_shutdown', reason: 'new' }, context); + const beforeQuit = ensureReady; + await handlers.get('session_shutdown')({ type: 'session_shutdown', reason: 'quit' }, context); + + assert.deepEqual(bodies.map((body) => body.hook_event_name), [ + 'session_start', 'turn_start', 'turn_start', 'session_shutdown', 'session_shutdown', + ]); + assert.deepEqual(bodies.filter((body) => body.hook_event_name === 'turn_start').map((body) => body.turn_index), [0, 1]); + assert.equal(ensureReady - beforeQuit, 0, 'quit forwards its final hook without restarting MCP'); + assert.equal(releases, 1); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test('custom Pi providers redirect only when every sibling API is supported', () => { + const selected = { + id: 'custom-response-model', + api: 'openai-responses', + provider: 'custom-enterprise', + baseUrl: 'https://ignored-by-managed-policy.example/v1', + }; + const serviceableCatalog = [ + selected, + { + id: 'custom-messages-model', + api: 'anthropic-messages', + provider: selected.provider, + baseUrl: 'https://ignored-by-managed-policy.example/v1/', + }, + ]; + + assert.deepEqual(decideManagedProviderRedirect(selected, serviceableCatalog), { + kind: 'redirect', + upstream: selected.baseUrl, + reason: 'provider uses only daemon-supported APIs and every model shares its endpoint', + }); + assert.equal(decideManagedProviderRedirect(selected, undefined).kind, 'skip'); + assert.deepEqual( + decideManagedProviderRedirect(selected, [ + ...serviceableCatalog, + { + id: 'custom-google-model', + api: 'google-generative-ai', + provider: selected.provider, + baseUrl: 'https://google.example', + }, + ]), + { + kind: 'skip', + code: 'provider-mixed-apis', + reason: + 'redirecting custom-enterprise would also move its unsupported ' + + 'google-generative-ai model custom-google-model', + }, + ); + assert.deepEqual( + decideManagedProviderRedirect(selected, [ + selected, + { + id: 'different-endpoint-model', + api: 'openai-completions', + provider: selected.provider, + baseUrl: 'https://different.example/v1', + }, + ]), + { + kind: 'skip', + code: 'provider-mixed-endpoints', + reason: + 'redirecting custom-enterprise would also move different-endpoint-model, which targets ' + + 'https://different.example/v1 rather than https://ignored-by-managed-policy.example/v1', + }, + ); +}); + +test('managed Pi tool rewrites require the exact call ID and recursively preserve shape', () => { + const current = { path: '/before', flags: [true, { retries: 2 }] }; + assert.deepEqual( + decideManagedToolTransform( + { + tool_call: { + tool_call_id: 'call-1', + input: { path: '/after', flags: [false, { retries: 3 }] }, + }, + }, + 'call-1', + current, + ), + { + kind: 'replace', + input: { path: '/after', flags: [false, { retries: 3 }] }, + }, + ); + assert.equal( + decideManagedToolTransform({ tool_call: { input: { path: '/after', flags: current.flags } } }, 'call-1', current) + .kind, + 'invalid', + ); + assert.equal( + decideManagedToolTransform( + { + tool_call: { + tool_call_id: 'call-1', + input: { path: '/after', flags: [false, { retries: 3 }], extra: true }, + }, + }, + 'call-1', + current, + ).kind, + 'invalid', + ); +}); + +test('managed Pi summaries preserve ordered text blocks and Unicode boundaries', () => { + assert.deepEqual( + summarizeManagedToolResult( + { + content: [ + { type: 'text', text: 'first' }, + { type: 'image', data: 'not-forwarded' }, + { type: 'text', text: 'second' }, + ], + }, + false, + ), + { content: 'first\nsecond', result_keys: ['content'] }, + ); + + const summary = summarizeManagedToolResult('x'.repeat(1_999) + '😀tail', false).content; + assert.equal(typeof summary, 'string'); + assert.equal(summary.includes('�'), false); + assert.equal(summary.includes('... [truncated 6 chars]'), true); +}); diff --git a/docs/nemo-relay-cli/about.mdx b/docs/nemo-relay-cli/about.mdx index 68e3fe70e..e9ae26955 100644 --- a/docs/nemo-relay-cli/about.mdx +++ b/docs/nemo-relay-cli/about.mdx @@ -69,6 +69,9 @@ Use these guide links to move from CLI setup into agent-specific instructions. - [Basic Usage](/nemo-relay-cli/basic-usage) explains gateway routes, transparent runs, shared configuration, hook forwarding, and runtime mapping. +- [Managed Daemon](/nemo-relay-cli/daemon) covers the brokered multi-user + daemon, managed hook and MCP commands, worker networking, immutable settings, + and lossless streaming transport. - [Coding Agent Installation](/nemo-relay-cli/plugin-installation) covers persistent Claude Code, Codex, and pi setup. - [Claude Code](/nemo-relay-cli/claude-code) covers transparent Claude Code @@ -78,6 +81,7 @@ Use these guide links to move from CLI setup into agent-specific instructions. GUI/app caveats, model provider routing, and remote-task limits. - [pi](/nemo-relay-cli/pi) covers the extension-based pi integration, tool gating, agent-run attribution, and the model-redirection gap. + Start with [Basic Usage](/nemo-relay-cli/basic-usage), then use the guide for the coding agent that you want to observe. diff --git a/docs/nemo-relay-cli/daemon.mdx b/docs/nemo-relay-cli/daemon.mdx new file mode 100644 index 000000000..f5c3ec17d --- /dev/null +++ b/docs/nemo-relay-cli/daemon.mdx @@ -0,0 +1,537 @@ +--- +title: 'Managed Daemon' +description: 'Deploy the brokered NeMo Relay daemon, MCP lifecycle client, managed hook forwarder, and worker.' +position: 7 +--- + +{/* SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 */} + +Use the managed daemon when an administrator must give multiple users the same +coding-agent configuration while preserving per-user NeMo Relay runtime state. +The daemon is the only public endpoint. It authenticates a machine-user route, +then either sends traffic through that route's worker or forwards it directly +to the configured provider in pass-through mode. + +This deployment model is separate from a personal installation. The existing +top-level `nemo-relay mcp` command and the hidden `nemo-relay hook-forward` +command remain available for personal integrations. + +## Process Topology + +The managed topology has the following processes: + +- `nemo-relay daemon` owns the public LLM and hook endpoints and the + authoritative route broker. +- `nemo-relay daemon mcp` registers one MCP client reference for the current + machine-user identity. It advertises no MCP tools. +- `nemo-relay daemon hook` forwards one native hook payload using the managed + route credential. +- `nemo-relay daemon worker` runs the per-machine-user Relay configuration and + remains attached to the daemon for its useful lifetime. + +All LLM and hook paths remain at the daemon root. Managed settings do not use +per-user route URLs. Requests use one Relay-specific public credential header: + +```text +x-nemo-relay-client-token: +``` + +Pi provider registrations also preserve Pi's runtime-selected endpoint in +`x-nemo-relay-upstream-base-url`. The authenticated daemon consumes that +routing header and removes it before contacting the provider. It is runtime +model metadata, not part of the managed settings artifact. + +Enterprise bootstrap must provide the credential through this environment +variable: + +```bash +export NEMO_RELAY_CLIENT_TOKEN='' +``` + +Claude Code reads custom provider headers from its native environment variable, +so the same bootstrap must derive exactly one header from that credential: + +```bash +export ANTHROPIC_CUSTOM_HEADERS="x-nemo-relay-client-token: ${NEMO_RELAY_CLIENT_TOKEN}" +``` + +Do not append another `x-nemo-relay-client-token` entry to an existing custom +header value. Managed diagnostics reject a missing, duplicate, or mismatched +credential header. + +The same credential identifies MCP registration, Codex and Claude LLM +requests, Pi provider requests, and managed hook requests. The daemon does +not require a token file or a token environment variable. During signed MCP +registration, it binds the credential digest to the user-machine fingerprint. +An existing binding cannot be reassigned to another fingerprint. The daemon +stores only credential digests, removes the public header before forwarding +a request, and preserves the caller's provider authentication. + +Enrollment is open to anyone who can reach the daemon and prove possession +of their machine identity. This proof establishes identity, not organizational +authorization. Restrict daemon access to trusted users through your network +or authenticated reverse proxy; do not expose open enrollment to an untrusted +network. TLS remains required for non-loopback communication. + +Challenge requests are limited to 16 per transport-peer IP in a 15-second +window, before the request body is read. MCP and worker challenges retain +separate capacity limits. Forwarded IP headers do not override this limit: +clients behind a reverse proxy share its peer budget. Apply additional +per-client admission controls at the proxy for larger deployments. + +## Start the Daemon + +Start the daemon with its loopback defaults: + +```bash +nemo-relay daemon +``` + +The default listener is `127.0.0.1:47632`. The full daemon-specific command +shape is: + +```text +nemo-relay daemon \ + [--bind <127.0.0.1|0.0.0.0>] \ + [--port ] \ + [--advertise-address ] \ + [--pass-through] \ + [--tls-cert ] \ + [--tls-key ] +``` + +Provision `NEMO_RELAY_CLIENT_TOKEN` in client environments only. Start and +keep an MCP connection active before sending LLM or hook requests, including +in `--pass-through` mode. Unknown credentials return `401`; registered routes +without a live MCP reference are unavailable and return `503`. Daemon restart +requires MCP re-registration to restore these in-memory route bindings. + +The shared `/models` and `/v1/models` GET endpoints return +`Cache-Control: no-store`, overriding provider cache policy to prevent +cross-credential catalog reuse. LLM streaming response headers are unchanged. + +Worker activation still requires the broker-issued, one-time grant delivered +through the protected inherited channel. Knowing the daemon address alone +does not allow a worker to register or replace an existing worker. + +The daemon accepts only `127.0.0.1` or `0.0.0.0` for `--bind`. A daemon bound +to `0.0.0.0` requires a concrete, reachable origin URL through +`--advertise-address`; `0.0.0.0` itself is never an advertised or target +address. + +For a non-loopback advertised origin, expose HTTPS either on the daemon +listener or through a trusted reverse proxy. For example, native TLS has the +following shape: + +```bash +nemo-relay daemon \ + --bind 0.0.0.0 \ + --port 8443 \ + --advertise-address https://relay.example.com:8443 \ + --tls-cert /etc/nemo-relay/tls.crt \ + --tls-key /etc/nemo-relay/tls.pk8 +``` + +Native TLS requires an `https` advertised URL. An `https` advertised URL may +also name a trusted reverse proxy that terminates TLS before forwarding to the +daemon listener. Every non-loopback `--daemon-address` target must use HTTPS. +If a trusted reverse proxy is in the data path, apply the streaming +requirements later on this page. + +### Use Explicit Pass-Through Mode + +Start a daemon that never creates workers with: + +```bash +nemo-relay daemon --pass-through +``` + +The daemon still authenticates the route credential. MCP registration can only +receive `UsePassThrough`, and the daemon cannot issue an activation grant or +accept worker registration. LLM requests use the same streaming transport +directly to the configured provider. Hook requests return the existing no-op +responses: + +| Agent | Response Body | +| ----------- | -------------------- | +| Codex | `{}` | +| Claude Code | `{"continue": true}` | +| Pi | `{}` | + +## Register the MCP Lifecycle Client + +Every managed MCP process must name its daemon explicitly: + +```bash +nemo-relay daemon mcp --daemon-address https://relay.example.com:443 +``` + +The address must be an HTTP or HTTPS origin with an explicit port. It cannot +contain credentials, a non-root path, a query, or a fragment. Plain HTTP is +accepted only for a loopback target. + +The MCP client authenticates and acquires its broker reference before it starts +the MCP protocol. The broker, not the MCP client, chooses one directive: + +| Directive | MCP Action | +| ---------------- | ----------------------------------------------------------------------------- | +| `ReuseWorker` | Use the worker already published for the fingerprint. | +| `WaitForWorker` | Wait while another MCP starts or drains the worker. | +| `LaunchWorker` | Start the current `nemo-relay` executable with the one-time activation grant. | +| `UsePassThrough` | Keep the MCP reference while the daemon forwards the route directly. | + +The MCP client never independently decides whether the route needs a worker. +When it receives `LaunchWorker`, it starts `daemon worker` on the same machine +and transfers the activation grant through a protected inherited standard-input +channel. The grant is not put in command arguments, environment variables, or +logs. + +## Forward Managed Hooks + +Use the agent-specific subcommand in immutable managed hook settings. The +supported command shapes are: + +```bash +nemo-relay daemon hook codex --daemon-address https://relay.example.com:443 --fail-closed +nemo-relay daemon hook claude --daemon-address https://relay.example.com:443 --fail-open +nemo-relay daemon hook pi --daemon-address https://relay.example.com:443 --fail-open +``` + +The hook process performs the following actions: + +1. Reads the agent's native hook payload from standard input. +2. Reads `NEMO_RELAY_CLIENT_TOKEN` from the environment. +3. Posts to the existing agent-specific hook path at the daemon root. +4. Adds `x-nemo-relay-client-token` exactly once. +5. Writes a successful, nonempty hook response to standard output without + changing its bytes. + +The hook process does not start a daemon or worker and does not perform an MCP +handshake. Use `--fail-open` for events that must not block the agent when the +daemon is unavailable. Use `--fail-closed` for policy events that must reject +the operation when delivery or evaluation fails. If neither option is present, +the existing event-specific failure policy applies. + +## Run a Daemon-Attached Worker + +The broker normally starts the worker. Its command shape is documented for +managed launchers and prescribed firewall, NAT-forwarding, and test scenarios: + +```text +nemo-relay daemon worker \ + --daemon-address \ + [--bind <127.0.0.1|0.0.0.0>] \ + [--port ] \ + [--advertise-address ] +``` + +Do not use the worker command as a standalone gateway. It requires the +one-time activation grant supplied by the MCP process and must authenticate to +the named daemon before it becomes routable. + +Worker network settings follow these rules: + +- The effective default is `127.0.0.1:0`; omitting `--port` lets the operating + system select an available port. +- `--bind` accepts only `127.0.0.1` or `0.0.0.0`. Hostnames, IPv6 addresses, + and other IPv4 addresses are rejected. +- An explicitly supplied port must be in `1..=65535`. Explicit `--port 0` is + rejected; omit the option for automatic allocation. +- Explicit ports are for documented firewall, NAT-forwarding, and test + deployments. +- A loopback worker advertises `127.0.0.1:` and does not accept + `--advertise-address`. +- `0.0.0.0` is bind-only. It requires a concrete daemon-reachable host or IP + through `--advertise-address`. +- `--daemon-address` is mandatory. A non-loopback daemon address must use + HTTPS. + +Only the daemon can send data-plane requests to a worker. The worker validates +the daemon-to-worker session credential from the request head before it reads +the request body. + +For a broker-launched worker, set the following environment variables on the +MCP process only when automatic network selection is insufficient: + +```bash +export NEMO_RELAY_WORKER_ADVERTISE_ADDRESS='worker.example.com' +export NEMO_RELAY_WORKER_PORT='9443' +``` + +`NEMO_RELAY_WORKER_ADVERTISE_ADDRESS` must be a concrete IPv4 address or +hostname that the daemon can reach. `NEMO_RELAY_WORKER_PORT` must be in +`1..=65535`; leave it unset for operating-system allocation. These overrides +are intended for prescribed firewall, NAT-forwarding, and test deployments, +and the MCP signs them into the broker registration before the daemon chooses +the worker launch directive. + +## Understand Broker Identity and Lifecycle + +Each daemon component proves a signed `nemo-relay` service identity and role, +and advertises its supported protocol range and capabilities. The binary +version is diagnostic metadata, not a route key or a compatibility decision. +Different compatible Relay binary versions can therefore share a worker. + +Each machine-user identity has an owner-private Ed25519 key. The public-key +digest is the route fingerprint. MCP registration also binds the digest of +`NEMO_RELAY_CLIENT_TOKEN` to that fingerprint. The daemon indexes requests by +the credential digest but does not persist or log the raw credential. The MCP +client trust-on-first-use pins the daemon identity to the normalized daemon +origin. + +The broker accepts a worker only when it proves both the single-use activation +grant and the same machine-user identity as the MCP client. Challenges expire, +are single use, and are replay protected. Later control messages use scoped +session credentials, request IDs, sequence numbers, and body hashes. + +One fingerprint moves through these broker states: + +| State | Behavior | +| ------------- | --------------------------------------------------------------------------- | +| `Empty` | The first MCP reference receives `LaunchWorker`. | +| `Activating` | Concurrent MCP references receive `WaitForWorker`. | +| `Ready` | Requests reuse the published worker. | +| `Draining` | New requests wait while accepted requests finish and the worker terminates. | +| `PassThrough` | The daemon forwards authenticated requests directly to providers. | +| `Recovering` | One connected MCP is nominated to replace a failed worker. | + +MCP clients renew their references every 10 seconds, and the daemon expires a +reference after 30 seconds without renewal. Workers send a heartbeat every five +seconds and expire after 20 seconds without one. There is no independent worker +idle timeout. + +If a worker loses its authenticated control relationship, it immediately stops +accepting new requests. Already accepted streams can finish while the worker +reconnects or re-registers after a daemon restart. The worker exits if it cannot +restore control within two minutes. + +When the last MCP reference leaves, the route enters a non-revivable drain. +Already accepted requests have up to two minutes to finish before the worker +terminates. A new MCP must wait for that termination, then begins a fresh +activation. If a ready worker fails while references remain, the broker +nominates one connected MCP to relaunch it. + +After MCP authentication succeeds, a worker activation, bind, registration, +readiness, or activation-channel failure moves the whole fingerprint route to +transient pass-through until all MCP references leave. A connection or identity +failure before authentication instead makes the MCP process log the error and +exit with a nonzero status. + +## Preserve Streaming Responses + +The daemon, worker, and pass-through paths use a pull-driven Hyper body from +provider to client. They do not collect a successful LLM response before +forwarding it. This transport preserves: + +- The response status. +- Ordered, multivalue end-to-end headers. +- The exact concatenated response-body byte sequence. +- HTTP trailers. +- SSE comments, heartbeats, `event`, `id`, `retry`, multiline `data`, and + `[DONE]` fields. +- Empty and non-UTF-8 data frames that the HTTP protocol accepts. + +HTTP implementations can split or combine DATA frames. The guarantee is the +same ordered body bytes and immediate availability, not matching TCP packets or +HTTP frame boundaries. + +The shared connection pools support HTTP/1.1 persistence and HTTP/2 +multiplexing. Delivery remains demand driven: a slow client applies bounded +backpressure upstream, and dropping the client cancels the corresponding +upstream work. Relay applies separate connection and response-head deadlines; +it does not apply a total response-lifetime deadline after streaming starts. + +The worker also transfers provider request bodies directly to the pooled +Hyper client when no LLM request guardrail, request interceptor, or +request-sanitization guardrail is registered. If one of those middleware types +needs the complete JSON request, the worker performs one bounded decode before +provider dispatch so the middleware can make its decision. Response bodies are +never collected for that purpose. + +Caller-visible stream fidelity takes precedence over response rewriting. Do +not use daemon-worker mode with middleware that must mutate, suppress, or +replace successful streaming response events. Any semantic observer or cache +must consume a bounded, nonblocking side channel. Falling behind can truncate +that observer's capture, but it must never delay or change delivery. + +The worker captures at most 4 MiB per streamed response for semantic +observation by default. Set +`NEMO_RELAY_DAEMON_OBSERVATION_CAPTURE_BYTES` to a positive byte count before +starting the worker to choose another bound. This setting changes only +side-band observability; it does not cap, buffer, or truncate caller-visible +delivery. + +Semantic observation has a separate 15-minute completion deadline, independent +of the provider's 60-second response-head deadline. Exceeding the observation +deadline marks only the captured observation as truncated; it does not stop +caller-visible streaming. + +The managed worker deliberately ignores per-user Relay configuration, +per-user plugin directories, and user lifecycle state. It loads only the +administrator-managed system configuration and plugin manifest while still +allowing provider secrets to arrive through their documented authentication +environment variables. This keeps the executed plugin configuration uniform +across users as well as keeping the coding-agent artifacts byte-identical. + +### Configure a Trusted Reverse Proxy + +A reverse proxy in front of the daemon becomes part of the streaming path. +Configure it to: + +- Disable response buffering. +- Disable compression or other response transformations. +- Disable cache coalescing. +- Preserve streaming over HTTP/1.1 or HTTP/2. +- Preserve response trailers. + +Test these settings end to end. TLS termination alone does not guarantee raw +stream preservation. + +## Distribute Immutable Managed Settings + +A managed bundle is an administrator-owned deployment artifact, not an output +of personal `nemo-relay install`. For each agent, platform, and deployment, its +plugin configuration and settings must remain byte-for-byte identical across +users and Relay binary releases. + +The bundle follows these rules: + +- Daemon and provider URLs and hook command text are fixed for the deployment. +- Artifacts contain no user path, fingerprint, credential, generation ID, or + Relay binary version. +- A stable administrator-owned dispatcher command or path survives Relay + upgrades. +- Refresh validates immutable artifacts and never rewrites them. +- An incompatible settings change uses a separately named v2 artifact rather + than replacing v1 bytes. + +Enterprise bootstrap owns credential provisioning and login-environment +injection. Relay validates `NEMO_RELAY_CLIENT_TOKEN` but does not install an +operating-system-specific login agent. + +Create a bundle once from fixed deployment values. Repeat `--agent` to select +the artifacts that the administrator distributes: + +```bash +nemo-relay daemon managed-bundle \ + --output /srv/nemo-relay/nemo-relay-managed-v1 \ + --daemon-address https://relay.example.com:443 \ + --dispatcher-command /opt/nvidia/bin/nemo-relay-dispatch \ + --platform linux \ + --agent codex \ + --agent claude \ + --agent pi +``` + +The destination must be new or already byte-identical. The command never +rewrites a different v1 bundle and prints only the canonical full-bundle +SHA-256. Provision that digest separately from the bundle; a digest stored +inside the same artifact is not a trust root. The dispatcher path is validated +for the target platform and must be an absolute, stable administrator path +outside known user and temporary directories. + +Validate a distributed bundle and its current environment without modifying +the bundle: + +```bash +nemo-relay doctor \ + --managed-bundle /path/to/nemo-relay-managed-v1 \ + --managed-bundle-sha256 '' +``` + +Doctor compares the exact file set and bytes with the canonical manifest and +the separately provisioned digest, then reports missing, unexpected, or +changed artifacts. This is a managed-only diagnostic: it does not load or fail because +of personal configuration, plugins, or agent installations. + +### Deploy the Managed Pi Extension + +The Pi artifact is a fixed TypeScript extension under +`pi/extension-v1/index.ts`. Load that exact administrator-installed file while +disabling discovered extensions; do not copy it into a per-user extension +directory or rewrite it during upgrades. For example: + +```bash +pi --no-extensions -e /srv/nemo-relay/nemo-relay-managed-v1/pi/extension-v1/index.ts +``` + +`--no-extensions` is mandatory for managed launches. Pi still loads the +explicit `-e` extension, but does not also load user, project, or discovered +extensions that could alter provider registration, tool arguments, or shell +policy after Relay has authorized an operation. + +The extension starts the fixed dispatcher as +`daemon mcp --daemon-address `, waits for MCP initialization before +registering managed providers, and +keeps one process-wide broker reference across Pi reload, new, resume, and fork +transitions. It sends Pi's session, agent, turn, compaction, tool, and custom +shell events to `/hooks/pi`; policy responses are converted back into Pi's +native `tool_call` and `user_bash` blocking results. Any selected Pi provider +whose models use only OpenAI Completions, OpenAI Responses, or Anthropic +Messages and share one endpoint is registered against the daemon, including +custom provider names. The registration attaches `x-nemo-relay-client-token` +and Pi's exact runtime-selected endpoint in +`x-nemo-relay-upstream-base-url`. No per-user upstream, fingerprint, +generation, user path, or route value is written into the managed artifact. + +### Codex Responses Compatibility + +Both the personal gateway and managed daemon accept Codex's ChatGPT-shaped +`POST /backend-api/codex/responses` path and canonicalize it to the ordinary +upstream Responses path. A WebSocket upgrade probe sent with `GET` to that +path, `/responses`, or `/v1/responses` receives `426 Upgrade Required`, which +causes clients that support the fallback to use HTTP streaming. An ordinary +GET remains `405 Method Not Allowed`. + +## Measure Daemon Transport + +Run the deterministic smoke check with: + +```bash +just daemon-transport-benchmark-smoke +``` + +The smoke check exercises OpenAI and Anthropic streaming fixtures over +HTTP/1.1 and HTTP/2. Stream-integrity failures fail the command; its timing +results are informational. + +Before a sustained comparison, build the existing size-optimized release and +an isolated `opt-level=3` candidate: + +```bash +just daemon-transport-benchmark-build-candidates +``` + +Start the deterministic provider with: + +```bash +just daemon-transport-benchmark-provider --bind 127.0.0.1:48100 +``` + +Start the Relay topology processes separately, then run the full load driver +against their endpoints. For example: + +```bash +export NEMO_RELAY_CLIENT_TOKEN='' + +just daemon-transport-benchmark \ + --direct-url http://127.0.0.1:48100 \ + --target daemon-pass-through=http://127.0.0.1:47632 \ + --target daemon-worker=http://127.0.0.1:47633 \ + --header-env daemon-pass-through:x-nemo-relay-client-token=NEMO_RELAY_CLIENT_TOKEN \ + --header-env daemon-worker:x-nemo-relay-client-token=NEMO_RELAY_CLIENT_TOKEN +``` + +The full preset uses a 10-second warmup and a 60-second measured interval. It +covers persistent HTTP/1.1 and HTTP/2 connections, 16 KiB and 1 MiB responses, +128 streamed events, concurrency 1, 16, 64, and 256, plus a separate +1,000-slow-stream capacity scenario. It records latency distributions, +throughput, goodput, process resource use, pool reuse, stream integrity, +trailers, cancellation, and reconnection data. + +Refer to `scripts/latency_benchmark/daemon_transport/README.md` for the complete +topology, worker-only target, protected-header setup, process metadata, and +report schema. Run base and candidate builds on the same otherwise-idle host. +Timing and throughput remain non-gating until the project establishes stable +reference baselines. diff --git a/justfile b/justfile index eb78af422..8044bb193 100644 --- a/justfile +++ b/justfile @@ -1220,6 +1220,46 @@ latency-benchmark *benchmark_args: test-latency-benchmark: uv run --locked python -m pytest scripts/latency_benchmark/tests +# Checks the detached Rust benchmark workspace that workspace-wide Rust commands do not cover. +check-daemon-transport-benchmark: + cargo fmt --manifest-path scripts/latency_benchmark/daemon_transport/Cargo.toml -- --check + cargo clippy --manifest-path scripts/latency_benchmark/daemon_transport/Cargo.toml --all-targets -- -D warnings + +# Small deterministic daemon transport check. Informational timings; stream integrity is required. +daemon-transport-benchmark-smoke: + #!/usr/bin/env bash + set -euo pipefail + result_dir={{ quote(output_dir) }} + result_dir="${result_dir:-target/benchmark-results}" + cargo run --locked --release \ + --manifest-path scripts/latency_benchmark/daemon_transport/Cargo.toml \ + --target-dir target/daemon-transport-driver \ + -- smoke \ + --output "$result_dir/daemon-transport-smoke.json" + +# Opt-in sustained daemon transport benchmark against already-running topology endpoints. +[positional-arguments] +daemon-transport-benchmark *benchmark_args: + cargo run --locked --release \ + --manifest-path scripts/latency_benchmark/daemon_transport/Cargo.toml \ + --target-dir target/daemon-transport-driver \ + -- load "$@" + +# Deterministic provider used by the opt-in daemon transport benchmark. +[positional-arguments] +daemon-transport-benchmark-provider *provider_args: + cargo run --locked --release \ + --manifest-path scripts/latency_benchmark/daemon_transport/Cargo.toml \ + --target-dir target/daemon-transport-driver \ + -- provider "$@" + +# Build the current size-optimized release and an isolated opt-level=3 candidate. +daemon-transport-benchmark-build-candidates: + cargo build --locked --release -p nemo-relay-cli + CARGO_PROFILE_RELEASE_OPT_LEVEL=3 cargo build --locked --release \ + --target-dir target/daemon-benchmark-opt3 \ + -p nemo-relay-cli + # --set [output_dir=] [ci=true|false] test-rust: #!/usr/bin/env bash @@ -1627,6 +1667,8 @@ test-pi: # client and loads no native addon, so nothing here depends on it. npm run typecheck --workspace=nemo-relay-pi npm test --workspace=nemo-relay-pi + node_modules/.bin/tsc -p crates/cli/src/daemon/managed/pi_extension/tsconfig.json --noEmit + node --test crates/cli/tests/managed_pi_extension_tests.mjs # --set [output_dir=] [ci=true|false] test-all: test-rust test-python test-python-langchain test-go test-node test-openclaw test-pi diff --git a/scripts/latency_benchmark/README.md b/scripts/latency_benchmark/README.md index f15d68838..556d66214 100644 --- a/scripts/latency_benchmark/README.md +++ b/scripts/latency_benchmark/README.md @@ -10,6 +10,12 @@ around OpenAI Responses, Anthropic Messages, Codex hooks, Claude Code hooks, and Relay process startup. The fixture runs deterministic providers on loopback, so network and model-service latency do not hide Relay overhead. +The daemon architecture has a separate Rust/Hyper transport fixture for +HTTP/1.1, HTTP/2, stream-integrity, high-concurrency, and slow-stream capacity +measurements. Refer to +[`daemon_transport/README.md`](daemon_transport/README.md) for its CI smoke +check and opt-in full topology matrix. + ## Before You Run Run all commands from the repository root. Start with the smoke test unless you diff --git a/scripts/latency_benchmark/config/daemon-transport-full.toml b/scripts/latency_benchmark/config/daemon-transport-full.toml new file mode 100644 index 000000000..0781e433a --- /dev/null +++ b/scripts/latency_benchmark/config/daemon-transport-full.toml @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Informational until the project establishes a stable reference baseline. +protocols = ["http1", "http2"] +providers = ["openai", "anthropic"] +response_bytes = [16384, 1048576] +events = 128 +concurrency = [1, 16, 64, 256] +warmup_seconds = 10 +duration_seconds = 60 +event_delay_micros = 0 +cancel_every = 100 +slow_streams = 1000 +slow_event_delay_millis = 100 diff --git a/scripts/latency_benchmark/config/daemon-transport-smoke.toml b/scripts/latency_benchmark/config/daemon-transport-smoke.toml new file mode 100644 index 000000000..0073d59d3 --- /dev/null +++ b/scripts/latency_benchmark/config/daemon-transport-smoke.toml @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +protocols = ["http1", "http2"] +providers = ["openai", "anthropic"] +response_bytes = [16384] +events = 128 +concurrency = [1] +warmup_seconds = 0 +duration_seconds = 1 +event_delay_micros = 0 +cancel_every = 0 +slow_streams = 0 +slow_event_delay_millis = 0 diff --git a/scripts/latency_benchmark/daemon_transport/Cargo.lock b/scripts/latency_benchmark/daemon_transport/Cargo.lock new file mode 100644 index 000000000..236e3c423 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/Cargo.lock @@ -0,0 +1,1221 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[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 2.0.119", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cc" +version = "1.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "005ec2760ca554fae18df7a11195552ec576cd665632a881bc011d5bb2fd4d80" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "clap" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 3.0.5", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8498c871161e1742aaa9d52551b2d6ebdd4c3d45a3be423e3728f33b955be550" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d85363c37faeca707aef026efa9f3b34d077bce547e48f770770625c6013679e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + +[[package]] +name = "find-msvc-tools" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e0f1c7c3a72c66fd80abe965175f7523475c0489a87d3ff9d6e8c87d87a9d2d" + +[[package]] +name = "flate2" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e634e2e0ebac1ee034020da1ca582e17ffe4e0f5e985823721e168928136dcb" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", +] + +[[package]] +name = "h2" +version = "0.4.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef8e5e5a340588f4452631496976cf8636d4a7ecf600239fdc27615d2530bc16" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "hdrhistogram" +version = "7.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f49d1053f4708f0af3cf9fc5bffc7e68a914a3c45becb231c80068c9c3f78bea" +dependencies = [ + "base64", + "byteorder", + "crossbeam-channel", + "flate2", + "nom", + "num-traits", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23169fe34a5fbcdd3f3862e78fb9b6fccd5f02a6dc6f732547005d45631ce71c" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[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.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27b501faa50e7a26c3d3560ca625132f4078a17771f4810baf70475ae48cbe43" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "rustls-native-certs", + "tokio", + "tokio-rustls", + "tower-service", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "libc", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "indexmap" +version = "2.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07aa2048142242915a31d35844fb311e0e53fcca590c3a0a40dcf1b841fa09eb" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "miniz_oxide" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b63fbc4a50860e98e7b2aa7804ded1db5cbc3aff9193adaff57a6931bf7c4b4c" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "mio" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b18443e9c262bfe8fa82f51666e2642c53393f7e5c27b3e1aeab922cff5b9d8" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "nemo-relay-daemon-transport-benchmark" +version = "0.1.0" +dependencies = [ + "anyhow", + "async-stream", + "base64", + "bytes", + "clap", + "hdrhistogram", + "http", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "ring", + "serde", + "serde_json", + "sha2", + "tempfile", + "tokio", + "toml", + "tower-service", +] + +[[package]] +name = "nom" +version = "8.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" +dependencies = [ + "memchr", +] + +[[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 = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +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 = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[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 = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9be42f50aa861c555654aa3a37f52f4b1074bacf4e48fe0ef7fa584e80f1f0f" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12df2e0110f65b775f769bb17ef989067a1d931b2eb822bd4346631eeada89f9" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.5", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0c85f2c3ef0b1cd58b36682f4b17aaa995f0e5db534d85692b4903abce21f67" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "libc", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "toml" +version = "0.9.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow 0.7.15", +] + +[[package]] +name = "toml_datetime" +version = "0.7.5+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow 1.0.4", +] + +[[package]] +name = "toml_writer" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" + +[[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 = [ + "pin-project-lite", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[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 = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" + +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zlib-rs" +version = "0.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34b31d188d9d685a4f9c7b46d6e36631b07058d2cfe190267adce54dc230bf12" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/scripts/latency_benchmark/daemon_transport/Cargo.toml b/scripts/latency_benchmark/daemon_transport/Cargo.toml new file mode 100644 index 000000000..7a5d41944 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/Cargo.toml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "nemo-relay-daemon-transport-benchmark" +version = "0.1.0" +edition = "2024" +publish = false + +[workspace] + +[dependencies] +anyhow = "1" +async-stream = "0.3" +base64 = "0.22" +bytes = "1" +clap = { version = "4", features = ["derive"] } +hdrhistogram = "7" +http = "1" +http-body-util = "0.1" +hyper = { version = "1", features = ["client", "server", "http1", "http2"] } +hyper-rustls = { version = "0.27", default-features = false, features = ["http1", "http2", "native-tokio", "ring", "tls12"] } +hyper-util = { version = "0.1", features = ["client-legacy", "http1", "http2", "server-auto", "tokio"] } +ring = "0.17" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +sha2 = "0.10" +tempfile = "3" +tokio = { version = "1", features = ["macros", "net", "process", "rt-multi-thread", "signal", "sync", "time"] } +toml = "0.9" +tower-service = "0.3" + +[profile.release] +debug = "line-tables-only" +lto = "thin" + +[lints.rust] +unsafe_code = "forbid" diff --git a/scripts/latency_benchmark/daemon_transport/README.md b/scripts/latency_benchmark/daemon_transport/README.md new file mode 100644 index 000000000..240b5ec59 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/README.md @@ -0,0 +1,155 @@ + + +# Daemon Transport Benchmark + +This standalone Rust/Hyper fixture measures streaming transport independently +of the existing Python latency suite. It supports HTTP/1.1 persistent +connections, cleartext HTTP/2 prior knowledge, remote HTTP/2 through ALPN, +and one client pool per protocol for the entire run. + +The benchmark is informational. Integrity failures make the command fail, but +latency and throughput values have no CI threshold until a stable baseline is +established. + +## CI Smoke Check + +Run the same short, direct-provider check used by CI: + +```bash +just daemon-transport-benchmark-smoke +``` + +The command starts an ephemeral deterministic provider in-process and checks +the load driver against that provider for OpenAI and Anthropic streams over +HTTP/1.1 and HTTP/2. It sends 128 events and 16 KiB per response, verifies the +body hash and trailers, and writes +`target/benchmark-results/daemon-transport-smoke.json`. This fast check does +not launch Relay and therefore does not exercise a daemon or worker hop. + +## Authenticated Worker-Only Runs + +A production worker accepts only broker-authenticated requests. Its endpoint +and credential are intentionally private, so do not scrape a daemon's state or +invent a worker header value. Supply a Relay binary to let the load driver +create an isolated worker target through the real activation flow: + +```bash +just daemon-transport-benchmark-provider --bind 127.0.0.1:48100 + +# In another terminal: +just daemon-transport-benchmark \ + --config scripts/latency_benchmark/config/daemon-transport-smoke.toml \ + --direct-url http://127.0.0.1:48100 \ + --worker-binary target/release/nemo-relay \ + --output target/benchmark-results/daemon-transport-worker-smoke.json +``` + +`--worker-binary` adds a `worker-only` target. The benchmark starts a real +daemon, MCP, and worker with isolated temporary state. A benchmark-only control +proxy binds an ephemeral loopback port, permits only Relay control-plane POST +paths, and verifies the authenticated worker registration and readiness flow. +It captures the worker endpoint and credential, forwards load directly to the +worker, and releases the MCP reference during cleanup. Harness credentials +remain in process memory and are never printed or serialized; the endpoint +appears in the report as ordinary target metadata. The provider URL must be an +HTTP origin with a numeric loopback address and explicit port because every +process and connection created by this orchestration is local and ephemeral. + +This run exercises the direct baseline and the worker data hop for both API +shapes and both HTTP protocols in the selected matrix. The helper daemon is +used only for worker activation and lifecycle control; it is not in the +`worker-only` request path. + +## Full Topology Run + +Build both Relay candidates without changing the workspace release profile: + +```bash +just daemon-transport-benchmark-build-candidates +``` + +The normal build is `target/release/nemo-relay`. The `opt-level=3` build is +`target/daemon-benchmark-opt3/release/nemo-relay`. The second build uses the +`CARGO_PROFILE_RELEASE_OPT_LEVEL` environment override and a separate target +directory, so it does not modify `Cargo.toml` or replace the normal release +binary. + +Start the deterministic provider on a fixed port before starting the Relay +processes being measured: + +```bash +just daemon-transport-benchmark-provider --bind 127.0.0.1:48100 +``` + +Configure each public Relay topology to use `http://127.0.0.1:48100` as its +provider. Keep those processes running and pass their public daemon URLs to the +load driver. A live MCP registration must own each public route for the entire +run. Use a pass-through daemon for `daemon-pass-through` and a normal daemon +with its activated worker for `daemon-worker`. + +Add `--worker-binary` to create the broker-authenticated `worker-only` target; +there is no supported direct `--target worker-only=...` form: + +```bash +export NEMO_RELAY_CLIENT_TOKEN='...' + +just daemon-transport-benchmark \ + --direct-url http://127.0.0.1:48100 \ + --worker-binary target/release/nemo-relay \ + --target daemon-pass-through=http://127.0.0.1:47632 \ + --target daemon-worker=http://127.0.0.1:47633 \ + --header-env daemon-pass-through:x-nemo-relay-client-token=NEMO_RELAY_CLIENT_TOKEN \ + --header-env daemon-worker:x-nemo-relay-client-token=NEMO_RELAY_CLIENT_TOKEN \ + --pid pass-through-daemon=1234 \ + --pid worker-daemon=1235 \ + --binary-metadata release=target/release/nemo-relay \ + --binary-metadata opt-level-3=target/daemon-benchmark-opt3/release/nemo-relay +``` + +Header values are read from environment variables and are neither serialized +into the report nor printed. Only public daemon targets accept header bindings; +the worker credential comes from the verified activation flow and never enters +the environment or command arguments. The report records configured header +names, not their values. + +Run the command once against processes from the normal release build and once +against equivalent processes from the `opt-level=3` build. Use a distinct +`--output` path for each run. Compare runs made on the same otherwise-idle +host. + +The full preset performs a 10-second warmup followed by a 60-second measured +interval for each configured topology and each combination of: + +- direct provider (always), daemon pass-through (when supplied), worker-only + (with `--worker-binary`), and daemon-plus-worker (when supplied); +- OpenAI Responses and Anthropic Messages; +- HTTP/1.1 and HTTP/2; +- 16 KiB and 1 MiB streamed responses with 128 events; +- concurrency 1, 16, 64, and 256. + +It also starts a separate 1,000-slow-stream capacity scenario for each +configured topology and protocol. One percent of the sustained requests are +cancelled after first content to exercise cancellation propagation. Remove +public targets or omit `--worker-binary` to benchmark a smaller subset; the +direct provider baseline is always present. + +## Result Schema + +The JSON report records response-head, first-content, per-event forwarding, +and total latency distributions; requests per second; MiB/s goodput; process +CPU and peak RSS samples; RSS growth per active stream; connection attempts; +estimated pool reuse; active HTTP/2 streams; cancellation and reconnect +counts; and missing, duplicate, reordered, cross-stream, corrupt, hash, status, +and trailer errors. + +Hyper does not expose queue depth or backpressure-stall counters from outside +the daemon and worker. Those fields are present as `null` with an explicit +reason. Supply daemon/worker telemetry when that instrumentation becomes +available instead of deriving misleading estimates from client timings. + +Body delivery is consumed frame-by-frame. The benchmark parser observes SSE +semantics only to calculate integrity and event-delay metrics; it is not part +of Relay's delivery path. diff --git a/scripts/latency_benchmark/daemon_transport/src/client.rs b/scripts/latency_benchmark/daemon_transport/src/client.rs new file mode 100644 index 000000000..0554c81fb --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/src/client.rs @@ -0,0 +1,932 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{BTreeMap, HashSet}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +use anyhow::{Context as _, Result, anyhow, ensure}; +use bytes::Bytes; +use hdrhistogram::Histogram; +use http::header::{ACCEPT, AUTHORIZATION, CONTENT_TYPE, TE}; +use http::{Request, Uri}; +use http_body_util::{BodyExt, Full}; +use hyper_rustls::{HttpsConnector, HttpsConnectorBuilder}; +use hyper_util::client::legacy::Client; +use hyper_util::client::legacy::connect::HttpConnector; +use hyper_util::rt::{TokioExecutor, TokioTimer}; +use serde::Serialize; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use tower_service::Service; + +use crate::config::{LoadOptions, MatrixConfig, Protocol, Provider, Target, Topology}; +use crate::metadata::{self, EnvironmentRecord}; +use crate::provider; +use crate::resources::{ResourceRecord, ResourceSampler}; + +const RESPONSE_BYTES: &str = "x-benchmark-response-bytes"; +const EVENT_COUNT: &str = "x-benchmark-event-count"; +const EVENT_DELAY_MICROS: &str = "x-benchmark-event-delay-micros"; +const STREAM_ID: &str = "x-benchmark-stream-id"; +const BODY_SHA256: &str = "x-benchmark-body-sha256"; +const MAX_HISTOGRAM_NANOS: u64 = 3_600_000_000_000; + +type BaseConnector = HttpsConnector; + +#[derive(Clone)] +struct CountingConnector { + inner: BaseConnector, + connections: Arc, +} + +impl Service for CountingConnector { + type Response = >::Response; + type Error = >::Error; + type Future = >::Future; + + fn poll_ready(&mut self, context: &mut Context<'_>) -> Poll> { + self.inner.poll_ready(context) + } + + fn call(&mut self, uri: Uri) -> Self::Future { + self.connections.fetch_add(1, Ordering::Relaxed); + self.inner.call(uri) + } +} + +#[derive(Clone)] +struct BenchmarkClient { + client: Client>, + connections: Arc, +} + +impl BenchmarkClient { + fn new(protocol: Protocol) -> Result { + let mut http = HttpConnector::new(); + http.enforce_http(false); + http.set_nodelay(true); + let connector = match protocol { + Protocol::Http1 => HttpsConnectorBuilder::new() + .with_native_roots() + .context("failed to load native TLS roots")? + .https_or_http() + .enable_http1() + .wrap_connector(http), + Protocol::Http2 => HttpsConnectorBuilder::new() + .with_native_roots() + .context("failed to load native TLS roots")? + .https_or_http() + .enable_http2() + .wrap_connector(http), + }; + let connections = Arc::new(AtomicU64::new(0)); + let connector = CountingConnector { + inner: connector, + connections: Arc::clone(&connections), + }; + let mut builder = Client::builder(TokioExecutor::new()); + builder.timer(TokioTimer::new()); + builder.pool_idle_timeout(Duration::from_secs(120)); + builder.pool_max_idle_per_host(usize::MAX); + if matches!(protocol, Protocol::Http2) { + builder.http2_only(true); + builder.http2_keep_alive_interval(Duration::from_secs(15)); + builder.http2_keep_alive_timeout(Duration::from_secs(5)); + builder.http2_keep_alive_while_idle(true); + } + Ok(Self { + client: builder.build(connector), + connections, + }) + } + + fn connections(&self) -> u64 { + self.connections.load(Ordering::Relaxed) + } +} + +struct Clients { + http1: BenchmarkClient, + http2: BenchmarkClient, +} + +impl Clients { + fn new() -> Result { + Ok(Self { + http1: BenchmarkClient::new(Protocol::Http1)?, + http2: BenchmarkClient::new(Protocol::Http2)?, + }) + } + + fn get(&self, protocol: Protocol) -> &BenchmarkClient { + match protocol { + Protocol::Http1 => &self.http1, + Protocol::Http2 => &self.http2, + } + } +} + +#[derive(Debug, Serialize)] +struct Report { + schema_version: u8, + informational_only: bool, + environment: EnvironmentRecord, + parameters: MatrixConfig, + targets: Vec, + scenarios: Vec, + validation_errors: Vec, +} + +#[derive(Debug, Serialize)] +struct TargetRecord { + topology: Topology, + url: String, + configured_header_names: Vec, +} + +#[derive(Debug, Serialize)] +struct ScenarioRecord { + kind: &'static str, + topology: Topology, + protocol: Protocol, + provider: Provider, + response_bytes: usize, + events: usize, + concurrency: usize, + configured_duration_seconds: Option, + measured_duration_seconds: f64, + requests_started: u64, + requests_completed: u64, + requests_cancelled: u64, + requests_per_second: f64, + goodput_mib_per_second: f64, + response_head: HistogramRecord, + first_content: HistogramRecord, + per_event_forwarding_delay: HistogramRecord, + total: HistogramRecord, + integrity: IntegrityRecord, + transport_errors: u64, + connections_opened_during_measurement: u64, + estimated_pool_reuses: u64, + reconnect_count: Option, + max_active_requests: usize, + queued_bytes_peak: Option, + backpressure_stalls: Option, + resources: BTreeMap, + unavailable_metrics: Vec<&'static str>, +} + +#[derive(Debug, Serialize)] +struct HistogramRecord { + samples: u64, + p50_ms: Option, + p95_ms: Option, + p99_ms: Option, + min_ms: Option, + max_ms: Option, +} + +impl HistogramRecord { + fn from_histogram(histogram: &Histogram) -> Self { + if histogram.is_empty() { + return Self { + samples: 0, + p50_ms: None, + p95_ms: None, + p99_ms: None, + min_ms: None, + max_ms: None, + }; + } + Self { + samples: histogram.len(), + p50_ms: Some(nanos_to_millis(histogram.value_at_quantile(0.50))), + p95_ms: Some(nanos_to_millis(histogram.value_at_quantile(0.95))), + p99_ms: Some(nanos_to_millis(histogram.value_at_quantile(0.99))), + min_ms: Some(nanos_to_millis(histogram.min())), + max_ms: Some(nanos_to_millis(histogram.max())), + } + } +} + +#[derive(Debug, Default, Serialize)] +struct IntegrityRecord { + missing_events: u64, + duplicate_events: u64, + reordered_events: u64, + corrupt_events: u64, + cross_stream_events: u64, + body_hash_mismatches: u64, + trailer_mismatches: u64, + status_errors: u64, +} + +struct Accumulator { + head: Histogram, + first: Histogram, + event_delay: Histogram, + total: Histogram, + started: u64, + completed: u64, + cancelled: u64, + bytes: u64, + integrity: IntegrityRecord, + transport_errors: u64, +} + +impl Accumulator { + fn new() -> Self { + Self { + head: latency_histogram(), + first: latency_histogram(), + event_delay: latency_histogram(), + total: latency_histogram(), + started: 0, + completed: 0, + cancelled: 0, + bytes: 0, + integrity: IntegrityRecord::default(), + transport_errors: 0, + } + } + + fn record(&mut self, result: Result) { + self.started += 1; + match result { + Ok(observation) => { + record_duration(&mut self.head, observation.head); + record_duration(&mut self.first, observation.first); + if observation.cancelled { + self.cancelled += 1; + return; + } + record_duration(&mut self.total, observation.total); + for delay in observation.event_delays { + let _ = self.event_delay.record(delay.clamp(1, MAX_HISTOGRAM_NANOS)); + } + self.completed += 1; + self.bytes += observation.bytes as u64; + self.integrity.add(observation.integrity); + } + Err(_) => self.transport_errors += 1, + } + } + + fn add(&mut self, other: Self) -> Result<()> { + self.head.add(&other.head)?; + self.first.add(&other.first)?; + self.event_delay.add(&other.event_delay)?; + self.total.add(&other.total)?; + self.started += other.started; + self.completed += other.completed; + self.cancelled += other.cancelled; + self.bytes += other.bytes; + self.integrity.add(other.integrity); + self.transport_errors += other.transport_errors; + Ok(()) + } +} + +impl IntegrityRecord { + fn add(&mut self, other: Self) { + self.missing_events += other.missing_events; + self.duplicate_events += other.duplicate_events; + self.reordered_events += other.reordered_events; + self.corrupt_events += other.corrupt_events; + self.cross_stream_events += other.cross_stream_events; + self.body_hash_mismatches += other.body_hash_mismatches; + self.trailer_mismatches += other.trailer_mismatches; + self.status_errors += other.status_errors; + } + + fn error_count(&self) -> u64 { + self.missing_events + + self.duplicate_events + + self.reordered_events + + self.corrupt_events + + self.cross_stream_events + + self.body_hash_mismatches + + self.trailer_mismatches + + self.status_errors + } +} + +struct Observation { + head: Duration, + first: Duration, + total: Duration, + event_delays: Vec, + bytes: usize, + integrity: IntegrityRecord, + cancelled: bool, +} + +pub async fn run(options: LoadOptions) -> Result<()> { + let environment = metadata::collect(&options.binaries)?; + let clients = Clients::new()?; + let request_sequence = Arc::new(AtomicU64::new(0)); + let mut scenarios = Vec::new(); + let mut validation_errors = Vec::new(); + + for protocol in &options.matrix.protocols { + for provider in &options.matrix.providers { + for response_bytes in &options.matrix.response_bytes { + for concurrency in &options.matrix.concurrency { + for target in &options.targets { + println!( + "benchmarking {} {} {} bytes={} concurrency={}", + target.topology, protocol, provider, response_bytes, concurrency + ); + let scenario = run_sustained_scenario( + clients.get(*protocol), + target, + *protocol, + *provider, + *response_bytes, + *concurrency, + &options, + Arc::clone(&request_sequence), + ) + .await?; + collect_validation_errors(&scenario, &mut validation_errors); + scenarios.push(scenario); + } + } + } + } + } + + if options.matrix.slow_streams > 0 { + for protocol in &options.matrix.protocols { + for target in &options.targets { + println!( + "benchmarking slow capacity {} {} streams={}", + target.topology, protocol, options.matrix.slow_streams + ); + let scenario = run_slow_scenario( + clients.get(*protocol), + target, + *protocol, + &options, + Arc::clone(&request_sequence), + ) + .await?; + collect_validation_errors(&scenario, &mut validation_errors); + scenarios.push(scenario); + } + } + } + + let report = Report { + schema_version: 1, + informational_only: true, + environment, + parameters: options.matrix, + targets: options + .targets + .iter() + .map(|target| TargetRecord { + topology: target.topology, + url: target.url.clone(), + configured_header_names: target + .headers + .iter() + .map(|(name, _)| name.to_string()) + .collect(), + }) + .collect(), + scenarios, + validation_errors, + }; + metadata::write_json(&options.output, &report)?; + println!("daemon transport report: {}", options.output.display()); + ensure!( + report.validation_errors.is_empty(), + "transport correctness validation failed: {}", + report.validation_errors.join("; ") + ); + Ok(()) +} + +pub async fn run_smoke(output: std::path::PathBuf) -> Result<()> { + let (url, stop, provider_task) = provider::spawn_ephemeral().await?; + let result = run(LoadOptions::smoke(output, url)).await; + let _ = stop.send(()); + provider_task + .await + .context("smoke provider task failed")??; + result +} + +#[allow(clippy::too_many_arguments)] +async fn run_sustained_scenario( + client: &BenchmarkClient, + target: &Target, + protocol: Protocol, + provider: Provider, + response_bytes: usize, + concurrency: usize, + options: &LoadOptions, + request_sequence: Arc, +) -> Result { + if options.matrix.warmup_seconds > 0 { + let _ = run_timed_phase( + client, + target, + provider, + response_bytes, + options.matrix.events, + concurrency, + Duration::from_secs(options.matrix.warmup_seconds), + options.matrix.event_delay_micros, + 0, + Arc::clone(&request_sequence), + Arc::new(AtomicUsize::new(0)), + Arc::new(AtomicUsize::new(0)), + ) + .await?; + } + + let connections_before = client.connections(); + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let sampler = ResourceSampler::start(options.processes.clone()).await; + let started = Instant::now(); + let accumulator = run_timed_phase( + client, + target, + provider, + response_bytes, + options.matrix.events, + concurrency, + Duration::from_secs(options.matrix.duration_seconds), + options.matrix.event_delay_micros, + options.matrix.cancel_every, + request_sequence, + Arc::clone(&active), + Arc::clone(&max_active), + ) + .await?; + let elapsed = started.elapsed(); + let resources = sampler.finish(max_active.load(Ordering::Relaxed)).await; + let connections = client.connections().saturating_sub(connections_before); + Ok(make_scenario( + "sustained", + target.topology, + protocol, + provider, + response_bytes, + options.matrix.events, + concurrency, + Some(options.matrix.duration_seconds), + elapsed, + accumulator, + connections, + max_active.load(Ordering::Relaxed), + resources, + )) +} + +async fn run_slow_scenario( + client: &BenchmarkClient, + target: &Target, + protocol: Protocol, + options: &LoadOptions, + request_sequence: Arc, +) -> Result { + let concurrency = options.matrix.slow_streams; + let response_bytes = options.matrix.response_bytes[0]; + let connections_before = client.connections(); + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let sampler = ResourceSampler::start(options.processes.clone()).await; + let started = Instant::now(); + let mut tasks = Vec::with_capacity(concurrency); + for _ in 0..concurrency { + let client = client.clone(); + let target = target.clone(); + let request_sequence = Arc::clone(&request_sequence); + let active = Arc::clone(&active); + let max_active = Arc::clone(&max_active); + let events = options.matrix.events; + let delay = options.matrix.slow_event_delay_millis.saturating_mul(1_000); + tasks.push(tokio::spawn(async move { + let mut accumulator = Accumulator::new(); + let ordinal = request_sequence.fetch_add(1, Ordering::Relaxed); + accumulator.record( + perform_request( + &client, + &target, + Provider::Openai, + response_bytes, + events, + delay, + false, + ordinal, + active, + max_active, + ) + .await, + ); + accumulator + })); + } + let mut accumulator = Accumulator::new(); + for task in tasks { + accumulator.add(task.await.context("slow-stream load task failed")?)?; + } + let elapsed = started.elapsed(); + let resources = sampler.finish(max_active.load(Ordering::Relaxed)).await; + let connections = client.connections().saturating_sub(connections_before); + Ok(make_scenario( + "slow-capacity", + target.topology, + protocol, + Provider::Openai, + response_bytes, + options.matrix.events, + concurrency, + None, + elapsed, + accumulator, + connections, + max_active.load(Ordering::Relaxed), + resources, + )) +} + +#[allow(clippy::too_many_arguments)] +async fn run_timed_phase( + client: &BenchmarkClient, + target: &Target, + provider: Provider, + response_bytes: usize, + events: usize, + concurrency: usize, + duration: Duration, + event_delay_micros: u64, + cancel_every: usize, + request_sequence: Arc, + active: Arc, + max_active: Arc, +) -> Result { + let barrier = Arc::new(tokio::sync::Barrier::new(concurrency)); + let mut tasks = Vec::with_capacity(concurrency); + for _ in 0..concurrency { + let client = client.clone(); + let target = target.clone(); + let barrier = Arc::clone(&barrier); + let request_sequence = Arc::clone(&request_sequence); + let active = Arc::clone(&active); + let max_active = Arc::clone(&max_active); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + let deadline = Instant::now() + duration; + let mut accumulator = Accumulator::new(); + while Instant::now() < deadline { + let ordinal = request_sequence.fetch_add(1, Ordering::Relaxed); + let cancel = cancel_every > 0 && ordinal.is_multiple_of(cancel_every as u64); + accumulator.record( + perform_request( + &client, + &target, + provider, + response_bytes, + events, + event_delay_micros, + cancel, + ordinal, + Arc::clone(&active), + Arc::clone(&max_active), + ) + .await, + ); + } + accumulator + })); + } + let mut accumulator = Accumulator::new(); + for task in tasks { + accumulator.add(task.await.context("load task failed")?)?; + } + Ok(accumulator) +} + +#[allow(clippy::too_many_arguments)] +async fn perform_request( + client: &BenchmarkClient, + target: &Target, + provider: Provider, + response_bytes: usize, + events: usize, + event_delay_micros: u64, + cancel_after_first: bool, + request_ordinal: u64, + active: Arc, + max_active: Arc, +) -> Result { + let stream_id = format!("{request_ordinal:016x}"); + let uri: Uri = format!("{}{}", target.url, provider.path()) + .parse() + .context("failed to construct target URI")?; + let mut builder = Request::post(uri) + .header(ACCEPT, "text/event-stream") + .header(TE, "trailers") + .header(CONTENT_TYPE, "application/json") + .header(AUTHORIZATION, "Bearer benchmark-provider-token") + .header(RESPONSE_BYTES, response_bytes) + .header(EVENT_COUNT, events) + .header(EVENT_DELAY_MICROS, event_delay_micros) + .header(STREAM_ID, &stream_id); + for (name, value) in &target.headers { + builder = builder.header(name, value); + } + let body = match provider { + Provider::Openai => Bytes::from_static(b"{\"model\":\"benchmark\",\"stream\":true,\"input\":\"ping\"}"), + Provider::Anthropic => Bytes::from_static( + b"{\"model\":\"benchmark\",\"stream\":true,\"max_tokens\":1024,\"messages\":[{\"role\":\"user\",\"content\":\"ping\"}]}", + ), + }; + let request = builder + .body(Full::new(body)) + .context("failed to build benchmark request")?; + let started = Instant::now(); + let active_now = active.fetch_add(1, Ordering::Relaxed) + 1; + max_active.fetch_max(active_now, Ordering::Relaxed); + let _active_guard = ActiveGuard(active); + let response = client + .client + .request(request) + .await + .context("request failed")?; + let head = started.elapsed(); + let mut integrity = IntegrityRecord::default(); + if !response.status().is_success() { + integrity.status_errors += 1; + } + let expected_events = response + .headers() + .get(EVENT_COUNT) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse().ok()) + .unwrap_or(events); + let expected_stream = response + .headers() + .get(STREAM_ID) + .and_then(|value| value.to_str().ok()) + .unwrap_or(&stream_id) + .to_owned(); + let mut body = response.into_body(); + let mut tracker = EventTracker::new(expected_stream, expected_events); + let mut hasher = Sha256::new(); + let mut bytes = 0; + let mut first = None; + let mut trailer_hash = None; + let mut trailer_events = None; + while let Some(frame) = body.frame().await { + let frame = frame.context("response body failed")?; + match frame.into_data() { + Ok(data) => { + bytes += data.len(); + hasher.update(&data); + let before = tracker.unique_events(); + tracker.push(&data); + if tracker.unique_events() > before && first.is_none() { + first = Some(started.elapsed()); + if cancel_after_first { + return Ok(Observation { + head, + first: first.expect("first content recorded"), + total: started.elapsed(), + event_delays: Vec::new(), + bytes, + integrity: IntegrityRecord::default(), + cancelled: true, + }); + } + } + } + Err(frame) => { + if let Ok(trailers) = frame.into_trailers() { + trailer_hash = trailers + .get(BODY_SHA256) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + trailer_events = trailers + .get(EVENT_COUNT) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()); + } + } + } + } + tracker.finish(); + integrity.add(tracker.integrity); + let actual_hash = format!("{:x}", hasher.finalize()); + if trailer_hash.as_deref() != Some(&actual_hash) { + integrity.body_hash_mismatches += 1; + } + if trailer_events != Some(expected_events) { + integrity.trailer_mismatches += 1; + } + Ok(Observation { + head, + first: first.ok_or_else(|| anyhow!("response contained no content event"))?, + total: started.elapsed(), + event_delays: tracker.event_delays, + bytes, + integrity, + cancelled: false, + }) +} + +struct ActiveGuard(Arc); + +impl Drop for ActiveGuard { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Relaxed); + } +} + +struct EventTracker { + buffer: Vec, + expected_stream: String, + expected_events: usize, + seen: HashSet, + last_sequence: Option, + event_delays: Vec, + integrity: IntegrityRecord, +} + +impl EventTracker { + fn new(expected_stream: String, expected_events: usize) -> Self { + Self { + buffer: Vec::new(), + expected_stream, + expected_events, + seen: HashSet::with_capacity(expected_events), + last_sequence: None, + event_delays: Vec::with_capacity(expected_events), + integrity: IntegrityRecord::default(), + } + } + + fn unique_events(&self) -> usize { + self.seen.len() + } + + fn push(&mut self, bytes: &[u8]) { + self.buffer.extend_from_slice(bytes); + while let Some(end) = self.buffer.windows(2).position(|window| window == b"\n\n") { + let event = self.buffer.drain(..end + 2).collect::>(); + self.parse_event(&event); + } + } + + fn parse_event(&mut self, event: &[u8]) { + for line in event.split(|byte| *byte == b'\n') { + let Some(data) = line.strip_prefix(b"data: ") else { + continue; + }; + if data == b"[DONE]" { + continue; + } + let parsed = serde_json::from_slice::(data); + let Ok(parsed) = parsed else { + self.integrity.corrupt_events += 1; + continue; + }; + let Some(sequence) = parsed + .get("s") + .and_then(Value::as_u64) + .map(|value| value as usize) + else { + self.integrity.corrupt_events += 1; + continue; + }; + if parsed.get("i").and_then(Value::as_str) != Some(&self.expected_stream) { + self.integrity.cross_stream_events += 1; + } + if !self.seen.insert(sequence) { + self.integrity.duplicate_events += 1; + } + if self.last_sequence.is_some_and(|last| sequence <= last) { + self.integrity.reordered_events += 1; + } + self.last_sequence = Some(sequence); + let emitted = parsed + .get("t") + .and_then(Value::as_str) + .and_then(|value| value.parse::().ok()); + if let Some(emitted) = emitted { + let delay = unix_time_nanos() + .saturating_sub(emitted) + .min(u64::MAX as u128) as u64; + self.event_delays.push(delay.max(1)); + } else { + self.integrity.corrupt_events += 1; + } + } + } + + fn finish(&mut self) { + if self.buffer.iter().any(|byte| !byte.is_ascii_whitespace()) { + self.integrity.corrupt_events += 1; + } + self.integrity.missing_events += + self.expected_events.saturating_sub(self.seen.len()) as u64; + } +} + +#[allow(clippy::too_many_arguments)] +fn make_scenario( + kind: &'static str, + topology: Topology, + protocol: Protocol, + provider: Provider, + response_bytes: usize, + events: usize, + concurrency: usize, + configured_duration_seconds: Option, + elapsed: Duration, + accumulator: Accumulator, + connections: u64, + max_active: usize, + resources: BTreeMap, +) -> ScenarioRecord { + let seconds = elapsed.as_secs_f64().max(f64::EPSILON); + let pool_reuses = accumulator.started.saturating_sub(connections); + ScenarioRecord { + kind, + topology, + protocol, + provider, + response_bytes, + events, + concurrency, + configured_duration_seconds, + measured_duration_seconds: seconds, + requests_started: accumulator.started, + requests_completed: accumulator.completed, + requests_cancelled: accumulator.cancelled, + requests_per_second: accumulator.completed as f64 / seconds, + goodput_mib_per_second: accumulator.bytes as f64 / (1024.0 * 1024.0) / seconds, + response_head: HistogramRecord::from_histogram(&accumulator.head), + first_content: HistogramRecord::from_histogram(&accumulator.first), + per_event_forwarding_delay: HistogramRecord::from_histogram(&accumulator.event_delay), + total: HistogramRecord::from_histogram(&accumulator.total), + integrity: accumulator.integrity, + transport_errors: accumulator.transport_errors, + connections_opened_during_measurement: connections, + estimated_pool_reuses: pool_reuses, + reconnect_count: None, + max_active_requests: max_active, + queued_bytes_peak: None, + backpressure_stalls: None, + resources, + unavailable_metrics: vec![ + "reconnect count requires connection lifecycle instrumentation", + "per-connection active HTTP/2 streams require connection-level instrumentation", + "queued byte depth requires daemon/worker instrumentation", + "backpressure stall count requires daemon/worker instrumentation", + ], + } +} + +fn collect_validation_errors(scenario: &ScenarioRecord, errors: &mut Vec) { + if scenario.integrity.error_count() > 0 || scenario.transport_errors > 0 { + errors.push(format!( + "{} {} {} bytes={} concurrency={} had {} integrity and {} transport errors", + scenario.topology, + scenario.protocol, + scenario.provider, + scenario.response_bytes, + scenario.concurrency, + scenario.integrity.error_count(), + scenario.transport_errors + )); + } +} + +fn latency_histogram() -> Histogram { + Histogram::new_with_bounds(1, MAX_HISTOGRAM_NANOS, 3).expect("valid latency histogram") +} + +fn record_duration(histogram: &mut Histogram, duration: Duration) { + let value = duration.as_nanos().min(MAX_HISTOGRAM_NANOS as u128) as u64; + let _ = histogram.record(value.max(1)); +} + +fn nanos_to_millis(value: u64) -> f64 { + value as f64 / 1_000_000.0 +} + +fn unix_time_nanos() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() +} diff --git a/scripts/latency_benchmark/daemon_transport/src/config.rs b/scripts/latency_benchmark/daemon_transport/src/config.rs new file mode 100644 index 000000000..d0a086fb2 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/src/config.rs @@ -0,0 +1,365 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{BTreeMap, BTreeSet}; +use std::env; +use std::fmt; +use std::path::{Path, PathBuf}; +use std::str::FromStr; + +use anyhow::{Context, Result, bail, ensure}; +use http::{HeaderName, HeaderValue}; +use serde::{Deserialize, Serialize}; + +use crate::metadata::BinarySpec; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize)] +#[serde(rename_all = "kebab-case")] +pub enum Protocol { + Http1, + Http2, +} + +impl fmt::Display for Protocol { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Http1 => formatter.write_str("http1"), + Self::Http2 => formatter.write_str("http2"), + } + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum Provider { + Openai, + Anthropic, +} + +impl fmt::Display for Provider { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Openai => formatter.write_str("openai"), + Self::Anthropic => formatter.write_str("anthropic"), + } + } +} + +impl Provider { + pub fn path(self) -> &'static str { + match self { + Self::Openai => "/v1/responses", + Self::Anthropic => "/v1/messages", + } + } +} + +#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)] +#[serde(rename_all = "kebab-case")] +pub enum Topology { + Direct, + DaemonPassThrough, + WorkerOnly, + DaemonWorker, +} + +impl fmt::Display for Topology { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Direct => formatter.write_str("direct"), + Self::DaemonPassThrough => formatter.write_str("daemon-pass-through"), + Self::WorkerOnly => formatter.write_str("worker-only"), + Self::DaemonWorker => formatter.write_str("daemon-worker"), + } + } +} + +impl FromStr for Topology { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + match value { + "direct" => Ok(Self::Direct), + "daemon-pass-through" => Ok(Self::DaemonPassThrough), + "worker-only" => Ok(Self::WorkerOnly), + "daemon-worker" => Ok(Self::DaemonWorker), + _ => bail!("unknown topology {value:?}"), + } + } +} + +#[derive(Clone, Debug)] +pub struct TargetSpec { + pub topology: Topology, + pub url: String, +} + +impl FromStr for TargetSpec { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + let (name, url) = value.split_once('=').context("target must use NAME=URL")?; + let topology = Topology::from_str(name)?; + ensure!( + topology != Topology::Direct, + "direct is supplied with --direct-url" + ); + validate_url(url)?; + Ok(Self { + topology, + url: url.trim_end_matches('/').to_owned(), + }) + } +} + +#[derive(Clone, Debug)] +pub struct TargetHeader { + pub topology: Topology, + pub name: HeaderName, + pub value: HeaderValue, +} + +impl FromStr for TargetHeader { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + let (target_and_header, environment_name) = value + .split_once('=') + .context("header binding must use TARGET:HEADER=ENV_NAME")?; + let (target, header) = target_and_header + .split_once(':') + .context("header binding must use TARGET:HEADER=ENV_NAME")?; + ensure!( + !environment_name.is_empty(), + "header environment name is empty" + ); + let topology = Topology::from_str(target)?; + let name = HeaderName::from_str(header).context("invalid HTTP header name")?; + let raw_value = env::var(environment_name) + .with_context(|| format!("environment variable {environment_name} is not set"))?; + let value = HeaderValue::from_str(&raw_value).context("invalid HTTP header value")?; + Ok(Self { + topology, + name, + value, + }) + } +} + +#[derive(Clone, Debug)] +pub struct ProcessSpec { + pub name: String, + pub pid: u32, +} + +impl FromStr for ProcessSpec { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + let (name, pid) = value.split_once('=').context("process must use NAME=PID")?; + ensure!(!name.is_empty(), "process name is empty"); + Ok(Self { + name: name.to_owned(), + pid: pid.parse().context("invalid process ID")?, + }) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct MatrixConfig { + pub protocols: Vec, + pub providers: Vec, + pub response_bytes: Vec, + pub events: usize, + pub concurrency: Vec, + pub warmup_seconds: u64, + pub duration_seconds: u64, + pub event_delay_micros: u64, + pub cancel_every: usize, + pub slow_streams: usize, + pub slow_event_delay_millis: u64, +} + +impl MatrixConfig { + pub fn validate(&self) -> Result<()> { + ensure!(!self.protocols.is_empty(), "protocols must not be empty"); + ensure!(!self.providers.is_empty(), "providers must not be empty"); + ensure!( + !self.response_bytes.is_empty(), + "response_bytes must not be empty" + ); + ensure!( + !self.concurrency.is_empty(), + "concurrency must not be empty" + ); + ensure!(self.events >= 128, "events must be at least 128"); + ensure!( + self.duration_seconds > 0, + "duration_seconds must be positive" + ); + ensure!( + self.response_bytes.iter().all(|size| *size > 0), + "response sizes must be positive" + ); + ensure!( + self.concurrency.iter().all(|value| *value > 0), + "concurrency must be positive" + ); + Ok(()) + } +} + +#[derive(Clone, Debug)] +pub struct Target { + pub topology: Topology, + pub url: String, + pub headers: Vec<(HeaderName, HeaderValue)>, +} + +#[derive(Clone, Debug)] +pub struct LoadOptions { + pub matrix: MatrixConfig, + pub targets: Vec, + pub processes: Vec, + pub binaries: Vec, + pub output: PathBuf, +} + +impl LoadOptions { + #[allow(clippy::too_many_arguments)] + pub fn from_file( + path: &Path, + direct_url: String, + targets: Vec, + headers: Vec, + processes: Vec, + binaries: Vec, + output: PathBuf, + ) -> Result { + validate_url(&direct_url)?; + let contents = std::fs::read_to_string(path) + .with_context(|| format!("failed to read {}", path.display()))?; + let matrix: MatrixConfig = toml::from_str(&contents) + .with_context(|| format!("failed to parse {}", path.display()))?; + matrix.validate()?; + + let mut seen = BTreeSet::from([Topology::Direct]); + let mut target_map = BTreeMap::from([( + Topology::Direct, + Target { + topology: Topology::Direct, + url: direct_url.trim_end_matches('/').to_owned(), + headers: Vec::new(), + }, + )]); + for target in targets { + ensure!( + target.topology != Topology::WorkerOnly, + "worker-only cannot be supplied as a direct target because production workers require a broker-private credential; use --worker-binary" + ); + ensure!( + seen.insert(target.topology), + "duplicate target {}", + target.topology + ); + target_map.insert( + target.topology, + Target { + topology: target.topology, + url: target.url, + headers: Vec::new(), + }, + ); + } + for header in headers { + let target = target_map + .get_mut(&header.topology) + .with_context(|| format!("header refers to absent target {}", header.topology))?; + ensure!( + !target.headers.iter().any(|(name, _)| name == header.name), + "duplicate header {} for {}", + header.name, + header.topology + ); + target.headers.push((header.name, header.value)); + } + for target in target_map.values() { + if !target.headers.is_empty() { + validate_credential_target(&target.url)?; + } + } + + Ok(Self { + matrix, + targets: target_map.into_values().collect(), + processes, + binaries, + output, + }) + } + + pub fn smoke(output: PathBuf, direct_url: String) -> Self { + Self { + matrix: MatrixConfig { + protocols: vec![Protocol::Http1, Protocol::Http2], + providers: vec![Provider::Openai, Provider::Anthropic], + response_bytes: vec![16 * 1024], + events: 128, + concurrency: vec![1], + warmup_seconds: 0, + duration_seconds: 1, + event_delay_micros: 0, + cancel_every: 0, + slow_streams: 0, + slow_event_delay_millis: 0, + }, + targets: vec![Target { + topology: Topology::Direct, + url: direct_url, + headers: Vec::new(), + }], + processes: Vec::new(), + binaries: Vec::new(), + output, + } + } +} + +fn validate_url(url: &str) -> Result<()> { + let uri = http::Uri::from_str(url).context("invalid target URL")?; + ensure!( + matches!(uri.scheme_str(), Some("http" | "https")), + "URL must use http or https" + ); + ensure!(uri.authority().is_some(), "URL must contain an authority"); + ensure!( + !uri.authority() + .is_some_and(|authority| authority.as_str().contains('@')), + "target URL must not contain userinfo" + ); + ensure!( + uri.path_and_query() + .is_none_or(|value| value.as_str() == "/"), + "target URL must not contain a path or query" + ); + Ok(()) +} + +fn validate_credential_target(url: &str) -> Result<()> { + let uri = http::Uri::from_str(url).context("invalid target URL")?; + if uri.scheme_str() != Some("http") { + return Ok(()); + } + let host = uri.host().context("target URL must contain a host")?; + let loopback = host.eq_ignore_ascii_case("localhost") + || host + .trim_matches(['[', ']']) + .parse::() + .is_ok_and(|address| address.is_loopback()); + ensure!( + loopback, + "credential-bearing benchmark targets must use HTTPS unless they are loopback" + ); + Ok(()) +} diff --git a/scripts/latency_benchmark/daemon_transport/src/main.rs b/scripts/latency_benchmark/daemon_transport/src/main.rs new file mode 100644 index 000000000..22039ff9f --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/src/main.rs @@ -0,0 +1,137 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +mod client; +mod config; +mod metadata; +mod orchestrate; +mod provider; +mod resources; + +use std::net::SocketAddr; +use std::path::PathBuf; + +use anyhow::Result; +use clap::{Args, Parser, Subcommand}; +use config::{LoadOptions, ProcessSpec, TargetHeader, TargetSpec}; + +#[derive(Debug, Parser)] +#[command(about = "Hyper-based transport benchmark for the NeMo Relay daemon")] +struct Cli { + #[command(subcommand)] + command: Command, +} + +#[derive(Debug, Subcommand)] +enum Command { + /// Run the deterministic OpenAI/Anthropic streaming provider. + Provider(ProviderArgs), + /// Benchmark already-running direct and Relay topology endpoints. + Load(LoadArgs), + /// Run the short, direct-provider HTTP/1.1 and HTTP/2 CI check. + Smoke(SmokeArgs), +} + +#[derive(Debug, Args)] +struct ProviderArgs { + #[arg(long, default_value = "127.0.0.1:48100")] + bind: SocketAddr, + + /// Write the selected URL after the listener is bound. + #[arg(long)] + ready_file: Option, +} + +#[derive(Debug, Args)] +struct LoadArgs { + #[arg( + long, + default_value = "scripts/latency_benchmark/config/daemon-transport-full.toml" + )] + config: PathBuf, + + /// URL of the deterministic provider, used as the direct baseline. + #[arg(long)] + direct_url: String, + + /// Additional public topology endpoint as NAME=URL. Names are daemon-pass-through and daemon-worker. + #[arg(long = "target")] + targets: Vec, + + /// Relay binary used to securely orchestrate a directly measured worker-only target. + #[arg(long)] + worker_binary: Option, + + /// Read a target-specific request header from the environment as TARGET:HEADER=ENV_NAME. + #[arg(long = "header-env")] + headers: Vec, + + /// Process sampled for CPU/RSS metadata as NAME=PID. The load driver is always sampled. + #[arg(long = "pid")] + processes: Vec, + + /// Relay binary recorded in metadata as PROFILE=PATH. Values are hashed but never executed. + #[arg(long = "binary-metadata")] + binaries: Vec, + + #[arg(long, default_value = "target/benchmark-results/daemon-transport.json")] + output: PathBuf, +} + +#[derive(Debug, Args)] +struct SmokeArgs { + #[arg( + long, + default_value = "target/benchmark-results/daemon-transport-smoke.json" + )] + output: PathBuf, +} + +#[tokio::main] +async fn main() -> Result<()> { + match Cli::parse().command { + Command::Provider(args) => provider::run(args.bind, args.ready_file).await, + Command::Load(args) => { + let mut options = LoadOptions::from_file( + &args.config, + args.direct_url, + args.targets, + args.headers, + args.processes, + args.binaries, + args.output, + )?; + let harness = match args.worker_binary.as_deref() { + Some(binary) => { + let harness = orchestrate::WorkerHarness::start( + binary, + &options + .targets + .iter() + .find(|target| target.topology == config::Topology::Direct) + .expect("LoadOptions always contains a direct target") + .url, + ) + .await?; + if let Err(error) = harness.add_to(&mut options) { + let _ = harness.shutdown().await; + return Err(error); + } + Some(harness) + } + None => None, + }; + let result = client::run(options).await; + let cleanup = match harness { + Some(harness) => harness.shutdown().await, + None => Ok(()), + }; + match (result, cleanup) { + (Err(error), _) => Err(error), + (Ok(()), Err(error)) => Err(error), + (Ok(()), Ok(())) => Ok(()), + } + } + Command::Smoke(args) => client::run_smoke(args.output).await, + } +} diff --git a/scripts/latency_benchmark/daemon_transport/src/metadata.rs b/scripts/latency_benchmark/daemon_transport/src/metadata.rs new file mode 100644 index 000000000..0e9b2e758 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/src/metadata.rs @@ -0,0 +1,106 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::str::FromStr; +use std::time::{SystemTime, UNIX_EPOCH}; + +use anyhow::{Context, Result, ensure}; +use serde::Serialize; +use sha2::{Digest, Sha256}; + +#[derive(Clone, Debug)] +pub struct BinarySpec { + pub profile: String, + pub path: PathBuf, +} + +impl FromStr for BinarySpec { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + let (profile, path) = value + .split_once('=') + .context("binary metadata must use PROFILE=PATH")?; + ensure!(!profile.is_empty(), "binary profile is empty"); + let path = PathBuf::from(path); + ensure!(path.is_file(), "binary does not exist: {}", path.display()); + Ok(Self { + profile: profile.to_owned(), + path, + }) + } +} + +#[derive(Debug, Serialize)] +pub struct EnvironmentRecord { + generated_unix_seconds: u64, + git_commit: String, + git_dirty: bool, + operating_system: String, + architecture: String, + logical_cpus: usize, + rustc: String, + benchmark_binary: String, + relay_binaries: Vec, +} + +#[derive(Debug, Serialize)] +struct BinaryRecord { + profile: String, + path: String, + size_bytes: u64, + sha256: String, +} + +pub fn collect(binaries: &[BinarySpec]) -> Result { + let benchmark_binary = + std::env::current_exe().context("failed to resolve benchmark executable")?; + Ok(EnvironmentRecord { + generated_unix_seconds: SystemTime::now() + .duration_since(UNIX_EPOCH) + .context("system clock precedes Unix epoch")? + .as_secs(), + git_commit: command_output("git", &["rev-parse", "HEAD"]) + .unwrap_or_else(|| "unknown".to_owned()), + git_dirty: command_output("git", &["status", "--porcelain"]) + .is_some_and(|value| !value.is_empty()), + operating_system: std::env::consts::OS.to_owned(), + architecture: std::env::consts::ARCH.to_owned(), + logical_cpus: std::thread::available_parallelism().map_or(1, usize::from), + rustc: command_output("rustc", &["--version"]).unwrap_or_else(|| "unknown".to_owned()), + benchmark_binary: benchmark_binary.display().to_string(), + relay_binaries: binaries.iter().map(binary_record).collect::>()?, + }) +} + +fn binary_record(spec: &BinarySpec) -> Result { + let contents = std::fs::read(&spec.path) + .with_context(|| format!("failed to read binary {}", spec.path.display()))?; + Ok(BinaryRecord { + profile: spec.profile.clone(), + path: spec.path.display().to_string(), + size_bytes: contents.len() as u64, + sha256: format!("{:x}", Sha256::digest(&contents)), + }) +} + +fn command_output(command: &str, arguments: &[&str]) -> Option { + let output = Command::new(command).args(arguments).output().ok()?; + output + .status + .success() + .then(|| String::from_utf8_lossy(&output.stdout).trim().to_owned()) +} + +pub fn write_json(path: &Path, value: &impl Serialize) -> Result<()> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + } + let serialized = + serde_json::to_vec_pretty(value).context("failed to serialize benchmark report")?; + std::fs::write(path, [serialized.as_slice(), b"\n"].concat()) + .with_context(|| format!("failed to write {}", path.display())) +} diff --git a/scripts/latency_benchmark/daemon_transport/src/orchestrate.rs b/scripts/latency_benchmark/daemon_transport/src/orchestrate.rs new file mode 100644 index 000000000..7030c7519 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/src/orchestrate.rs @@ -0,0 +1,727 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Benchmark-only orchestration for an authenticated, directly measured worker hop. +//! +//! Production workers deliberately expose neither their endpoint nor their daemon-to-worker +//! credential. This module starts an isolated real daemon and MCP, observes their loopback control +//! exchange without persisting either value, and adds the resulting worker target to the load +//! driver in memory. It never weakens the production worker's request authentication. + +use std::convert::Infallible; +use std::net::{IpAddr, Ipv4Addr}; +use std::path::Path; +use std::process::Stdio; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use anyhow::{Context, Result, anyhow, bail, ensure}; +use base64::Engine as _; +use bytes::Bytes; +use http::header::{CONNECTION, CONTENT_LENGTH, HOST, TRANSFER_ENCODING}; +use http::{HeaderName, HeaderValue, Method, Request, Response, StatusCode, Uri}; +use http_body_util::{BodyExt as _, Full, Limited}; +use hyper::body::Incoming; +use hyper::service::service_fn; +use hyper_util::client::legacy::Client; +use hyper_util::client::legacy::connect::HttpConnector; +use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer}; +use hyper_util::server::conn::auto::Builder as ServerBuilder; +use ring::rand::{SecureRandom as _, SystemRandom}; +use serde_json::Value; +use tempfile::TempDir; +use tokio::net::{TcpListener, TcpStream}; +use tokio::process::{Child, Command}; +use tokio::sync::oneshot; + +use crate::config::{LoadOptions, ProcessSpec, Target, Topology}; + +const WORKER_TOKEN_HEADER: &str = "x-nemo-relay-worker-token"; +const CHALLENGE_PATH: &str = "/_nemo-relay/control/v1/challenge"; +const MCP_REGISTER_PATH: &str = "/_nemo-relay/control/v1/mcp/register"; +const MCP_HEARTBEAT_PATH: &str = "/_nemo-relay/control/v1/mcp/heartbeat"; +const MCP_RELEASE_PATH: &str = "/_nemo-relay/control/v1/mcp/release"; +const MCP_ACTIVATION_FAILED_PATH: &str = "/_nemo-relay/control/v1/mcp/activation-failed"; +const WORKER_REGISTER_PATH: &str = "/_nemo-relay/control/v1/worker/register"; +const WORKER_RECOVER_PATH: &str = "/_nemo-relay/control/v1/worker/recover"; +const WORKER_READY_PATH: &str = "/_nemo-relay/control/v1/worker/ready"; +const WORKER_HEARTBEAT_PATH: &str = "/_nemo-relay/control/v1/worker/heartbeat"; +const MAX_CONTROL_BODY_BYTES: usize = 256 * 1024; +const STARTUP_TIMEOUT: Duration = Duration::from_secs(30); +const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10); + +type ProxyClient = Client>; + +struct PendingWorker { + worker_id: String, + endpoint: String, + data_token: String, +} + +struct WorkerAccess { + endpoint: String, + data_token: String, +} + +#[derive(Default)] +struct CaptureState { + pending: Option, + terminal: bool, +} + +struct ProxyState { + backend_origin: String, + expected_worker_endpoint: String, + client: ProxyClient, + capture: Mutex, + ready: Mutex>>>, +} + +/// Owns every process and temporary identity file created for one worker-only benchmark target. +pub struct WorkerHarness { + daemon: Child, + mcp: Child, + proxy: tokio::task::JoinHandle<()>, + state_dir: TempDir, + endpoint: String, + data_token: String, + worker_port: u16, + process_specs: Vec, +} + +impl WorkerHarness { + /// Starts a real worker through the normal daemon/MCP activation path. + pub async fn start(relay_binary: &Path, provider_url: &str) -> Result { + ensure!( + relay_binary.is_file(), + "Relay binary does not exist: {}", + relay_binary.display() + ); + ensure_loopback_origin(provider_url, "provider")?; + + let daemon_port = reserve_loopback_port().await?; + let worker_port = reserve_loopback_port().await?; + ensure!( + daemon_port != worker_port, + "ephemeral port allocator returned a duplicate port" + ); + let proxy_listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .context("failed to bind benchmark control proxy")?; + let proxy_address = proxy_listener + .local_addr() + .context("failed to read benchmark control proxy address")?; + let proxy_origin = format!("http://{proxy_address}"); + let backend_origin = format!("http://127.0.0.1:{daemon_port}"); + let expected_worker_endpoint = format!("http://127.0.0.1:{worker_port}"); + let state_dir = tempfile::Builder::new() + .prefix("nemo-relay-daemon-benchmark-") + .tempdir() + .context("failed to create isolated benchmark state directory")?; + let route_token = random_route_token()?; + + let (ready, wait_for_ready) = oneshot::channel(); + let proxy_state = Arc::new(ProxyState { + backend_origin: backend_origin.clone(), + expected_worker_endpoint: expected_worker_endpoint.clone(), + client: proxy_client(), + capture: Mutex::new(CaptureState::default()), + ready: Mutex::new(Some(ready)), + }); + let proxy = tokio::spawn(serve_control_proxy( + proxy_listener, + Arc::clone(&proxy_state), + )); + + let mut daemon = relay_command(relay_binary, state_dir.path(), provider_url); + daemon + .arg("daemon") + .arg("--bind") + .arg("127.0.0.1") + .arg("--port") + .arg(daemon_port.to_string()) + .arg("--advertise-address") + .arg(&proxy_origin) + .env_remove("NEMO_RELAY_CLIENT_TOKEN") + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::inherit()) + .kill_on_drop(true); + let mut daemon = daemon + .spawn() + .with_context(|| format!("failed to launch {} daemon", relay_binary.display()))?; + if let Err(error) = wait_for_listener(daemon_port, &mut daemon).await { + proxy.abort(); + let _ = daemon.start_kill(); + return Err(error); + } + + let mut mcp = relay_command(relay_binary, state_dir.path(), provider_url); + mcp.arg("daemon") + .arg("mcp") + .arg("--daemon-address") + .arg(&proxy_origin) + .env("NEMO_RELAY_CLIENT_TOKEN", &route_token) + .env("NEMO_RELAY_WORKER_PORT", worker_port.to_string()) + .stdin(Stdio::piped()) + .stdout(Stdio::null()) + .stderr(Stdio::inherit()) + .kill_on_drop(true); + let mut mcp = match mcp.spawn() { + Ok(child) => child, + Err(error) => { + proxy.abort(); + let _ = daemon.start_kill(); + return Err(error).context("failed to launch benchmark MCP"); + } + }; + + let access = match tokio::time::timeout(STARTUP_TIMEOUT, wait_for_ready).await { + Ok(Ok(Ok(access))) => access, + Ok(Ok(Err(error))) => { + stop_child(&mut mcp).await; + stop_child(&mut daemon).await; + proxy.abort(); + bail!("benchmark control proxy rejected worker activation: {error}"); + } + Ok(Err(_)) => { + stop_child(&mut mcp).await; + stop_child(&mut daemon).await; + proxy.abort(); + bail!("benchmark control proxy stopped before worker readiness"); + } + Err(_) => { + let mcp_status = mcp.try_wait().context("failed to inspect benchmark MCP")?; + stop_child(&mut mcp).await; + stop_child(&mut daemon).await; + proxy.abort(); + bail!( + "timed out waiting for authenticated worker readiness; MCP status: {mcp_status:?}" + ); + } + }; + ensure!( + access.endpoint == expected_worker_endpoint, + "worker registered an unexpected endpoint" + ); + if let Err(error) = wait_for_worker_acceptance(&access).await { + let mcp_cleanup = release_mcp(&mut mcp).await; + let worker_cleanup = wait_for_worker_exit(worker_port).await; + stop_child(&mut daemon).await; + proxy.abort(); + if let Err(cleanup_error) = mcp_cleanup.and(worker_cleanup) { + return Err(error.context(format!( + "benchmark harness cleanup also failed: {cleanup_error:#}" + ))); + } + return Err(error); + } + + let mut process_specs = vec![ + ProcessSpec { + name: "worker-harness-daemon".into(), + pid: daemon.id().context("benchmark daemon has no process ID")?, + }, + ProcessSpec { + name: "worker-harness-mcp".into(), + pid: mcp.id().context("benchmark MCP has no process ID")?, + }, + ]; + if let Some(pid) = find_worker_pid(worker_port, &proxy_origin).await { + process_specs.push(ProcessSpec { + name: "worker".into(), + pid, + }); + } + + Ok(Self { + daemon, + mcp, + proxy, + state_dir, + endpoint: access.endpoint, + data_token: access.data_token, + worker_port, + process_specs, + }) + } + + /// Adds the authenticated direct-worker target without exposing its credential to the caller. + pub fn add_to(&self, options: &mut LoadOptions) -> Result<()> { + ensure!( + !options + .targets + .iter() + .any(|target| target.topology == Topology::WorkerOnly), + "worker-only target is already configured" + ); + options.targets.push(Target { + topology: Topology::WorkerOnly, + url: self.endpoint.clone(), + headers: vec![( + HeaderName::from_static(WORKER_TOKEN_HEADER), + HeaderValue::from_str(&self.data_token) + .context("daemon returned an invalid worker credential")?, + )], + }); + options.processes.extend(self.process_specs.clone()); + Ok(()) + } + + /// Gracefully releases the MCP reference, verifies the worker exits, and removes state. + pub async fn shutdown(mut self) -> Result<()> { + let _state_dir_lifetime = &self.state_dir; + let mcp_result = release_mcp(&mut self.mcp).await; + let worker_result = wait_for_worker_exit(self.worker_port).await; + stop_child(&mut self.daemon).await; + self.proxy.abort(); + mcp_result.and(worker_result) + } +} + +fn relay_command(binary: &Path, state_dir: &Path, provider_url: &str) -> Command { + let mut command = Command::new(binary); + command + .arg("--openai-base-url") + .arg(provider_url) + .arg("--anthropic-base-url") + .arg(provider_url) + .env("XDG_CONFIG_HOME", state_dir) + .env("XDG_CONFIG_DIRS", state_dir) + .env("NEMO_RELAY_OPENAI_BASE_URL", provider_url) + .env("NEMO_RELAY_ANTHROPIC_BASE_URL", provider_url); + command +} + +async fn reserve_loopback_port() -> Result { + let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)) + .await + .context("failed to reserve a loopback port")?; + listener + .local_addr() + .map(|address| address.port()) + .context("failed to read reserved loopback port") +} + +fn random_route_token() -> Result { + let mut bytes = [0_u8; 32]; + SystemRandom::new() + .fill(&mut bytes) + .map_err(|_| anyhow!("failed to generate benchmark route credential"))?; + Ok(base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes)) +} + +fn ensure_loopback_origin(origin: &str, label: &str) -> Result<()> { + let uri = origin + .parse::() + .with_context(|| format!("invalid {label} URL"))?; + ensure!( + uri.scheme_str() == Some("http"), + "benchmark {label} must use loopback HTTP" + ); + let authority = uri + .authority() + .with_context(|| format!("{label} URL has no authority"))?; + let host = authority.host(); + let loopback = host + .parse::() + .is_ok_and(|address| address.is_loopback()); + ensure!(loopback, "benchmark {label} must be loopback-only"); + ensure!( + authority.port_u16().is_some(), + "benchmark {label} needs an explicit port" + ); + ensure!( + uri.path_and_query() + .is_none_or(|value| value.as_str() == "/"), + "benchmark {label} URL cannot contain a path or query" + ); + Ok(()) +} + +fn proxy_client() -> ProxyClient { + let mut connector = HttpConnector::new(); + connector.enforce_http(true); + connector.set_nodelay(true); + let mut builder = Client::builder(TokioExecutor::new()); + builder.timer(TokioTimer::new()); + builder.pool_idle_timeout(Duration::from_secs(30)); + builder.build(connector) +} + +async fn serve_control_proxy(listener: TcpListener, state: Arc) { + let mut connections = tokio::task::JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let Ok((stream, peer)) = accepted else { + fail_capture(&state, "control proxy listener failed".into()); + return; + }; + if !peer.ip().is_loopback() { + continue; + } + let service_state = Arc::clone(&state); + connections.spawn(async move { + let service = service_fn(move |request| { + proxy_control_request(request, Arc::clone(&service_state)) + }); + let builder = ServerBuilder::new(TokioExecutor::new()); + let _ = builder + .serve_connection(TokioIo::new(stream), service) + .await; + }); + } + Some(completed) = connections.join_next(), if !connections.is_empty() => { + if completed.is_err() { + fail_capture(&state, "control proxy connection task failed".into()); + return; + } + } + } + } +} + +async fn proxy_control_request( + request: Request, + state: Arc, +) -> Result>, Infallible> { + if request.method() != Method::POST { + return Ok(Response::builder() + .status(StatusCode::METHOD_NOT_ALLOWED) + .body(Full::new(Bytes::from_static(b"method not allowed"))) + .expect("static proxy rejection")); + } + if !allowed_control_path(request.uri().path()) { + return Ok(Response::builder() + .status(StatusCode::NOT_FOUND) + .body(Full::new(Bytes::from_static(b"not found"))) + .expect("static proxy rejection")); + } + let response = match proxy_control_request_inner(request, &state).await { + Ok(response) => response, + Err(error) => { + fail_capture(&state, error.to_string()); + Response::builder() + .status(StatusCode::BAD_GATEWAY) + .body(Full::new(Bytes::from_static( + b"benchmark control proxy failure", + ))) + .expect("static proxy error response") + } + }; + Ok(response) +} + +async fn proxy_control_request_inner( + request: Request, + state: &ProxyState, +) -> Result>> { + let path = request.uri().path().to_owned(); + let path_and_query = request + .uri() + .path_and_query() + .context("control request has no path")? + .as_str() + .to_owned(); + let (parts, body) = request.into_parts(); + let body = Limited::new(body, MAX_CONTROL_BODY_BYTES) + .collect() + .await + .map_err(|error| { + anyhow!("control request body failed or exceeded the benchmark proxy limit: {error}") + })? + .to_bytes(); + let request_json = if matches!(path.as_str(), WORKER_REGISTER_PATH | WORKER_READY_PATH) { + Some(serde_json::from_slice::(&body).context("invalid worker control JSON")?) + } else { + None + }; + + let uri = format!("{}{path_and_query}", state.backend_origin) + .parse::() + .context("failed to construct control backend URI")?; + let mut forwarded = Request::builder() + .method(parts.method) + .version(parts.version) + .uri(uri) + .body(Full::new(body)) + .context("failed to build forwarded control request")?; + copy_end_to_end_headers(&parts.headers, forwarded.headers_mut()); + let response = state + .client + .request(forwarded) + .await + .context("control backend request failed")?; + let (response_parts, response_body) = response.into_parts(); + let response_body = Limited::new(response_body, MAX_CONTROL_BODY_BYTES) + .collect() + .await + .map_err(|error| { + anyhow!("control response body failed or exceeded the benchmark proxy limit: {error}") + })? + .to_bytes(); + + if response_parts.status.is_success() { + match path.as_str() { + WORKER_REGISTER_PATH => capture_registration( + state, + request_json + .as_ref() + .context("missing registration request")?, + &serde_json::from_slice(&response_body) + .context("invalid worker registration response")?, + )?, + WORKER_READY_PATH => capture_readiness( + state, + request_json.as_ref().context("missing readiness request")?, + )?, + _ => {} + } + } + + let mut rebuilt = Response::builder() + .status(response_parts.status) + .version(response_parts.version) + .body(Full::new(response_body)) + .context("failed to rebuild control response")?; + copy_end_to_end_headers(&response_parts.headers, rebuilt.headers_mut()); + Ok(rebuilt) +} + +fn allowed_control_path(path: &str) -> bool { + matches!( + path, + CHALLENGE_PATH + | MCP_REGISTER_PATH + | MCP_HEARTBEAT_PATH + | MCP_RELEASE_PATH + | MCP_ACTIVATION_FAILED_PATH + | WORKER_REGISTER_PATH + | WORKER_RECOVER_PATH + | WORKER_READY_PATH + | WORKER_HEARTBEAT_PATH + ) +} + +fn copy_end_to_end_headers(source: &http::HeaderMap, destination: &mut http::HeaderMap) { + for (name, value) in source { + if !matches!( + name, + &CONNECTION | &CONTENT_LENGTH | &HOST | &TRANSFER_ENCODING + ) { + destination.append(name, value.clone()); + } + } +} + +fn capture_registration(state: &ProxyState, request: &Value, response: &Value) -> Result<()> { + let worker_id = json_string(request, "/worker_id")?; + let endpoint = json_string(request, "/endpoint")?; + ensure!( + endpoint == state.expected_worker_endpoint, + "worker registration endpoint did not match the prescribed loopback port" + ); + let data_token = json_string(response, "/data_token")?; + let mut capture = lock(&state.capture); + ensure!( + !capture.terminal, + "worker registered after terminal readiness" + ); + if let Some(existing) = capture.pending.as_ref() { + ensure!( + existing.worker_id == worker_id + && existing.endpoint == endpoint + && existing.data_token == data_token, + "worker registration retry changed authenticated values" + ); + } else { + capture.pending = Some(PendingWorker { + worker_id, + endpoint, + data_token, + }); + } + Ok(()) +} + +fn capture_readiness(state: &ProxyState, request: &Value) -> Result<()> { + let worker_id = json_string(request, "/payload/worker_id")?; + let access = { + let mut capture = lock(&state.capture); + ensure!(!capture.terminal, "duplicate terminal worker readiness"); + let pending = capture + .pending + .take() + .context("worker became ready before an authenticated registration response")?; + ensure!( + pending.worker_id == worker_id, + "worker readiness ID did not match registration" + ); + capture.terminal = true; + WorkerAccess { + endpoint: pending.endpoint, + data_token: pending.data_token, + } + }; + let sender = lock(&state.ready) + .take() + .context("worker readiness was already reported")?; + sender + .send(Ok(access)) + .map_err(|_| anyhow!("worker readiness receiver was dropped")) +} + +fn fail_capture(state: &ProxyState, error: String) { + let should_send = { + let mut capture = lock(&state.capture); + if capture.terminal { + false + } else { + capture.terminal = true; + true + } + }; + if should_send && let Some(sender) = lock(&state.ready).take() { + let _ = sender.send(Err(error)); + } +} + +fn json_string(value: &Value, pointer: &str) -> Result { + value + .pointer(pointer) + .and_then(Value::as_str) + .filter(|value| !value.is_empty()) + .map(str::to_owned) + .with_context(|| format!("control JSON omitted {pointer}")) +} + +async fn wait_for_listener(port: u16, child: &mut Child) -> Result<()> { + let deadline = Instant::now() + STARTUP_TIMEOUT; + loop { + if let Some(status) = child + .try_wait() + .context("failed to inspect benchmark daemon")? + { + bail!("benchmark daemon exited before listening: {status}"); + } + if TcpStream::connect((Ipv4Addr::LOCALHOST, port)) + .await + .is_ok() + { + return Ok(()); + } + ensure!( + Instant::now() < deadline, + "timed out waiting for benchmark daemon listener" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + +async fn wait_for_worker_acceptance(access: &WorkerAccess) -> Result<()> { + let client = proxy_client(); + let deadline = Instant::now() + STARTUP_TIMEOUT; + loop { + let uri = format!("{}/v1/responses", access.endpoint) + .parse::() + .context("invalid captured worker endpoint")?; + let request = Request::post(uri) + .header(WORKER_TOKEN_HEADER, &access.data_token) + .header("content-type", "application/json") + .header("authorization", "Bearer benchmark-provider-token") + .header("x-benchmark-response-bytes", 16 * 1024) + .header("x-benchmark-event-count", 128) + .body(Full::new(Bytes::from_static( + b"{\"model\":\"benchmark\",\"stream\":true,\"input\":\"readiness\"}", + ))) + .context("failed to build worker readiness request")?; + match client.request(request).await { + Ok(response) if response.status().is_success() => { + response + .into_body() + .collect() + .await + .context("worker readiness response body failed")?; + return Ok(()); + } + Ok(response) if response.status() == StatusCode::SERVICE_UNAVAILABLE => {} + Ok(response) => bail!( + "worker readiness request returned HTTP {}", + response.status() + ), + Err(_) => {} + } + ensure!( + Instant::now() < deadline, + "timed out waiting for worker request admission" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + +async fn wait_for_worker_exit(port: u16) -> Result<()> { + let deadline = Instant::now() + SHUTDOWN_TIMEOUT; + loop { + if TcpStream::connect((Ipv4Addr::LOCALHOST, port)) + .await + .is_err() + { + return Ok(()); + } + ensure!( + Instant::now() < deadline, + "benchmark worker did not exit after MCP release" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + +async fn release_mcp(mcp: &mut Child) -> Result<()> { + drop(mcp.stdin.take()); + match tokio::time::timeout(SHUTDOWN_TIMEOUT, mcp.wait()).await { + Ok(Ok(_)) => Ok(()), + Ok(Err(error)) => Err(error).context("failed to wait for benchmark MCP"), + Err(_) => { + stop_child(mcp).await; + bail!("benchmark MCP did not exit after its input was closed") + } + } +} + +async fn stop_child(child: &mut Child) { + if child.try_wait().ok().flatten().is_none() { + let _ = child.start_kill(); + } + let _ = tokio::time::timeout(Duration::from_secs(2), child.wait()).await; +} + +#[cfg(unix)] +async fn find_worker_pid(port: u16, daemon_origin: &str) -> Option { + let output = Command::new("ps") + .args(["-axo", "pid=,command="]) + .output() + .await + .ok()?; + let port = format!("--port {port}"); + String::from_utf8_lossy(&output.stdout) + .lines() + .find_map(|line| { + let trimmed = line.trim(); + let (pid, command) = trimmed.split_once(char::is_whitespace)?; + (command.contains("daemon worker") + && command.contains(&port) + && command.contains(daemon_origin)) + .then(|| pid.parse().ok()) + .flatten() + }) +} + +#[cfg(not(unix))] +async fn find_worker_pid(_port: u16, _daemon_origin: &str) -> Option { + None +} + +fn lock(mutex: &Mutex) -> std::sync::MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(|error| error.into_inner()) +} diff --git a/scripts/latency_benchmark/daemon_transport/src/provider.rs b/scripts/latency_benchmark/daemon_transport/src/provider.rs new file mode 100644 index 000000000..6868ef8c4 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/src/provider.rs @@ -0,0 +1,349 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::convert::Infallible; +use std::fmt::Write as _; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use anyhow::{Context, Result, ensure}; +use bytes::Bytes; +use http::header::{CONTENT_TYPE, TRAILER}; +use http::{HeaderMap, HeaderValue, Method, Request, Response, StatusCode}; +use http_body_util::{BodyExt, Full, StreamBody, combinators::UnsyncBoxBody}; +use hyper::body::{Frame, Incoming}; +use hyper::service::service_fn; +use hyper_util::rt::{TokioExecutor, TokioIo}; +use hyper_util::server::conn::auto::Builder; +use serde::Serialize; +use sha2::{Digest, Sha256}; +use tokio::net::TcpListener; +use tokio::sync::oneshot; + +const RESPONSE_BYTES: &str = "x-benchmark-response-bytes"; +const EVENT_COUNT: &str = "x-benchmark-event-count"; +const EVENT_DELAY_MICROS: &str = "x-benchmark-event-delay-micros"; +const STREAM_ID: &str = "x-benchmark-stream-id"; +const BODY_SHA256: &str = "x-benchmark-body-sha256"; + +type ResponseBody = UnsyncBoxBody; + +#[derive(Default)] +struct ProviderStats { + next_stream: AtomicU64, + accepted: AtomicU64, + completed: AtomicU64, + cancelled: AtomicU64, +} + +#[derive(Serialize)] +struct ProviderSnapshot { + accepted: u64, + completed: u64, + cancelled: u64, +} + +impl ProviderStats { + fn snapshot(&self) -> ProviderSnapshot { + ProviderSnapshot { + accepted: self.accepted.load(Ordering::Relaxed), + completed: self.completed.load(Ordering::Relaxed), + cancelled: self.cancelled.load(Ordering::Relaxed), + } + } +} + +struct CompletionGuard { + stats: Arc, + complete: bool, +} + +impl Drop for CompletionGuard { + fn drop(&mut self) { + if !self.complete { + self.stats.cancelled.fetch_add(1, Ordering::Relaxed); + } + } +} + +pub async fn run(bind: SocketAddr, ready_file: Option) -> Result<()> { + let listener = TcpListener::bind(bind) + .await + .with_context(|| format!("failed to bind provider on {bind}"))?; + let address = listener + .local_addr() + .context("failed to read provider address")?; + let url = format!("http://{address}"); + if let Some(path) = ready_file { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent) + .with_context(|| format!("failed to create {}", parent.display()))?; + } + std::fs::write(&path, format!("{url}\n")) + .with_context(|| format!("failed to write {}", path.display()))?; + } + println!("{{\"provider_url\":\"{url}\"}}"); + serve(listener, async { + let _ = tokio::signal::ctrl_c().await; + }) + .await +} + +pub async fn spawn_ephemeral() -> Result<( + String, + oneshot::Sender<()>, + tokio::task::JoinHandle>, +)> { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .context("failed to bind smoke provider")?; + let url = format!( + "http://{}", + listener + .local_addr() + .context("failed to read smoke provider address")? + ); + let (stop, stopped) = oneshot::channel(); + let task = tokio::spawn(async move { + serve(listener, async { + let _ = stopped.await; + }) + .await + }); + Ok((url, stop, task)) +} + +async fn serve(listener: TcpListener, shutdown: impl Future) -> Result<()> { + let stats = Arc::new(ProviderStats::default()); + tokio::pin!(shutdown); + loop { + let (stream, _) = tokio::select! { + accepted = listener.accept() => accepted.context("provider accept failed")?, + _ = &mut shutdown => break, + }; + stream + .set_nodelay(true) + .context("failed to set TCP_NODELAY")?; + let connection_stats = Arc::clone(&stats); + tokio::spawn(async move { + let service = service_fn(move |request| handle(request, Arc::clone(&connection_stats))); + let builder = Builder::new(TokioExecutor::new()); + if let Err(error) = builder + .serve_connection(TokioIo::new(stream), service) + .await + { + eprintln!("provider connection failed: {error}"); + } + }); + } + Ok(()) +} + +async fn handle( + request: Request, + stats: Arc, +) -> Result, Infallible> { + let response = match (request.method(), request.uri().path()) { + (&Method::GET, "/healthz") => { + full_response(StatusCode::OK, Bytes::from_static(b"ok"), "text/plain") + } + (&Method::GET, "/metrics") => { + let body = serde_json::to_vec(&stats.snapshot()).expect("provider metrics serialize"); + full_response(StatusCode::OK, Bytes::from(body), "application/json") + } + (&Method::POST, "/v1/responses") => { + stream_response(&request, "response.output_text.delta", stats) + } + (&Method::POST, "/v1/messages") => stream_response(&request, "content_block_delta", stats), + _ => full_response( + StatusCode::NOT_FOUND, + Bytes::from_static(b"not found"), + "text/plain", + ), + }; + Ok(response) +} + +fn stream_response( + request: &Request, + event_type: &'static str, + stats: Arc, +) -> Response { + let parameters = parse_parameters(request.headers(), &stats); + let Ok((response_bytes, event_count, delay, stream_id)) = parameters else { + return full_response( + StatusCode::BAD_REQUEST, + Bytes::from(parameters.unwrap_err().to_string()), + "text/plain", + ); + }; + let mut event_buffer = String::new(); + let mut base_size = done_event().len(); + for sequence in 0..event_count { + event_buffer.clear(); + write_event(&mut event_buffer, event_type, &stream_id, sequence, 0, 0); + base_size += event_buffer.len(); + } + if response_bytes < base_size { + return full_response( + StatusCode::BAD_REQUEST, + Bytes::from(format!( + "response size {response_bytes} is smaller than minimum {base_size} for {event_count} events" + )), + "text/plain", + ); + } + + stats.accepted.fetch_add(1, Ordering::Relaxed); + let remaining = response_bytes - base_size; + let body_stats = Arc::clone(&stats); + let response_stream_id = stream_id.clone(); + let body = async_stream::stream! { + let mut guard = CompletionGuard { stats: Arc::clone(&body_stats), complete: false }; + let mut hasher = Sha256::new(); + for sequence in 0..event_count { + if sequence > 0 && !delay.is_zero() { + tokio::time::sleep(delay).await; + } + let padding = remaining / event_count + usize::from(sequence < remaining % event_count); + let emitted = unix_time_nanos(); + let event = make_event(event_type, &stream_id, sequence, emitted, padding); + hasher.update(&event); + yield Ok(Frame::data(event)); + } + let done = done_event(); + hasher.update(&done); + yield Ok(Frame::data(done)); + let mut trailers = HeaderMap::new(); + trailers.insert(BODY_SHA256, HeaderValue::from_str(&format!("{:x}", hasher.finalize())).expect("SHA-256 header")); + trailers.insert(EVENT_COUNT, HeaderValue::from_str(&event_count.to_string()).expect("event count header")); + yield Ok(Frame::trailers(trailers)); + guard.complete = true; + body_stats.completed.fetch_add(1, Ordering::Relaxed); + }; + + Response::builder() + .status(StatusCode::OK) + .header(CONTENT_TYPE, "text/event-stream") + .header(TRAILER, format!("{BODY_SHA256}, {EVENT_COUNT}")) + .header(STREAM_ID, response_stream_id) + .header(EVENT_COUNT, event_count) + .body(StreamBody::new(body).boxed_unsync()) + .expect("valid benchmark response") +} + +fn parse_parameters( + headers: &HeaderMap, + stats: &ProviderStats, +) -> Result<(usize, usize, Duration, String)> { + let response_bytes = header_number(headers, RESPONSE_BYTES, 16 * 1024)?; + ensure!( + response_bytes <= 1024 * 1024 * 1024, + "response size is too large" + ); + let event_count = header_number(headers, EVENT_COUNT, 128)?; + ensure!(event_count >= 128, "event count must be at least 128"); + ensure!(event_count <= 1_000_000, "event count is too large"); + let delay_micros: u64 = header_number(headers, EVENT_DELAY_MICROS, 0)?; + let stream_id = headers + .get(STREAM_ID) + .map(|value| value.to_str().context("stream ID is not ASCII")) + .transpose()? + .map(str::to_owned) + .unwrap_or_else(|| format!("{:016x}", stats.next_stream.fetch_add(1, Ordering::Relaxed))); + ensure!( + stream_id + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'-'), + "stream ID contains unsupported characters" + ); + Ok(( + response_bytes, + event_count, + Duration::from_micros(delay_micros), + stream_id, + )) +} + +fn header_number(headers: &HeaderMap, name: &'static str, default: T) -> Result +where + T: std::str::FromStr, + T::Err: std::error::Error + Send + Sync + 'static, +{ + headers + .get(name) + .map(|value| { + value + .to_str() + .context("numeric benchmark header is not ASCII")? + .parse() + .context("invalid numeric benchmark header") + }) + .transpose() + .map(|value| value.unwrap_or(default)) +} + +fn make_event( + event_type: &str, + stream_id: &str, + sequence: usize, + emitted_unix_nanos: u128, + padding: usize, +) -> Bytes { + let mut event = String::new(); + write_event( + &mut event, + event_type, + stream_id, + sequence, + emitted_unix_nanos, + padding, + ); + Bytes::from(event) +} + +fn write_event( + event: &mut String, + event_type: &str, + stream_id: &str, + sequence: usize, + emitted_unix_nanos: u128, + padding: usize, +) { + let metadata = if sequence == 0 { + format!(": benchmark-heartbeat\nevent: {event_type}\nid: {stream_id}-0\nretry: 1000\n") + } else { + String::new() + }; + write!(event, + "{metadata}data: {{\"type\":\"{event_type}\",\"i\":\"{stream_id}\",\"s\":{sequence},\"t\":\"{emitted_unix_nanos:020}\",\"d\":\"{}\"}}\n\n", + "x".repeat(padding) + ) + .expect("writing to a String cannot fail"); +} + +fn done_event() -> Bytes { + Bytes::from_static(b"data: [DONE]\n\n") +} + +fn unix_time_nanos() -> u128 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos() +} + +fn full_response( + status: StatusCode, + body: Bytes, + content_type: &'static str, +) -> Response { + Response::builder() + .status(status) + .header(CONTENT_TYPE, content_type) + .body(Full::new(body).boxed_unsync()) + .expect("valid benchmark response") +} diff --git a/scripts/latency_benchmark/daemon_transport/src/resources.rs b/scripts/latency_benchmark/daemon_transport/src/resources.rs new file mode 100644 index 000000000..5ae8f96a7 --- /dev/null +++ b/scripts/latency_benchmark/daemon_transport/src/resources.rs @@ -0,0 +1,151 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeMap; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use serde::Serialize; +use tokio::sync::oneshot; +use tokio::time::MissedTickBehavior; + +use crate::config::ProcessSpec; + +#[derive(Clone, Debug, Default)] +struct Aggregate { + baseline_rss_bytes: Option, + peak_rss_bytes: Option, + cpu_total: f64, + cpu_samples: u64, +} + +#[derive(Debug, Serialize)] +pub struct ResourceRecord { + pid: u32, + baseline_rss_bytes: Option, + peak_rss_bytes: Option, + rss_growth_bytes: Option, + rss_growth_per_active_stream_bytes: Option, + /// Average of `ps %cpu` samples, which is lifetime-average process CPU rather than + /// benchmark-window CPU for processes that existed before the sampler started. + average_cpu_percent: Option, +} + +pub struct ResourceSampler { + processes: Vec, + state: Arc>>, + stop: Option>, + task: tokio::task::JoinHandle<()>, +} + +impl ResourceSampler { + pub async fn start(mut processes: Vec) -> Self { + processes.push(ProcessSpec { + name: "load-driver".to_owned(), + pid: std::process::id(), + }); + let state = Arc::new(Mutex::new(BTreeMap::new())); + update(&processes, &state).await; + let (stop, mut stopped) = oneshot::channel(); + let task_processes = processes.clone(); + let task_state = Arc::clone(&state); + let task = tokio::spawn(async move { + let mut interval = tokio::time::interval(Duration::from_secs(1)); + interval.set_missed_tick_behavior(MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = interval.tick() => update(&task_processes, &task_state).await, + _ = &mut stopped => break, + } + } + }); + Self { + processes, + state, + stop: Some(stop), + task, + } + } + + pub async fn finish(mut self, active_streams: usize) -> BTreeMap { + update(&self.processes, &self.state).await; + if let Some(stop) = self.stop.take() { + let _ = stop.send(()); + } + let _ = self.task.await; + let state = self.state.lock().expect("resource state mutex poisoned"); + self.processes + .iter() + .map(|process| { + let aggregate = state.get(&process.pid).cloned().unwrap_or_default(); + let growth = aggregate + .baseline_rss_bytes + .zip(aggregate.peak_rss_bytes) + .map(|(baseline, peak)| peak.saturating_sub(baseline)); + let cpu = (aggregate.cpu_samples > 0) + .then(|| aggregate.cpu_total / aggregate.cpu_samples as f64); + ( + process.name.clone(), + ResourceRecord { + pid: process.pid, + baseline_rss_bytes: aggregate.baseline_rss_bytes, + peak_rss_bytes: aggregate.peak_rss_bytes, + rss_growth_bytes: growth, + rss_growth_per_active_stream_bytes: growth + .map(|value| value as f64 / active_streams.max(1) as f64), + average_cpu_percent: cpu, + }, + ) + }) + .collect() + } +} + +async fn update(processes: &[ProcessSpec], state: &Arc>>) { + let samples = sample_processes(processes).await; + let mut state = state.lock().expect("resource state mutex poisoned"); + for (pid, (rss, cpu)) in samples { + let aggregate = state.entry(pid).or_default(); + aggregate.baseline_rss_bytes.get_or_insert(rss); + aggregate.peak_rss_bytes = Some(aggregate.peak_rss_bytes.unwrap_or(0).max(rss)); + aggregate.cpu_total += cpu; + aggregate.cpu_samples += 1; + } +} + +#[cfg(unix)] +async fn sample_processes(processes: &[ProcessSpec]) -> BTreeMap { + if processes.is_empty() { + return BTreeMap::new(); + } + let pids = processes + .iter() + .map(|process| process.pid.to_string()) + .collect::>() + .join(","); + let output = tokio::time::timeout( + Duration::from_secs(2), + tokio::process::Command::new("ps") + .args(["-o", "pid=", "-o", "rss=", "-o", "%cpu=", "-p", &pids]) + .output(), + ) + .await; + let Ok(Ok(output)) = output else { + return BTreeMap::new(); + }; + String::from_utf8_lossy(&output.stdout) + .lines() + .filter_map(|line| { + let mut fields = line.split_whitespace(); + let pid = fields.next()?.parse().ok()?; + let rss_kib: u64 = fields.next()?.parse().ok()?; + let cpu = fields.next()?.replace(',', ".").parse().ok()?; + Some((pid, (rss_kib.saturating_mul(1024), cpu))) + }) + .collect() +} + +#[cfg(not(unix))] +async fn sample_processes(_processes: &[ProcessSpec]) -> BTreeMap { + BTreeMap::new() +} diff --git a/scripts/licensing/attributions_lockfile_md.py b/scripts/licensing/attributions_lockfile_md.py index fb17de566..7a0f13ff7 100755 --- a/scripts/licensing/attributions_lockfile_md.py +++ b/scripts/licensing/attributions_lockfile_md.py @@ -392,7 +392,7 @@ def _render_rust_metadata_fallback_attribution(crate: dict[str, Any]) -> tuple[s license_name = _normalize_license_name(str(crate.get("license") or "UNKNOWN")) parts = [ - f"## {name} - {version}\n", + f"## {name} - {version}\n\n", f"**Repository URL**: {repo}\n", f"**License Type(s)**: {license_name}\n", f"### License: {spdx_url(license_name, fallback='https://spdx.org/licenses/')}\n", @@ -442,7 +442,7 @@ def _render_rust_crate_attribution( rendered = "".join( [ - f"## {name} - {version}\n", + f"## {name} - {version}\n\n", f"**Repository URL**: {repo}\n", f"**License Type(s)**: {license_id}\n", f"### License: https://spdx.org/licenses/{license_id}.html\n",