From d8b93e4130c9e1343b12ed43eb913c5434bb7d08 Mon Sep 17 00:00:00 2001 From: tapiocaboy Date: Mon, 6 Jul 2026 17:40:58 +1000 Subject: [PATCH] Added and fixed mcp streaming --- .github/workflows/ci.yml | 4 +- Cargo.lock | 777 ++++++++++++++++++++++++++++++++------ Cargo.toml | 46 ++- README.md | 40 +- docs/TUTORIAL.md | 415 ++++++++++++++++++++ examples/tool_calling.rs | 59 +++ src/agent/builder.rs | 28 +- src/agent/manifest.rs | 3 + src/agent/runtime.rs | 212 ++++++++++- src/auth.rs | 664 ++++++++++++++++++++++++++++++++ src/error.rs | 8 + src/grpc/mod.rs | 84 ++++- src/grpc/pb.rs | 42 +-- src/identity.rs | 12 + src/lib.rs | 37 +- src/llm/mod.rs | 8 +- src/llm/nvidia.rs | 270 +++++++++++-- src/llm/types.rs | 89 ++++- src/llm/usage.rs | 89 +++++ src/mcp/client.rs | 318 +++++++++++++--- src/mcp/mod.rs | 83 +++- src/server/mod.rs | 139 +++++-- src/server/rest.rs | 120 +++++- src/session/sqlite.rs | 87 +++-- src/skills/mod.rs | 124 ++++++ src/tls.rs | 77 ++++ src/trust.rs | 1 + src/vector/chunk.rs | 94 +++++ src/vector/memory.rs | 44 ++- src/vector/mod.rs | 96 +++++ src/vector/pg.rs | 258 +++++++++++++ src/vector/pinecone.rs | 108 ++++-- src/vector/qdrant.rs | 25 +- src/x509.rs | 147 ++++++++ tests/a2a_and_sessions.rs | 1 + tests/tools_and_policy.rs | 274 ++++++++++++++ 36 files changed, 4496 insertions(+), 387 deletions(-) create mode 100644 docs/TUTORIAL.md create mode 100644 examples/tool_calling.rs create mode 100644 src/auth.rs create mode 100644 src/llm/usage.rs create mode 100644 src/tls.rs create mode 100644 src/vector/chunk.rs create mode 100644 src/vector/pg.rs create mode 100644 src/x509.rs create mode 100644 tests/tools_and_policy.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4a850e6..73252c5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -69,11 +69,11 @@ jobs: - run: cargo doc --no-deps --features full msrv: - name: MSRV (1.86) + name: MSRV (1.88) runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - - uses: dtolnay/rust-toolchain@1.86 + - uses: dtolnay/rust-toolchain@1.88 - uses: Swatinem/rust-cache@v2 - run: cargo check --all-targets --features full diff --git a/Cargo.lock b/Cargo.lock index dc3d02b..57302f5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -29,6 +29,45 @@ dependencies = [ "rustversion", ] +[[package]] +name = "asn1-rs" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5493c3bedbacf7fd7382c6346bbd66687d12bbaad3a89a2d2c303ee6cf20b048" +dependencies = [ + "asn1-rs-derive", + "asn1-rs-impl", + "displaydoc", + "nom", + "num-traits", + "rusticata-macros", + "thiserror 1.0.69", + "time", +] + +[[package]] +name = "asn1-rs-derive" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "965c2d33e53cb6b267e148a4cb0760bc01f4904c1cd4bb4002a085bb016d1490" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "asn1-rs-impl" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -75,30 +114,26 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] -name = "axum" -version = "0.7.9" +name = "aws-lc-rs" +version = "1.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" +checksum = "4342d8937fc7e5dd9b1c60292261c0670c882a2cd1719cfc11b1af41731e32ad" dependencies = [ - "async-trait", - "axum-core 0.4.5", - "bytes", - "futures-util", - "http", - "http-body", - "http-body-util", - "itoa", - "matchit 0.7.3", - "memchr", - "mime", - "percent-encoding", - "pin-project-lite", - "rustversion", - "serde", - "sync_wrapper", - "tower 0.5.3", - "tower-layer", - "tower-service", + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d9ceb1da931507a12f4fccea479dccd00da1943e1b4ae72d8e502d707361444" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", ] [[package]] @@ -107,7 +142,7 @@ version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" dependencies = [ - "axum-core 0.5.6", + "axum-core", "base64", "bytes", "form_urlencoded", @@ -118,7 +153,7 @@ dependencies = [ "hyper", "hyper-util", "itoa", - "matchit 0.8.4", + "matchit", "memchr", "mime", "percent-encoding", @@ -131,7 +166,7 @@ dependencies = [ "sync_wrapper", "tokio", "tokio-tungstenite", - "tower 0.5.3", + "tower", "tower-layer", "tower-service", "tracing", @@ -139,41 +174,43 @@ dependencies = [ [[package]] name = "axum-core" -version = "0.4.5" +version = "0.5.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" dependencies = [ - "async-trait", "bytes", - "futures-util", + "futures-core", "http", "http-body", "http-body-util", "mime", "pin-project-lite", - "rustversion", "sync_wrapper", "tower-layer", "tower-service", + "tracing", ] [[package]] -name = "axum-core" -version = "0.5.6" +name = "axum-server" +version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +checksum = "b1df331683d982a0b9492b38127151e6453639cd34926eb9c07d4cd8c6d22bfc" dependencies = [ + "arc-swap", "bytes", - "futures-core", + "either", + "fs-err", "http", "http-body", - "http-body-util", - "mime", + "hyper", + "hyper-util", "pin-project-lite", - "sync_wrapper", - "tower-layer", + "rustls", + "rustls-pki-types", + "tokio", + "tokio-rustls", "tower-service", - "tracing", ] [[package]] @@ -212,6 +249,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + [[package]] name = "bs58" version = "0.5.1" @@ -227,6 +273,12 @@ version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[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.0" @@ -240,6 +292,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e228eec9be7c17ccb640b59b36a5cd805ea2a564a4c5e162c2f659fea30d3b96" dependencies = [ "find-msvc-tools", + "jobserver", + "libc", "shlex", ] @@ -266,6 +320,21 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + [[package]] name = "combine" version = "4.6.7" @@ -286,32 +355,48 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "corrosive_agents" version = "0.1.0" dependencies = [ "async-stream", "async-trait", - "axum 0.8.9", + "axum", + "axum-server", "base64", "bs58", "dotenvy", "ed25519-dalek", "futures-util", + "hmac 0.12.1", + "pgvector", "prost", "rand 0.8.6", + "rcgen", "redis", "reqwest", + "ring", "rusqlite", "semver", "serde", "serde_json", - "thiserror", + "sha2 0.10.9", + "thiserror 2.0.18", "tokio", + "tokio-postgres", "tonic", + "tonic-prost", "tracing", "tracing-subscriber", + "utoipa", "uuid", + "x509-parser", ] [[package]] @@ -342,6 +427,24 @@ dependencies = [ "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + [[package]] name = "curve25519-dalek" version = "4.1.3" @@ -351,7 +454,7 @@ dependencies = [ "cfg-if", "cpufeatures 0.2.17", "curve25519-dalek-derive", - "digest", + "digest 0.10.7", "fiat-crypto", "rustc_version", "subtle", @@ -381,18 +484,51 @@ version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" dependencies = [ - "const-oid", + "const-oid 0.9.6", "zeroize", ] +[[package]] +name = "der-parser" +version = "9.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cd0a5c643689626bec213c4d8bd4d96acc8ffdb4ad4bb6bc16abf27d5f4b553" +dependencies = [ + "asn1-rs", + "displaydoc", + "nom", + "num-bigint", + "num-traits", + "rusticata-macros", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" + [[package]] name = "digest" version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "crypto-common", + "block-buffer 0.10.4", + "crypto-common 0.1.7", + "subtle", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid 0.10.2", + "crypto-common 0.2.2", + "ctutils", ] [[package]] @@ -412,6 +548,12 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "ed25519" version = "2.2.3" @@ -432,7 +574,7 @@ dependencies = [ "ed25519", "rand_core 0.6.4", "serde", - "sha2", + "sha2 0.10.9", "subtle", "zeroize", ] @@ -459,6 +601,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -504,6 +652,22 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs-err" +version = "3.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b91aa448ca50d7e79433bdf3ee8d99215430d2ec02ade5aefab2a073a1822e8a" +dependencies = [ + "autocfg", + "tokio", +] + +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures" version = "0.3.32" @@ -611,7 +775,7 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "wasm-bindgen", ] @@ -653,19 +817,13 @@ dependencies = [ "futures-core", "futures-sink", "http", - "indexmap 2.14.0", + "indexmap", "slab", "tokio", "tokio-util", "tracing", ] -[[package]] -name = "hashbrown" -version = "0.12.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" - [[package]] name = "hashbrown" version = "0.14.5" @@ -690,6 +848,24 @@ dependencies = [ "hashbrown 0.14.5", ] +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest 0.10.7", +] + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", +] + [[package]] name = "http" version = "1.4.2" @@ -735,6 +911,15 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "hybrid-array" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "818356c5132c1fede50f837ca96afbe78ff42413047f4abb886217845e1b6c8c" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.10.1" @@ -803,7 +988,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.4", + "socket2", "tokio", "tower-service", "tracing", @@ -912,16 +1097,6 @@ dependencies = [ "icu_properties", ] -[[package]] -name = "indexmap" -version = "1.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" -dependencies = [ - "autocfg", - "hashbrown 0.12.3", -] - [[package]] name = "indexmap" version = "2.14.0" @@ -930,6 +1105,8 @@ checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", "hashbrown 0.17.1", + "serde", + "serde_core", ] [[package]] @@ -962,6 +1139,16 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + [[package]] name = "js-sys" version = "0.3.103" @@ -985,6 +1172,15 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libredox" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c943259e342f1e06ff2da7a83eabdfe7f92ce10262688dbf1895ff0b3e6e4652" +dependencies = [ + "libc", +] + [[package]] name = "libsqlite3-sys" version = "0.30.1" @@ -1025,15 +1221,19 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "matchit" -version = "0.7.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" [[package]] -name = "matchit" -version = "0.8.4" +name = "md-5" +version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] [[package]] name = "memchr" @@ -1047,6 +1247,12 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "mio" version = "1.2.1" @@ -1054,10 +1260,20 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02bd0af71c67b473010cbbc60715ee815645a4dc942899111f494b4b737d6fda" dependencies = [ "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1077,6 +1293,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + [[package]] name = "num-integer" version = "0.1.46" @@ -1095,6 +1317,33 @@ dependencies = [ "autocfg", ] +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", +] + +[[package]] +name = "objc2-system-configuration" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7216bd11cbda54ccabcab84d523dc93b858ec75ecfb3a7d89513fa22464da396" +dependencies = [ + "objc2-core-foundation", +] + +[[package]] +name = "oid-registry" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8d8034d9489cdaf79228eb9f6a3b8d7bb32ba00d6645ebd48eef4077ceb5bd9" +dependencies = [ + "asn1-rs", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1124,12 +1373,51 @@ dependencies = [ "windows-link", ] +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64", + "serde_core", +] + [[package]] name = "percent-encoding" version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pgvector" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3673cba5b9a124916096a423b806a9f29620972c6c97b08db5f2053e9428b481" +dependencies = [ + "bytes", + "postgres-types", +] + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared", + "serde", +] + +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.13" @@ -1172,6 +1460,37 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "postgres-protocol" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" +dependencies = [ + "base64", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "hmac 0.13.0", + "md-5", + "memchr", + "rand 0.10.2", + "sha2 0.11.0", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" +dependencies = [ + "bytes", + "fallible-iterator 0.2.0", + "postgres-protocol", + "serde_core", + "serde_json", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -1181,6 +1500,12 @@ dependencies = [ "zerovec", ] +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + [[package]] name = "ppv-lite86" version = "0.2.21" @@ -1201,9 +1526,9 @@ dependencies = [ [[package]] name = "prost" -version = "0.13.5" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" dependencies = [ "bytes", "prost-derive", @@ -1211,9 +1536,9 @@ dependencies = [ [[package]] name = "prost-derive" -version = "0.13.5" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", "itertools 0.14.0", @@ -1235,8 +1560,8 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2 0.6.4", - "thiserror", + "socket2", + "thiserror 2.0.18", "tokio", "tracing", "web-time", @@ -1258,7 +1583,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror", + "thiserror 2.0.18", "tinyvec", "tracing", "web-time", @@ -1273,7 +1598,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.4", + "socket2", "tracing", "windows-sys 0.61.2", ] @@ -1384,6 +1709,19 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rcgen" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75e669e5202259b5314d1ea5397316ad400819437857b90861765f24c4cf80a2" +dependencies = [ + "pem", + "ring", + "rustls-pki-types", + "time", + "yasna", +] + [[package]] name = "redis" version = "0.27.6" @@ -1447,7 +1785,7 @@ dependencies = [ "tokio", "tokio-rustls", "tokio-util", - "tower 0.5.3", + "tower", "tower-http", "tower-service", "url", @@ -1479,7 +1817,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e" dependencies = [ "bitflags", - "fallible-iterator", + "fallible-iterator 0.3.0", "fallible-streaming-iterator", "hashlink", "libsqlite3-sys", @@ -1501,12 +1839,23 @@ dependencies = [ "semver", ] +[[package]] +name = "rusticata-macros" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632" +dependencies = [ + "nom", +] + [[package]] name = "rustls" version = "0.23.41" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b92b125634d9b795e7beca796cc790df15a7fb38323bf3196fda83292d06b1f" dependencies = [ + "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -1531,6 +1880,7 @@ version = "0.103.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -1634,7 +1984,7 @@ checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", "cpufeatures 0.2.17", - "digest", + "digest 0.10.7", ] [[package]] @@ -1651,7 +2001,18 @@ checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", "cpufeatures 0.2.17", - "digest", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -1688,6 +2049,12 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "siphasher" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" + [[package]] name = "slab" version = "0.4.12" @@ -1700,16 +2067,6 @@ version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" -[[package]] -name = "socket2" -version = "0.5.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" -dependencies = [ - "libc", - "windows-sys 0.52.0", -] - [[package]] name = "socket2" version = "0.6.4" @@ -1736,6 +2093,17 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "stringprep" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4df3d392d81bd458a8a621b8bffbd2302a12ffe288a9d931670948749463b1" +dependencies = [ + "unicode-bidi", + "unicode-normalization", + "unicode-properties", +] + [[package]] name = "subtle" version = "2.6.1" @@ -1773,13 +2141,33 @@ dependencies = [ "syn", ] +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", ] [[package]] @@ -1802,6 +2190,36 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "time" +version = "0.3.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18dfaaeddcb932337b5e7866ee7d0ce9b76d2fd092997146f187ec09b4558a50" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" + +[[package]] +name = "time-macros" +version = "0.2.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c431b87111666e491a90baa837f914fb45cd5dc3c268591b0220ff5057f2085f" +dependencies = [ + "num-conv", + "time-core", +] + [[package]] name = "tinystr" version = "0.8.3" @@ -1839,7 +2257,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "signal-hook-registry", - "socket2 0.6.4", + "socket2", "tokio-macros", "windows-sys 0.61.2", ] @@ -1855,6 +2273,32 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-postgres" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand 0.10.2", + "socket2", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-rustls" version = "0.26.4" @@ -1903,13 +2347,12 @@ dependencies = [ [[package]] name = "tonic" -version = "0.12.3" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "877c5b330756d856ffcc4553ab34a5684481ade925ecc54bcd1bf02b1d0d4d52" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ - "async-stream", "async-trait", - "axum 0.7.9", + "axum", "base64", "bytes", "h2", @@ -1921,34 +2364,26 @@ dependencies = [ "hyper-util", "percent-encoding", "pin-project", - "prost", - "socket2 0.5.10", + "socket2", + "sync_wrapper", "tokio", + "tokio-rustls", "tokio-stream", - "tower 0.4.13", + "tower", "tower-layer", "tower-service", "tracing", ] [[package]] -name = "tower" -version = "0.4.13" +name = "tonic-prost" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ - "futures-core", - "futures-util", - "indexmap 1.9.3", - "pin-project", - "pin-project-lite", - "rand 0.8.6", - "slab", - "tokio", - "tokio-util", - "tower-layer", - "tower-service", - "tracing", + "bytes", + "prost", + "tonic", ] [[package]] @@ -1959,9 +2394,12 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", "tracing", @@ -1979,7 +2417,7 @@ dependencies = [ "http", "http-body", "pin-project-lite", - "tower 0.5.3", + "tower", "tower-layer", "tower-service", "url", @@ -2074,7 +2512,7 @@ dependencies = [ "log", "rand 0.9.4", "sha1", - "thiserror", + "thiserror 2.0.18", ] [[package]] @@ -2083,12 +2521,33 @@ version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +[[package]] +name = "unicode-bidi" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5" + [[package]] name = "unicode-ident" version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-properties" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" + [[package]] name = "untrusted" version = "0.9.0" @@ -2113,6 +2572,29 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "utoipa" +version = "5.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bde15df68e80b16c7d16b9616e80770ad158988daa56a27dccd1e55558b0160" +dependencies = [ + "indexmap", + "serde", + "serde_json", + "utoipa-gen", +] + +[[package]] +name = "utoipa-gen" +version = "5.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba0b99ee52df3028635d93840c797102da61f8a7bb3cf751032455895b52ef8" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "uuid" version = "1.23.4" @@ -2158,6 +2640,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.4+wasi-0.2.12" @@ -2167,6 +2658,15 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.126" @@ -2264,6 +2764,19 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "whoami" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "998767ef88740d1f5b0682a9c53c24431453923962269c2db68ee43788c5a40d" +dependencies = [ + "libc", + "libredox", + "objc2-system-configuration", + "wasite", + "web-sys", +] + [[package]] name = "windows-link" version = "0.2.1" @@ -2364,6 +2877,32 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "x509-parser" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fcbc162f30700d6f3f82a24bf7cc62ffe7caea42c0b2cba8bf7f3ae50cf51f69" +dependencies = [ + "asn1-rs", + "data-encoding", + "der-parser", + "lazy_static", + "nom", + "oid-registry", + "rusticata-macros", + "thiserror 1.0.69", + "time", +] + +[[package]] +name = "yasna" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" +dependencies = [ + "time", +] + [[package]] name = "yoke" version = "0.8.3" diff --git a/Cargo.toml b/Cargo.toml index 3c6238b..7ab0a57 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ name = "corrosive_agents" version = "0.1.0" edition = "2021" -rust-version = "1.86" +rust-version = "1.88" authors = ["Kalanamith "] description = "Build verifiable, interactive AI agents powered by NVIDIA Nemotron free LLM models — MCP, skills, Ed25519 identity, REST/WebSocket/gRPC transports, and pluggable vector stores (Pinecone, Qdrant, custom)." license = "MIT OR Apache-2.0" @@ -15,14 +15,22 @@ exclude = [".env", ".github/", "TODO.md", "tools/"] [features] default = ["server"] -## REST + WebSocket serving via axum. -server = ["dep:axum"] +## REST + WebSocket serving via axum (includes API-key/JWT auth middleware). +server = ["dep:axum", "dep:hmac", "dep:sha2", "dep:ring"] ## gRPC serving via tonic (protobuf code is vendored; no protoc needed). -grpc = ["dep:tonic", "dep:prost"] +grpc = ["dep:tonic", "dep:tonic-prost", "dep:prost", "dep:hmac", "dep:sha2", "dep:ring"] +## TLS serving helpers for REST (axum-server) and gRPC (tonic/tls-ring). +tls = ["dep:axum-server", "tonic?/tls-ring"] +## OpenAPI 3 document served at /openapi.json. +openapi = ["server", "dep:utoipa"] +## X.509 certificate-based identity (self-signed certs from agent keys). +x509 = ["dep:rcgen", "dep:x509-parser", "ed25519-dalek/pkcs8"] ## Pinecone vector store backend (REST API). pinecone = [] ## Qdrant vector store backend (REST API). qdrant = ["uuid/v5"] +## pgvector (PostgreSQL) vector store backend. +pgvector = ["dep:tokio-postgres", "dep:pgvector"] ## SQLite-backed persistent session store. sqlite-sessions = ["dep:rusqlite"] ## Redis-backed persistent session store. @@ -31,14 +39,18 @@ redis-sessions = ["dep:redis"] full = [ "server", "grpc", + "tls", + "openapi", + "x509", "pinecone", "qdrant", + "pgvector", "sqlite-sessions", "redis-sessions", ] [dependencies] -tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time", "io-util", "process", "net"] } +tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time", "io-util", "process", "net", "signal"] } serde = { version = "1", features = ["derive"] } serde_json = "1" thiserror = "2" @@ -54,12 +66,30 @@ semver = "1" tracing = "0.1" uuid = { version = "1", features = ["v4"] } -# feature = "server" +# feature = "server" / "grpc" (JWT HS256 + RS256/JWKS verification) axum = { version = "0.8", features = ["ws"], optional = true } +hmac = { version = "0.12", optional = true } +sha2 = { version = "0.10", optional = true } +ring = { version = "0.17", optional = true } # feature = "grpc" -tonic = { version = "0.12", optional = true } -prost = { version = "0.13", optional = true } +tonic = { version = "0.14", optional = true } +tonic-prost = { version = "0.14", optional = true } +prost = { version = "0.14", optional = true } + +# feature = "tls" +axum-server = { version = "0.8", features = ["tls-rustls"], optional = true } + +# feature = "openapi" +utoipa = { version = "5", optional = true } + +# feature = "x509" +rcgen = { version = "0.13", optional = true } +x509-parser = { version = "0.16", optional = true } + +# feature = "pgvector" +tokio-postgres = { version = "0.7", features = ["with-serde_json-1"], optional = true } +pgvector = { version = "0.4", features = ["postgres"], optional = true } # feature = "sqlite-sessions" rusqlite = { version = "0.32", features = ["bundled"], optional = true } diff --git a/README.md b/README.md index 456f5f2..e98221c 100644 --- a/README.md +++ b/README.md @@ -23,14 +23,19 @@ or any **custom vector store**. |---|---| | 🏗️ **Builder pattern** | Fluent `Agent::builder()…build()` construction with semver + config validation | | 📄 **JSON manifests** | Load name/version/capabilities/skills/MCP servers from a file | -| 🔐 **Verifiable identity** | Ed25519-signed manifests; `did:key` DIDs; key rotation chains & revocations (`TrustStore`) | -| 🧠 **NVIDIA Nemotron** | Chat, streaming (SSE), and embeddings via the free NIM API | -| 🛠️ **Skills** | Async JSON-in/JSON-out abilities, invocable locally or over any transport | -| 🔌 **MCP** | stdio JSON-RPC client: `initialize`, `tools/list`, `tools/call` | -| 🌐 **Transports** | REST + WebSocket (axum) and gRPC (tonic, vendored protos — no `protoc` needed) | +| 🔐 **Verifiable identity** | Ed25519 manifests; `did:key` DIDs; X.509 certs; key rotation & revocation (`TrustStore`) | +| 🧠 **NVIDIA Nemotron** | Chat, streaming, tool calling, embeddings — with retry/backoff + rate-limit handling | +| 🛠️ **Skills** | Async JSON abilities with a sandbox: allowlists, permissions, timeouts, panic isolation | +| 🔁 **Tool loop** | `chat_with_tools`: the model auto-invokes skills; usage accounting hooks built in | +| 🔌 **MCP** | stdio + streamable-HTTP/SSE transports; tools, resources, and prompts | +| 🌐 **Transports** | REST + WebSocket + gRPC, with API-key/JWT auth, TLS, graceful shutdown, `/ready`, OpenAPI | | 🤝 **A2A delegation** | `RemoteAgent` peers with pinned-key/DID verification; delegate chat & skills | | 💾 **Sessions** | Pluggable `SessionStore`: in-memory, SQLite, or Redis persistence | -| 📚 **Vector stores** | One trait; in-memory, Qdrant, Pinecone backends; `remember`/`recall` RAG helpers | +| 📚 **Vector stores** | In-memory, Qdrant, Pinecone, pgvector; metadata filters, chunking, `remember`/`recall` | + +📖 **New to the library? Read the [tutorial](docs/TUTORIAL.md)** — it walks +from an empty project to a production-shaped agent, and ships on +[docs.rs](https://docs.rs/corrosive_agents) as the `tutorial` module. ## Installation @@ -46,15 +51,19 @@ Optional features: corrosive_agents = { version = "0.1", features = ["full"] } # grpc + qdrant + pinecone ``` -| Feature | Default | Enables | -|-------------------|---------|----------------------------------------| -| `server` | ✅ | REST + WebSocket serving (axum) | -| `grpc` | — | gRPC serving + generated client (tonic)| -| `pinecone` | — | Pinecone vector store backend | -| `qdrant` | — | Qdrant vector store backend | -| `sqlite-sessions` | — | SQLite-persisted conversation history | -| `redis-sessions` | — | Redis-persisted conversation history | -| `full` | — | All of the above | +| Feature | Default | Enables | +|-------------------|---------|------------------------------------------------| +| `server` | ✅ | REST + WebSocket serving + auth middleware | +| `grpc` | — | gRPC serving + generated client (tonic) | +| `tls` | — | TLS helpers for REST and gRPC | +| `openapi` | — | OpenAPI 3 document at `/openapi.json` | +| `x509` | — | X.509 certificate-based identity | +| `pinecone` | — | Pinecone vector store backend | +| `qdrant` | — | Qdrant vector store backend | +| `pgvector` | — | PostgreSQL/pgvector vector store backend | +| `sqlite-sessions` | — | SQLite-persisted conversation history | +| `redis-sessions` | — | Redis-persisted conversation history | +| `full` | — | All of the above | ## Quickstart @@ -198,6 +207,7 @@ methods: `upsert`, `search`, `delete`). | `build_agent` | Builder pattern, capabilities, skills, signed manifest, chat | `cargo run --example build_agent` | | `agent_from_json` | Loading an agent + MCP config from `examples/agent.json` | `cargo run --example agent_from_json` | | `interactive_chat` | Terminal REPL with streamed tokens | `cargo run --example interactive_chat` | +| `tool_calling` | Model auto-invokes skills (function calling) + usage hooks | `cargo run --example tool_calling` | | `sign_and_verify` | Public-key verification end to end (offline) | `cargo run --example sign_and_verify` | | `vector_rag` | Embeddings + vector store + retrieval-augmented answers | `cargo run --example vector_rag` | | `serve` | REST + WebSocket server | `cargo run --example serve` | diff --git a/docs/TUTORIAL.md b/docs/TUTORIAL.md new file mode 100644 index 0000000..d551d33 --- /dev/null +++ b/docs/TUTORIAL.md @@ -0,0 +1,415 @@ +# Building agents with `corrosive_agents` — the tutorial + +This walks you from an empty project to a production-shaped agent: +chat → skills → tool calling → memory (RAG) → identity & trust → serving +(REST/WebSocket/gRPC with auth and TLS) → agent-to-agent delegation. + +Code blocks marked `rust,no_run` compile against the default features; blocks +that need optional features are marked `rust,ignore` and name the feature. + +## 0. Setup + +Add the crate (REST + WebSocket serving is on by default): + +```toml +[dependencies] +corrosive_agents = "0.1" +tokio = { version = "1", features = ["full"] } +serde_json = "1" +``` + +Get a **free** NVIDIA API key: visit , open any +model page, press *Get API Key*. Export it: + +```sh +export NVIDIA_API_KEY=nvapi-... +``` + +## 1. Your first agent + +Everything starts with the builder. `name` and `version` (valid SemVer) are +required; `build()` validates and returns the immutable agent. + +```rust,no_run +use corrosive_agents::prelude::*; + +#[tokio::main] +async fn main() -> Result<()> { + let agent = Agent::builder() + .name("tutor") + .version("0.1.0") + .description("Explains Rust concepts briefly") + .system_prompt("You are a Rust tutor. Answer in at most four sentences.") + .model(models::NEMOTRON_3_NANO_30B) + .llm(NvidiaClient::from_env()?) + .build()?; + + // History is kept per session id — the second question can say "it". + let a1 = agent.chat("lesson-1", "What is an Arc?").await?; + let a2 = agent.chat("lesson-1", "When would I combine it with a Mutex?").await?; + println!("{a1}\n---\n{a2}"); + Ok(()) +} +``` + +Streaming variant — print tokens as they arrive: + +```rust,no_run +use corrosive_agents::prelude::*; +use futures_util::StreamExt; + +# async fn run(agent: Agent) -> Result<()> { +let mut stream = agent.chat_stream("lesson-1", "Explain lifetimes").await?; +while let Some(chunk) = stream.next().await { + let chunk = chunk?; + if chunk.done { break } + print!("{}", chunk.delta); +} +# Ok(()) +# } +``` + +The NVIDIA client retries 429/5xx with exponential backoff automatically; +tune it with `.with_retry_policy(RetryPolicy { .. })` or disable via +`RetryPolicy::none()`. + +## 2. Define the agent in JSON + +Ship the agent's *shape* as data. Save as `agent.json`: + +```json +{ + "name": "tutor", + "version": "1.0.0", + "description": "Explains Rust concepts briefly", + "model": "nvidia/nemotron-3-nano-30b-a3b", + "system_prompt": "You are a Rust tutor.", + "capabilities": [{ "name": "chat", "description": "Q&A" }], + "skills": ["lookup_doc"], + "mcp_servers": [ + { "name": "fs", "command": "npx", + "args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"] } + ] +} +``` + +Load it, then bind code to the declared skill names: + +```rust,no_run +use corrosive_agents::prelude::*; +use serde_json::json; + +# async fn run() -> Result<()> { +let agent = AgentBuilder::from_json_file("agent.json")? + .skill(FnSkill::new("lookup_doc", "Looks up a std type", |input| async move { + Ok(json!({ "url": format!("https://doc.rust-lang.org/std/?search={}", + input["query"].as_str().unwrap_or_default()) })) + })) + .llm(NvidiaClient::from_env()?) + .build()?; +# Ok(()) +# } +``` + +## 3. Skills and the sandbox + +A skill is a named async JSON-in/JSON-out function. Implement the `Skill` +trait for full control, or wrap a closure with `FnSkill`. Give the model a +real schema — it dramatically improves call accuracy: + +```rust,no_run +use corrosive_agents::prelude::*; +use serde_json::json; + +let weather = FnSkill::new("weather", "Current weather for a city", |input| async move { + let city = input["city"].as_str().unwrap_or("nowhere").to_string(); + Ok(json!({ "city": city, "temp_c": 21 })) +}) +.with_schema(json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"] +})) +.with_permissions(["net"]); // declared requirement +``` + +The **`SkillPolicy`** is the sandbox: an allowlist of skill names, a set of +granted permissions, and an execution timeout (default 30 s). Skills run on +their own task, so a panic becomes an error instead of killing the agent. + +```rust,no_run +use corrosive_agents::prelude::*; +use std::time::Duration; + +# fn f(builder: AgentBuilder) -> AgentBuilder { +builder.skill_policy( + SkillPolicy::new() + .allow_only(["weather", "lookup_doc"]) + .grant("net") + .with_timeout(Duration::from_secs(5)), +) +# } +``` + +## 4. Tool calling — let the model use your skills + +`chat_with_tools` offers every registered skill to the model as a function. +When the model asks for a tool, the agent executes it (through the policy!), +feeds the JSON result back, and loops until the model produces text: + +```rust,no_run +use corrosive_agents::prelude::*; + +# async fn run(agent: Agent) -> Result<()> { +let reply = agent + .chat_with_tools("session", "What's the weather in Colombo?", 8) + .await?; // ≤ 8 tool rounds +# Ok(()) +# } +``` + +Over REST the same thing is `POST /chat` with `{"use_tools": true}`. + +Track spend with usage hooks — a closure works as an observer: + +```rust,no_run +use corrosive_agents::prelude::*; + +# fn f(builder: AgentBuilder) -> AgentBuilder { +builder.usage_observer(|event: &UsageEvent| { + println!("[{}] {} used {} tokens", event.session_id, event.model, event.usage.total_tokens); +}) +# } +``` + +`agent.usage()` returns cumulative totals at any time. + +## 5. Memory: sessions and RAG + +**Sessions** live in a `SessionStore`. The default is in-memory; swap in +SQLite or Redis for persistence across restarts (features `sqlite-sessions` +/ `redis-sessions`): + +```rust,ignore +// --features sqlite-sessions +let agent = Agent::builder() + .name("persistent").version("1.0.0") + .llm(NvidiaClient::from_env()?) + .session_store(SqliteSessionStore::open("sessions.db")?) + .build()?; +// After a restart, session ids resume where they left off. +``` + +**RAG** needs an embedding provider (the `NvidiaClient` doubles as one) and a +`VectorStore` — in-memory, Qdrant (`qdrant`), Pinecone (`pinecone`), +Postgres/pgvector (`pgvector`), or your own trait impl: + +```rust,no_run +use corrosive_agents::prelude::*; +use serde_json::json; + +# async fn run() -> Result<()> { +let nvidia = NvidiaClient::from_env()?; +let agent = Agent::builder() + .name("librarian").version("0.1.0") + .llm(nvidia.clone()) + .embeddings(nvidia) + .vector_store(InMemoryVectorStore::new()) + .build()?; + +// Index — one fact, a batch, or a whole document (auto-chunked): +agent.remember("Nemotron 3 Nano is fast.", json!({"topic": "models"})).await?; +agent.remember_document(&std::fs::read_to_string("guide.md")?, + json!({"source": "guide"}), 1200, 200).await?; + +// Retrieve — optionally filtered by metadata: +let hits = agent.recall("which model is fast?", 3).await?; +let guide_only = agent + .recall_filtered("chunking", 3, &MetadataFilter::new().eq("source", json!("guide"))) + .await?; +# Ok(()) +# } +``` + +## 6. Identity, verification, and trust + +Give the agent an Ed25519 identity and the manifest is signed at `build()`: + +```rust,no_run +use corrosive_agents::prelude::*; + +# fn main() -> Result<()> { +let agent = Agent::builder() + .name("trusted").version("1.0.0") + .generate_identity() + .build()?; + +let manifest_json = agent.manifest().to_json()?; // ship this anywhere +let did = agent.identity().unwrap().did_key(); // did:key:z6Mk… + +// Any consumer, offline: +let received = AgentManifest::from_json(&manifest_json)?; +received.verify()?; // embedded key +received.verify_with(&did)?; // or pinned key / DID +# Ok(()) +# } +``` + +Store the secret (`identity.secret_key_base64()`) somewhere safe and restore +with `AgentIdentity::from_secret_base64` to keep the same identity across +releases. + +**Rotation & revocation** — replace a key without losing consumer trust: + +```rust,no_run +use corrosive_agents::prelude::*; + +# fn main() -> Result<()> { +let old = AgentIdentity::generate(); +let new = AgentIdentity::generate(); +let mut manifest = AgentManifest::new("trusted", "1.1.0"); +manifest.sign(&old)?; +manifest.rotate_identity(&old, &new)?; // old key endorses new; re-signed + +let mut trust = TrustStore::new(); +trust.trust(&old.public_key_base64())?; // consumer pinned the OLD key +trust.verify_manifest(&manifest)?; // still verifies via the chain +trust.revoke(Revocation::create(&old, "rotated away"))?; // kills post-revocation use +# Ok(()) +# } +``` + +**X.509** (feature `x509`) — for PKI-speaking integrations: + +```rust,ignore +// --features x509 +let cert_pem = corrosive_agents::x509::generate_certificate_pem(&identity, "trusted")?; +corrosive_agents::x509::verify_manifest_with_certificate(&manifest, &cert_pem)?; +``` + +## 7. MCP: tools, resources, prompts + +Declare servers in the manifest (stdio) or connect over streamable HTTP: + +```rust,no_run +use corrosive_agents::prelude::*; +use serde_json::json; + +# async fn run(agent: Agent) -> Result<()> { +agent.connect_mcp_servers().await?; // manifest-declared servers + +let tools = agent.mcp_tools("fs").await?; +let listing = agent.call_mcp_tool("fs", "list_directory", json!({"path": "/tmp"})).await?; + +// Or an HTTP/SSE server, ad hoc: +let remote = McpClient::connect( + &McpServerConfig::http("docs", "https://mcp.example.com/mcp") + .with_header("Authorization", "Bearer token"), +).await?; +let resources = remote.list_resources().await?; +let readme = remote.read_resource("file:///README.md").await?; +let prompts = remote.list_prompts().await?; +# Ok(()) +# } +``` + +## 8. Serving: REST, WebSocket, gRPC + +```rust,no_run +use std::sync::Arc; +use corrosive_agents::auth::AuthScheme; +use corrosive_agents::prelude::*; +use corrosive_agents::server; + +# async fn run(agent: Agent) -> Result<()> { +let agent = Arc::new(agent); + +// Open server with graceful shutdown on Ctrl-C/SIGTERM: +server::serve_with_shutdown( + agent.clone(), + "0.0.0.0:8080".parse().unwrap(), + server::shutdown_signal(), +).await?; + +// API-key or JWT protected (probes stay open): +server::serve_with_auth( + agent, + "0.0.0.0:8080".parse().unwrap(), + AuthScheme::jwt_hs256("shared-secret").with_issuer("my-platform"), +).await?; +# Ok(()) +# } +``` + +For an identity provider (Auth0, Keycloak, …), verify RS256 tokens against +its JWKS instead of sharing a secret: + +```rust,no_run +use corrosive_agents::auth::{AuthScheme, JwksStore}; + +# async fn run() -> corrosive_agents::Result<()> { +let jwks = JwksStore::from_url("https://idp.example.com/.well-known/jwks.json").await?; +let auth = AuthScheme::jwt_rs256(jwks) + .with_issuer("https://idp.example.com/") + .with_audience("corrosive-agents"); +# Ok(()) +# } +``` + +(`JwksStore::refresh()` re-fetches keys — call it from a periodic task if +your provider rotates them.) + +- `/health` is liveness; `/ready` is readiness — flip it with + `agent.set_ready(false)` while warming up or draining. +- With feature `openapi`, the full spec is served at `/openapi.json`. +- With feature `tls`: `server::serve_tls(agent, addr, &TlsConfig::from_pem_files("cert.pem", "key.pem"))`. +- gRPC (feature `grpc`) mirrors all of it: `grpc::serve`, `grpc::serve_with_auth`, + `grpc::serve_with_shutdown`, `grpc::serve_tls` — clients authenticate with + `authorization` / `x-api-key` metadata. + +WebSocket protocol (at `/ws`): send +`{"type":"chat","message":"hi","stream":true}`, receive `chunk` events and a +final `done`. + +## 9. Agent-to-agent (A2A) delegation + +Compose agents across processes/machines. Pin the peer's DID and the client +verifies its signed manifest before the first request: + +```rust,no_run +use corrosive_agents::prelude::*; +use serde_json::json; + +# async fn run() -> Result<()> { +let orchestrator = Agent::builder() + .name("orchestrator").version("1.0.0") + .peer("cruncher", + RemoteAgent::new("http://worker:8080") + .with_pinned_key("did:key:z6Mk...worker...")) + .build()?; + +let reply = orchestrator.delegate_chat("cruncher", "job-7", "crunch this").await?; +let out = orchestrator.delegate_skill("cruncher", "fibonacci", json!({"n": 42})).await?; +# Ok(()) +# } +``` + +A mismatched identity fails **before** any payload is sent. Run +`cargo run --example a2a_delegation` to see it end to end, offline. + +## 10. Production checklist + +- [ ] Persist the identity secret (vault/KMS), not in the manifest or repo. +- [ ] Pin peer keys or DIDs for every `RemoteAgent`. +- [ ] `SkillPolicy` with an explicit allowlist + timeouts on any deployment + that loads third-party skills. +- [ ] Persistent `SessionStore` (SQLite single node, Redis multi-node). +- [ ] `serve_with_auth` + `tls` (or terminate TLS at your ingress). +- [ ] Wire `usage_observer` into your metrics stack. +- [ ] `/ready` in your orchestrator's readiness probe; + `serve_with_shutdown(…, shutdown_signal())` for clean rollouts. +- [ ] Watch NVIDIA model availability — catalog entries can 404; probe with a + 1-token completion and keep model ids configurable via the manifest. + +Every topic above has a runnable example in `examples/` — start with +`build_agent`, `vector_rag`, `sign_and_verify`, and `a2a_delegation`. diff --git a/examples/tool_calling.rs b/examples/tool_calling.rs new file mode 100644 index 0000000..083a88b --- /dev/null +++ b/examples/tool_calling.rs @@ -0,0 +1,59 @@ +//! Tool calling: the model automatically invokes the agent's skills. +//! +//! ```sh +//! export NVIDIA_API_KEY=nvapi-... +//! cargo run --example tool_calling +//! ``` + +use corrosive_agents::prelude::*; +use serde_json::json; + +#[tokio::main] +async fn main() -> Result<()> { + let _ = dotenvy::dotenv(); + + let agent = Agent::builder() + .name("calculator-agent") + .version("0.1.0") + .system_prompt( + "You are a precise assistant. Use the provided tools for any arithmetic; \ + never compute it yourself.", + ) + .model(models::LLAMA_NEMOTRON_SUPER_49B) + .skill( + FnSkill::new( + "multiply", + "Multiplies two integers a and b", + |input| async move { + let a = input["a"].as_i64().unwrap_or(0); + let b = input["b"].as_i64().unwrap_or(0); + println!(" [skill] multiply({a}, {b})"); + Ok(json!({ "product": a * b })) + }, + ) + .with_schema(json!({ + "type": "object", + "properties": { + "a": { "type": "integer" }, + "b": { "type": "integer" } + }, + "required": ["a", "b"] + })), + ) + .usage_observer(|event: &UsageEvent| { + println!( + " [usage] {} tokens ({})", + event.usage.total_tokens, event.model + ); + }) + .llm(NvidiaClient::from_env()?) + .build()?; + + let question = "What is 1337 multiplied by 42? Use the tool."; + println!("user > {question}"); + let reply = agent.chat_with_tools("demo", question, 5).await?; + println!("agent> {reply}"); + + println!("\ntotal usage: {:?}", agent.usage()); + Ok(()) +} diff --git a/src/agent/builder.rs b/src/agent/builder.rs index 1ce0649..a38b1f2 100644 --- a/src/agent/builder.rs +++ b/src/agent/builder.rs @@ -10,10 +10,10 @@ use crate::a2a::RemoteAgent; use crate::agent::{Agent, AgentManifest, Capability}; use crate::error::{Error, Result}; use crate::identity::AgentIdentity; -use crate::llm::{EmbeddingProvider, LlmProvider}; +use crate::llm::{EmbeddingProvider, LlmProvider, UsageObserver, UsageTotals}; use crate::mcp::McpServerConfig; use crate::session::{InMemorySessionStore, SessionStore}; -use crate::skills::{Skill, SkillRegistry}; +use crate::skills::{Skill, SkillPolicy, SkillRegistry}; use crate::vector::VectorStore; /// Builds an [`Agent`] step by step. @@ -44,8 +44,10 @@ pub struct AgentBuilder { embeddings: Option>, vector_store: Option>, skills: Vec>, + skill_policy: Option, session_store: Option>, peers: HashMap>, + usage_observer: Option>, } impl AgentBuilder { @@ -176,6 +178,24 @@ impl AgentBuilder { self } + /// Set the sandbox policy for skill execution (allowlist, permissions, + /// timeout). Defaults to [`SkillPolicy::default`]: every skill allowed, + /// no permissions granted, 30-second timeout. + #[must_use] + pub fn skill_policy(mut self, policy: SkillPolicy) -> Self { + self.skill_policy = Some(policy); + self + } + + /// Receive a [`UsageEvent`](crate::llm::UsageEvent) after every + /// completion — for metrics, billing, or rate-limit dashboards. Plain + /// closures work: `.usage_observer(|e: &UsageEvent| println!("{e:?}"))`. + #[must_use] + pub fn usage_observer(mut self, observer: impl UsageObserver + 'static) -> Self { + self.usage_observer = Some(Arc::new(observer)); + self + } + /// Give the agent an existing identity. The manifest will be signed with /// it during [`build`](Self::build). #[must_use] @@ -237,11 +257,15 @@ impl AgentBuilder { embeddings: self.embeddings, vector_store: self.vector_store, skills: registry, + skill_policy: self.skill_policy.unwrap_or_default(), mcp_clients: RwLock::new(HashMap::new()), sessions: self .session_store .unwrap_or_else(|| Arc::new(InMemorySessionStore::new())), peers: RwLock::new(self.peers), + usage_totals: Arc::new(UsageTotals::default()), + usage_observer: self.usage_observer, + ready: std::sync::atomic::AtomicBool::new(true), }) } } diff --git a/src/agent/manifest.rs b/src/agent/manifest.rs index 81a96d5..41957e9 100644 --- a/src/agent/manifest.rs +++ b/src/agent/manifest.rs @@ -13,6 +13,7 @@ use crate::mcp::McpServerConfig; /// Capabilities are declarative metadata: they describe the agent to humans /// and other agents, and can carry free-form JSON configuration. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct Capability { /// Machine-readable capability name, e.g. `"chat"` or `"rag"`. pub name: String, @@ -24,6 +25,7 @@ pub struct Capability { pub enabled: bool, /// Optional free-form configuration for this capability. #[serde(default, skip_serializing_if = "Option::is_none")] + #[cfg_attr(feature = "openapi", schema(value_type = Option))] pub config: Option, } @@ -82,6 +84,7 @@ impl Capability { /// } /// ``` #[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct AgentManifest { /// Agent name (required, non-empty). pub name: String, diff --git a/src/agent/runtime.rs b/src/agent/runtime.rs index 6f28e39..419c820 100644 --- a/src/agent/runtime.rs +++ b/src/agent/runtime.rs @@ -14,14 +14,18 @@ use crate::a2a::RemoteAgent; use crate::agent::{AgentBuilder, AgentManifest, Capability}; use crate::error::{Error, Result}; use crate::identity::AgentIdentity; -use crate::llm::{ChatMessage, ChatRequest, EmbeddingProvider, LlmProvider, StreamChunk}; +use crate::llm::{ + ChatMessage, ChatRequest, ChatResponse, EmbeddingProvider, LlmProvider, StreamChunk, ToolSpec, + UsageEvent, UsageObserver, UsageSnapshot, UsageTotals, +}; use crate::mcp::{McpClient, McpTool}; use crate::session::SessionStore; -use crate::skills::SkillRegistry; -use crate::vector::{Document, SearchResult, VectorStore}; +use crate::skills::{SkillPolicy, SkillRegistry}; +use crate::vector::{chunk_text, Document, MetadataFilter, SearchResult, VectorStore}; /// Public, serializable snapshot of an agent — what `GET /agent` returns. #[derive(Debug, Clone, Serialize, Deserialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct AgentInfo { /// Agent name. pub name: String, @@ -55,9 +59,13 @@ pub struct Agent { pub(crate) embeddings: Option>, pub(crate) vector_store: Option>, pub(crate) skills: SkillRegistry, + pub(crate) skill_policy: SkillPolicy, pub(crate) mcp_clients: RwLock>>, pub(crate) sessions: Arc, pub(crate) peers: RwLock>>, + pub(crate) usage_totals: Arc, + pub(crate) usage_observer: Option>, + pub(crate) ready: std::sync::atomic::AtomicBool, } impl std::fmt::Debug for Agent { @@ -141,9 +149,46 @@ impl Agent { &self.skills } - /// Execute a registered skill by name. + /// Execute a registered skill by name, subject to the agent's + /// [`SkillPolicy`] (allowlist, permissions, timeout). Execution runs on + /// a separate task so a panicking skill cannot take the agent down. pub async fn execute_skill(&self, name: &str, input: Value) -> Result { - self.skills.execute(name, input).await + let skill = self + .skills + .get(name) + .ok_or_else(|| Error::SkillNotFound(name.to_string()))?; + self.skill_policy.check(skill.as_ref())?; + + let handle = tokio::spawn(async move { skill.execute(input).await }); + let abort = handle.abort_handle(); + let outcome = match self.skill_policy.timeout() { + Some(limit) => tokio::time::timeout(limit, handle).await.map_err(|_| { + abort.abort(); // don't leave the runaway skill running + Error::Skill(format!("skill '{name}' timed out after {limit:?}")) + })?, + None => handle.await, + }; + outcome.map_err(|e| Error::Skill(format!("skill '{name}' panicked: {e}")))? + } + + /// Record token usage from a completed response. + fn record_usage(&self, session_id: &str, response: &ChatResponse) { + if let Some(usage) = &response.usage { + let event = UsageEvent { + session_id: session_id.to_string(), + model: response.model.clone(), + usage: usage.clone(), + }; + self.usage_totals.record(&event); + if let Some(observer) = &self.usage_observer { + observer.on_usage(&event); + } + } + } + + /// Cumulative token usage across all sessions since the agent started. + pub fn usage(&self) -> UsageSnapshot { + self.usage_totals.snapshot() } fn require_llm(&self) -> Result> { @@ -178,6 +223,7 @@ impl Agent { let llm = self.require_llm()?; let request = self.request_for(self.conversation(session_id, message).await?); let response = llm.chat(request).await?; + self.record_usage(session_id, &response); self.sessions .append( session_id, @@ -190,6 +236,70 @@ impl Agent { Ok(response.content) } + /// Like [`chat`](Self::chat), but the model may call the agent's + /// registered skills as tools (function calling). + /// + /// The loop runs until the model answers with plain text (or + /// `max_rounds` tool rounds elapse): tool calls are executed through + /// [`execute_skill`](Self::execute_skill) — so the [`SkillPolicy`] is + /// enforced — and their JSON results are fed back to the model. Skill + /// failures are reported to the model as `{"error": …}` results rather + /// than aborting the turn. + pub async fn chat_with_tools( + &self, + session_id: &str, + message: impl AsRef, + max_rounds: usize, + ) -> Result { + let message = message.as_ref(); + let llm = self.require_llm()?; + let tools: Vec = self + .skills + .list() + .iter() + .map(|s| ToolSpec::from_skill(s.as_ref())) + .collect(); + + let mut messages = self.conversation(session_id, message).await?; + // Everything after (and including) the user message gets persisted. + let mut transcript: Vec = vec![ChatMessage::user(message)]; + + for _ in 0..max_rounds.max(1) { + let mut request = self.request_for(messages.clone()); + if !tools.is_empty() { + request.tools = Some(tools.clone()); + } + let response = llm.chat(request).await?; + self.record_usage(session_id, &response); + + if response.tool_calls.is_empty() { + transcript.push(ChatMessage::assistant(&response.content)); + self.sessions.append(session_id, &transcript).await?; + return Ok(response.content); + } + + let assistant = ChatMessage::assistant_tool_calls( + response.content.clone(), + response.tool_calls.clone(), + ); + messages.push(assistant.clone()); + transcript.push(assistant); + + for call in response.tool_calls { + let output = match self.execute_skill(&call.name, call.arguments.clone()).await { + Ok(value) => value.to_string(), + Err(e) => serde_json::json!({ "error": e.to_string() }).to_string(), + }; + let result = ChatMessage::tool_result(&call.id, output); + messages.push(result.clone()); + transcript.push(result); + } + } + Err(Error::Llm(format!( + "tool-calling loop did not converge within {max_rounds} rounds" + ))) + } + /// Streaming variant of [`chat`](Self::chat): yields incremental /// [`StreamChunk`]s and records the full reply in the session once the /// stream completes. @@ -356,8 +466,85 @@ impl Agent { Ok(id) } + /// Embed and store several texts at once (batched embedding + batched + /// upsert). Returns the generated document ids, in input order. + pub async fn remember_batch(&self, texts: &[String], metadata: Value) -> Result> { + let embeddings = self + .embeddings + .clone() + .ok_or(Error::NotConfigured("embedding provider"))?; + let store = self + .vector_store + .clone() + .ok_or(Error::NotConfigured("vector store"))?; + if texts.is_empty() { + return Ok(Vec::new()); + } + + let vectors = embeddings.embed_documents(texts).await?; + if vectors.len() != texts.len() { + return Err(Error::Llm(format!( + "embedding provider returned {} vectors for {} texts", + vectors.len(), + texts.len() + ))); + } + let mut ids = Vec::with_capacity(texts.len()); + let documents: Vec = texts + .iter() + .zip(vectors) + .map(|(text, vector)| { + let id = uuid::Uuid::new_v4().to_string(); + ids.push(id.clone()); + Document::new(id, vector) + .with_text(text) + .with_metadata(metadata.clone()) + }) + .collect(); + store.upsert_batched(documents, 64).await?; + Ok(ids) + } + + /// Chunk a long document (word-boundary chunks of `max_chars` with + /// `overlap` characters of carried context — see + /// [`chunk_text`](crate::vector::chunk_text)), then embed and store every + /// chunk. Each chunk's metadata gains a `_chunk` index. + pub async fn remember_document( + &self, + text: &str, + metadata: Value, + max_chars: usize, + overlap: usize, + ) -> Result> { + let chunks = chunk_text(text, max_chars, overlap); + let mut ids = Vec::with_capacity(chunks.len()); + for (index, chunk) in chunks.iter().enumerate() { + let mut chunk_metadata = metadata.clone(); + if let Value::Object(map) = &mut chunk_metadata { + map.insert("_chunk".into(), Value::from(index)); + } + ids.extend( + self.remember_batch(std::slice::from_ref(chunk), chunk_metadata) + .await?, + ); + } + Ok(ids) + } + /// Embed `query` and return the `top_k` most similar remembered documents. pub async fn recall(&self, query: &str, top_k: usize) -> Result> { + self.recall_filtered(query, top_k, &MetadataFilter::new()) + .await + } + + /// Like [`recall`](Self::recall), restricted to documents whose metadata + /// matches `filter`. + pub async fn recall_filtered( + &self, + query: &str, + top_k: usize, + filter: &MetadataFilter, + ) -> Result> { let embeddings = self .embeddings .clone() @@ -368,11 +555,24 @@ impl Agent { .ok_or(Error::NotConfigured("vector store"))?; let vector = embeddings.embed_query(query).await?; - store.search(vector, top_k).await + store.search_filtered(vector, top_k, filter).await } /// The configured vector store, when present. pub fn vector_store(&self) -> Option> { self.vector_store.clone() } + + /// Readiness flag served by `GET /ready` (liveness at `/health` is + /// always `ok`). Starts `true`; flip it while warming caches or draining + /// before shutdown. + pub fn set_ready(&self, ready: bool) { + self.ready + .store(ready, std::sync::atomic::Ordering::Relaxed); + } + + /// Current readiness state. + pub fn is_ready(&self) -> bool { + self.ready.load(std::sync::atomic::Ordering::Relaxed) + } } diff --git a/src/auth.rs b/src/auth.rs new file mode 100644 index 0000000..c81a143 --- /dev/null +++ b/src/auth.rs @@ -0,0 +1,664 @@ +//! Authentication for agent transports (features `server` / `grpc`). +//! +//! One [`AuthScheme`] protects both REST/WebSocket +//! ([`server::router_with_auth`](crate::server::router_with_auth)) and gRPC +//! ([`grpc::serve_with_auth`](crate::grpc::serve_with_auth)): +//! +//! - **API keys** — clients send `Authorization: Bearer ` or +//! `X-Api-Key: `. +//! - **JWT (HS256)** — shared-secret HMAC-SHA256 tokens. +//! - **JWT (RS256 + JWKS)** — RSA tokens verified against a JSON Web Key +//! Set, e.g. from your identity provider's +//! `/.well-known/jwks.json` ([`JwksStore`]). +//! +//! All JWTs must carry a future `exp` and, when configured, matching +//! `iss`/`aud` claims. Health/readiness endpoints stay unauthenticated so +//! probes keep working. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use base64::engine::general_purpose::URL_SAFE_NO_PAD; +use base64::Engine as _; +use hmac::{Hmac, Mac}; +use sha2::Sha256; + +use crate::error::{Error, Result}; + +/// Allowed clock skew when checking `exp`/`nbf` (seconds). +const LEEWAY_SECS: u64 = 30; + +/// How clients must authenticate. +#[derive(Clone)] +pub enum AuthScheme { + /// A fixed set of accepted API keys. + ApiKeys(HashSet), + /// HS256-signed JWTs verified with a shared secret. + JwtHs256 { + /// Shared HMAC secret. + secret: String, + /// Required `iss` claim, when set. + issuer: Option, + /// Required `aud` claim, when set. + audience: Option, + }, + /// RS256-signed JWTs verified against a JSON Web Key Set. + JwtRs256 { + /// The RSA public keys (kid → key), loadable from JWKS JSON or URL. + keys: Arc, + /// Required `iss` claim, when set. + issuer: Option, + /// Required `aud` claim, when set. + audience: Option, + }, +} + +impl std::fmt::Debug for AuthScheme { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Never print key material. + match self { + Self::ApiKeys(keys) => f + .debug_struct("AuthScheme::ApiKeys") + .field("count", &keys.len()) + .finish(), + Self::JwtHs256 { + issuer, audience, .. + } => f + .debug_struct("AuthScheme::JwtHs256") + .field("issuer", issuer) + .field("audience", audience) + .finish_non_exhaustive(), + Self::JwtRs256 { + keys, + issuer, + audience, + } => f + .debug_struct("AuthScheme::JwtRs256") + .field("keys", &keys.len()) + .field("issuer", issuer) + .field("audience", audience) + .finish_non_exhaustive(), + } + } +} + +/// One RSA public key from a JWKS document (raw big-endian components). +#[derive(Clone)] +struct RsaJwk { + n: Vec, + e: Vec, +} + +/// A set of RSA public keys for RS256 verification, keyed by `kid`. +/// +/// Load once from static JWKS JSON ([`from_json`](Self::from_json)) or from +/// a URL ([`from_url`](Self::from_url)); call [`refresh`](Self::refresh) on +/// key rotation (e.g. from a periodic task). Verification itself is +/// synchronous against the cached keys. +pub struct JwksStore { + url: Option, + keys: std::sync::RwLock>, +} + +impl std::fmt::Debug for JwksStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("JwksStore") + .field("url", &self.url) + .field("keys", &self.len()) + .finish() + } +} + +impl JwksStore { + /// Parse a static JWKS document (`{"keys":[{"kty":"RSA",…}]}`). + /// Non-RSA keys and keys marked for encryption (`"use":"enc"`) are + /// skipped. + pub fn from_json(jwks: &str) -> Result { + let keys = Self::parse(jwks)?; + if keys.is_empty() { + return Err(Error::Auth("JWKS contains no usable RSA keys".into())); + } + Ok(Self { + url: None, + keys: std::sync::RwLock::new(keys), + }) + } + + /// Fetch a JWKS document from `url` (typically + /// `https:///.well-known/jwks.json`). + pub async fn from_url(url: impl Into) -> Result { + let url = url.into(); + let store = Self { + url: Some(url), + keys: std::sync::RwLock::new(HashMap::new()), + }; + store.refresh().await?; + Ok(store) + } + + /// Re-fetch the JWKS from the configured URL, replacing the cached keys. + /// Returns the number of keys loaded. No-op error for JSON-loaded stores. + pub async fn refresh(&self) -> Result { + let url = self + .url + .as_deref() + .ok_or_else(|| Error::Auth("this JWKS store was loaded from static JSON".into()))?; + let body = reqwest::get(url) + .await + .map_err(|e| Error::Auth(format!("JWKS fetch from {url} failed: {e}")))? + .text() + .await + .map_err(|e| Error::Auth(format!("JWKS fetch from {url} failed: {e}")))?; + let keys = Self::parse(&body)?; + if keys.is_empty() { + return Err(Error::Auth(format!( + "JWKS at {url} contains no usable RSA keys" + ))); + } + let count = keys.len(); + *self.keys.write().expect("jwks lock poisoned") = keys; + Ok(count) + } + + /// Number of cached keys. + pub fn len(&self) -> usize { + self.keys.read().expect("jwks lock poisoned").len() + } + + /// `true` when no keys are cached. + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + fn parse(jwks: &str) -> Result> { + let document: serde_json::Value = serde_json::from_str(jwks) + .map_err(|e| Error::Auth(format!("invalid JWKS JSON: {e}")))?; + let mut keys = HashMap::new(); + for (index, key) in document + .get("keys") + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .enumerate() + { + let get = |field: &str| key.get(field).and_then(serde_json::Value::as_str); + if get("kty") != Some("RSA") || get("use").is_some_and(|u| u != "sig") { + continue; + } + if get("alg").is_some_and(|a| a != "RS256") { + continue; + } + let (Some(n), Some(e)) = (get("n"), get("e")) else { + continue; + }; + let n = URL_SAFE_NO_PAD + .decode(n) + .map_err(|_| Error::Auth("invalid base64url modulus in JWKS".into()))?; + let e = URL_SAFE_NO_PAD + .decode(e) + .map_err(|_| Error::Auth("invalid base64url exponent in JWKS".into()))?; + let kid = get("kid") + .map(String::from) + .unwrap_or_else(|| index.to_string()); + keys.insert(kid, RsaJwk { n, e }); + } + Ok(keys) + } + + /// Look up by `kid`; a token without `kid` matches only a single-key set. + fn get(&self, kid: Option<&str>) -> Option { + let keys = self.keys.read().expect("jwks lock poisoned"); + match kid { + Some(kid) => keys.get(kid).cloned(), + None if keys.len() == 1 => keys.values().next().cloned(), + None => None, + } + } +} + +impl AuthScheme { + /// Accept any of the given API keys. + pub fn api_keys(keys: I) -> Self + where + I: IntoIterator, + S: Into, + { + Self::ApiKeys(keys.into_iter().map(Into::into).collect()) + } + + /// Accept HS256 JWTs signed with `secret` (tokens must carry `exp`). + pub fn jwt_hs256(secret: impl Into) -> Self { + Self::JwtHs256 { + secret: secret.into(), + issuer: None, + audience: None, + } + } + + /// Accept RS256 JWTs verified against a [`JwksStore`] + /// (tokens must carry `exp`; `kid` selects the key). + pub fn jwt_rs256(keys: JwksStore) -> Self { + Self::JwtRs256 { + keys: Arc::new(keys), + issuer: None, + audience: None, + } + } + + /// Additionally require the `iss` claim to equal `issuer`. + #[must_use] + pub fn with_issuer(self, issuer: impl Into) -> Self { + let issuer = Some(issuer.into()); + match self { + Self::JwtHs256 { + secret, audience, .. + } => Self::JwtHs256 { + secret, + issuer, + audience, + }, + Self::JwtRs256 { keys, audience, .. } => Self::JwtRs256 { + keys, + issuer, + audience, + }, + other => other, + } + } + + /// Additionally require the `aud` claim to equal `audience`. + #[must_use] + pub fn with_audience(self, audience: impl Into) -> Self { + let audience = Some(audience.into()); + match self { + Self::JwtHs256 { secret, issuer, .. } => Self::JwtHs256 { + secret, + issuer, + audience, + }, + Self::JwtRs256 { keys, issuer, .. } => Self::JwtRs256 { + keys, + issuer, + audience, + }, + other => other, + } + } + + /// Authorize a request given its raw `Authorization` header value and/or + /// `X-Api-Key` header value. + pub fn authorize(&self, authorization: Option<&str>, api_key: Option<&str>) -> Result<()> { + let bearer = authorization + .and_then(|v| { + v.strip_prefix("Bearer ") + .or_else(|| v.strip_prefix("bearer ")) + }) + .map(str::trim); + + match self { + Self::ApiKeys(keys) => match bearer.or(api_key) { + Some(candidate) if keys.contains(candidate) => Ok(()), + Some(_) => Err(Error::Auth("invalid API key".into())), + None => Err(Error::Auth( + "missing credentials: send 'Authorization: Bearer ' or 'X-Api-Key'".into(), + )), + }, + Self::JwtHs256 { + secret, + issuer, + audience, + } => { + let token = bearer.ok_or_else(|| { + Error::Auth("missing credentials: send 'Authorization: Bearer '".into()) + })?; + verify_jwt_hs256(token, secret, issuer.as_deref(), audience.as_deref()) + } + Self::JwtRs256 { + keys, + issuer, + audience, + } => { + let token = bearer.ok_or_else(|| { + Error::Auth("missing credentials: send 'Authorization: Bearer '".into()) + })?; + verify_jwt_rs256(token, keys, issuer.as_deref(), audience.as_deref()) + } + } + } +} + +/// Split a compact JWT and parse its header; returns +/// `(header_json, header_b64, payload_b64, signature_bytes)`. +fn split_jwt(token: &str) -> Result<(serde_json::Value, &str, &str, Vec)> { + let mut parts = token.split('.'); + let (Some(header), Some(payload), Some(signature), None) = + (parts.next(), parts.next(), parts.next(), parts.next()) + else { + return Err(Error::Auth("invalid JWT: expected three segments".into())); + }; + let header_json: serde_json::Value = serde_json::from_slice(&b64url(header, "header")?) + .map_err(|_| Error::Auth("invalid JWT: malformed header".into()))?; + let signature = b64url(signature, "signature")?; + Ok((header_json, header, payload, signature)) +} + +fn unix_now() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_secs()) + .unwrap_or(0) +} + +fn b64url(part: &str, what: &str) -> Result> { + URL_SAFE_NO_PAD + .decode(part) + .map_err(|_| Error::Auth(format!("invalid JWT: {what} is not base64url"))) +} + +/// Validate an HS256 JWT: signature, `exp` (required), `nbf`, and optional +/// `iss`/`aud` claims. +fn verify_jwt_hs256( + token: &str, + secret: &str, + issuer: Option<&str>, + audience: Option<&str>, +) -> Result<()> { + let (header_json, header, payload, signature) = split_jwt(token)?; + if header_json.get("alg").and_then(|v| v.as_str()) != Some("HS256") { + return Err(Error::Auth("invalid JWT: only HS256 is accepted".into())); + } + + // Constant-time signature check over "
.". + let mut mac = Hmac::::new_from_slice(secret.as_bytes()) + .map_err(|_| Error::Auth("invalid JWT secret".into()))?; + mac.update(header.as_bytes()); + mac.update(b"."); + mac.update(payload.as_bytes()); + mac.verify_slice(&signature) + .map_err(|_| Error::Auth("invalid JWT: signature mismatch".into()))?; + + validate_claims(payload, issuer, audience) +} + +/// Validate an RS256 JWT against a JWKS: signature (RSA PKCS#1 v1.5, +/// SHA-256), `exp` (required), `nbf`, and optional `iss`/`aud` claims. +fn verify_jwt_rs256( + token: &str, + keys: &JwksStore, + issuer: Option<&str>, + audience: Option<&str>, +) -> Result<()> { + let (header_json, header, payload, signature) = split_jwt(token)?; + if header_json.get("alg").and_then(|v| v.as_str()) != Some("RS256") { + return Err(Error::Auth("invalid JWT: only RS256 is accepted".into())); + } + + let kid = header_json.get("kid").and_then(|v| v.as_str()); + let key = keys.get(kid).ok_or_else(|| { + Error::Auth(match kid { + Some(kid) => format!("invalid JWT: no JWKS key with kid '{kid}' (try refreshing)"), + None => "invalid JWT: token has no 'kid' and the JWKS is ambiguous".into(), + }) + })?; + + let message = format!("{header}.{payload}"); + ring::signature::RsaPublicKeyComponents { + n: &key.n, + e: &key.e, + } + .verify( + &ring::signature::RSA_PKCS1_2048_8192_SHA256, + message.as_bytes(), + &signature, + ) + .map_err(|_| Error::Auth("invalid JWT: signature mismatch".into()))?; + + validate_claims(payload, issuer, audience) +} + +/// Shared registered-claims validation (`exp`, `nbf`, `iss`, `aud`). +fn validate_claims(payload: &str, issuer: Option<&str>, audience: Option<&str>) -> Result<()> { + let claims: serde_json::Value = serde_json::from_slice(&b64url(payload, "payload")?) + .map_err(|_| Error::Auth("invalid JWT: malformed claims".into()))?; + + let now = unix_now(); + let exp = claims + .get("exp") + .and_then(|v| v.as_u64()) + .ok_or_else(|| Error::Auth("invalid JWT: missing 'exp' claim".into()))?; + if exp + LEEWAY_SECS < now { + return Err(Error::Auth("invalid JWT: token is expired".into())); + } + if let Some(nbf) = claims.get("nbf").and_then(|v| v.as_u64()) { + if nbf > now + LEEWAY_SECS { + return Err(Error::Auth("invalid JWT: token is not yet valid".into())); + } + } + if let Some(expected) = issuer { + if claims.get("iss").and_then(|v| v.as_str()) != Some(expected) { + return Err(Error::Auth("invalid JWT: wrong issuer".into())); + } + } + if let Some(expected) = audience { + let aud = claims.get("aud"); + let matches = match aud { + Some(serde_json::Value::String(s)) => s == expected, + Some(serde_json::Value::Array(list)) => { + list.iter().any(|v| v.as_str() == Some(expected)) + } + _ => false, + }; + if !matches { + return Err(Error::Auth("invalid JWT: wrong audience".into())); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn make_jwt(secret: &str, claims: serde_json::Value) -> String { + let header = URL_SAFE_NO_PAD.encode(br#"{"alg":"HS256","typ":"JWT"}"#); + let payload = URL_SAFE_NO_PAD.encode(claims.to_string().as_bytes()); + let mut mac = Hmac::::new_from_slice(secret.as_bytes()).unwrap(); + mac.update(format!("{header}.{payload}").as_bytes()); + let signature = URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes()); + format!("{header}.{payload}.{signature}") + } + + fn future_exp() -> u64 { + unix_now() + 3600 + } + + #[test] + fn api_keys_accept_bearer_and_header() { + let auth = AuthScheme::api_keys(["k1", "k2"]); + auth.authorize(Some("Bearer k1"), None).unwrap(); + auth.authorize(None, Some("k2")).unwrap(); + assert!(auth.authorize(Some("Bearer nope"), None).is_err()); + assert!(auth.authorize(None, None).is_err()); + } + + #[test] + fn jwt_validates_signature_and_claims() { + let auth = AuthScheme::jwt_hs256("s3cret").with_issuer("corrosive"); + let good = make_jwt("s3cret", json!({ "exp": future_exp(), "iss": "corrosive" })); + auth.authorize(Some(&format!("Bearer {good}")), None) + .unwrap(); + + let wrong_secret = make_jwt("other", json!({ "exp": future_exp(), "iss": "corrosive" })); + assert!(auth + .authorize(Some(&format!("Bearer {wrong_secret}")), None) + .is_err()); + + let wrong_issuer = make_jwt("s3cret", json!({ "exp": future_exp(), "iss": "evil" })); + assert!(auth + .authorize(Some(&format!("Bearer {wrong_issuer}")), None) + .is_err()); + + let expired = make_jwt("s3cret", json!({ "exp": 1000, "iss": "corrosive" })); + assert!(auth + .authorize(Some(&format!("Bearer {expired}")), None) + .is_err()); + + let missing_exp = make_jwt("s3cret", json!({ "iss": "corrosive" })); + assert!(auth + .authorize(Some(&format!("Bearer {missing_exp}")), None) + .is_err()); + } + + #[test] + fn jwt_audience_string_or_array() { + let auth = AuthScheme::jwt_hs256("s").with_audience("agents"); + let single = make_jwt("s", json!({ "exp": future_exp(), "aud": "agents" })); + auth.authorize(Some(&format!("Bearer {single}")), None) + .unwrap(); + let list = make_jwt( + "s", + json!({ "exp": future_exp(), "aud": ["other", "agents"] }), + ); + auth.authorize(Some(&format!("Bearer {list}")), None) + .unwrap(); + let wrong = make_jwt("s", json!({ "exp": future_exp(), "aud": "nope" })); + assert!(auth + .authorize(Some(&format!("Bearer {wrong}")), None) + .is_err()); + } + + #[test] + fn non_hs256_rejected() { + let auth = AuthScheme::jwt_hs256("s"); + // alg=none downgrade attempt + let header = URL_SAFE_NO_PAD.encode(br#"{"alg":"none"}"#); + let payload = URL_SAFE_NO_PAD.encode(json!({ "exp": future_exp() }).to_string()); + let forged = format!("{header}.{payload}."); + assert!(auth + .authorize(Some(&format!("Bearer {forged}")), None) + .is_err()); + } + + // ── RS256 / JWKS ───────────────────────────────────────────────────── + + /// A fixed RSA-2048 test key (PKCS#8 DER, base64). Test-only material. + const TEST_RSA_PKCS8_B64: &str = "MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQCmKPH2Su6Z49De2egkDfwTXMLAhRQebihlfgI8C9qO0Ij5FCw2v9vb+0KKgfR4sbe1dS02IJ8dBSOlgMka8DIKoz8d9KdKe+EI4traqf0xpD/YxPh7JxeyVqZyMMrrt+XW4q3wLPoApGnsdZJwJVPTN+TkPTVA9x4D4EV3cHgUS+Volb7XSeKD5035pNjkl1KMP0zhhwDyTp+2OJKY6PnnFZ+dlbbVgPbknSnOMlRlSKsM3PmbMX3ek69ylcSfchgxvS7ACCw23OyYv5IesV2ElFa9cqXA3Thlv6auoT8LETMcbOCDrzaVE4ezp4nJozQaWg17yArwuUJQrHhUtmpVAgMBAAECggEAD7/pz3Ki0yto7Payrlg1AJDWVPFISuoeIiCjjZWCDe9uTE3BMx9Uc7GJSR+wUJBn3WdR9cN50YJfMpzWcxs5YxC+NtSt1r2PQwxdukRKn531/1IMS4AVGu5jsIc3dMhlnMy3uABLYiwzbhpm2wJuw6KUq52xoorJ6YwkiYG7oBC6gJ1LM2TSIL+iCGN7QuK57Wi/wcZ7646mDy2NBrHG7svM5axQkeUkyyTEp1HwbDECPbQGpd1pXCLqVJq8yy3MEDDbZ71JTfwclxbsSyUiioWGbsnEvzFlfbCKjRtT+KYrvNykp2pqArgCw9wlPB0GGWZyABR8gX0gJOSdbwszgQKBgQDPDqGsoIyxVKk+3R0E3AuPJrSHjKSrtnUBIPX6Y78eonByIyEwZ8uT5mEhqWGrbtcBSoROvXWMnIHXopT1BMc6uIdxWWqnvjtRlSBJgcQ25d+7GvOkshbDT9Z9muvoGCg8rVwhSfZJ4AxoSKgvnt9xeFL2+xERGR0EJBPi9FYWjQKBgQDNb45fAIFKq6AwvQZG1LQduYBWdByeAotAtmjUXW8wI2Y4dki/3g9x0SsPuRD1hslsfa22KY8jJDjPGH1fwxyjM8X+Uvh3pMAo+eqkWEsKnGUdWklzFsOgeg0WDV3tx8xDnrBqqLMYrF9qwZWkqkJuyw0MA388MTefLt9eqbR06QKBgQDMnJ3h1GoUFyCEscaqdbSqisodpTtZQJ3RNrw86nMEF+vcrqBukDOZ/TCBLjwJSCgJ65Rhp1HRWRvqdoySsF0cxxt5RK5kA1XlIePdH/JBedoksNaSKzbZXT0NtJlpKu4gQARqFQfgKxq3tw0Uuf/+xrPdw28zIUkOPYS1Y1TrRQKBgENxKxEfLlLgMw+tDoF0VMkpW+uF5NsuxJ5zA8kr/1OTW3yPwGRUt0dLPtLDk8C3Bis6uyuBSz9jJc8/H/GvMRiW55oNjQpiKL+LBC/92GzcWQmg2VoSEBj/2Inzy3FDVVihoRLy3RDtjcmTUdgkGPkcaeUWxM9y7OqyTZxbJCX5AoGAVIBP27/zLsnZKFbkV/FLJMkAJUYCGYJyx6KrWExRiSTFJhQjeLKzzD0XdGHL1uX0AXdi+Qx3aEYJLI0lyru4Hnsokb4C+d3PepsVi73Wlt/PQZMQ8PLvVUlg9x4DPZR+bGCdCDh9BXXoMntcEfLofJO1FooUumZgD/U+CXfC4B0="; + const TEST_RSA_N_B64URL: &str = "pijx9krumePQ3tnoJA38E1zCwIUUHm4oZX4CPAvajtCI-RQsNr_b2_tCioH0eLG3tXUtNiCfHQUjpYDJGvAyCqM_HfSnSnvhCOLa2qn9MaQ_2MT4eycXslamcjDK67fl1uKt8Cz6AKRp7HWScCVT0zfk5D01QPceA-BFd3B4FEvlaJW-10nig-dN-aTY5JdSjD9M4YcA8k6ftjiSmOj55xWfnZW21YD25J0pzjJUZUirDNz5mzF93pOvcpXEn3IYMb0uwAgsNtzsmL-SHrFdhJRWvXKlwN04Zb-mrqE_CxEzHGzgg682lROHs6eJyaM0GloNe8gK8LlCUKx4VLZqVQ"; + + fn test_jwks() -> String { + json!({ + "keys": [{ + "kty": "RSA", + "kid": "test-key", + "use": "sig", + "alg": "RS256", + "n": TEST_RSA_N_B64URL, + "e": "AQAB", + }] + }) + .to_string() + } + + fn make_rs256_jwt(claims: serde_json::Value, kid: Option<&str>) -> String { + use base64::engine::general_purpose::STANDARD; + let der = STANDARD.decode(TEST_RSA_PKCS8_B64).unwrap(); + let key_pair = ring::signature::RsaKeyPair::from_pkcs8(&der).unwrap(); + + let header = match kid { + Some(kid) => json!({ "alg": "RS256", "typ": "JWT", "kid": kid }), + None => json!({ "alg": "RS256", "typ": "JWT" }), + }; + let header = URL_SAFE_NO_PAD.encode(header.to_string()); + let payload = URL_SAFE_NO_PAD.encode(claims.to_string()); + let message = format!("{header}.{payload}"); + + let mut signature = vec![0; key_pair.public().modulus_len()]; + key_pair + .sign( + &ring::signature::RSA_PKCS1_SHA256, + &ring::rand::SystemRandom::new(), + message.as_bytes(), + &mut signature, + ) + .unwrap(); + format!("{message}.{}", URL_SAFE_NO_PAD.encode(signature)) + } + + #[test] + fn rs256_jwks_roundtrip() { + let auth = + AuthScheme::jwt_rs256(JwksStore::from_json(&test_jwks()).unwrap()).with_issuer("idp"); + + let good = make_rs256_jwt( + json!({ "exp": future_exp(), "iss": "idp" }), + Some("test-key"), + ); + auth.authorize(Some(&format!("Bearer {good}")), None) + .unwrap(); + + // Single-key sets also match tokens without a kid. + let no_kid = make_rs256_jwt(json!({ "exp": future_exp(), "iss": "idp" }), None); + auth.authorize(Some(&format!("Bearer {no_kid}")), None) + .unwrap(); + } + + #[test] + fn rs256_rejects_tampering_and_wrong_kid() { + let auth = AuthScheme::jwt_rs256(JwksStore::from_json(&test_jwks()).unwrap()); + + let token = make_rs256_jwt( + json!({ "exp": future_exp(), "role": "user" }), + Some("test-key"), + ); + // Tamper with the payload (escalate role) keeping the signature. + let mut parts: Vec<&str> = token.split('.').collect(); + let forged_payload = + URL_SAFE_NO_PAD.encode(json!({ "exp": future_exp(), "role": "admin" }).to_string()); + parts[1] = &forged_payload; + let forged = parts.join("."); + assert!(auth + .authorize(Some(&format!("Bearer {forged}")), None) + .is_err()); + + let unknown_kid = make_rs256_jwt(json!({ "exp": future_exp() }), Some("other-key")); + let err = auth + .authorize(Some(&format!("Bearer {unknown_kid}")), None) + .unwrap_err(); + assert!(err.to_string().contains("kid"), "got: {err}"); + + let expired = make_rs256_jwt(json!({ "exp": 1000 }), Some("test-key")); + assert!(auth + .authorize(Some(&format!("Bearer {expired}")), None) + .is_err()); + } + + #[test] + fn hs256_token_rejected_by_rs256_scheme() { + // Algorithm-confusion guard: an HS256 token signed with the public + // modulus as the HMAC secret must not pass an RS256 scheme. + let auth = AuthScheme::jwt_rs256(JwksStore::from_json(&test_jwks()).unwrap()); + let confused = make_jwt(TEST_RSA_N_B64URL, json!({ "exp": future_exp() })); + assert!(auth + .authorize(Some(&format!("Bearer {confused}")), None) + .is_err()); + } + + #[test] + fn jwks_parsing_skips_unusable_keys() { + let jwks = json!({ + "keys": [ + { "kty": "EC", "kid": "ec-key", "crv": "P-256" }, + { "kty": "RSA", "kid": "enc-key", "use": "enc", + "n": TEST_RSA_N_B64URL, "e": "AQAB" }, + { "kty": "RSA", "kid": "good", "use": "sig", + "n": TEST_RSA_N_B64URL, "e": "AQAB" }, + ] + }) + .to_string(); + let store = JwksStore::from_json(&jwks).unwrap(); + assert_eq!(store.len(), 1); + assert!(store.get(Some("good")).is_some()); + assert!(store.get(Some("ec-key")).is_none()); + } +} diff --git a/src/error.rs b/src/error.rs index 42a044d..e83f1b0 100644 --- a/src/error.rs +++ b/src/error.rs @@ -48,6 +48,14 @@ pub enum Error { #[error("skill execution error: {0}")] Skill(String), + /// A skill (or other action) was refused by the active policy. + #[error("permission denied: {0}")] + PermissionDenied(String), + + /// Authentication or authorization failed. + #[error("authentication error: {0}")] + Auth(String), + /// Transport-layer serving failure (REST/WS/gRPC). #[error("server error: {0}")] Server(String), diff --git a/src/grpc/mod.rs b/src/grpc/mod.rs index 9a46847..6476d87 100644 --- a/src/grpc/mod.rs +++ b/src/grpc/mod.rs @@ -37,6 +37,8 @@ fn to_status(error: Error) -> Status { Error::SkillNotFound(_) => Status::not_found(error.to_string()), Error::Config(_) | Error::Json(_) => Status::invalid_argument(error.to_string()), Error::NotConfigured(_) => Status::unimplemented(error.to_string()), + Error::Auth(_) => Status::unauthenticated(error.to_string()), + Error::PermissionDenied(_) => Status::permission_denied(error.to_string()), Error::Verification(_) | Error::Identity(_) => { Status::failed_precondition(error.to_string()) } @@ -178,11 +180,91 @@ pub async fn serve(agent: Arc, addr: SocketAddr) -> Result<()> { .map_err(|e| Error::Server(format!("gRPC server failed: {e}"))) } +/// A tonic interceptor enforcing an [`AuthScheme`](crate::auth::AuthScheme) +/// from the `authorization` / `x-api-key` request metadata. +// tonic's Interceptor contract fixes the Result<_, Status> signature. +#[allow(clippy::result_large_err)] +fn auth_interceptor( + auth: Arc, +) -> impl FnMut(Request<()>) -> std::result::Result, Status> + Clone { + move |request: Request<()>| { + let metadata = request.metadata(); + let authorization = metadata + .get("authorization") + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + let api_key = metadata + .get("x-api-key") + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + auth.authorize(authorization.as_deref(), api_key.as_deref()) + .map_err(|e| Status::unauthenticated(e.to_string()))?; + Ok(request) + } +} + +/// [`serve`], with every RPC protected by `auth` (clients send an +/// `authorization: Bearer …` or `x-api-key` metadata entry). +pub async fn serve_with_auth( + agent: Arc, + addr: SocketAddr, + auth: crate::auth::AuthScheme, +) -> Result<()> { + tracing::info!("gRPC API (authenticated) listening on {addr}"); + let service = pb::agent_service_server::AgentServiceServer::with_interceptor( + AgentGrpcService::new(agent), + auth_interceptor(Arc::new(auth)), + ); + tonic::transport::Server::builder() + .add_service(service) + .serve(addr) + .await + .map_err(|e| Error::Server(format!("gRPC server failed: {e}"))) +} + +/// [`serve`], shutting down gracefully when `signal` resolves. Pair with +/// [`server::shutdown_signal`](crate::server::shutdown_signal) (feature +/// `server`) or your own future. +pub async fn serve_with_shutdown( + agent: Arc, + addr: SocketAddr, + signal: impl std::future::Future + Send + 'static, +) -> Result<()> { + tracing::info!("gRPC API listening on {addr} (graceful shutdown armed)"); + tonic::transport::Server::builder() + .add_service(AgentGrpcService::new(agent).into_server()) + .serve_with_shutdown(addr, signal) + .await + .map_err(|e| Error::Server(format!("gRPC server failed: {e}"))) +} + +/// Serve gRPC over TLS (features `grpc` + `tls`). +#[cfg(feature = "tls")] +#[cfg_attr(docsrs, doc(cfg(feature = "tls")))] +pub async fn serve_tls( + agent: Arc, + addr: SocketAddr, + tls: &crate::tls::TlsConfig, +) -> Result<()> { + let (cert, key) = tls.pem_pair()?; + let identity = tonic::transport::Identity::from_pem(cert, key); + let tls_config = tonic::transport::ServerTlsConfig::new().identity(identity); + tracing::info!("gRPC API listening on {addr} (TLS)"); + tonic::transport::Server::builder() + .tls_config(tls_config) + .map_err(|e| Error::Server(format!("invalid TLS material: {e}")))? + .add_service(AgentGrpcService::new(agent).into_server()) + .serve(addr) + .await + .map_err(|e| Error::Server(format!("gRPC server failed: {e}"))) +} + impl Agent { /// Serve this agent's gRPC API on `addr`. /// /// Convenience for [`grpc::serve`](serve); requires the agent to be - /// wrapped in an [`Arc`]. + /// wrapped in an [`Arc`]. See also [`serve_with_auth`], + /// [`serve_with_shutdown`], and [`serve_tls`]. pub async fn serve_grpc(self: Arc, addr: SocketAddr) -> Result<()> { serve(self, addr).await } diff --git a/src/grpc/pb.rs b/src/grpc/pb.rs index b4fba27..30dfcd1 100644 --- a/src/grpc/pb.rs +++ b/src/grpc/pb.rs @@ -1,7 +1,7 @@ // This file is @generated by prost-build. -#[derive(Clone, Copy, PartialEq, ::prost::Message)] +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct GetInfoRequest {} -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct CapabilityInfo { #[prost(string, tag = "1")] pub name: ::prost::alloc::string::String, @@ -28,7 +28,7 @@ pub struct AgentInfo { #[prost(bool, tag = "7")] pub signed: bool, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ChatRequest { /// Session to continue; server generates one when empty. #[prost(string, tag = "1")] @@ -36,14 +36,14 @@ pub struct ChatRequest { #[prost(string, tag = "2")] pub message: ::prost::alloc::string::String, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ChatReply { #[prost(string, tag = "1")] pub session_id: ::prost::alloc::string::String, #[prost(string, tag = "2")] pub reply: ::prost::alloc::string::String, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct ChatChunk { #[prost(string, tag = "1")] pub session_id: ::prost::alloc::string::String, @@ -52,7 +52,7 @@ pub struct ChatChunk { #[prost(bool, tag = "3")] pub done: bool, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct SkillRequest { #[prost(string, tag = "1")] pub name: ::prost::alloc::string::String, @@ -60,7 +60,7 @@ pub struct SkillRequest { #[prost(string, tag = "2")] pub input_json: ::prost::alloc::string::String, } -#[derive(Clone, PartialEq, ::prost::Message)] +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct SkillReply { /// JSON-encoded output. #[prost(string, tag = "1")] @@ -95,7 +95,7 @@ pub mod agent_service_client { } impl AgentServiceClient where - T: tonic::client::GrpcService, + T: tonic::client::GrpcService, T::Error: Into, T::ResponseBody: Body + std::marker::Send + 'static, ::Error: Into + std::marker::Send, @@ -116,12 +116,12 @@ pub mod agent_service_client { F: tonic::service::Interceptor, T::ResponseBody: Default, T: tonic::codegen::Service< - http::Request, + http::Request, Response = http::Response< - >::ResponseBody, + >::ResponseBody, >, >, - >>::Error: + >>::Error: Into + std::marker::Send + std::marker::Sync, { AgentServiceClient::new(InterceptedService::new(inner, interceptor)) @@ -165,7 +165,7 @@ pub mod agent_service_client { self.inner.ready().await.map_err(|e| { tonic::Status::unknown(format!("Service was not ready: {}", e.into())) })?; - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static("/corrosive.agent.v1.AgentService/GetInfo"); let mut req = request.into_request(); @@ -183,7 +183,7 @@ pub mod agent_service_client { self.inner.ready().await.map_err(|e| { tonic::Status::unknown(format!("Service was not ready: {}", e.into())) })?; - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static("/corrosive.agent.v1.AgentService/Chat"); let mut req = request.into_request(); @@ -202,7 +202,7 @@ pub mod agent_service_client { self.inner.ready().await.map_err(|e| { tonic::Status::unknown(format!("Service was not ready: {}", e.into())) })?; - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static("/corrosive.agent.v1.AgentService/ChatStream"); let mut req = request.into_request(); @@ -220,7 +220,7 @@ pub mod agent_service_client { self.inner.ready().await.map_err(|e| { tonic::Status::unknown(format!("Service was not ready: {}", e.into())) })?; - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let path = http::uri::PathAndQuery::from_static( "/corrosive.agent.v1.AgentService/ExecuteSkill", ); @@ -335,7 +335,7 @@ pub mod agent_service_server { B: Body + std::marker::Send + 'static, B::Error: Into + std::marker::Send + 'static, { - type Response = http::Response; + type Response = http::Response; type Error = std::convert::Infallible; type Future = BoxFuture; fn poll_ready( @@ -369,7 +369,7 @@ pub mod agent_service_server { let inner = self.inner.clone(); let fut = async move { let method = GetInfoSvc(inner); - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let mut grpc = tonic::server::Grpc::new(codec) .apply_compression_config( accept_compression_encodings, @@ -407,7 +407,7 @@ pub mod agent_service_server { let inner = self.inner.clone(); let fut = async move { let method = ChatSvc(inner); - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let mut grpc = tonic::server::Grpc::new(codec) .apply_compression_config( accept_compression_encodings, @@ -450,7 +450,7 @@ pub mod agent_service_server { let inner = self.inner.clone(); let fut = async move { let method = ChatStreamSvc(inner); - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let mut grpc = tonic::server::Grpc::new(codec) .apply_compression_config( accept_compression_encodings, @@ -489,7 +489,7 @@ pub mod agent_service_server { let inner = self.inner.clone(); let fut = async move { let method = ExecuteSkillSvc(inner); - let codec = tonic::codec::ProstCodec::default(); + let codec = tonic_prost::ProstCodec::default(); let mut grpc = tonic::server::Grpc::new(codec) .apply_compression_config( accept_compression_encodings, @@ -505,7 +505,7 @@ pub mod agent_service_server { Box::pin(fut) } _ => Box::pin(async move { - let mut response = http::Response::new(empty_body()); + let mut response = http::Response::new(tonic::body::Body::default()); let headers = response.headers_mut(); headers.insert( tonic::Status::GRPC_STATUS, diff --git a/src/identity.rs b/src/identity.rs index dcf7f89..af33bb2 100644 --- a/src/identity.rs +++ b/src/identity.rs @@ -73,6 +73,18 @@ impl AgentIdentity { pub fn sign(&self, message: &[u8]) -> String { B64.encode(self.signing_key.sign(message).to_bytes()) } + + /// The private key as PKCS#8 DER — used by the X.509 module to build + /// certificates. **Handle like the secret key itself.** + #[cfg(feature = "x509")] + pub(crate) fn pkcs8_der(&self) -> Result> { + use ed25519_dalek::pkcs8::EncodePrivateKey; + let document = self + .signing_key + .to_pkcs8_der() + .map_err(|e| Error::Identity(format!("PKCS#8 encoding failed: {e}")))?; + Ok(document.as_bytes().to_vec()) + } } /// Multicodec prefix identifying an Ed25519 public key inside a `did:key`. diff --git a/src/lib.rs b/src/lib.rs index befb8fb..001ec27 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -67,6 +67,10 @@ pub mod skills; pub mod trust; pub mod vector; +#[cfg(any(feature = "server", feature = "grpc"))] +#[cfg_attr(docsrs, doc(cfg(any(feature = "server", feature = "grpc"))))] +pub mod auth; + #[cfg(feature = "server")] #[cfg_attr(docsrs, doc(cfg(feature = "server")))] pub mod server; @@ -75,6 +79,21 @@ pub mod server; #[cfg_attr(docsrs, doc(cfg(feature = "grpc")))] pub mod grpc; +#[cfg(feature = "tls")] +#[cfg_attr(docsrs, doc(cfg(feature = "tls")))] +pub mod tls; + +#[cfg(feature = "x509")] +#[cfg_attr(docsrs, doc(cfg(feature = "x509")))] +pub mod x509; + +/// The step-by-step tutorial: from an empty project to a production-shaped +/// agent. Rendered from `docs/TUTORIAL.md`; its core code blocks are +/// compile-checked as doctests. +pub mod tutorial { + #![doc = include_str!("../docs/TUTORIAL.md")] +} + pub use agent::{Agent, AgentBuilder, AgentInfo, AgentManifest, Capability}; pub use error::{Error, Result}; @@ -90,13 +109,23 @@ pub mod prelude { pub use crate::identity::AgentIdentity; pub use crate::llm::{ models, ChatMessage, ChatRequest, ChatResponse, EmbeddingProvider, LlmProvider, - NvidiaClient, StreamChunk, + NvidiaClient, RetryPolicy, StreamChunk, ToolCall, ToolSpec, UsageEvent, UsageObserver, + UsageSnapshot, }; - pub use crate::mcp::{McpClient, McpServerConfig, McpTool}; + pub use crate::mcp::{McpClient, McpPrompt, McpResource, McpServerConfig, McpTool}; pub use crate::session::{InMemorySessionStore, SessionStore}; - pub use crate::skills::{FnSkill, Skill, SkillRegistry}; + pub use crate::skills::{FnSkill, Skill, SkillPolicy, SkillRegistry}; pub use crate::trust::{Revocation, RotationProof, TrustStore}; - pub use crate::vector::{Document, InMemoryVectorStore, SearchResult, VectorStore}; + pub use crate::vector::{ + chunk_text, Document, InMemoryVectorStore, MetadataFilter, SearchResult, VectorStore, + }; + + #[cfg(any(feature = "server", feature = "grpc"))] + pub use crate::auth::{AuthScheme, JwksStore}; + #[cfg(feature = "tls")] + pub use crate::tls::TlsConfig; + #[cfg(feature = "pgvector")] + pub use crate::vector::PgVectorStore; #[cfg(feature = "redis-sessions")] pub use crate::session::RedisSessionStore; diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 022fdac..a019332 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -2,9 +2,13 @@ mod nvidia; mod types; +mod usage; -pub use nvidia::{models, NvidiaClient}; -pub use types::{ChatMessage, ChatRequest, ChatResponse, Role, StreamChunk, Usage}; +pub use nvidia::{models, NvidiaClient, RetryPolicy}; +pub use types::{ + ChatMessage, ChatRequest, ChatResponse, Role, StreamChunk, ToolCall, ToolSpec, Usage, +}; +pub use usage::{UsageEvent, UsageObserver, UsageSnapshot, UsageTotals}; use futures_util::stream::BoxStream; diff --git a/src/llm/nvidia.rs b/src/llm/nvidia.rs index eed4e8b..2d90fd3 100644 --- a/src/llm/nvidia.rs +++ b/src/llm/nvidia.rs @@ -1,17 +1,27 @@ -//! NVIDIA NIM client — chat, streaming, and embeddings for Nemotron models. +//! NVIDIA NIM client — chat, streaming, tools, and embeddings for Nemotron +//! models. //! //! Talks to the OpenAI-compatible endpoint at //! `https://integrate.api.nvidia.com/v1`. Free API keys are available from //! (every model page has a "Get API Key" button). +//! +//! Requests are retried with exponential backoff and jitter on `429` (rate +//! limit, honoring `Retry-After`), `5xx`, and transport errors — see +//! [`RetryPolicy`]. + +use std::time::Duration; use async_stream::try_stream; use futures_util::stream::BoxStream; use futures_util::StreamExt; +use rand::Rng; use serde::Deserialize; -use serde_json::json; +use serde_json::{json, Value}; use crate::error::{Error, Result}; -use crate::llm::types::{ChatRequest, ChatResponse, StreamChunk, Usage}; +use crate::llm::types::{ + ChatMessage, ChatRequest, ChatResponse, Role, StreamChunk, ToolCall, Usage, +}; use crate::llm::{EmbeddingProvider, LlmProvider}; /// Well-known NVIDIA NIM model ids usable with a free build.nvidia.com key. @@ -44,6 +54,53 @@ pub mod models { const DEFAULT_BASE_URL: &str = "https://integrate.api.nvidia.com/v1"; const DEFAULT_TIMEOUT_SECS: u64 = 300; +/// How the client retries failed requests. +/// +/// Retries fire on HTTP 429 / 5xx and on transport errors. Delays grow +/// exponentially from `base_delay` (with up to 20% jitter) and are capped at +/// `max_delay`; a `Retry-After` response header, when present, overrides the +/// computed delay. +#[derive(Debug, Clone)] +pub struct RetryPolicy { + /// Maximum number of retries after the initial attempt. + pub max_retries: u32, + /// Delay before the first retry. + pub base_delay: Duration, + /// Upper bound for any single delay. + pub max_delay: Duration, +} + +impl Default for RetryPolicy { + fn default() -> Self { + Self { + max_retries: 3, + base_delay: Duration::from_millis(500), + max_delay: Duration::from_secs(30), + } + } +} + +impl RetryPolicy { + /// Disable retries entirely. + pub fn none() -> Self { + Self { + max_retries: 0, + ..Default::default() + } + } + + fn delay_for(&self, attempt: u32, retry_after: Option) -> Duration { + if let Some(after) = retry_after { + return after.min(self.max_delay); + } + let exponential = self.base_delay.saturating_mul(2u32.saturating_pow(attempt)); + let capped = exponential.min(self.max_delay); + // Up to 20% jitter to avoid thundering herds. + let jitter = rand::thread_rng().gen_range(0.0..=0.2); + capped.mul_f64(1.0 + jitter).min(self.max_delay) + } +} + /// Client for NVIDIA NIM inference endpoints (Nemotron chat + embeddings). #[derive(Clone)] pub struct NvidiaClient { @@ -52,6 +109,7 @@ pub struct NvidiaClient { base_url: String, model: String, embedding_model: String, + retry: RetryPolicy, } impl std::fmt::Debug for NvidiaClient { @@ -60,6 +118,7 @@ impl std::fmt::Debug for NvidiaClient { .field("base_url", &self.base_url) .field("model", &self.model) .field("embedding_model", &self.embedding_model) + .field("retry", &self.retry) .finish_non_exhaustive() } } @@ -77,6 +136,7 @@ impl NvidiaClient { base_url: DEFAULT_BASE_URL.to_string(), model: models::NEMOTRON_3_NANO_30B.to_string(), embedding_model: models::EMBED_QA_E5_V5.to_string(), + retry: RetryPolicy::default(), } } @@ -115,47 +175,157 @@ impl NvidiaClient { self } - fn chat_body(&self, request: &ChatRequest, stream: bool) -> serde_json::Value { - json!({ + /// Override the retry/backoff behavior (see [`RetryPolicy`]). + #[must_use] + pub fn with_retry_policy(mut self, retry: RetryPolicy) -> Self { + self.retry = retry; + self + } + + /// Convert a [`ChatMessage`] to the OpenAI wire format. + fn wire_message(message: &ChatMessage) -> Value { + let role = match message.role { + Role::System => "system", + Role::User => "user", + Role::Assistant => "assistant", + Role::Tool => "tool", + }; + let mut wire = json!({ "role": role, "content": message.content }); + if let Some(calls) = &message.tool_calls { + wire["tool_calls"] = Value::Array( + calls + .iter() + .map(|c| { + json!({ + "id": c.id, + "type": "function", + "function": { + "name": c.name, + "arguments": c.arguments.to_string(), + }, + }) + }) + .collect(), + ); + } + if let Some(id) = &message.tool_call_id { + wire["tool_call_id"] = Value::String(id.clone()); + } + wire + } + + fn chat_body(&self, request: &ChatRequest, stream: bool) -> Value { + let messages: Vec = request.messages.iter().map(Self::wire_message).collect(); + let mut body = json!({ "model": request.model.as_deref().unwrap_or(&self.model), - "messages": request.messages, + "messages": messages, "temperature": request.temperature.unwrap_or(0.6), "top_p": request.top_p.unwrap_or(0.95), "max_tokens": request.max_tokens.unwrap_or(2048), "stream": stream, - }) + }); + if let Some(tools) = &request.tools { + if !tools.is_empty() { + body["tools"] = Value::Array( + tools + .iter() + .map(|t| { + json!({ + "type": "function", + "function": { + "name": t.name, + "description": t.description, + "parameters": t.parameters, + }, + }) + }) + .collect(), + ); + body["tool_choice"] = Value::String("auto".into()); + } + } + body } - async fn post_json(&self, path: &str, body: serde_json::Value) -> Result { - let response = self - .http - .post(format!("{}{path}", self.base_url)) - .bearer_auth(&self.api_key) - .json(&body) - .send() - .await?; - let status = response.status(); - if !status.is_success() { - let detail = response.text().await.unwrap_or_default(); - return Err(Error::Llm(format!( - "NVIDIA API returned {status}: {detail}" - ))); + fn retry_after(response: &reqwest::Response) -> Option { + response + .headers() + .get(reqwest::header::RETRY_AFTER)? + .to_str() + .ok()? + .parse::() + .ok() + .map(Duration::from_secs) + } + + /// POST with retry on 429/5xx and transport errors. + async fn post_json(&self, path: &str, body: Value) -> Result { + let url = format!("{}{path}", self.base_url); + let mut attempt: u32 = 0; + loop { + let outcome = self + .http + .post(&url) + .bearer_auth(&self.api_key) + .json(&body) + .send() + .await; + + let (retryable, retry_after, error) = match outcome { + Ok(response) => { + let status = response.status(); + if status.is_success() { + return Ok(response); + } + let retry_after = Self::retry_after(&response); + let retryable = status.as_u16() == 429 || status.is_server_error(); + let detail = response.text().await.unwrap_or_default(); + ( + retryable, + retry_after, + Error::Llm(format!("NVIDIA API returned {status}: {detail}")), + ) + } + Err(e) => { + // Connection/timeout problems are worth retrying; anything + // else (e.g. request building) is not. + let retryable = e.is_timeout() || e.is_connect() || e.is_request(); + (retryable, None, Error::Http(e)) + } + }; + + if !retryable || attempt >= self.retry.max_retries { + return Err(error); + } + let delay = self.retry.delay_for(attempt, retry_after); + tracing::warn!( + "NVIDIA request failed (attempt {}/{}), retrying in {delay:?}: {error}", + attempt + 1, + self.retry.max_retries, + ); + tokio::time::sleep(delay).await; + attempt += 1; } - Ok(response) } async fn embed(&self, texts: &[String], input_type: &str) -> Result>> { - let body = json!({ - "model": self.embedding_model, - "input": texts, - "input_type": input_type, - "encoding_format": "float", - }); - let response = self.post_json("/embeddings", body).await?; - let parsed: EmbeddingsResponse = response.json().await?; - let mut data = parsed.data; - data.sort_by_key(|d| d.index); - Ok(data.into_iter().map(|d| d.embedding).collect()) + // The embeddings endpoint caps batch sizes; chunk transparently. + const BATCH: usize = 32; + let mut all = Vec::with_capacity(texts.len()); + for batch in texts.chunks(BATCH) { + let body = json!({ + "model": self.embedding_model, + "input": batch, + "input_type": input_type, + "encoding_format": "float", + }); + let response = self.post_json("/embeddings", body).await?; + let parsed: EmbeddingsResponse = response.json().await?; + let mut data = parsed.data; + data.sort_by_key(|d| d.index); + all.extend(data.into_iter().map(|d| d.embedding)); + } + Ok(all) } } @@ -177,6 +347,34 @@ struct ChatChoice { struct ChoiceMessage { #[serde(default)] content: Option, + #[serde(default)] + tool_calls: Vec, +} + +#[derive(Deserialize)] +struct WireToolCall { + #[serde(default)] + id: String, + function: WireFunction, +} + +#[derive(Deserialize)] +struct WireFunction { + name: String, + #[serde(default)] + arguments: String, +} + +impl WireToolCall { + fn into_tool_call(self) -> ToolCall { + let arguments = serde_json::from_str(&self.function.arguments) + .unwrap_or(Value::String(self.function.arguments)); + ToolCall { + id: self.id, + name: self.function.name, + arguments, + } + } } #[derive(Deserialize)] @@ -223,6 +421,12 @@ impl LlmProvider for NvidiaClient { Ok(ChatResponse { content: choice.message.content.unwrap_or_default(), model: completion.model, + tool_calls: choice + .message + .tool_calls + .into_iter() + .map(WireToolCall::into_tool_call) + .collect(), usage: completion.usage, }) } diff --git a/src/llm/types.rs b/src/llm/types.rs index 363bd50..19aac07 100644 --- a/src/llm/types.rs +++ b/src/llm/types.rs @@ -1,6 +1,7 @@ //! Chat request/response types shared by all LLM providers. use serde::{Deserialize, Serialize}; +use serde_json::Value; /// Who authored a chat message. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] @@ -16,6 +17,39 @@ pub enum Role { Tool, } +/// A tool invocation requested by the model. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct ToolCall { + /// Provider-assigned call id (echo it back in the tool result message). + pub id: String, + /// Name of the tool/skill to invoke. + pub name: String, + /// Parsed JSON arguments. + pub arguments: Value, +} + +/// A tool the model is allowed to call, in provider-neutral form. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ToolSpec { + /// Tool name (matches a registered skill name). + pub name: String, + /// What the tool does — shown to the model. + pub description: String, + /// JSON Schema of the tool's arguments. + pub parameters: Value, +} + +impl ToolSpec { + /// Build a spec from a registered [`Skill`](crate::skills::Skill). + pub fn from_skill(skill: &dyn crate::skills::Skill) -> Self { + Self { + name: skill.name().to_string(), + description: skill.description().to_string(), + parameters: skill.input_schema(), + } + } +} + /// One message in a conversation. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatMessage { @@ -23,30 +57,56 @@ pub struct ChatMessage { pub role: Role, /// Message text. pub content: String, + /// Tool invocations attached to an assistant message. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + /// For [`Role::Tool`] messages: the id of the call being answered. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_call_id: Option, } impl ChatMessage { - /// Build a system message. - pub fn system(content: impl Into) -> Self { + fn plain(role: Role, content: impl Into) -> Self { Self { - role: Role::System, + role, content: content.into(), + tool_calls: None, + tool_call_id: None, } } + /// Build a system message. + pub fn system(content: impl Into) -> Self { + Self::plain(Role::System, content) + } + /// Build a user message. pub fn user(content: impl Into) -> Self { - Self { - role: Role::User, - content: content.into(), - } + Self::plain(Role::User, content) } /// Build an assistant message. pub fn assistant(content: impl Into) -> Self { + Self::plain(Role::Assistant, content) + } + + /// Build an assistant message that requests tool invocations. + pub fn assistant_tool_calls(content: impl Into, calls: Vec) -> Self { Self { role: Role::Assistant, content: content.into(), + tool_calls: Some(calls), + tool_call_id: None, + } + } + + /// Build a tool-result message answering `tool_call_id`. + pub fn tool_result(tool_call_id: impl Into, content: impl Into) -> Self { + Self { + role: Role::Tool, + content: content.into(), + tool_calls: None, + tool_call_id: Some(tool_call_id.into()), } } } @@ -68,6 +128,9 @@ pub struct ChatRequest { /// Maximum number of tokens to generate. #[serde(skip_serializing_if = "Option::is_none")] pub max_tokens: Option, + /// Tools the model may call (enables function calling when non-empty). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tools: Option>, } impl ChatRequest { @@ -99,6 +162,13 @@ impl ChatRequest { self.max_tokens = Some(max_tokens); self } + + /// Offer tools to the model. + #[must_use] + pub fn with_tools(mut self, tools: Vec) -> Self { + self.tools = Some(tools); + self + } } /// Token accounting reported by the provider. @@ -118,10 +188,13 @@ pub struct Usage { /// A completed (non-streaming) chat response. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ChatResponse { - /// The assistant's reply text. + /// The assistant's reply text (may be empty when tools are called). pub content: String, /// The model that produced the reply. pub model: String, + /// Tool invocations the model requested (empty when none). + #[serde(default)] + pub tool_calls: Vec, /// Token usage, when reported. #[serde(skip_serializing_if = "Option::is_none")] pub usage: Option, diff --git a/src/llm/usage.rs b/src/llm/usage.rs new file mode 100644 index 0000000..aa9e2ea --- /dev/null +++ b/src/llm/usage.rs @@ -0,0 +1,89 @@ +//! Token-usage accounting hooks. +//! +//! Every non-streaming completion an [`Agent`](crate::agent::Agent) makes is +//! reported to an optional [`UsageObserver`] +//! ([`AgentBuilder::usage_observer`](crate::agent::AgentBuilder::usage_observer)) +//! and always accumulated into the agent's built-in totals +//! ([`Agent::usage`](crate::agent::Agent::usage)). + +use std::sync::atomic::{AtomicU64, Ordering}; + +use serde::{Deserialize, Serialize}; + +use crate::llm::Usage; + +/// One completion's worth of token accounting. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UsageEvent { + /// Session that triggered the completion. + pub session_id: String, + /// Model that served it. + pub model: String, + /// Tokens consumed. + pub usage: Usage, +} + +/// Receives a [`UsageEvent`] after every completion — implement to feed +/// metrics, billing, or rate-limit dashboards. +pub trait UsageObserver: Send + Sync { + /// Called after each successful completion that reported usage. + fn on_usage(&self, event: &UsageEvent); +} + +/// Call an `Fn(&UsageEvent)` closure as an observer. +impl UsageObserver for F { + fn on_usage(&self, event: &UsageEvent) { + self(event) + } +} + +/// Lock-free running totals across all sessions (the agent keeps one). +#[derive(Debug, Default)] +pub struct UsageTotals { + requests: AtomicU64, + prompt_tokens: AtomicU64, + completion_tokens: AtomicU64, + total_tokens: AtomicU64, +} + +/// A point-in-time copy of [`UsageTotals`]. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct UsageSnapshot { + /// Number of completions that reported usage. + pub requests: u64, + /// Cumulative prompt tokens. + pub prompt_tokens: u64, + /// Cumulative completion tokens. + pub completion_tokens: u64, + /// Cumulative total tokens. + pub total_tokens: u64, +} + +impl UsageTotals { + /// Add one event to the totals. + pub fn record(&self, event: &UsageEvent) { + self.requests.fetch_add(1, Ordering::Relaxed); + self.prompt_tokens + .fetch_add(u64::from(event.usage.prompt_tokens), Ordering::Relaxed); + self.completion_tokens + .fetch_add(u64::from(event.usage.completion_tokens), Ordering::Relaxed); + self.total_tokens + .fetch_add(u64::from(event.usage.total_tokens), Ordering::Relaxed); + } + + /// A consistent-enough copy of the counters. + pub fn snapshot(&self) -> UsageSnapshot { + UsageSnapshot { + requests: self.requests.load(Ordering::Relaxed), + prompt_tokens: self.prompt_tokens.load(Ordering::Relaxed), + completion_tokens: self.completion_tokens.load(Ordering::Relaxed), + total_tokens: self.total_tokens.load(Ordering::Relaxed), + } + } +} + +impl UsageObserver for UsageTotals { + fn on_usage(&self, event: &UsageEvent) { + self.record(event); + } +} diff --git a/src/mcp/client.rs b/src/mcp/client.rs index 78d520f..6e462d3 100644 --- a/src/mcp/client.rs +++ b/src/mcp/client.rs @@ -1,4 +1,5 @@ -//! A minimal MCP client over stdio (JSON-RPC 2.0, newline-delimited). +//! MCP client over stdio (JSON-RPC, newline-delimited) or streamable +//! HTTP/SSE. use std::process::Stdio; use std::sync::atomic::{AtomicI64, Ordering}; @@ -6,10 +7,10 @@ use std::sync::atomic::{AtomicI64, Ordering}; use serde_json::{json, Value}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; use tokio::process::{Child, ChildStdin, ChildStdout, Command}; -use tokio::sync::Mutex; +use tokio::sync::{Mutex, RwLock}; use crate::error::{Error, Result}; -use crate::mcp::{McpServerConfig, McpTool}; +use crate::mcp::{McpPrompt, McpResource, McpServerConfig, McpTool}; const PROTOCOL_VERSION: &str = "2024-11-05"; @@ -18,11 +19,23 @@ struct McpIo { stdout: BufReader, } -/// A connected MCP server (spawned child process, stdio transport). +enum Transport { + Stdio { + io: Mutex, + child: Mutex, + }, + Http { + http: reqwest::Client, + url: String, + headers: std::collections::HashMap, + session_id: RwLock>, + }, +} + +/// A connected MCP server (stdio child process or streamable-HTTP endpoint). pub struct McpClient { name: String, - io: Mutex, - child: Mutex, + transport: Transport, next_id: AtomicI64, server_info: Value, } @@ -31,43 +44,68 @@ impl std::fmt::Debug for McpClient { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("McpClient") .field("name", &self.name) + .field( + "transport", + &match &self.transport { + Transport::Stdio { .. } => "stdio", + Transport::Http { .. } => "http", + }, + ) .field("server_info", &self.server_info) .finish_non_exhaustive() } } impl McpClient { - /// Spawn the configured server process and perform the MCP `initialize` - /// handshake. + /// Connect to the configured server (spawn + handshake for stdio, POST + /// handshake for HTTP) and perform the MCP `initialize` exchange. pub async fn connect(config: &McpServerConfig) -> Result { - let mut command = Command::new(&config.command); - command - .args(&config.args) - .envs(&config.env) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::null()) - .kill_on_drop(true); - - let mut child = command - .spawn() - .map_err(|e| Error::Mcp(format!("failed to spawn '{}': {e}", config.command)))?; - let stdin = child - .stdin - .take() - .ok_or_else(|| Error::Mcp("child stdin unavailable".into()))?; - let stdout = child - .stdout - .take() - .ok_or_else(|| Error::Mcp("child stdout unavailable".into()))?; + let transport = if let Some(url) = &config.url { + Transport::Http { + http: reqwest::Client::new(), + url: url.clone(), + headers: config.headers.clone(), + session_id: RwLock::new(None), + } + } else { + if config.command.trim().is_empty() { + return Err(Error::Mcp(format!( + "MCP server '{}' has neither a command nor a url", + config.name + ))); + } + let mut command = Command::new(&config.command); + command + .args(&config.args) + .envs(&config.env) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .kill_on_drop(true); + + let mut child = command + .spawn() + .map_err(|e| Error::Mcp(format!("failed to spawn '{}': {e}", config.command)))?; + let stdin = child + .stdin + .take() + .ok_or_else(|| Error::Mcp("child stdin unavailable".into()))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| Error::Mcp("child stdout unavailable".into()))?; + Transport::Stdio { + io: Mutex::new(McpIo { + stdin, + stdout: BufReader::new(stdout), + }), + child: Mutex::new(child), + } + }; let client = Self { name: config.name.clone(), - io: Mutex::new(McpIo { - stdin, - stdout: BufReader::new(stdout), - }), - child: Mutex::new(child), + transport, next_id: AtomicI64::new(1), server_info: Value::Null, }; @@ -106,6 +144,8 @@ impl McpClient { &self.server_info } + // ── Tools ──────────────────────────────────────────────────────────── + /// List the tools this server offers. pub async fn list_tools(&self) -> Result> { let result = self.request("tools/list", json!({})).await?; @@ -137,42 +177,221 @@ impl McpClient { Ok(result.get("content").cloned().unwrap_or(result)) } - /// Terminate the server process. + // ── Resources ──────────────────────────────────────────────────────── + + /// List the resources this server offers. + pub async fn list_resources(&self) -> Result> { + let result = self.request("resources/list", json!({})).await?; + let resources = result + .get("resources") + .cloned() + .ok_or_else(|| Error::Mcp("resources/list response missing 'resources'".into()))?; + Ok(serde_json::from_value(resources)?) + } + + /// Read a resource by URI; returns the `contents` array (text or blob + /// entries). + pub async fn read_resource(&self, uri: &str) -> Result { + let result = self + .request("resources/read", json!({ "uri": uri })) + .await?; + Ok(result.get("contents").cloned().unwrap_or(result)) + } + + // ── Prompts ────────────────────────────────────────────────────────── + + /// List the prompt templates this server offers. + pub async fn list_prompts(&self) -> Result> { + let result = self.request("prompts/list", json!({})).await?; + let prompts = result + .get("prompts") + .cloned() + .ok_or_else(|| Error::Mcp("prompts/list response missing 'prompts'".into()))?; + Ok(serde_json::from_value(prompts)?) + } + + /// Expand a prompt template with arguments; returns the rendered + /// `messages` array. + pub async fn get_prompt(&self, name: &str, arguments: Value) -> Result { + let result = self + .request( + "prompts/get", + json!({ "name": name, "arguments": arguments }), + ) + .await?; + Ok(result.get("messages").cloned().unwrap_or(result)) + } + + /// Terminate the connection (kills the child process for stdio; ends the + /// HTTP session best-effort). pub async fn shutdown(&self) -> Result<()> { - let mut child = self.child.lock().await; - child - .kill() - .await - .map_err(|e| Error::Mcp(format!("failed to kill server: {e}"))) + match &self.transport { + Transport::Stdio { child, .. } => { + let mut child = child.lock().await; + child + .kill() + .await + .map_err(|e| Error::Mcp(format!("failed to kill server: {e}"))) + } + Transport::Http { + http, + url, + headers, + session_id, + } => { + if let Some(sid) = session_id.read().await.clone() { + let mut request = http.delete(url).header("Mcp-Session-Id", sid); + for (name, value) in headers { + request = request.header(name, value); + } + let _ = request.send().await; // best-effort per spec + } + Ok(()) + } + } } + // ── JSON-RPC plumbing ──────────────────────────────────────────────── + async fn notify(&self, method: &str, params: Value) -> Result<()> { let message = json!({ "jsonrpc": "2.0", "method": method, "params": params }); - let mut io = self.io.lock().await; - Self::write_message(&mut io.stdin, &message).await + match &self.transport { + Transport::Stdio { io, .. } => { + let mut io = io.lock().await; + Self::write_message(&mut io.stdin, &message).await + } + Transport::Http { .. } => { + // Notifications over HTTP get a 202 with no body. + self.http_post(&message, None).await.map(|_| ()) + } + } } async fn request(&self, method: &str, params: Value) -> Result { let id = self.next_id.fetch_add(1, Ordering::Relaxed); let message = json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params }); - let mut io = self.io.lock().await; - Self::write_message(&mut io.stdin, &message).await?; + let value = match &self.transport { + Transport::Stdio { io, .. } => { + let mut io = io.lock().await; + Self::write_message(&mut io.stdin, &message).await?; + Self::read_response_stdio(&mut io.stdout, id, method, &self.name).await? + } + Transport::Http { .. } => self + .http_post(&message, Some(id)) + .await? + .ok_or_else(|| Error::Mcp(format!("'{method}' returned no response")))?, + }; + if let Some(error) = value.get("error") { + return Err(Error::Mcp(format!("'{method}' failed: {error}"))); + } + Ok(value.get("result").cloned().unwrap_or(Value::Null)) + } + + /// POST one JSON-RPC message over the streamable-HTTP transport. Returns + /// the matching response envelope (or `None` for notifications). + async fn http_post(&self, message: &Value, expect_id: Option) -> Result> { + let Transport::Http { + http, + url, + headers, + session_id, + } = &self.transport + else { + unreachable!("http_post called on stdio transport"); + }; + + let mut request = http + .post(url) + .header("Content-Type", "application/json") + .header("Accept", "application/json, text/event-stream") + .json(message); + for (name, value) in headers { + request = request.header(name, value); + } + if let Some(sid) = session_id.read().await.clone() { + request = request.header("Mcp-Session-Id", sid); + } + + let response = request + .send() + .await + .map_err(|e| Error::Mcp(format!("HTTP request to '{}' failed: {e}", self.name)))?; + + let status = response.status(); + // The server assigns a session id on initialize; echo it afterwards. + if let Some(sid) = response + .headers() + .get("mcp-session-id") + .and_then(|v| v.to_str().ok()) + { + *session_id.write().await = Some(sid.to_string()); + } + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(Error::Mcp(format!( + "server '{}' returned {status}: {body}", + self.name + ))); + } + let Some(expect_id) = expect_id else { + return Ok(None); // notification — 202/200 with ignorable body + }; + + let content_type = response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .unwrap_or("") + .to_string(); + let body = response + .text() + .await + .map_err(|e| Error::Mcp(format!("failed to read response body: {e}")))?; + + if content_type.starts_with("text/event-stream") { + // Scan SSE events for the JSON-RPC response with our id. + for line in body.lines() { + let Some(data) = line.trim().strip_prefix("data:") else { + continue; + }; + let Ok(value) = serde_json::from_str::(data.trim()) else { + continue; + }; + if value.get("id").and_then(Value::as_i64) == Some(expect_id) { + return Ok(Some(value)); + } + } + Err(Error::Mcp(format!( + "SSE stream from '{}' ended without a response for id {expect_id}", + self.name + ))) + } else { + let value: Value = serde_json::from_str(&body) + .map_err(|e| Error::Mcp(format!("invalid JSON from '{}': {e}", self.name)))?; + Ok(Some(value)) + } + } + + async fn read_response_stdio( + stdout: &mut BufReader, + id: i64, + method: &str, + name: &str, + ) -> Result { // Read newline-delimited JSON until our response id shows up, // skipping notifications and unrelated messages. let mut line = String::new(); loop { line.clear(); - let read = io - .stdout + let read = stdout .read_line(&mut line) .await .map_err(|e| Error::Mcp(format!("read from server failed: {e}")))?; if read == 0 { return Err(Error::Mcp(format!( - "server '{}' closed the connection during '{method}'", - self.name + "server '{name}' closed the connection during '{method}'" ))); } let trimmed = line.trim(); @@ -185,10 +404,7 @@ impl McpClient { if value.get("id").and_then(Value::as_i64) != Some(id) { continue; } - if let Some(error) = value.get("error") { - return Err(Error::Mcp(format!("'{method}' failed: {error}"))); - } - return Ok(value.get("result").cloned().unwrap_or(Value::Null)); + return Ok(value); } } diff --git a/src/mcp/mod.rs b/src/mcp/mod.rs index 9411e24..81baebd 100644 --- a/src/mcp/mod.rs +++ b/src/mcp/mod.rs @@ -1,9 +1,17 @@ //! Model Context Protocol (MCP) support. //! //! Agents can declare MCP servers in their JSON manifest and connect to them -//! over stdio using JSON-RPC 2.0. The client speaks the `initialize` -//! handshake and the `tools/list` / `tools/call` methods, so any standard -//! MCP tool server (filesystem, git, fetch, …) plugs straight in. +//! over two transports: +//! +//! - **stdio** — the server is spawned as a child process and spoken to over +//! newline-delimited JSON-RPC (the classic transport). +//! - **streamable HTTP / SSE** — JSON-RPC messages are POSTed to a URL; the +//! server answers with JSON or a `text/event-stream` body, and an +//! `Mcp-Session-Id` header carries session affinity. +//! +//! The client speaks the `initialize` handshake plus **tools** +//! (`tools/list`, `tools/call`), **resources** (`resources/list`, +//! `resources/read`), and **prompts** (`prompts/list`, `prompts/get`). mod client; @@ -13,12 +21,18 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; -/// How to launch an MCP server, as declared in the agent manifest. +/// How to reach an MCP server, as declared in the agent manifest. +/// +/// Set `command` (+ `args`/`env`) for a stdio server, **or** `url` +/// (+ `headers`) for a streamable-HTTP/SSE server. When both are present, +/// `url` wins. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct McpServerConfig { /// Local name used to address the server (e.g. `"fs"`). pub name: String, - /// Executable to spawn (e.g. `"npx"` or `"uvx"`). + /// Executable to spawn for the stdio transport (e.g. `"npx"`). + #[serde(default)] pub command: String, /// Arguments passed to the executable. #[serde(default)] @@ -26,10 +40,17 @@ pub struct McpServerConfig { /// Extra environment variables for the child process. #[serde(default)] pub env: HashMap, + /// Endpoint URL for the streamable-HTTP transport + /// (e.g. `"https://mcp.example.com/mcp"`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub url: Option, + /// Extra HTTP headers (e.g. `Authorization`) for the HTTP transport. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub headers: HashMap, } impl McpServerConfig { - /// Create a config from a name, command, and arguments. + /// Create a stdio-transport config from a name, command, and arguments. pub fn new( name: impl Into, command: impl Into, @@ -40,8 +61,29 @@ impl McpServerConfig { command: command.into(), args: args.into_iter().map(Into::into).collect(), env: HashMap::new(), + url: None, + headers: HashMap::new(), } } + + /// Create a streamable-HTTP transport config from a name and URL. + pub fn http(name: impl Into, url: impl Into) -> Self { + Self { + name: name.into(), + command: String::new(), + args: Vec::new(), + env: HashMap::new(), + url: Some(url.into()), + headers: HashMap::new(), + } + } + + /// Add an HTTP header (HTTP transport only). + #[must_use] + pub fn with_header(mut self, name: impl Into, value: impl Into) -> Self { + self.headers.insert(name.into(), value.into()); + self + } } /// A tool advertised by an MCP server. @@ -56,3 +98,32 @@ pub struct McpTool { #[serde(default, rename = "inputSchema")] pub input_schema: serde_json::Value, } + +/// A resource advertised by an MCP server. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpResource { + /// Resource URI (pass to [`McpClient::read_resource`]). + pub uri: String, + /// Human-readable name. + #[serde(default)] + pub name: String, + /// Description, when provided. + #[serde(default)] + pub description: String, + /// MIME type, when provided. + #[serde(default, rename = "mimeType", skip_serializing_if = "Option::is_none")] + pub mime_type: Option, +} + +/// A prompt template advertised by an MCP server. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct McpPrompt { + /// Prompt name (pass to [`McpClient::get_prompt`]). + pub name: String, + /// Description, when provided. + #[serde(default)] + pub description: String, + /// Declared arguments (name/description/required triples). + #[serde(default)] + pub arguments: serde_json::Value, +} diff --git a/src/server/mod.rs b/src/server/mod.rs index d610a74..5a1a20a 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -3,25 +3,47 @@ //! //! # Endpoints //! -//! | Method | Path | Description | -//! |--------|-------------------|--------------------------------------------| -//! | GET | `/health` | Liveness probe | -//! | GET | `/agent` | [`AgentInfo`](crate::agent::AgentInfo) | -//! | GET | `/agent/manifest` | The (signed) JSON manifest | -//! | GET | `/capabilities` | Declared capabilities | -//! | GET | `/skills` | Registered skills with input schemas | -//! | POST | `/skills/{name}` | Execute a skill (JSON body → JSON result) | -//! | POST | `/chat` | One chat turn | -//! | POST | `/verify` | Verify a posted manifest's signature | -//! | GET | `/ws` | WebSocket for interactive/streaming chat | +//! | Method | Path | Auth | Description | +//! |--------|-------------------|------|--------------------------------------------| +//! | GET | `/health` | no | Liveness probe (always `ok` while serving) | +//! | GET | `/ready` | no | Readiness probe ([`Agent::set_ready`]) | +//! | GET | `/agent` | yes | [`AgentInfo`](crate::agent::AgentInfo) | +//! | GET | `/agent/manifest` | yes | The (signed) JSON manifest | +//! | GET | `/capabilities` | yes | Declared capabilities | +//! | GET | `/skills` | yes | Registered skills with input schemas | +//! | POST | `/skills/{name}` | yes | Execute a skill (JSON body → JSON result) | +//! | POST | `/chat` | yes | One chat turn | +//! | POST | `/verify` | yes | Verify a posted manifest's signature | +//! | GET | `/ws` | yes | WebSocket for interactive/streaming chat | +//! | GET | `/openapi.json` | no | OpenAPI 3 document (feature `openapi`) | +//! +//! "Auth: yes" applies only when the router is built with +//! [`router_with_auth`] / [`serve_with_auth`]; the plain [`router`] is open. //! //! ```no_run //! use std::sync::Arc; +//! use corrosive_agents::auth::AuthScheme; //! use corrosive_agents::prelude::*; +//! use corrosive_agents::server; //! //! # async fn run() -> corrosive_agents::Result<()> { //! let agent = Arc::new(Agent::builder().name("svc").version("0.1.0").build()?); -//! agent.serve("127.0.0.1:8080".parse().unwrap()).await?; +//! +//! // Open, until Ctrl-C/SIGTERM: +//! server::serve_with_shutdown( +//! agent.clone(), +//! "127.0.0.1:8080".parse().unwrap(), +//! server::shutdown_signal(), +//! ) +//! .await?; +//! +//! // Or API-key protected: +//! server::serve_with_auth( +//! agent, +//! "127.0.0.1:8080".parse().unwrap(), +//! AuthScheme::api_keys(["super-secret"]), +//! ) +//! .await?; //! # Ok(()) //! # } //! ``` @@ -29,29 +51,103 @@ mod rest; mod ws; +use std::future::Future; use std::net::SocketAddr; use std::sync::Arc; use axum::Router; use crate::agent::Agent; +use crate::auth::AuthScheme; use crate::error::{Error, Result}; -/// Build the axum [`Router`] for an agent — compose it into a larger app or -/// serve it directly with [`serve`]. +/// Build the axum [`Router`] for an agent (no authentication) — compose it +/// into a larger app or serve it directly with [`serve`]. pub fn router(agent: Arc) -> Router { - rest::router(agent) + rest::router(agent, None) +} + +/// Build the router with every endpoint except `/health`, `/ready`, and +/// `/openapi.json` protected by `auth`. +pub fn router_with_auth(agent: Arc, auth: AuthScheme) -> Router { + rest::router(agent, Some(Arc::new(auth))) +} + +async fn bind(addr: SocketAddr) -> Result { + tokio::net::TcpListener::bind(addr) + .await + .map_err(|e| Error::Server(format!("failed to bind {addr}: {e}"))) } /// Bind `addr` and serve the agent's REST + WebSocket API until the task is /// cancelled. pub async fn serve(agent: Arc, addr: SocketAddr) -> Result<()> { - let app = router(agent); - let listener = tokio::net::TcpListener::bind(addr) - .await - .map_err(|e| Error::Server(format!("failed to bind {addr}: {e}")))?; + let listener = bind(addr).await?; tracing::info!("REST/WebSocket API listening on http://{addr}"); - axum::serve(listener, app) + axum::serve(listener, router(agent)) + .await + .map_err(|e| Error::Server(e.to_string())) +} + +/// [`serve`], with all non-probe endpoints protected by `auth`. +pub async fn serve_with_auth(agent: Arc, addr: SocketAddr, auth: AuthScheme) -> Result<()> { + let listener = bind(addr).await?; + tracing::info!("REST/WebSocket API (authenticated) listening on http://{addr}"); + axum::serve(listener, router_with_auth(agent, auth)) + .await + .map_err(|e| Error::Server(e.to_string())) +} + +/// [`serve`], shutting down gracefully (in-flight requests drain) when +/// `signal` resolves. Pair with [`shutdown_signal`] for Ctrl-C/SIGTERM. +pub async fn serve_with_shutdown( + agent: Arc, + addr: SocketAddr, + signal: impl Future + Send + 'static, +) -> Result<()> { + let listener = bind(addr).await?; + tracing::info!("REST/WebSocket API listening on http://{addr} (graceful shutdown armed)"); + axum::serve(listener, router(agent)) + .with_graceful_shutdown(signal) + .await + .map_err(|e| Error::Server(e.to_string())) +} + +/// Resolves on Ctrl-C or SIGTERM — the conventional shutdown trigger for +/// containerized deployments. +pub async fn shutdown_signal() { + let ctrl_c = async { + let _ = tokio::signal::ctrl_c().await; + }; + #[cfg(unix)] + { + let mut sigterm = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("failed to install SIGTERM handler"); + tokio::select! { + _ = ctrl_c => {} + _ = sigterm.recv() => {} + } + } + #[cfg(not(unix))] + ctrl_c.await; + tracing::info!("shutdown signal received"); +} + +/// Serve REST + WebSocket over TLS (feature `tls`). +#[cfg(feature = "tls")] +#[cfg_attr(docsrs, doc(cfg(feature = "tls")))] +pub async fn serve_tls( + agent: Arc, + addr: SocketAddr, + tls: &crate::tls::TlsConfig, +) -> Result<()> { + let (cert, key) = tls.pem_pair()?; + let config = axum_server::tls_rustls::RustlsConfig::from_pem(cert, key) + .await + .map_err(|e| Error::Server(format!("invalid TLS material: {e}")))?; + tracing::info!("REST/WebSocket API listening on https://{addr}"); + axum_server::bind_rustls(addr, config) + .serve(router(agent).into_make_service()) .await .map_err(|e| Error::Server(e.to_string())) } @@ -60,7 +156,8 @@ impl Agent { /// Serve this agent's REST + WebSocket API on `addr`. /// /// Convenience for [`server::serve`](serve); requires the agent to be - /// wrapped in an [`Arc`]. + /// wrapped in an [`Arc`]. See also [`serve_with_auth`], + /// [`serve_with_shutdown`], and [`serve_tls`]. pub async fn serve(self: Arc, addr: SocketAddr) -> Result<()> { serve(self, addr).await } diff --git a/src/server/rest.rs b/src/server/rest.rs index 8dede03..8db80e5 100644 --- a/src/server/rest.rs +++ b/src/server/rest.rs @@ -2,8 +2,9 @@ use std::sync::Arc; -use axum::extract::{Path, State}; +use axum::extract::{Path, Request, State}; use axum::http::StatusCode; +use axum::middleware::Next; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; @@ -11,11 +12,19 @@ use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use crate::agent::{Agent, AgentInfo, AgentManifest, Capability}; +use crate::auth::AuthScheme; use crate::error::Error; -pub(crate) fn router(agent: Arc) -> Router { - Router::new() +pub(crate) fn router(agent: Arc, auth: Option>) -> Router { + // Probes (and the API spec) stay open so orchestrators can always reach + // them; everything else is optionally auth-gated. + let public = Router::new() .route("/health", get(health)) + .route("/ready", get(ready)); + #[cfg(feature = "openapi")] + let public = public.route("/openapi.json", get(openapi)); + + let mut protected = Router::new() .route("/agent", get(agent_info)) .route("/agent/manifest", get(manifest)) .route("/capabilities", get(capabilities)) @@ -23,8 +32,30 @@ pub(crate) fn router(agent: Arc) -> Router { .route("/skills/{name}", post(execute_skill)) .route("/chat", post(chat)) .route("/verify", post(verify)) - .route("/ws", get(super::ws::upgrade)) - .with_state(agent) + .route("/ws", get(super::ws::upgrade)); + + if let Some(auth) = auth { + protected = protected.route_layer(axum::middleware::from_fn( + move |request: Request, next: Next| { + let auth = Arc::clone(&auth); + async move { authenticate(&auth, request, next).await } + }, + )); + } + + public.merge(protected).with_state(agent) +} + +async fn authenticate(auth: &AuthScheme, request: Request, next: Next) -> Response { + let headers = request.headers(); + let authorization = headers + .get(axum::http::header::AUTHORIZATION) + .and_then(|v| v.to_str().ok()); + let api_key = headers.get("x-api-key").and_then(|v| v.to_str().ok()); + match auth.authorize(authorization, api_key) { + Ok(()) => next.run(request).await, + Err(e) => e.into_response(), + } } impl IntoResponse for Error { @@ -32,6 +63,8 @@ impl IntoResponse for Error { let status = match &self { Error::SkillNotFound(_) => StatusCode::NOT_FOUND, Error::Config(_) | Error::Json(_) => StatusCode::BAD_REQUEST, + Error::Auth(_) => StatusCode::UNAUTHORIZED, + Error::PermissionDenied(_) => StatusCode::FORBIDDEN, Error::Verification(_) | Error::Identity(_) => StatusCode::UNPROCESSABLE_ENTITY, Error::NotConfigured(_) => StatusCode::NOT_IMPLEMENTED, _ => StatusCode::INTERNAL_SERVER_ERROR, @@ -40,22 +73,47 @@ impl IntoResponse for Error { } } +#[cfg_attr(feature = "openapi", utoipa::path(get, path = "/health", + responses((status = 200, description = "Liveness probe"))))] async fn health() -> Json { Json(json!({ "status": "ok" })) } +#[cfg_attr(feature = "openapi", utoipa::path(get, path = "/ready", responses( + (status = 200, description = "Agent is ready for traffic"), + (status = 503, description = "Agent is not ready"))))] +async fn ready(State(agent): State>) -> Response { + if agent.is_ready() { + Json(json!({ "ready": true })).into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({ "ready": false })), + ) + .into_response() + } +} + +#[cfg_attr(feature = "openapi", utoipa::path(get, path = "/agent", + responses((status = 200, description = "Public agent info", body = AgentInfo))))] async fn agent_info(State(agent): State>) -> Json { Json(agent.info()) } +#[cfg_attr(feature = "openapi", utoipa::path(get, path = "/agent/manifest", + responses((status = 200, description = "The signed agent manifest", body = AgentManifest))))] async fn manifest(State(agent): State>) -> Json { Json(agent.manifest().clone()) } +#[cfg_attr(feature = "openapi", utoipa::path(get, path = "/capabilities", + responses((status = 200, description = "Declared capabilities", body = [Capability]))))] async fn capabilities(State(agent): State>) -> Json> { Json(agent.manifest().capabilities.clone()) } +#[cfg_attr(feature = "openapi", utoipa::path(get, path = "/skills", + responses((status = 200, description = "Registered skills with input schemas"))))] async fn list_skills(State(agent): State>) -> Json { let skills: Vec = agent .skills() @@ -66,12 +124,19 @@ async fn list_skills(State(agent): State>) -> Json { "name": skill.name(), "description": skill.description(), "input_schema": skill.input_schema(), + "required_permissions": skill.required_permissions(), }) }) .collect(); Json(json!({ "skills": skills })) } +#[cfg_attr(feature = "openapi", utoipa::path(post, path = "/skills/{name}", + params(("name" = String, Path, description = "Skill name")), + responses( + (status = 200, description = "Skill output (JSON)"), + (status = 403, description = "Refused by the skill policy"), + (status = 404, description = "No such skill"))))] async fn execute_skill( State(agent): State>, Path(name): Path, @@ -82,20 +147,28 @@ async fn execute_skill( /// Body for `POST /chat`. #[derive(Debug, Deserialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] struct ChatBody { /// Session to continue; a fresh one is created when omitted. session_id: Option, /// The user message. message: String, + /// When `true`, the model may call the agent's skills (tool loop). + #[serde(default)] + use_tools: bool, } /// Response for `POST /chat`. #[derive(Debug, Serialize)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] struct ChatReply { session_id: String, reply: String, } +#[cfg_attr(feature = "openapi", utoipa::path(post, path = "/chat", + request_body = ChatBody, + responses((status = 200, description = "Assistant reply", body = ChatReply))))] async fn chat( State(agent): State>, Json(body): Json, @@ -103,13 +176,48 @@ async fn chat( let session_id = body .session_id .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let reply = agent.chat(&session_id, &body.message).await?; + let reply = if body.use_tools { + agent.chat_with_tools(&session_id, &body.message, 8).await? + } else { + agent.chat(&session_id, &body.message).await? + }; Ok(Json(ChatReply { session_id, reply })) } +#[cfg_attr(feature = "openapi", utoipa::path(post, path = "/verify", + request_body = AgentManifest, + responses((status = 200, description = "Verification verdict"))))] async fn verify(Json(manifest): Json) -> Json { match manifest.verify() { Ok(()) => Json(json!({ "valid": true })), Err(e) => Json(json!({ "valid": false, "reason": e.to_string() })), } } + +#[cfg(feature = "openapi")] +#[derive(utoipa::OpenApi)] +#[openapi( + info( + title = "corrosive_agents", + description = "REST API of a corrosive agent (see also /ws for WebSocket chat)." + ), + paths( + health, + ready, + agent_info, + manifest, + capabilities, + list_skills, + execute_skill, + chat, + verify + ), + components(schemas(AgentInfo, AgentManifest, Capability, ChatBody, ChatReply)) +)] +struct ApiDoc; + +#[cfg(feature = "openapi")] +async fn openapi() -> Json { + use utoipa::OpenApi as _; + Json(serde_json::to_value(ApiDoc::openapi()).unwrap_or(Value::Null)) +} diff --git a/src/session/sqlite.rs b/src/session/sqlite.rs index a585da7..c108904 100644 --- a/src/session/sqlite.rs +++ b/src/session/sqlite.rs @@ -5,12 +5,15 @@ use std::sync::{Arc, Mutex}; use rusqlite::{params, Connection}; use crate::error::{Error, Result}; -use crate::llm::{ChatMessage, Role}; +use crate::llm::ChatMessage; use crate::session::SessionStore; /// Persists conversation history in a SQLite database (bundled, no external /// service required). Safe to share across an application via `Arc`. /// +/// Messages are stored as JSON, so tool calls and future message fields +/// round-trip losslessly. +/// /// ```no_run /// use corrosive_agents::session::SqliteSessionStore; /// @@ -30,27 +33,6 @@ impl std::fmt::Debug for SqliteSessionStore { } } -fn role_to_str(role: Role) -> &'static str { - match role { - Role::System => "system", - Role::User => "user", - Role::Assistant => "assistant", - Role::Tool => "tool", - } -} - -fn role_from_str(s: &str) -> Result { - match s { - "system" => Ok(Role::System), - "user" => Ok(Role::User), - "assistant" => Ok(Role::Assistant), - "tool" => Ok(Role::Tool), - other => Err(Error::Config(format!( - "unknown role '{other}' in session store" - ))), - } -} - fn db_err(e: rusqlite::Error) -> Error { Error::Config(format!("sqlite session store: {e}")) } @@ -71,8 +53,7 @@ impl SqliteSessionStore { "CREATE TABLE IF NOT EXISTS messages ( session_id TEXT NOT NULL, seq INTEGER NOT NULL, - role TEXT NOT NULL, - content TEXT NOT NULL, + message TEXT NOT NULL, PRIMARY KEY (session_id, seq) );", ) @@ -104,44 +85,36 @@ impl SessionStore for SqliteSessionStore { let session_id = session_id.to_string(); let rows = self .with_conn(move |conn| { - let mut statement = conn.prepare( - "SELECT role, content FROM messages WHERE session_id = ?1 ORDER BY seq", - )?; + let mut statement = conn + .prepare("SELECT message FROM messages WHERE session_id = ?1 ORDER BY seq")?; let rows = statement - .query_map(params![session_id], |row| { - Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) - })? + .query_map(params![session_id], |row| row.get::<_, String>(0))? .collect::, _>>()?; Ok(rows) }) .await?; - rows.into_iter() - .map(|(role, content)| { - Ok(ChatMessage { - role: role_from_str(&role)?, - content, - }) - }) + rows.iter() + .map(|json| serde_json::from_str(json).map_err(Error::from)) .collect() } async fn append(&self, session_id: &str, messages: &[ChatMessage]) -> Result<()> { let session_id = session_id.to_string(); - let rows: Vec<(String, String)> = messages + let rows: Vec = messages .iter() - .map(|m| (role_to_str(m.role).to_string(), m.content.clone())) - .collect(); + .map(serde_json::to_string) + .collect::>()?; self.with_conn(move |conn| { - for (role, content) in rows { + for message in rows { conn.execute( - "INSERT INTO messages (session_id, seq, role, content) + "INSERT INTO messages (session_id, seq, message) VALUES ( ?1, (SELECT COALESCE(MAX(seq), -1) + 1 FROM messages WHERE session_id = ?1), - ?2, ?3 + ?2 )", - params![session_id, role, content], + params![session_id, message], )?; } Ok(()) @@ -177,6 +150,7 @@ impl SessionStore for SqliteSessionStore { #[cfg(test)] mod tests { use super::*; + use crate::llm::Role; #[tokio::test] async fn roundtrip_in_memory() { @@ -210,6 +184,31 @@ mod tests { assert_eq!(store.list_sessions().await.unwrap(), vec!["s2"]); } + #[tokio::test] + async fn tool_messages_roundtrip_losslessly() { + use crate::llm::ToolCall; + let store = SqliteSessionStore::in_memory().unwrap(); + let call = ToolCall { + id: "call-1".into(), + name: "lookup".into(), + arguments: serde_json::json!({ "q": "rust" }), + }; + store + .append( + "s", + &[ + ChatMessage::assistant_tool_calls("", vec![call.clone()]), + ChatMessage::tool_result("call-1", r#"{"answer":42}"#), + ], + ) + .await + .unwrap(); + + let history = store.load("s").await.unwrap(); + assert_eq!(history[0].tool_calls.as_ref().unwrap()[0], call); + assert_eq!(history[1].tool_call_id.as_deref(), Some("call-1")); + } + #[tokio::test] async fn persists_across_handles_to_same_file() { let dir = std::env::temp_dir().join(format!("corrosive-test-{}", uuid::Uuid::new_v4())); diff --git a/src/skills/mod.rs b/src/skills/mod.rs index b0d53de..13c3285 100644 --- a/src/skills/mod.rs +++ b/src/skills/mod.rs @@ -28,10 +28,116 @@ pub trait Skill: Send + Sync { serde_json::json!({ "type": "object" }) } + /// Permissions this skill needs (free-form labels such as `"net"`, + /// `"fs:read"`). The agent's [`SkillPolicy`] must grant all of them or + /// execution is refused. Defaults to none. + fn required_permissions(&self) -> Vec { + Vec::new() + } + /// Execute the skill. async fn execute(&self, input: Value) -> Result; } +/// The agent-level sandbox for skill execution: which skills may run, which +/// permissions are granted, and how long a skill may take. +/// +/// The default policy allows every registered skill, grants no permissions +/// (so skills that declare [`Skill::required_permissions`] are refused until +/// granted), and applies a 30-second timeout. +/// +/// ``` +/// use corrosive_agents::skills::SkillPolicy; +/// use std::time::Duration; +/// +/// let policy = SkillPolicy::new() +/// .allow_only(["fetch", "summarize"]) // everything else is refused +/// .grant("net") // satisfy `required_permissions` +/// .with_timeout(Duration::from_secs(5)); +/// ``` +#[derive(Debug, Clone)] +pub struct SkillPolicy { + allowed_skills: Option>, + granted_permissions: std::collections::HashSet, + timeout: Option, +} + +impl Default for SkillPolicy { + fn default() -> Self { + Self { + allowed_skills: None, + granted_permissions: std::collections::HashSet::new(), + timeout: Some(std::time::Duration::from_secs(30)), + } + } +} + +impl SkillPolicy { + /// The default policy (all skills allowed, no permissions granted, + /// 30-second timeout). + pub fn new() -> Self { + Self::default() + } + + /// Restrict execution to an explicit allowlist of skill names. + #[must_use] + pub fn allow_only(mut self, names: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.allowed_skills = Some(names.into_iter().map(Into::into).collect()); + self + } + + /// Grant a permission label (see [`Skill::required_permissions`]). + #[must_use] + pub fn grant(mut self, permission: impl Into) -> Self { + self.granted_permissions.insert(permission.into()); + self + } + + /// Cap how long a single skill execution may run. + #[must_use] + pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self { + self.timeout = Some(timeout); + self + } + + /// Remove the execution timeout. + #[must_use] + pub fn without_timeout(mut self) -> Self { + self.timeout = None; + self + } + + /// The configured timeout, if any. + pub fn timeout(&self) -> Option { + self.timeout + } + + /// Check whether `skill` may run under this policy. + pub fn check(&self, skill: &dyn Skill) -> Result<()> { + if let Some(allowed) = &self.allowed_skills { + if !allowed.contains(skill.name()) { + return Err(Error::PermissionDenied(format!( + "skill '{}' is not on the allowlist", + skill.name() + ))); + } + } + for permission in skill.required_permissions() { + if !self.granted_permissions.contains(&permission) { + return Err(Error::PermissionDenied(format!( + "skill '{}' requires permission '{permission}' which is not granted", + skill.name() + ))); + } + } + Ok(()) + } +} + type SkillFuture = Pin> + Send>>; type SkillFn = dyn Fn(Value) -> SkillFuture + Send + Sync; @@ -50,6 +156,7 @@ pub struct FnSkill { name: String, description: String, schema: Value, + permissions: Vec, handler: Arc, } @@ -64,6 +171,7 @@ impl FnSkill { name: name.into(), description: description.into(), schema: serde_json::json!({ "type": "object" }), + permissions: Vec::new(), handler: Arc::new(move |input| Box::pin(handler(input))), } } @@ -74,6 +182,18 @@ impl FnSkill { self.schema = schema; self } + + /// Declare permissions this skill requires (must be granted by the + /// agent's [`SkillPolicy`] before the skill may run). + #[must_use] + pub fn with_permissions(mut self, permissions: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.permissions = permissions.into_iter().map(Into::into).collect(); + self + } } #[async_trait::async_trait] @@ -90,6 +210,10 @@ impl Skill for FnSkill { self.schema.clone() } + fn required_permissions(&self) -> Vec { + self.permissions.clone() + } + async fn execute(&self, input: Value) -> Result { (self.handler)(input).await } diff --git a/src/tls.rs b/src/tls.rs new file mode 100644 index 0000000..efbdbd3 --- /dev/null +++ b/src/tls.rs @@ -0,0 +1,77 @@ +//! TLS configuration helpers (feature `tls`). +//! +//! One [`TlsConfig`] drives both transports: +//! [`server::serve_tls`](crate::server::serve_tls) for REST/WebSocket and +//! [`grpc::serve_tls`](crate::grpc::serve_tls) for gRPC. +//! +//! ```no_run +//! use corrosive_agents::tls::TlsConfig; +//! +//! let tls = TlsConfig::from_pem_files("cert.pem", "key.pem"); +//! ``` + +use std::path::PathBuf; + +use crate::error::{Error, Result}; + +enum Source { + Files { cert: PathBuf, key: PathBuf }, + Pem { cert: Vec, key: Vec }, +} + +/// A certificate/key pair for serving TLS. +pub struct TlsConfig { + source: Source, +} + +impl std::fmt::Debug for TlsConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match &self.source { + Source::Files { cert, key } => f + .debug_struct("TlsConfig") + .field("cert", cert) + .field("key", key) + .finish(), + Source::Pem { .. } => f.debug_struct("TlsConfig").field("source", &"pem").finish(), + } + } +} + +impl TlsConfig { + /// Load the certificate chain and private key from PEM files at serve + /// time. + pub fn from_pem_files(cert: impl Into, key: impl Into) -> Self { + Self { + source: Source::Files { + cert: cert.into(), + key: key.into(), + }, + } + } + + /// Use in-memory PEM data. + pub fn from_pem(cert_pem: impl Into>, key_pem: impl Into>) -> Self { + Self { + source: Source::Pem { + cert: cert_pem.into(), + key: key_pem.into(), + }, + } + } + + /// The (certificate, key) PEM bytes. + pub(crate) fn pem_pair(&self) -> Result<(Vec, Vec)> { + match &self.source { + Source::Files { cert, key } => { + let cert = std::fs::read(cert).map_err(|e| { + Error::Config(format!("cannot read TLS cert {}: {e}", cert.display())) + })?; + let key = std::fs::read(key).map_err(|e| { + Error::Config(format!("cannot read TLS key {}: {e}", key.display())) + })?; + Ok((cert, key)) + } + Source::Pem { cert, key } => Ok((cert.clone(), key.clone())), + } + } +} diff --git a/src/trust.rs b/src/trust.rs index eda6069..f56b04e 100644 --- a/src/trust.rs +++ b/src/trust.rs @@ -53,6 +53,7 @@ fn unix_now() -> u64 { /// A signed endorsement of a new key by the previous key. #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))] pub struct RotationProof { /// Base64 public key being rotated away from. pub previous_key: String, diff --git a/src/vector/chunk.rs b/src/vector/chunk.rs new file mode 100644 index 0000000..2d16ba7 --- /dev/null +++ b/src/vector/chunk.rs @@ -0,0 +1,94 @@ +//! Text chunking helpers for indexing large documents. + +/// Split `text` into chunks of at most `max_chars` characters, breaking on +/// word boundaries, with roughly `overlap` characters of trailing context +/// repeated at the start of the next chunk. +/// +/// Useful before embedding: retrieval works best when each chunk is a +/// self-contained passage (a few hundred to ~2000 characters, depending on +/// the embedding model). +/// +/// ``` +/// use corrosive_agents::vector::chunk_text; +/// +/// let chunks = chunk_text(&"lorem ipsum ".repeat(100), 200, 40); +/// assert!(chunks.iter().all(|c| c.chars().count() <= 200)); +/// assert!(chunks.len() > 1); +/// ``` +pub fn chunk_text(text: &str, max_chars: usize, overlap: usize) -> Vec { + let max_chars = max_chars.max(1); + let overlap = overlap.min(max_chars / 2); + let words: Vec<&str> = text.split_whitespace().collect(); + if words.is_empty() { + return Vec::new(); + } + + let mut chunks = Vec::new(); + let mut current: Vec<&str> = Vec::new(); + let mut current_len = 0usize; + + for word in words { + let word_len = word.chars().count(); + let sep = usize::from(!current.is_empty()); + if current_len + sep + word_len > max_chars && !current.is_empty() { + chunks.push(current.join(" ")); + // Seed the next chunk with ~`overlap` chars of trailing words. + let mut carried: Vec<&str> = Vec::new(); + let mut carried_len = 0usize; + for prev in current.iter().rev() { + let prev_len = prev.chars().count(); + if carried_len + prev_len > overlap { + break; + } + carried_len += prev_len + 1; + carried.push(prev); + } + carried.reverse(); + current = carried; + current_len = current.iter().map(|w| w.chars().count()).sum::() + + current.len().saturating_sub(1); + } + current_len += usize::from(!current.is_empty()) + word_len; + current.push(word); + } + if !current.is_empty() { + chunks.push(current.join(" ")); + } + chunks +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn respects_max_chars() { + let text = "alpha beta gamma delta epsilon zeta eta theta iota kappa".repeat(20); + for chunk in chunk_text(&text, 80, 20) { + assert!(chunk.chars().count() <= 80, "chunk too long: {chunk}"); + } + } + + #[test] + fn overlap_carries_context() { + let text = "one two three four five six seven eight nine ten"; + let chunks = chunk_text(text, 20, 8); + assert!(chunks.len() >= 2); + // Some trailing words of chunk N reappear at the start of chunk N+1. + let first_tail = chunks[0].split_whitespace().last().unwrap(); + assert!(chunks[1].contains(first_tail)); + } + + #[test] + fn short_text_is_one_chunk() { + assert_eq!(chunk_text("hello world", 100, 10), vec!["hello world"]); + assert!(chunk_text(" ", 100, 10).is_empty()); + } + + #[test] + fn single_oversized_word_still_emits() { + let long_word = "x".repeat(50); + let chunks = chunk_text(&long_word, 10, 2); + assert_eq!(chunks.len(), 1); + } +} diff --git a/src/vector/memory.rs b/src/vector/memory.rs index a64b5cc..4b1caa5 100644 --- a/src/vector/memory.rs +++ b/src/vector/memory.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; use tokio::sync::RwLock; use crate::error::{Error, Result}; -use crate::vector::{Document, SearchResult, VectorStore}; +use crate::vector::{Document, MetadataFilter, SearchResult, VectorStore}; /// A thread-safe in-memory vector store using cosine similarity. /// @@ -66,9 +66,20 @@ impl VectorStore for InMemoryVectorStore { } async fn search(&self, vector: Vec, top_k: usize) -> Result> { + self.search_filtered(vector, top_k, &MetadataFilter::new()) + .await + } + + async fn search_filtered( + &self, + vector: Vec, + top_k: usize, + filter: &MetadataFilter, + ) -> Result> { let store = self.documents.read().await; let mut results: Vec = store .values() + .filter(|doc| filter.is_empty() || filter.matches(&doc.metadata)) .map(|doc| SearchResult { id: doc.id.clone(), score: cosine_similarity(&vector, &doc.vector), @@ -149,4 +160,35 @@ mod tests { .await .is_err()); } + + #[tokio::test] + async fn metadata_filter_narrows_results() { + let store = InMemoryVectorStore::new(); + store + .upsert(vec![ + Document::new("r1", vec![1.0, 0.0]).with_metadata(json!({"lang": "rust"})), + Document::new("g1", vec![1.0, 0.0]).with_metadata(json!({"lang": "go"})), + Document::new("r2", vec![0.9, 0.1]).with_metadata(json!({"lang": "rust"})), + ]) + .await + .unwrap(); + + let filter = MetadataFilter::new().eq("lang", json!("rust")); + let hits = store + .search_filtered(vec![1.0, 0.0], 10, &filter) + .await + .unwrap(); + assert_eq!(hits.len(), 2); + assert!(hits.iter().all(|h| h.metadata["lang"] == "rust")); + } + + #[tokio::test] + async fn batched_upsert_stores_everything() { + let store = InMemoryVectorStore::new(); + let docs: Vec = (0..25) + .map(|i| Document::new(format!("d{i}"), vec![i as f32, 1.0])) + .collect(); + store.upsert_batched(docs, 10).await.unwrap(); + assert_eq!(store.len().await, 25); + } } diff --git a/src/vector/mod.rs b/src/vector/mod.rs index 333404f..17ad417 100644 --- a/src/vector/mod.rs +++ b/src/vector/mod.rs @@ -9,13 +9,20 @@ //! //! Implement the trait yourself to plug in any other backend. +mod chunk; mod memory; +#[cfg(feature = "pgvector")] +mod pg; #[cfg(feature = "pinecone")] mod pinecone; #[cfg(feature = "qdrant")] mod qdrant; +pub use chunk::chunk_text; pub use memory::InMemoryVectorStore; +#[cfg(feature = "pgvector")] +#[cfg_attr(docsrs, doc(cfg(feature = "pgvector")))] +pub use pg::PgVectorStore; #[cfg(feature = "pinecone")] #[cfg_attr(docsrs, doc(cfg(feature = "pinecone")))] pub use pinecone::PineconeStore; @@ -23,11 +30,62 @@ pub use pinecone::PineconeStore; #[cfg_attr(docsrs, doc(cfg(feature = "qdrant")))] pub use qdrant::QdrantStore; +use std::collections::BTreeMap; + use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::error::Result; +/// An equality filter over document metadata. +/// +/// Backends translate it natively where possible (Qdrant payload filters, +/// Pinecone `$eq` filters, pgvector `@>` containment); the trait's default +/// implementation over-fetches and filters client-side, so custom stores get +/// filtering for free. +/// +/// ``` +/// use corrosive_agents::vector::MetadataFilter; +/// use serde_json::json; +/// +/// let filter = MetadataFilter::new() +/// .eq("topic", json!("rust")) +/// .eq("year", json!(2026)); +/// assert!(filter.matches(&json!({ "topic": "rust", "year": 2026, "extra": 1 }))); +/// assert!(!filter.matches(&json!({ "topic": "go" }))); +/// ``` +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)] +pub struct MetadataFilter { + /// Field → required value (all must match). + pub equals: BTreeMap, +} + +impl MetadataFilter { + /// An empty filter (matches everything). + pub fn new() -> Self { + Self::default() + } + + /// Require `metadata[key] == value`. + #[must_use] + pub fn eq(mut self, key: impl Into, value: Value) -> Self { + self.equals.insert(key.into(), value); + self + } + + /// `true` when no conditions are set. + pub fn is_empty(&self) -> bool { + self.equals.is_empty() + } + + /// Does `metadata` satisfy every condition? + pub fn matches(&self, metadata: &Value) -> bool { + self.equals + .iter() + .all(|(key, expected)| metadata.get(key) == Some(expected)) + } +} + /// A document stored in (or destined for) a vector store. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Document { @@ -96,6 +154,44 @@ pub trait VectorStore: Send + Sync { /// Return the `top_k` most similar documents to `vector`. async fn search(&self, vector: Vec, top_k: usize) -> Result>; + /// Like [`search`](Self::search), keeping only documents whose metadata + /// matches `filter`. + /// + /// The default implementation over-fetches (4×) and filters client-side; + /// backends override it with native filters where the database supports + /// them. + async fn search_filtered( + &self, + vector: Vec, + top_k: usize, + filter: &MetadataFilter, + ) -> Result> { + if filter.is_empty() { + return self.search(vector, top_k).await; + } + let fetch = top_k.saturating_mul(4).max(top_k); + let results = self.search(vector, fetch).await?; + Ok(results + .into_iter() + .filter(|r| filter.matches(&r.metadata)) + .take(top_k) + .collect()) + } + + /// Upsert in fixed-size batches — use for large corpora to keep request + /// sizes bounded. `batch_size` of 0 is treated as 1. + async fn upsert_batched(&self, documents: Vec, batch_size: usize) -> Result<()> { + let batch_size = batch_size.max(1); + // Consume the Vec in chunks without cloning documents. + let mut documents = documents; + while !documents.is_empty() { + let rest = documents.split_off(documents.len().min(batch_size)); + self.upsert(documents).await?; + documents = rest; + } + Ok(()) + } + /// Delete documents by id. async fn delete(&self, ids: &[String]) -> Result<()>; } diff --git a/src/vector/pg.rs b/src/vector/pg.rs new file mode 100644 index 0000000..899a92c --- /dev/null +++ b/src/vector/pg.rs @@ -0,0 +1,258 @@ +//! pgvector (PostgreSQL) vector store backend, enabled with feature +//! `pgvector`. + +use std::sync::Arc; + +use pgvector::Vector; +use serde_json::Value; +use tokio_postgres::NoTls; + +use crate::error::{Error, Result}; +use crate::vector::{Document, MetadataFilter, SearchResult, VectorStore}; + +fn pg_err(e: tokio_postgres::Error) -> Error { + Error::VectorStore(format!("pgvector: {e}")) +} + +/// A [`VectorStore`] backed by PostgreSQL with the +/// [pgvector](https://github.com/pgvector/pgvector) extension. +/// +/// Documents live in a table `(id TEXT PRIMARY KEY, embedding vector(n), +/// text TEXT, metadata JSONB)`; similarity is cosine (`<=>` operator) and +/// metadata filters use JSONB containment (`@>`), so they can be served by a +/// GIN index. +/// +/// ```no_run +/// use corrosive_agents::vector::PgVectorStore; +/// +/// # async fn run() -> corrosive_agents::Result<()> { +/// let store = PgVectorStore::connect( +/// "host=localhost user=postgres password=secret dbname=agents", +/// "documents", +/// ) +/// .await?; +/// store.ensure_table(1024).await?; // installs the extension + table if missing +/// # Ok(()) +/// # } +/// ``` +/// +/// Connections are unencrypted (`NoTls`); front Postgres with TLS at the +/// network layer or a proxy if needed. +#[derive(Clone)] +pub struct PgVectorStore { + client: Arc, + table: String, +} + +impl std::fmt::Debug for PgVectorStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PgVectorStore") + .field("table", &self.table) + .finish_non_exhaustive() + } +} + +fn validate_table_name(table: &str) -> Result<()> { + let valid = !table.is_empty() && table.chars().all(|c| c.is_ascii_alphanumeric() || c == '_'); + if valid { + Ok(()) + } else { + Err(Error::Config(format!( + "invalid table name '{table}': use only ASCII letters, digits, and underscores" + ))) + } +} + +impl PgVectorStore { + /// Connect with a `tokio_postgres` connection string and target table. + /// + /// The connection driver is spawned onto the current Tokio runtime. + pub async fn connect(connection_string: &str, table: impl Into) -> Result { + let table = table.into(); + validate_table_name(&table)?; + let (client, connection) = tokio_postgres::connect(connection_string, NoTls) + .await + .map_err(pg_err)?; + tokio::spawn(async move { + if let Err(e) = connection.await { + tracing::error!("pgvector connection error: {e}"); + } + }); + Ok(Self { + client: Arc::new(client), + table, + }) + } + + /// Install the `vector` extension and create the table + cosine index if + /// they do not exist. `dimensions` must match your embedding model + /// (e.g. 1024 for `nv-embedqa-e5-v5`). + pub async fn ensure_table(&self, dimensions: usize) -> Result<()> { + let table = &self.table; + let ddl = format!( + "CREATE EXTENSION IF NOT EXISTS vector; + CREATE TABLE IF NOT EXISTS {table} ( + id TEXT PRIMARY KEY, + embedding vector({dimensions}) NOT NULL, + text TEXT, + metadata JSONB NOT NULL DEFAULT 'null'::jsonb + ); + CREATE INDEX IF NOT EXISTS {table}_metadata_idx ON {table} USING GIN (metadata);" + ); + self.client.batch_execute(&ddl).await.map_err(pg_err) + } + + async fn run_query( + &self, + vector: Vec, + top_k: usize, + filter: Option<&MetadataFilter>, + ) -> Result> { + let embedding = Vector::from(vector); + let limit = top_k as i64; + let table = &self.table; + + let rows = if let Some(filter) = filter { + let conditions = Value::Object( + filter + .equals + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + ); + let sql = format!( + "SELECT id, text, metadata, 1 - (embedding <=> $1) AS score + FROM {table} WHERE metadata @> $3 ORDER BY embedding <=> $1 LIMIT $2" + ); + self.client + .query(&sql, &[&embedding, &limit, &conditions]) + .await + .map_err(pg_err)? + } else { + let sql = format!( + "SELECT id, text, metadata, 1 - (embedding <=> $1) AS score + FROM {table} ORDER BY embedding <=> $1 LIMIT $2" + ); + self.client + .query(&sql, &[&embedding, &limit]) + .await + .map_err(pg_err)? + }; + + Ok(rows + .into_iter() + .map(|row| SearchResult { + id: row.get::<_, String>(0), + text: row.get::<_, Option>(1), + metadata: row.get::<_, Value>(2), + score: row.get::<_, f64>(3) as f32, + }) + .collect()) + } +} + +#[async_trait::async_trait] +impl VectorStore for PgVectorStore { + async fn upsert(&self, documents: Vec) -> Result<()> { + let table = &self.table; + let sql = format!( + "INSERT INTO {table} (id, embedding, text, metadata) VALUES ($1, $2, $3, $4) + ON CONFLICT (id) DO UPDATE + SET embedding = EXCLUDED.embedding, + text = EXCLUDED.text, + metadata = EXCLUDED.metadata" + ); + let statement = self.client.prepare(&sql).await.map_err(pg_err)?; + for doc in documents { + let embedding = Vector::from(doc.vector); + self.client + .execute(&statement, &[&doc.id, &embedding, &doc.text, &doc.metadata]) + .await + .map_err(pg_err)?; + } + Ok(()) + } + + async fn search(&self, vector: Vec, top_k: usize) -> Result> { + self.run_query(vector, top_k, None).await + } + + async fn search_filtered( + &self, + vector: Vec, + top_k: usize, + filter: &MetadataFilter, + ) -> Result> { + if filter.is_empty() { + return self.run_query(vector, top_k, None).await; + } + self.run_query(vector, top_k, Some(filter)).await + } + + async fn delete(&self, ids: &[String]) -> Result<()> { + let table = &self.table; + let sql = format!("DELETE FROM {table} WHERE id = ANY($1)"); + let ids: Vec<&str> = ids.iter().map(String::as_str).collect(); + self.client.execute(&sql, &[&ids]).await.map_err(pg_err)?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn table_names_are_validated() { + assert!(validate_table_name("documents").is_ok()); + assert!(validate_table_name("docs_v2").is_ok()); + assert!(validate_table_name("docs; DROP TABLE users").is_err()); + assert!(validate_table_name("").is_err()); + } + + /// Full roundtrip against a real Postgres with pgvector. Skipped unless + /// `PG_URL` is set (e.g. `PG_URL="host=localhost user=postgres" cargo + /// test --features pgvector`). + #[tokio::test] + async fn roundtrip_against_real_postgres() { + let Ok(url) = std::env::var("PG_URL") else { + eprintln!("PG_URL not set — skipping pgvector integration test"); + return; + }; + let table = format!("corrosive_test_{}", uuid::Uuid::new_v4().simple()); + let store = PgVectorStore::connect(&url, &table).await.unwrap(); + store.ensure_table(2).await.unwrap(); + + store + .upsert(vec![ + Document::new("a", vec![1.0, 0.0]) + .with_text("alpha") + .with_metadata(json!({"lang": "rust"})), + Document::new("b", vec![0.0, 1.0]).with_metadata(json!({"lang": "go"})), + ]) + .await + .unwrap(); + + let hits = store.search(vec![1.0, 0.1], 1).await.unwrap(); + assert_eq!(hits[0].id, "a"); + assert_eq!(hits[0].text.as_deref(), Some("alpha")); + + let filter = MetadataFilter::new().eq("lang", json!("go")); + let hits = store + .search_filtered(vec![1.0, 0.0], 5, &filter) + .await + .unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].id, "b"); + + store.delete(&["a".into(), "b".into()]).await.unwrap(); + assert!(store.search(vec![1.0, 0.0], 5).await.unwrap().is_empty()); + + store + .client + .batch_execute(&format!("DROP TABLE {table}")) + .await + .unwrap(); + } +} diff --git a/src/vector/pinecone.rs b/src/vector/pinecone.rs index d938364..ac18203 100644 --- a/src/vector/pinecone.rs +++ b/src/vector/pinecone.rs @@ -3,7 +3,7 @@ use serde_json::{json, Map, Value}; use crate::error::{Error, Result}; -use crate::vector::{Document, SearchResult, VectorStore}; +use crate::vector::{Document, MetadataFilter, SearchResult, VectorStore}; /// A [`VectorStore`] backed by a [Pinecone](https://www.pinecone.io) index. /// @@ -55,6 +55,52 @@ impl PineconeStore { } Ok(body) } + + async fn query( + &self, + vector: Vec, + top_k: usize, + filter: Option, + ) -> Result> { + let mut body = json!({ + "vector": vector, + "topK": top_k, + "includeMetadata": true, + "namespace": self.namespace, + }); + if let Some(filter) = filter { + body["filter"] = filter; + } + let response = self.post("/query", body).await?; + let matches = response + .get("matches") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(); + Ok(matches + .into_iter() + .map(|hit| { + let mut metadata = hit.get("metadata").cloned().unwrap_or(Value::Null); + let text = metadata + .get("_text") + .and_then(Value::as_str) + .map(String::from); + if let Value::Object(map) = &mut metadata { + map.remove("_text"); + } + SearchResult { + id: hit + .get("id") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + score: hit.get("score").and_then(Value::as_f64).unwrap_or(0.0) as f32, + text, + metadata, + } + }) + .collect()) + } } fn to_pinecone_metadata(doc: &Document) -> Value { @@ -105,41 +151,33 @@ impl VectorStore for PineconeStore { } async fn search(&self, vector: Vec, top_k: usize) -> Result> { - let body = json!({ - "vector": vector, - "topK": top_k, - "includeMetadata": true, - "namespace": self.namespace, - }); - let response = self.post("/query", body).await?; - let matches = response - .get("matches") - .and_then(Value::as_array) - .cloned() - .unwrap_or_default(); - Ok(matches - .into_iter() - .map(|hit| { - let mut metadata = hit.get("metadata").cloned().unwrap_or(Value::Null); - let text = metadata - .get("_text") - .and_then(Value::as_str) - .map(String::from); - if let Value::Object(map) = &mut metadata { - map.remove("_text"); - } - SearchResult { - id: hit - .get("id") - .and_then(Value::as_str) - .unwrap_or_default() - .to_string(), - score: hit.get("score").and_then(Value::as_f64).unwrap_or(0.0) as f32, - text, - metadata, - } + self.query(vector, top_k, None).await + } + + async fn search_filtered( + &self, + vector: Vec, + top_k: usize, + filter: &MetadataFilter, + ) -> Result> { + if filter.is_empty() { + return self.query(vector, top_k, None).await; + } + // Native Pinecone filter. Metadata is stored flat, so only scalar + // values match; nested values were stringified at upsert time. + let conditions: serde_json::Map = filter + .equals + .iter() + .map(|(key, value)| { + let comparable = match value { + Value::String(_) | Value::Number(_) | Value::Bool(_) => value.clone(), + other => Value::String(other.to_string()), + }; + (key.clone(), json!({ "$eq": comparable })) }) - .collect()) + .collect(); + self.query(vector, top_k, Some(Value::Object(conditions))) + .await } async fn delete(&self, ids: &[String]) -> Result<()> { diff --git a/src/vector/qdrant.rs b/src/vector/qdrant.rs index 580a237..1411917 100644 --- a/src/vector/qdrant.rs +++ b/src/vector/qdrant.rs @@ -3,7 +3,7 @@ use serde_json::{json, Value}; use crate::error::{Error, Result}; -use crate::vector::{Document, SearchResult, VectorStore}; +use crate::vector::{Document, MetadataFilter, SearchResult, VectorStore}; /// UUID v5 namespace for deterministically mapping document ids to Qdrant /// point ids (Qdrant only accepts unsigned integers or UUIDs as point ids). @@ -110,11 +110,32 @@ impl VectorStore for QdrantStore { } async fn search(&self, vector: Vec, top_k: usize) -> Result> { + self.search_filtered(vector, top_k, &MetadataFilter::new()) + .await + } + + async fn search_filtered( + &self, + vector: Vec, + top_k: usize, + filter: &MetadataFilter, + ) -> Result> { let url = format!( "{}/collections/{}/points/search", self.base_url, self.collection ); - let body = json!({ "vector": vector, "limit": top_k, "with_payload": true }); + let mut body = json!({ "vector": vector, "limit": top_k, "with_payload": true }); + if !filter.is_empty() { + // Native Qdrant payload filter (metadata lives under "metadata.*"). + let must: Vec = filter + .equals + .iter() + .map(|(key, value)| { + json!({ "key": format!("metadata.{key}"), "match": { "value": value } }) + }) + .collect(); + body["filter"] = json!({ "must": must }); + } let response = self.send(self.http.post(&url).json(&body)).await?; let hits = response .get("result") diff --git a/src/x509.rs b/src/x509.rs new file mode 100644 index 0000000..44168ff --- /dev/null +++ b/src/x509.rs @@ -0,0 +1,147 @@ +//! X.509 certificate-based identity (feature `x509`). +//! +//! An agent's Ed25519 identity can be wrapped in a **self-signed X.509 +//! certificate** — useful when the party you integrate with speaks PKI +//! rather than raw keys or DIDs. The certificate carries the same public key +//! that signs the agent's manifest, so verifying a manifest against a +//! certificate is: extract the key from the cert (checking the cert's own +//! self-signature), then verify the manifest with it. +//! +//! ``` +//! use corrosive_agents::identity::AgentIdentity; +//! use corrosive_agents::agent::AgentManifest; +//! use corrosive_agents::x509; +//! +//! # fn main() -> corrosive_agents::Result<()> { +//! let identity = AgentIdentity::generate(); +//! let cert_pem = x509::generate_certificate_pem(&identity, "research-agent")?; +//! +//! let mut manifest = AgentManifest::new("research-agent", "1.0.0"); +//! manifest.sign(&identity)?; +//! +//! // A consumer holding only the certificate and the manifest: +//! x509::verify_manifest_with_certificate(&manifest, &cert_pem)?; +//! # Ok(()) +//! # } +//! ``` + +use base64::engine::general_purpose::STANDARD as B64; +use base64::Engine as _; +use ed25519_dalek::{Signature, Verifier, VerifyingKey}; + +use crate::agent::AgentManifest; +use crate::error::{Error, Result}; +use crate::identity::AgentIdentity; + +/// OID of the Ed25519 signature/key algorithm (RFC 8410). +const ED25519_OID: &str = "1.3.101.112"; + +/// Generate a self-signed X.509 certificate (PEM) for the identity, with +/// `common_name` as both subject CN and DNS SAN. +pub fn generate_certificate_pem(identity: &AgentIdentity, common_name: &str) -> Result { + let key_pair = rcgen::KeyPair::try_from(identity.pkcs8_der()?.as_slice()) + .map_err(|e| Error::Identity(format!("rcgen rejected the key: {e}")))?; + let mut params = rcgen::CertificateParams::new(vec![common_name.to_string()]) + .map_err(|e| Error::Identity(format!("invalid certificate params: {e}")))?; + params + .distinguished_name + .push(rcgen::DnType::CommonName, common_name); + let certificate = params + .self_signed(&key_pair) + .map_err(|e| Error::Identity(format!("certificate generation failed: {e}")))?; + Ok(certificate.pem()) +} + +/// Extract the base64 Ed25519 public key from a PEM certificate, after +/// checking that the certificate is Ed25519-signed and its self-signature is +/// valid. +pub fn public_key_from_certificate_pem(cert_pem: &str) -> Result { + let (_, pem) = x509_parser::pem::parse_x509_pem(cert_pem.as_bytes()) + .map_err(|e| Error::Identity(format!("invalid PEM: {e}")))?; + let certificate = pem + .parse_x509() + .map_err(|e| Error::Identity(format!("invalid X.509 certificate: {e}")))?; + + let spki = certificate.public_key(); + if spki.algorithm.algorithm.to_id_string() != ED25519_OID { + return Err(Error::Identity( + "certificate does not contain an Ed25519 key".into(), + )); + } + let key_bytes: [u8; 32] = spki + .subject_public_key + .data + .as_ref() + .try_into() + .map_err(|_| Error::Identity("certificate key must be 32 bytes".into()))?; + let verifying_key = VerifyingKey::from_bytes(&key_bytes) + .map_err(|e| Error::Identity(format!("invalid Ed25519 key in certificate: {e}")))?; + + // Verify the certificate's self-signature over its TBS section. + if certificate.signature_algorithm.algorithm.to_id_string() != ED25519_OID { + return Err(Error::Identity( + "certificate is not Ed25519-self-signed".into(), + )); + } + let signature_bytes: [u8; 64] = certificate + .signature_value + .data + .as_ref() + .try_into() + .map_err(|_| Error::Identity("certificate signature must be 64 bytes".into()))?; + let signature = Signature::from_bytes(&signature_bytes); + verifying_key + .verify(certificate.tbs_certificate.as_ref(), &signature) + .map_err(|_| Error::Verification("certificate self-signature is invalid".into()))?; + + Ok(B64.encode(key_bytes)) +} + +/// Verify a manifest against the identity in an X.509 certificate: the +/// certificate must be valid and self-consistent, and its key must have +/// signed the manifest. +pub fn verify_manifest_with_certificate(manifest: &AgentManifest, cert_pem: &str) -> Result<()> { + let public_key = public_key_from_certificate_pem(cert_pem)?; + manifest.verify_with(&public_key) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn certificate_roundtrip() { + let identity = AgentIdentity::generate(); + let pem = generate_certificate_pem(&identity, "test-agent").unwrap(); + assert!(pem.contains("BEGIN CERTIFICATE")); + + let extracted = public_key_from_certificate_pem(&pem).unwrap(); + assert_eq!(extracted, identity.public_key_base64()); + } + + #[test] + fn manifest_verifies_against_certificate() { + let identity = AgentIdentity::generate(); + let pem = generate_certificate_pem(&identity, "signed-agent").unwrap(); + + let mut manifest = AgentManifest::new("signed-agent", "1.0.0"); + manifest.sign(&identity).unwrap(); + verify_manifest_with_certificate(&manifest, &pem).unwrap(); + } + + #[test] + fn foreign_certificate_is_rejected() { + let identity = AgentIdentity::generate(); + let other = AgentIdentity::generate(); + let other_pem = generate_certificate_pem(&other, "imposter").unwrap(); + + let mut manifest = AgentManifest::new("victim", "1.0.0"); + manifest.sign(&identity).unwrap(); + assert!(verify_manifest_with_certificate(&manifest, &other_pem).is_err()); + } + + #[test] + fn garbage_pem_is_rejected() { + assert!(public_key_from_certificate_pem("not a pem").is_err()); + } +} diff --git a/tests/a2a_and_sessions.rs b/tests/a2a_and_sessions.rs index 3325583..4a54823 100644 --- a/tests/a2a_and_sessions.rs +++ b/tests/a2a_and_sessions.rs @@ -23,6 +23,7 @@ impl LlmProvider for EchoLlm { Ok(ChatResponse { content: format!("echo:{last_user}"), model: "echo".into(), + tool_calls: Vec::new(), usage: None, }) } diff --git a/tests/tools_and_policy.rs b/tests/tools_and_policy.rs new file mode 100644 index 0000000..fd5047e --- /dev/null +++ b/tests/tools_and_policy.rs @@ -0,0 +1,274 @@ +//! Integration tests for the tool-calling loop, usage accounting, and the +//! skill sandbox — all offline via a scripted mock LLM. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use corrosive_agents::llm::Role; +use corrosive_agents::prelude::*; +use futures_util::stream::BoxStream; +use serde_json::json; + +/// A scripted LLM: first turn requests the `add` tool, second turn answers +/// with the tool's result. Reports fixed token usage on every call. +struct ToolScriptLlm { + calls: AtomicUsize, +} + +impl ToolScriptLlm { + fn new() -> Self { + Self { + calls: AtomicUsize::new(0), + } + } +} + +#[async_trait::async_trait] +impl LlmProvider for ToolScriptLlm { + async fn chat(&self, request: ChatRequest) -> Result { + let call_number = self.calls.fetch_add(1, Ordering::SeqCst); + let usage = Some(corrosive_agents::llm::Usage { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }); + if call_number == 0 { + assert!( + request.tools.as_ref().is_some_and(|t| !t.is_empty()), + "tools should be offered to the model" + ); + Ok(ChatResponse { + content: String::new(), + model: "scripted".into(), + tool_calls: vec![ToolCall { + id: "call-1".into(), + name: "add".into(), + arguments: json!({ "a": 2, "b": 40 }), + }], + usage, + }) + } else { + // The tool result must have been fed back. + let tool_message = request + .messages + .iter() + .find(|m| matches!(m.role, Role::Tool)) + .expect("tool result message present"); + assert_eq!(tool_message.tool_call_id.as_deref(), Some("call-1")); + assert!(tool_message.content.contains("42")); + Ok(ChatResponse { + content: "the answer is 42".into(), + model: "scripted".into(), + tool_calls: Vec::new(), + usage, + }) + } + } + + async fn chat_stream( + &self, + _request: ChatRequest, + ) -> Result>> { + unimplemented!("not used in these tests") + } + + fn default_model(&self) -> &str { + "scripted" + } +} + +fn add_skill() -> FnSkill { + FnSkill::new("add", "Adds a and b", |input| async move { + let sum = input["a"].as_i64().unwrap_or(0) + input["b"].as_i64().unwrap_or(0); + Ok(json!({ "sum": sum })) + }) +} + +#[tokio::test] +async fn tool_loop_executes_skills_and_records_history() { + let agent = Agent::builder() + .name("tools") + .version("1.0.0") + .llm(ToolScriptLlm::new()) + .skill(add_skill()) + .build() + .unwrap(); + + let reply = agent + .chat_with_tools("s", "what is 2 + 40?", 5) + .await + .unwrap(); + assert_eq!(reply, "the answer is 42"); + + // Session: user, assistant(tool_calls), tool result, final assistant. + let history = agent.session_history("s").await.unwrap(); + assert_eq!(history.len(), 4); + assert!(history[1].tool_calls.is_some()); + assert!(matches!(history[2].role, Role::Tool)); + assert_eq!(history[3].content, "the answer is 42"); +} + +#[tokio::test] +async fn usage_totals_and_observer_fire() { + let seen = Arc::new(AtomicUsize::new(0)); + let seen_clone = Arc::clone(&seen); + + let agent = Agent::builder() + .name("metered") + .version("1.0.0") + .llm(ToolScriptLlm::new()) + .skill(add_skill()) + .usage_observer(move |event: &UsageEvent| { + assert_eq!(event.model, "scripted"); + seen_clone.fetch_add(1, Ordering::SeqCst); + }) + .build() + .unwrap(); + + agent.chat_with_tools("s", "sum please", 5).await.unwrap(); + + // Two completions happened (tool round + final answer). + assert_eq!(seen.load(Ordering::SeqCst), 2); + let totals = agent.usage(); + assert_eq!(totals.requests, 2); + assert_eq!(totals.total_tokens, 30); + assert_eq!(totals.prompt_tokens, 20); +} + +#[tokio::test] +async fn allowlist_blocks_unlisted_skills() { + let agent = Agent::builder() + .name("locked") + .version("1.0.0") + .skill(add_skill()) + .skill(FnSkill::new("noop", "does nothing", |_| async move { + Ok(json!({})) + })) + .skill_policy(SkillPolicy::new().allow_only(["noop"])) + .build() + .unwrap(); + + agent.execute_skill("noop", json!({})).await.unwrap(); + let err = agent.execute_skill("add", json!({})).await.unwrap_err(); + assert!(matches!(err, Error::PermissionDenied(_)), "got: {err}"); +} + +#[tokio::test] +async fn required_permissions_are_enforced() { + let net_skill = || { + FnSkill::new("fetch", "Fetches a URL", |_| async move { Ok(json!({})) }) + .with_permissions(["net"]) + }; + + // Not granted → refused. + let denied = Agent::builder() + .name("sandboxed") + .version("1.0.0") + .skill(net_skill()) + .build() + .unwrap(); + let err = denied.execute_skill("fetch", json!({})).await.unwrap_err(); + assert!(matches!(err, Error::PermissionDenied(_))); + + // Granted → runs. + let granted = Agent::builder() + .name("sandboxed") + .version("1.0.0") + .skill(net_skill()) + .skill_policy(SkillPolicy::new().grant("net")) + .build() + .unwrap(); + granted.execute_skill("fetch", json!({})).await.unwrap(); +} + +#[tokio::test] +async fn skill_timeout_fires() { + let agent = Agent::builder() + .name("slowpoke") + .version("1.0.0") + .skill(FnSkill::new("sleepy", "sleeps forever", |_| async move { + tokio::time::sleep(Duration::from_secs(3600)).await; + Ok(json!({})) + })) + .skill_policy(SkillPolicy::new().with_timeout(Duration::from_millis(50))) + .build() + .unwrap(); + + let err = agent.execute_skill("sleepy", json!({})).await.unwrap_err(); + assert!( + matches!(&err, Error::Skill(msg) if msg.contains("timed out")), + "got: {err}" + ); +} + +#[tokio::test] +async fn panicking_skill_does_not_take_down_the_agent() { + let agent = Agent::builder() + .name("resilient") + .version("1.0.0") + .skill(FnSkill::new("boom", "panics", |_| async move { + panic!("kaboom"); + })) + .skill(add_skill()) + .build() + .unwrap(); + + let err = agent.execute_skill("boom", json!({})).await.unwrap_err(); + assert!( + matches!(&err, Error::Skill(msg) if msg.contains("panicked")), + "got: {err}" + ); + + // The agent still works afterwards. + let out = agent + .execute_skill("add", json!({ "a": 1, "b": 1 })) + .await + .unwrap(); + assert_eq!(out["sum"], 2); +} + +#[tokio::test] +async fn tool_loop_gives_up_after_max_rounds() { + /// Always demands another tool call — never converges. + struct LoopingLlm; + + #[async_trait::async_trait] + impl LlmProvider for LoopingLlm { + async fn chat(&self, _request: ChatRequest) -> Result { + Ok(ChatResponse { + content: String::new(), + model: "looper".into(), + tool_calls: vec![ToolCall { + id: "again".into(), + name: "add".into(), + arguments: json!({ "a": 1, "b": 1 }), + }], + usage: None, + }) + } + async fn chat_stream( + &self, + _request: ChatRequest, + ) -> Result>> { + unimplemented!() + } + fn default_model(&self) -> &str { + "looper" + } + } + + let agent = Agent::builder() + .name("bounded") + .version("1.0.0") + .llm(LoopingLlm) + .skill(add_skill()) + .build() + .unwrap(); + + let err = agent.chat_with_tools("s", "loop", 3).await.unwrap_err(); + assert!( + matches!(&err, Error::Llm(msg) if msg.contains("3 rounds")), + "got: {err}" + ); +}